Merge branch 'vertex-imagen' into ah/vertex-imagen-preview-2

andrewheard · andrewheard · commit 3e7a3b83ee89 · 2025-01-24T17:26:09.000-05:00
# Conflicts:
#	FirebaseVertexAI/Sources/VertexAI.swift
diff --git a/FirebaseVertexAI/Sources/Types/Internal/Imagen/ImagenImageRepresentable.swift b/FirebaseVertexAI/Sources/Types/Internal/Imagen/ImagenImageRepresentable.swift
@@ -14,7 +14,12 @@
 
 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, *)
-public protocol ImagenImageRepresentable {
-  var _imagenImage: _ImagenImage { get }
+protocol ImagenImageRepresentable {
+  /// Internal representation of the image for use with the Imagen model.
+  ///
+  /// - Important: Not needed by SDK users.
+  var _internalImagenImage: _InternalImagenImage { get }
 }
diff --git a/FirebaseVertexAI/Sources/Types/Internal/Imagen/InternalImagenImage.swift b/FirebaseVertexAI/Sources/Types/Internal/Imagen/InternalImagenImage.swift
@@ -14,8 +14,15 @@
 
 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, *)
-public struct _ImagenImage {
+struct _InternalImagenImage {
   let mimeType: String
   let bytesBase64Encoded: String?
   let gcsURI: String?
diff --git a/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenFileDataImage.swift b/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenFileDataImage.swift
@@ -27,8 +27,10 @@ public struct ImagenFileDataImage {
 
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 extension ImagenFileDataImage: ImagenImageRepresentable {
-  public var _imagenImage: _ImagenImage {
-    _ImagenImage(mimeType: mimeType, bytesBase64Encoded: nil, gcsURI: gcsURI)
+  // TODO(andrewheard): Make this public when the SDK supports Imagen operations that take images as
+  // input (upscaling / editing).
+  var _internalImagenImage: _InternalImagenImage {
+    _InternalImagenImage(mimeType: mimeType, bytesBase64Encoded: nil, gcsURI: gcsURI)
   }
 }
 
diff --git a/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenGenerationConfig.swift b/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenGenerationConfig.swift
@@ -16,12 +16,17 @@
 public struct ImagenGenerationConfig {
   public var numberOfImages: Int?
   public var negativePrompt: String?
+  public var imageFormat: ImagenImageFormat?
   public var aspectRatio: ImagenAspectRatio?
+  public var addWatermark: Bool?
 
   public init(numberOfImages: Int? = nil, negativePrompt: String? = nil,
-              aspectRatio: ImagenAspectRatio? = nil) {
+              imageFormat: ImagenImageFormat? = nil, aspectRatio: ImagenAspectRatio? = nil,
+              addWatermark: Bool? = nil) {
     self.numberOfImages = numberOfImages
     self.negativePrompt = negativePrompt
+    self.imageFormat = imageFormat
     self.aspectRatio = aspectRatio
+    self.addWatermark = addWatermark
   }
 }
diff --git a/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenGenerationResponse.swift b/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenGenerationResponse.swift
@@ -15,7 +15,7 @@
 import Foundation
 
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
-public struct ImagenGenerationResponse<T: ImagenImageRepresentable> {
+public struct ImagenGenerationResponse<T> {
   public let images: [T]
   public let filteredReason: String?
 }
diff --git a/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenInlineDataImage.swift b/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenInlineDataImage.swift
@@ -31,8 +31,10 @@ public struct ImagenInlineDataImage {
 
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 extension ImagenInlineDataImage: ImagenImageRepresentable {
-  public var _imagenImage: _ImagenImage {
-    _ImagenImage(
+  // TODO(andrewheard): Make this public when the SDK supports Imagen operations that take images as
+  // input (upscaling / editing).
+  var _internalImagenImage: _InternalImagenImage {
+    _InternalImagenImage(
       mimeType: mimeType,
       bytesBase64Encoded: data.base64EncodedString(),
       gcsURI: nil
diff --git a/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenModel.swift b/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenModel.swift
@@ -24,8 +24,6 @@ public final class ImagenModel {
   /// The backing service responsible for sending and receiving model requests to the backend.
   let generativeAIService: GenerativeAIService
 
-  let modelConfig: ImagenModelConfig?
-
   let safetySettings: ImagenSafetySettings?
 
   /// Configuration parameters for sending requests to the backend.
@@ -34,7 +32,6 @@ public final class ImagenModel {
   init(name: String,
        projectID: String,
        apiKey: String,
-       modelConfig: ImagenModelConfig?,
        safetySettings: ImagenSafetySettings?,
        requestOptions: RequestOptions,
        appCheck: AppCheckInterop?,
@@ -48,7 +45,6 @@ public final class ImagenModel {
       auth: auth,
       urlSession: urlSession
     )
-    self.modelConfig = modelConfig
     self.safetySettings = safetySettings
     self.requestOptions = requestOptions
   }
@@ -61,7 +57,6 @@ public final class ImagenModel {
       parameters: ImagenModel.imageGenerationParameters(
         storageURI: nil,
         generationConfig: generationConfig,
-        modelConfig: modelConfig,
         safetySettings: safetySettings
       )
     )
@@ -75,15 +70,14 @@ public final class ImagenModel {
       parameters: ImagenModel.imageGenerationParameters(
         storageURI: storageURI,
         generationConfig: generationConfig,
-        modelConfig: modelConfig,
         safetySettings: safetySettings
       )
     )
   }
 
-  func generateImages<T: Decodable>(prompt: String,
-                                    parameters: ImageGenerationParameters) async throws
-    -> ImagenGenerationResponse<T> {
+  func generateImages<T>(prompt: String,
+                         parameters: ImageGenerationParameters) async throws
+    -> ImagenGenerationResponse<T> where T: Decodable, T: ImagenImageRepresentable {
     let request = ImagenGenerationRequest<T>(
       model: modelResourceName,
       options: requestOptions,
@@ -96,7 +90,6 @@ public final class ImagenModel {
 
   static func imageGenerationParameters(storageURI: String?,
                                         generationConfig: ImagenGenerationConfig?,
-                                        modelConfig: ImagenModelConfig?,
                                         safetySettings: ImagenSafetySettings?)
     -> ImageGenerationParameters {
     return ImageGenerationParameters(
@@ -106,13 +99,13 @@ public final class ImagenModel {
       aspectRatio: generationConfig?.aspectRatio?.rawValue,
       safetyFilterLevel: safetySettings?.safetyFilterLevel?.rawValue,
       personGeneration: safetySettings?.personFilterLevel?.rawValue,
-      outputOptions: modelConfig?.imageFormat.map {
+      outputOptions: generationConfig?.imageFormat.map {
         ImageGenerationOutputOptions(
           mimeType: $0.mimeType,
           compressionQuality: $0.compressionQuality
         )
       },
-      addWatermark: modelConfig?.addWatermark,
+      addWatermark: generationConfig?.addWatermark,
       includeResponsibleAIFilterReason: true
     )
   }
diff --git a/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenModelConfig.swift b/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenModelConfig.swift
diff --git a/FirebaseVertexAI/Sources/VertexAI.swift b/FirebaseVertexAI/Sources/VertexAI.swift
@@ -113,7 +113,6 @@ public class VertexAI {
       name: modelResourceName(modelName: modelName),
       projectID: projectID,
       apiKey: apiKey,
-      modelConfig: modelConfig,
       safetySettings: safetySettings,
       requestOptions: requestOptions,
       appCheck: appCheck,
diff --git a/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift b/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift
@@ -64,7 +64,6 @@ final class IntegrationTests: XCTestCase {
     )
     imagenModel = vertex.imagenModel(
       modelName: "imagen-3.0-fast-generate-001",
-      modelConfig: ImagenModelConfig(imageFormat: .jpeg(compressionQuality: 70)),
       safetySettings: ImagenSafetySettings(
         safetyFilterLevel: .blockLowAndAbove,
         personFilterLevel: .blockAll
@@ -255,7 +254,10 @@ final class IntegrationTests: XCTestCase {
     overlooking a vast African savanna at sunset. Golden hour light, long shadows, sharp focus on
     the lion, shallow depth of field, detailed fur texture, DSLR, 85mm lens.
     """
-    let generationConfig = ImagenGenerationConfig(aspectRatio: .landscape16x9)
+    let generationConfig = ImagenGenerationConfig(
+      imageFormat: .jpeg(compressionQuality: 70),
+      aspectRatio: .landscape16x9
+    )
 
     let response = try await imagenModel.generateImages(
       prompt: imagePrompt,
diff --git a/FirebaseVertexAI/Tests/Unit/Types/Imagen/ImageGenerationParametersTests.swift b/FirebaseVertexAI/Tests/Unit/Types/Imagen/ImageGenerationParametersTests.swift
@@ -40,7 +40,6 @@ final class ImageGenerationParametersTests: XCTestCase {
     let parameters = ImagenModel.imageGenerationParameters(
       storageURI: nil,
       generationConfig: nil,
-      modelConfig: nil,
       safetySettings: nil
     )
 
@@ -64,37 +63,6 @@ final class ImageGenerationParametersTests: XCTestCase {
     let parameters = ImagenModel.imageGenerationParameters(
       storageURI: storageURI,
       generationConfig: nil,
-      modelConfig: nil,
-      safetySettings: nil
-    )
-
-    XCTAssertEqual(parameters, expectedParameters)
-  }
-
-  func testParameters_includeModelConfig() throws {
-    let compressionQuality = 80
-    let imageFormat = ImagenImageFormat.jpeg(compressionQuality: compressionQuality)
-    let addWatermark = true
-    let modelConfig = ImagenModelConfig(imageFormat: imageFormat, addWatermark: addWatermark)
-    let expectedParameters = ImageGenerationParameters(
-      sampleCount: 1,
-      storageURI: nil,
-      negativePrompt: nil,
-      aspectRatio: nil,
-      safetyFilterLevel: nil,
-      personGeneration: nil,
-      outputOptions: ImageGenerationOutputOptions(
-        mimeType: imageFormat.mimeType,
-        compressionQuality: imageFormat.compressionQuality
-      ),
-      addWatermark: addWatermark,
-      includeResponsibleAIFilterReason: true
-    )
-
-    let parameters = ImagenModel.imageGenerationParameters(
-      storageURI: nil,
-      generationConfig: nil,
-      modelConfig: modelConfig,
       safetySettings: nil
     )
 
@@ -104,11 +72,16 @@ final class ImageGenerationParametersTests: XCTestCase {
   func testParameters_includeGenerationConfig() throws {
     let sampleCount = 2
     let negativePrompt = "test-negative-prompt"
+    let compressionQuality = 80
+    let imageFormat = ImagenImageFormat.jpeg(compressionQuality: compressionQuality)
     let aspectRatio = ImagenAspectRatio.landscape16x9
+    let addWatermark = true
     let generationConfig = ImagenGenerationConfig(
       numberOfImages: sampleCount,
       negativePrompt: negativePrompt,
-      aspectRatio: aspectRatio
+      imageFormat: imageFormat,
+      aspectRatio: aspectRatio,
+      addWatermark: addWatermark
     )
     let expectedParameters = ImageGenerationParameters(
       sampleCount: sampleCount,
@@ -117,15 +90,17 @@ final class ImageGenerationParametersTests: XCTestCase {
       aspectRatio: aspectRatio.rawValue,
       safetyFilterLevel: nil,
       personGeneration: nil,
-      outputOptions: nil,
-      addWatermark: nil,
+      outputOptions: ImageGenerationOutputOptions(
+        mimeType: imageFormat.mimeType,
+        compressionQuality: imageFormat.compressionQuality
+      ),
+      addWatermark: addWatermark,
       includeResponsibleAIFilterReason: true
     )
 
     let parameters = ImagenModel.imageGenerationParameters(
       storageURI: nil,
       generationConfig: generationConfig,
-      modelConfig: nil,
       safetySettings: nil
     )
 
@@ -155,7 +130,6 @@ final class ImageGenerationParametersTests: XCTestCase {
     let parameters = ImagenModel.imageGenerationParameters(
       storageURI: nil,
       generationConfig: nil,
-      modelConfig: nil,
       safetySettings: safetySettings
     )
 
@@ -168,15 +142,16 @@ final class ImageGenerationParametersTests: XCTestCase {
     let storageURI = "gs://test-bucket/path"
     let sampleCount = 4
     let negativePrompt = "test-negative-prompt"
+    let imageFormat = ImagenImageFormat.png()
     let aspectRatio = ImagenAspectRatio.portrait3x4
+    let addWatermark = false
     let generationConfig = ImagenGenerationConfig(
       numberOfImages: sampleCount,
       negativePrompt: negativePrompt,
-      aspectRatio: aspectRatio
+      imageFormat: imageFormat,
+      aspectRatio: aspectRatio,
+      addWatermark: addWatermark
     )
-    let imageFormat = ImagenImageFormat.png()
-    let addWatermark = false
-    let modelConfig = ImagenModelConfig(imageFormat: imageFormat, addWatermark: addWatermark)
     let safetyFilterLevel = ImagenSafetyFilterLevel.blockNone
     let personFilterLevel = ImagenPersonFilterLevel.blockAll
     let safetySettings = ImagenSafetySettings(
@@ -201,7 +176,6 @@ final class ImageGenerationParametersTests: XCTestCase {
     let parameters = ImagenModel.imageGenerationParameters(
       storageURI: storageURI,
       generationConfig: generationConfig,
-      modelConfig: modelConfig,
       safetySettings: safetySettings
     )
 
diff --git a/FirebaseVertexAI/Tests/Unit/Types/Imagen/ImagenFileDataImageTests.swift b/FirebaseVertexAI/Tests/Unit/Types/Imagen/ImagenFileDataImageTests.swift
@@ -35,9 +35,9 @@ final class ImagenFileDataImageTests: XCTestCase {
 
     XCTAssertEqual(image.mimeType, mimeType)
     XCTAssertEqual(image.gcsURI, gcsURI)
-    XCTAssertEqual(image._imagenImage.mimeType, mimeType)
-    XCTAssertEqual(image._imagenImage.gcsURI, gcsURI)
-    XCTAssertNil(image._imagenImage.bytesBase64Encoded)
+    XCTAssertEqual(image._internalImagenImage.mimeType, mimeType)
+    XCTAssertEqual(image._internalImagenImage.gcsURI, gcsURI)
+    XCTAssertNil(image._internalImagenImage.bytesBase64Encoded)
   }
 
   func testDecodeImage_missingGCSURI_throws() throws {
diff --git a/FirebaseVertexAI/Tests/Unit/Types/Imagen/ImagenInlineDataImageTests.swift b/FirebaseVertexAI/Tests/Unit/Types/Imagen/ImagenInlineDataImageTests.swift
@@ -35,9 +35,9 @@ final class ImagenInlineDataImageTests: XCTestCase {
 
     XCTAssertEqual(image.mimeType, mimeType)
     XCTAssertEqual(image.data.base64EncodedString(), bytesBase64Encoded)
-    XCTAssertEqual(image._imagenImage.mimeType, mimeType)
-    XCTAssertEqual(image._imagenImage.bytesBase64Encoded, bytesBase64Encoded)
-    XCTAssertNil(image._imagenImage.gcsURI)
+    XCTAssertEqual(image._internalImagenImage.mimeType, mimeType)
+    XCTAssertEqual(image._internalImagenImage.bytesBase64Encoded, bytesBase64Encoded)
+    XCTAssertNil(image._internalImagenImage.gcsURI)
   }
 
   func testDecodeImage_missingBytesBase64Encoded_throws() throws {

Original file line number	Diff line number	Diff line change
`@@ -27,8 +27,10 @@ public struct ImagenFileDataImage {`
`27`	`27`
`28`	`28`	`@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)`
`29`	`29`	`extension ImagenFileDataImage: ImagenImageRepresentable {`
`30`		`- public var _imagenImage: _ImagenImage {`
`31`		`- _ImagenImage(mimeType: mimeType, bytesBase64Encoded: nil, gcsURI: gcsURI)`
	`30`	`+ // TODO(andrewheard): Make this public when the SDK supports Imagen operations that take images as`
	`31`	`+ // input (upscaling / editing).`
	`32`	`+ var _internalImagenImage: _InternalImagenImage {`
	`33`	`+ _InternalImagenImage(mimeType: mimeType, bytesBase64Encoded: nil, gcsURI: gcsURI)`
`32`	`34`	`}`
`33`	`35`	`}`
`34`	`36`
Original file line number	Diff line number	Diff line change
`@@ -15,7 +15,7 @@`
`15`	`15`	`import Foundation`
`16`	`16`
`17`	`17`	`@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)`
`18`		`-public struct ImagenGenerationResponse<T: ImagenImageRepresentable> {`
	`18`	`+public struct ImagenGenerationResponse<T> {`
`19`	`19`	`public let images: [T]`
`20`	`20`	`public let filteredReason: String?`
`21`	`21`	`}`