Add ImagenModel

Author: andrewheard

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/Public/Imagen/ImagenModel.swift

@@ -0,0 +1,78 @@
+import FirebaseAppCheckInterop
+import FirebaseAuthInterop
+import Foundation
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+public final class ImagenModel {
+  /// The resource name of the model in the backend; has the format "models/model-name".
+  let modelResourceName: String
+
+  /// The backing service responsible for sending and receiving model requests to the backend.
+  let generativeAIService: GenerativeAIService
+
+  /// Configuration parameters for sending requests to the backend.
+  let requestOptions: RequestOptions
+
+  init(name: String,
+       projectID: String,
+       apiKey: String,
+       requestOptions: RequestOptions,
+       appCheck: AppCheckInterop?,
+       auth: AuthInterop?,
+       urlSession: URLSession = .shared) {
+    modelResourceName = name
+    generativeAIService = GenerativeAIService(
+      projectID: projectID,
+      apiKey: apiKey,
+      appCheck: appCheck,
+      auth: auth,
+      urlSession: urlSession
+    )
+    self.requestOptions = requestOptions
+  }
+
+  public func generateImages(prompt: String) async throws
+    -> ImageGenerationResponse<ImagenInlineDataImage> {
+    return try await generateImages(
+      prompt: prompt,
+      parameters: imageGenerationParameters(storageURI: nil)
+    )
+  }
+
+  public func generateImages(prompt: String, storageURI: String) async throws
+    -> ImageGenerationResponse<ImagenFileDataImage> {
+    return try await generateImages(
+      prompt: prompt,
+      parameters: imageGenerationParameters(storageURI: storageURI)
+    )
+  }
+
+  func generateImages<T: Decodable>(prompt: String,
+                                    parameters: ImageGenerationParameters) async throws
+    -> ImageGenerationResponse<T> {
+    let request = ImageGenerationRequest<T>(
+      model: modelResourceName,
+      options: requestOptions,
+      instances: [ImageGenerationInstance(prompt: prompt)],
+      parameters: parameters
+    )
+
+    return try await generativeAIService.loadRequest(request: request)
+  }
+
+  func imageGenerationParameters(storageURI: String?) -> ImageGenerationParameters {
+    // TODO(#14221): Add support for configuring these parameters.
+    return ImageGenerationParameters(
+      sampleCount: 1,
+      storageURI: storageURI,
+      seed: nil,
+      negativePrompt: nil,
+      aspectRatio: nil,
+      safetyFilterLevel: nil,
+      personGeneration: nil,
+      outputOptions: nil,
+      addWatermark: nil,
+      includeResponsibleAIFilterReason: true
+    )
+  }
+}

FirebaseVertexAI/Sources/VertexAI.swift

@@ -104,6 +104,18 @@ public class VertexAI {
     )
   }
 
+  public func imagenModel(modelName: String, requestOptions: RequestOptions = RequestOptions())
+    -> ImagenModel {
+    return ImagenModel(
+      name: modelResourceName(modelName: modelName),
+      projectID: projectID,
+      apiKey: apiKey,
+      requestOptions: requestOptions,
+      appCheck: appCheck,
+      auth: auth
+    )
+  }
+
   /// Class to enable VertexAI to register via the Objective-C based Firebase component system
   /// to include VertexAI in the userAgent.
   @objc(FIRVertexAIComponent) class FirebaseVertexAIComponent: NSObject {}