[Vertex AI] Add image generation support using Imagen

.github/workflows/vertexai.yml

@@ -9,6 +9,7 @@ on:
   schedule:
     # Run every day at 11pm (PST) - cron uses UTC times
     - cron:  '0 7 * * *'
+  workflow_dispatch:
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
@@ -102,6 +103,7 @@ jobs:
     needs: spm-package-resolved
     env:
       TEST_RUNNER_FIRAAppCheckDebugToken: ${{ secrets.VERTEXAI_INTEGRATION_FAC_DEBUG_TOKEN }}
+      TEST_RUNNER_VTXIntegrationImagen: ${{ github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' }}
       FIREBASECI_USE_LATEST_GOOGLEAPPMEASUREMENT: 1
       secrets_passphrase: ${{ secrets.GHASecretsGPGPassphrase1 }}
     steps:

FirebaseVertexAI/CHANGELOG.md

@@ -1,3 +1,11 @@
+# Unreleased
+- [feature] **Public Preview**: Added support for generating images using the
+  Imagen 3 model.
+  <br /><br />
+  Note: This feature is in Public Preview, which means that the it is not
+  subject to any SLA or deprecation policy and could change in
+  backwards-incompatible ways.
+
 # 11.6.0
 - [changed] The token counts from `GenerativeModel.countTokens(...)` now include
   tokens from the schema for JSON output and function calling; reported token

FirebaseVertexAI/Sources/GenerationConfig.swift

@@ -162,4 +162,17 @@ public struct GenerationConfig {
 // MARK: - Codable Conformances
 
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
-extension GenerationConfig: Encodable {}
+extension GenerationConfig: Encodable {
+  enum CodingKeys: String, CodingKey {
+    case temperature
+    case topP
+    case topK
+    case candidateCount
+    case maxOutputTokens
+    case presencePenalty
+    case frequencyPenalty
+    case stopSequences
+    case responseMIMEType = "responseMimeType"
+    case responseSchema
+  }
+}

FirebaseVertexAI/Sources/GenerativeAIRequest.swift

@@ -52,3 +52,6 @@ public struct RequestOptions {
     self.init(timeout: timeout, apiVersion: .v1beta)
   }
 }
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+extension RequestOptions: Equatable {}

FirebaseVertexAI/Sources/GenerativeAIService.swift

@@ -203,7 +203,6 @@ struct GenerativeAIService {
     }
 
     let encoder = JSONEncoder()
-    encoder.keyEncodingStrategy = .convertToSnakeCase
     urlRequest.httpBody = try encoder.encode(request)
     urlRequest.timeoutInterval = request.options.timeout
 

FirebaseVertexAI/Sources/Types/Internal/Imagen/ImageGenerationInstance.swift

@@ -0,0 +1,26 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+struct ImageGenerationInstance {
+  let prompt: String
+}
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+extension ImageGenerationInstance: Equatable {}
+
+// MARK: - Codable Conformance
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+extension ImageGenerationInstance: Encodable {}

FirebaseVertexAI/Sources/Types/Internal/Imagen/ImageGenerationOutputOptions.swift

@@ -0,0 +1,29 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import Foundation
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+struct ImageGenerationOutputOptions {
+  let mimeType: String
+  let compressionQuality: Int?
+}
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+extension ImageGenerationOutputOptions: Equatable {}
+
+// MARK: - Codable Conformance
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+extension ImageGenerationOutputOptions: Encodable {}

FirebaseVertexAI/Sources/Types/Internal/Imagen/ImageGenerationParameters.swift

@@ -0,0 +1,62 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+struct ImageGenerationParameters {
+  let sampleCount: Int?
+  let storageURI: String?
+  let negativePrompt: String?
+  let aspectRatio: String?
+  let safetyFilterLevel: String?
+  let personGeneration: String?
+  let outputOptions: ImageGenerationOutputOptions?
+  let addWatermark: Bool?
+  let includeResponsibleAIFilterReason: Bool?
+}
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+extension ImageGenerationParameters: Equatable {}
+
+// MARK: - Codable Conformance
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+extension ImageGenerationParameters: Encodable {
+  enum CodingKeys: String, CodingKey {
+    case sampleCount
+    case storageURI = "storageUri"
+    case negativePrompt
+    case aspectRatio
+    case safetyFilterLevel = "safetySetting"
+    case personGeneration
+    case outputOptions
+    case addWatermark
+    case includeResponsibleAIFilterReason = "includeRaiReason"
+  }
+
+  func encode(to encoder: any Encoder) throws {
+    var container = encoder.container(keyedBy: CodingKeys.self)
+    try container.encodeIfPresent(sampleCount, forKey: .sampleCount)
+    try container.encodeIfPresent(storageURI, forKey: .storageURI)
+    try container.encodeIfPresent(negativePrompt, forKey: .negativePrompt)
+    try container.encodeIfPresent(aspectRatio, forKey: .aspectRatio)
+    try container.encodeIfPresent(safetyFilterLevel, forKey: .safetyFilterLevel)
+    try container.encodeIfPresent(personGeneration, forKey: .personGeneration)
+    try container.encodeIfPresent(outputOptions, forKey: .outputOptions)
+    try container.encodeIfPresent(addWatermark, forKey: .addWatermark)
+    try container.encodeIfPresent(
+      includeResponsibleAIFilterReason,
+      forKey: .includeResponsibleAIFilterReason
+    )
+  }
+}

FirebaseVertexAI/Sources/Types/Internal/Imagen/ImagenGenerationRequest.swift

@@ -0,0 +1,54 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import Foundation
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+struct ImagenGenerationRequest<ImageType: ImagenImageRepresentable> {
+  let model: String
+  let options: RequestOptions
+  let instances: [ImageGenerationInstance]
+  let parameters: ImageGenerationParameters
+
+  init(model: String, options: RequestOptions, instances: [ImageGenerationInstance],
+       parameters: ImageGenerationParameters) {
+    self.model = model
+    self.options = options
+    self.instances = instances
+    self.parameters = parameters
+  }
+}
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+extension ImagenGenerationRequest: GenerativeAIRequest where ImageType: Decodable {
+  typealias Response = ImagenGenerationResponse<ImageType>
+
+  var url: URL {
+    return URL(string: "\(Constants.baseURL)/\(options.apiVersion)/\(model):predict")!
+  }
+}
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+extension ImagenGenerationRequest: Encodable {
+  enum CodingKeys: CodingKey {
+    case instances
+    case parameters
+  }
+
+  func encode(to encoder: any Encoder) throws {
+    var container = encoder.container(keyedBy: CodingKeys.self)
+    try container.encode(instances, forKey: .instances)
+    try container.encode(parameters, forKey: .parameters)
+  }
+}

FirebaseVertexAI/Sources/Types/Internal/Imagen/ImagenImageRepresentable.swift

@@ -0,0 +1,25 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import Foundation
+
+// TODO(andrewheard): Make this public when the SDK supports Imagen operations that take images as
+// input (upscaling / editing).
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+protocol ImagenImageRepresentable {
+  /// Internal representation of the image for use with the Imagen model.
+  ///
+  /// - Important: Not needed by SDK users.
+  var _internalImagenImage: _InternalImagenImage { get }
+}

FirebaseVertexAI/Sources/Types/Internal/Imagen/InternalImagenImage.swift

@@ -0,0 +1,29 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import Foundation
+
+/// Internal representation of an image for the Imagen model.
+///
+/// - Important: For internal use by types conforming to ``ImagenImageRepresentable``; all
+/// properties are `internal` and are not needed by SDK users.
+///
+/// TODO(andrewheard): Make this public when the SDK supports Imagen operations that take images as
+/// input (upscaling / editing).
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+struct _InternalImagenImage {
+  let mimeType: String
+  let bytesBase64Encoded: String?
+  let gcsURI: String?
+}

FirebaseVertexAI/Sources/Types/Internal/Imagen/RAIFilteredReason.swift

@@ -0,0 +1,25 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+struct RAIFilteredReason {
+  let raiFilteredReason: String
+}
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+extension RAIFilteredReason: Decodable {
+  enum CodingKeys: CodingKey {
+    case raiFilteredReason
+  }
+}

FirebaseVertexAI/Sources/Types/Internal/InternalPart.swift

@@ -34,6 +34,11 @@ struct FileData: Codable, Equatable, Sendable {
     self.fileURI = fileURI
     self.mimeType = mimeType
   }
+
+  enum CodingKeys: String, CodingKey {
+    case fileURI = "fileUri"
+    case mimeType
+  }
 }
 
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)