[Vertex AI] Make ImagenImageRepresentable internal

firebase · Jan 13, 2025 · ff270f2 · ff270f2
1 parent a3a026a
commit ff270f2
Show file tree

Hide file tree

Showing 8 changed files with 33 additions and 17 deletions.
diff --git a/...lic/Imagen/ImagenImageRepresentable.swift → ...nal/Imagen/ImagenImageRepresentable.swift b/...lic/Imagen/ImagenImageRepresentable.swift → ...nal/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/...ces/Types/Public/Imagen/ImagenImage.swift → ...Internal/Imagen/InternalImagenImage.swift b/...ces/Types/Public/Imagen/ImagenImage.swift → ...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/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
@@ -81,9 +81,9 @@ public final class ImagenModel {
     )
   }
 
-  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,

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 {