Merge branch 'davidmotson.imagen_support' into davidmotson.imagen_docs

Author: David Motsonashvili

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/FirebaseVertexAI.kt

@@ -24,8 +24,8 @@ import com.google.firebase.auth.internal.InternalAuthProvider
 import com.google.firebase.inject.Provider
 import com.google.firebase.vertexai.type.Content
 import com.google.firebase.vertexai.type.GenerationConfig
-import com.google.firebase.vertexai.type.ImageSafetySettings
-import com.google.firebase.vertexai.type.ImagenModelConfig
+import com.google.firebase.vertexai.type.ImagenGenerationConfig
+import com.google.firebase.vertexai.type.ImagenSafetySettings
 import com.google.firebase.vertexai.type.InvalidLocationException
 import com.google.firebase.vertexai.type.RequestOptions
 import com.google.firebase.vertexai.type.SafetySetting
@@ -93,8 +93,8 @@ internal constructor(
   @JvmOverloads
   public fun imageModel(
     modelName: String,
-    generationConfig: ImagenModelConfig? = null,
-    safetySettings: ImageSafetySettings? = null,
+    generationConfig: ImagenGenerationConfig? = null,
+    safetySettings: ImagenSafetySettings? = null,
     requestOptions: RequestOptions = RequestOptions(),
   ): ImageModel {
     if (location.trim().isEmpty() || location.contains("/")) {

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/ImageModel.kt

@@ -28,13 +28,11 @@ import com.google.firebase.vertexai.internal.util.toInternal
 import com.google.firebase.vertexai.internal.util.toPublicGCS
 import com.google.firebase.vertexai.internal.util.toPublicInline
 import com.google.firebase.vertexai.type.FirebaseVertexAIException
-import com.google.firebase.vertexai.type.ImageSafetySettings
+import com.google.firebase.vertexai.type.ImagenSafetySettings
 import com.google.firebase.vertexai.type.ImagenGCSImage
 import com.google.firebase.vertexai.type.ImagenGenerationConfig
 import com.google.firebase.vertexai.type.ImagenGenerationResponse
-import com.google.firebase.vertexai.type.ImagenImageRepresentible
 import com.google.firebase.vertexai.type.ImagenInlineImage
-import com.google.firebase.vertexai.type.ImagenModelConfig
 import com.google.firebase.vertexai.type.PromptBlockedException
 import com.google.firebase.vertexai.type.RequestOptions
 import kotlin.time.Duration
@@ -48,16 +46,16 @@ import kotlinx.coroutines.tasks.await
 public class ImageModel
 internal constructor(
   private val modelName: String,
-  private val generationConfig: ImagenModelConfig? = null,
-  private val safetySettings: ImageSafetySettings? = null,
+  private val generationConfig: ImagenGenerationConfig? = null,
+  private val safetySettings: ImagenSafetySettings? = null,
   private val controller: APIController,
 ) {
   @JvmOverloads
   internal constructor(
     modelName: String,
     apiKey: String,
-    generationConfig: ImagenModelConfig? = null,
-    safetySettings: ImageSafetySettings? = null,
+    generationConfig: ImagenGenerationConfig? = null,
+    safetySettings: ImagenSafetySettings? = null,
     requestOptions: RequestOptions = RequestOptions(),
     appCheckTokenProvider: InteropAppCheckTokenProvider? = null,
     internalAuthProvider: InternalAuthProvider? = null,
@@ -115,10 +113,9 @@ internal constructor(
    */
   public suspend fun generateImage(
     prompt: String,
-    config: ImagenGenerationConfig?,
   ): ImagenGenerationResponse<ImagenInlineImage> =
     try {
-      controller.generateImage(constructRequest(prompt, null, config)).toPublicInline().validate()
+      controller.generateImage(constructRequest(prompt, null, generationConfig)).toPublicInline().validate()
     } catch (e: Throwable) {
       throw FirebaseVertexAIException.from(e)
     }
@@ -132,10 +129,9 @@ internal constructor(
   public suspend fun generateImage(
     prompt: String,
     gcsUri: String,
-    config: ImagenGenerationConfig?,
   ): ImagenGenerationResponse<ImagenGCSImage> =
     try {
-      controller.generateImage(constructRequest(prompt, gcsUri, config)).toPublicGCS().validate()
+      controller.generateImage(constructRequest(prompt, gcsUri, generationConfig)).toPublicGCS().validate()
     } catch (e: Throwable) {
       throw FirebaseVertexAIException.from(e)
     }
@@ -168,7 +164,7 @@ internal constructor(
   }
 }
 
-private fun <T : ImagenImageRepresentible> ImagenGenerationResponse<T>.validate():
+private fun <T> ImagenGenerationResponse<T>.validate():
   ImagenGenerationResponse<T> {
   if (images.isEmpty()) {
     throw PromptBlockedException(message = filteredReason ?: ImageModel.DEFAULT_FILTERED_ERROR)

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/internal/GenerateImageResponse.kt

@@ -23,7 +23,7 @@ internal data class GenerateImageResponse(val predictions: List<ImagenImageRespo
 
 @Serializable
 internal data class ImagenImageResponse(
-  val bytesBase64Encoded: String?,
-  val gcsUri: String?,
+  val bytesBase64Encoded: String? = null,
+  val gcsUri: String? = null,
   val mimeType: String,
 )

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/internal/util/conversions.kt

@@ -205,7 +205,7 @@ internal fun com.google.firebase.vertexai.type.Schema.toInternal(): Schema =
     items?.toInternal(),
   )
 
-internal fun com.google.firebase.vertexai.type.ImageFormat.toInternal(): ImageOutputOptions =
+internal fun com.google.firebase.vertexai.type.ImagenImageFormat.toInternal(): ImageOutputOptions =
   ImageOutputOptions(mimeType, compressionQuality)
 
 internal fun JSONObject.toInternal() = Json.decodeFromString<JsonObject>(toString())

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/ImagenGCSImage.kt

@@ -22,10 +22,4 @@ package com.google.firebase.vertexai.type
  * @param gcsUri Contains the gs:// uri for the image.
  * @param mimeType Contains the mime type of the image eg. "image/png"
  */
-public class ImagenGCSImage(public val gcsUri: String, public val mimeType: String) :
-  ImagenImageRepresentible {
-
-  override fun asImagenImage(): ImagenImage {
-    return ImagenImage(null, gcsUri, mimeType)
-  }
-}
+public class ImagenGCSImage(public val gcsUri: String, public val mimeType: String) {}

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/ImagenGenerationConfig.kt

@@ -22,9 +22,13 @@ package com.google.firebase.vertexai.type
  * generated images.
  * @param numberOfImages How many images should be generated.
  * @param aspectRatio The aspect ratio of the generated images.
+ * @param imageFormat The file format/compression of the generated images.
+ * @param addWatermark Adds an invisible watermark to mark the image as AI generated.
  */
 public class ImagenGenerationConfig(
-  public val negativePrompt: String? = null,
-  public val numberOfImages: Int = 1,
-  public val aspectRatio: ImagenAspectRatio? = null,
+    public val negativePrompt: String? = null,
+    public val numberOfImages: Int? = 1,
+    public val aspectRatio: ImagenAspectRatio? = null,
+    public val imageFormat: ImagenImageFormat? = null,
+    public val addWatermark: Boolean? = null
 ) {}

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/ImagenGenerationResponse.kt

@@ -23,7 +23,7 @@ package com.google.firebase.vertexai.type
  * @param filteredReason if fewer images were generated than were requested, this field will contain
  * the reason they were filtered out.
  */
-public class ImagenGenerationResponse<T : ImagenImageRepresentible>(
+public class ImagenGenerationResponse<T>(
   public val images: List<T>,
   public val filteredReason: String?,
 ) {}

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/ImagenImage.kt

@@ -24,12 +24,7 @@ package com.google.firebase.vertexai.type
  * @param mimeType Contains the mime type of the image eg. "image/png"
  */
 public class ImagenImage(
-  public val data: ByteArray?,
-  public val gcsUri: String?,
-  public val mimeType: String,
-) : ImagenImageRepresentible {
-
-  override fun asImagenImage(): ImagenImage {
-    return this
-  }
-}
+  internal val data: ByteArray?,
+  internal val gcsUri: String?,
+  internal val mimeType: String,
+) {}

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/ImageFormat.kt renamed to firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/ImagenImageFormat.kt

@@ -23,22 +23,22 @@ package com.google.firebase.vertexai.type
  * number meaning the image is permitted to be lower quality to reduce size. This parameter is not
  * relevant for every mimetype
  */
-public class ImageFormat
+public class ImagenImageFormat
 private constructor(public val mimeType: String, public val compressionQuality: Int?) {
 
   public companion object {
     /**
-     * An [ImageFormat] representing a JPEG image.
+     * An [ImagenImageFormat] representing a JPEG image.
      * @param compressionQuality an int (1-100) representing how the quality of the image, a lower
      * number meaning the image is permitted to be lower quality to reduce size.
      */
-    public fun jpeg(compressionQuality: Int?): ImageFormat {
-      return ImageFormat("image/jpeg", compressionQuality)
+    public fun jpeg(compressionQuality: Int? = null): ImagenImageFormat {
+      return ImagenImageFormat("image/jpeg", compressionQuality)
     }
 
-    /** An [ImageFormat] representing a PNG image */
-    public fun png(): ImageFormat {
-      return ImageFormat("image/png", null)
+    /** An [ImagenFormat] representing a PNG image */
+    public fun png(): ImagenImageFormat {
+      return ImagenImageFormat("image/png", null)
     }
   }
 }