add safety scores to image generation

Author: David Motsonashvili

firebase-ai/src/main/kotlin/com/google/firebase/ai/ImagenModel.kt

@@ -177,6 +177,7 @@ internal constructor(
       GenerateImageRequest.ImagenParameters(
         sampleCount = generationConfig?.numberOfImages ?: 1,
         includeRaiReason = true,
+        includeSafetyAttributes = true,
         addWatermark = generationConfig?.addWatermark,
         personGeneration = safetySettings?.personFilterLevel?.internalVal,
         negativePrompt = generationConfig?.negativePrompt,
@@ -206,6 +207,7 @@ internal constructor(
       GenerateImageRequest.ImagenParameters(
         sampleCount = generationConfig?.numberOfImages ?: 1,
         includeRaiReason = true,
+        includeSafetyAttributes = true,
         addWatermark = generationConfig?.addWatermark,
         personGeneration = safetySettings?.personFilterLevel?.internalVal,
         negativePrompt = generationConfig?.negativePrompt,

firebase-ai/src/main/kotlin/com/google/firebase/ai/common/Request.kt

@@ -93,6 +93,7 @@ internal data class GenerateImageRequest(
   internal data class ImagenParameters(
     val sampleCount: Int,
     val includeRaiReason: Boolean,
+    val includeSafetyAttributes: Boolean,
     val storageUri: String?,
     val negativePrompt: String?,
     val aspectRatio: String?,

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

@@ -42,7 +42,7 @@ internal constructor(public val images: List<T>, public val filteredReason: Stri
     internal fun toPublicInline() =
       ImagenGenerationResponse(
         images = predictions.filter { it.mimeType != null }.map { it.toPublicInline() },
-        null,
+          predictions.firstNotNullOfOrNull { it.raiFilteredReason },
       )
   }
 
@@ -52,10 +52,25 @@ internal constructor(public val images: List<T>, public val filteredReason: Stri
     val gcsUri: String? = null,
     val mimeType: String? = null,
     val raiFilteredReason: String? = null,
+    val safetyAttributes: ImagenSafetyAttributes? = null,
   ) {
     internal fun toPublicInline() =
-      ImagenInlineImage(Base64.decode(bytesBase64Encoded!!, Base64.NO_WRAP), mimeType!!)
+      ImagenInlineImage(
+        Base64.decode(bytesBase64Encoded!!, Base64.NO_WRAP),
+        mimeType!!,
+        safetyAttributes?.toPublic() ?: emptyMap()
+      )
 
     internal fun toPublicGCS() = ImagenGCSImage(gcsUri!!, mimeType!!)
   }
+
+  @Serializable
+  internal data class ImagenSafetyAttributes(
+    val categories: List<String>? = null,
+    val scores: List<Double>? = null
+  ){
+    internal fun toPublic(): Map<String, Double> {
+      return categories?.zip(scores!!)?.toMap() ?: emptyMap()
+    }
+  }
 }

firebase-ai/src/main/kotlin/com/google/firebase/ai/type/ImagenInlineImage.kt

@@ -28,10 +28,14 @@ import kotlinx.serialization.Serializable
  * @property data The raw image bytes in JPEG or PNG format, as specified by [mimeType].
  * @property mimeType The IANA standard MIME type of the image data; either `"image/png"` or
  * `"image/jpeg"`; to request a different format, see [ImagenGenerationConfig.imageFormat].
+ * @property safetyAttributes a set of safety attributes with their associated score.
  */
 @PublicPreviewAPI
 public class ImagenInlineImage
-internal constructor(public val data: ByteArray, public val mimeType: String) {
+internal constructor(
+  public val data: ByteArray,
+  public val mimeType: String,
+  public val safetyAttributes: Map<String, Double>) {
 
   /**
    * Returns the image as an Android OS native [Bitmap] so that it can be saved or sent to the UI.
@@ -53,5 +57,5 @@ public fun Bitmap.toImagenInlineImage(): ImagenInlineImage {
   val byteArrayOutputStream = ByteArrayOutputStream()
   this.compress(Bitmap.CompressFormat.JPEG, 80, byteArrayOutputStream)
   val byteArray = byteArrayOutputStream.toByteArray()
-  return ImagenInlineImage(data = byteArray, mimeType = "image/jpeg")
+  return ImagenInlineImage(data = byteArray, mimeType = "image/jpeg", safetyAttributes = emptyMap())
 }