[Vertex AI] Add ImagenSafetySettings type and param (#14237)

Author: andrewheard

FirebaseVertexAI/Sources/Types/Internal/Imagen/ImageGenerationParameters.swift

@@ -16,7 +16,6 @@
 struct ImageGenerationParameters {
   let sampleCount: Int?
   let storageURI: String?
-  let seed: Int32?
   let negativePrompt: String?
   let aspectRatio: String?
   let safetyFilterLevel: String?
@@ -36,7 +35,6 @@ extension ImageGenerationParameters: Encodable {
   enum CodingKeys: String, CodingKey {
     case sampleCount
     case storageURI = "storageUri"
-    case seed
     case negativePrompt
     case aspectRatio
     case safetyFilterLevel = "safetySetting"
@@ -50,7 +48,6 @@ extension ImageGenerationParameters: Encodable {
     var container = encoder.container(keyedBy: CodingKeys.self)
     try container.encodeIfPresent(sampleCount, forKey: .sampleCount)
     try container.encodeIfPresent(storageURI, forKey: .storageURI)
-    try container.encodeIfPresent(seed, forKey: .seed)
     try container.encodeIfPresent(negativePrompt, forKey: .negativePrompt)
     try container.encodeIfPresent(aspectRatio, forKey: .aspectRatio)
     try container.encodeIfPresent(safetyFilterLevel, forKey: .safetyFilterLevel)

FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenModel.swift

@@ -24,12 +24,15 @@ public final class ImagenModel {
   /// The backing service responsible for sending and receiving model requests to the backend.
   let generativeAIService: GenerativeAIService
 
+  let safetySettings: ImagenSafetySettings?
+
   /// Configuration parameters for sending requests to the backend.
   let requestOptions: RequestOptions
 
   init(name: String,
        projectID: String,
        apiKey: String,
+       safetySettings: ImagenSafetySettings?,
        requestOptions: RequestOptions,
        appCheck: AppCheckInterop?,
        auth: AuthInterop?,
@@ -42,6 +45,7 @@ public final class ImagenModel {
       auth: auth,
       urlSession: urlSession
     )
+    self.safetySettings = safetySettings
     self.requestOptions = requestOptions
   }
 
@@ -50,7 +54,11 @@ public final class ImagenModel {
     -> ImageGenerationResponse<ImagenInlineDataImage> {
     return try await generateImages(
       prompt: prompt,
-      parameters: imageGenerationParameters(storageURI: nil, generationConfig: generationConfig)
+      parameters: ImagenModel.imageGenerationParameters(
+        storageURI: nil,
+        generationConfig: generationConfig,
+        safetySettings: safetySettings
+      )
     )
   }
 
@@ -59,9 +67,10 @@ public final class ImagenModel {
     -> ImageGenerationResponse<ImagenFileDataImage> {
     return try await generateImages(
       prompt: prompt,
-      parameters: imageGenerationParameters(
+      parameters: ImagenModel.imageGenerationParameters(
         storageURI: storageURI,
-        generationConfig: generationConfig
+        generationConfig: generationConfig,
+        safetySettings: safetySettings
       )
     )
   }
@@ -79,26 +88,25 @@ public final class ImagenModel {
     return try await generativeAIService.loadRequest(request: request)
   }
 
-  func imageGenerationParameters(storageURI: String?,
-                                 generationConfig: ImagenGenerationConfig? = nil)
+  static func imageGenerationParameters(storageURI: String?,
+                                        generationConfig: ImagenGenerationConfig?,
+                                        safetySettings: ImagenSafetySettings?)
     -> ImageGenerationParameters {
-    // TODO(#14221): Add support for configuring remaining parameters.
     return ImageGenerationParameters(
       sampleCount: generationConfig?.numberOfImages ?? 1,
       storageURI: storageURI,
-      seed: nil,
       negativePrompt: generationConfig?.negativePrompt,
       aspectRatio: generationConfig?.aspectRatio?.rawValue,
-      safetyFilterLevel: nil,
-      personGeneration: nil,
+      safetyFilterLevel: safetySettings?.safetyFilterLevel?.rawValue,
+      personGeneration: safetySettings?.personGeneration?.rawValue,
       outputOptions: generationConfig?.imageFormat.map {
         ImageGenerationOutputOptions(
           mimeType: $0.mimeType,
           compressionQuality: $0.compressionQuality
         )
       },
       addWatermark: generationConfig?.addWatermark,
-      includeResponsibleAIFilterReason: true
+      includeResponsibleAIFilterReason: safetySettings?.includeFilterReason ?? true
     )
   }
 }

FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenSafetySettings.swift

@@ -0,0 +1,65 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import Foundation
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+public struct ImagenSafetySettings {
+  let safetyFilterLevel: SafetyFilterLevel?
+  let includeFilterReason: Bool?
+  let personGeneration: PersonGeneration?
+
+  public init(safetyFilterLevel: SafetyFilterLevel? = nil, includeFilterReason: Bool? = nil,
+              personGeneration: PersonGeneration? = nil) {
+    self.safetyFilterLevel = safetyFilterLevel
+    self.includeFilterReason = includeFilterReason
+    self.personGeneration = personGeneration
+  }
+}
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+public extension ImagenSafetySettings {
+  struct SafetyFilterLevel: ProtoEnum {
+    enum Kind: String {
+      case blockLowAndAbove = "block_low_and_above"
+      case blockMediumAndAbove = "block_medium_and_above"
+      case blockOnlyHigh = "block_only_high"
+      case blockNone = "block_none"
+    }
+
+    public static let blockLowAndAbove = SafetyFilterLevel(kind: .blockLowAndAbove)
+    public static let blockMediumAndAbove = SafetyFilterLevel(kind: .blockMediumAndAbove)
+    public static let blockOnlyHigh = SafetyFilterLevel(kind: .blockOnlyHigh)
+    public static let blockNone = SafetyFilterLevel(kind: .blockNone)
+
+    let rawValue: String
+  }
+}
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+public extension ImagenSafetySettings {
+  struct PersonGeneration: ProtoEnum {
+    enum Kind: String {
+      case blockAll = "dont_allow"
+      case allowAdult = "allow_adult"
+      case allowAll = "allow_all"
+    }
+
+    public static let blockAll = PersonGeneration(kind: .blockAll)
+    public static let allowAdult = PersonGeneration(kind: .allowAdult)
+    public static let allowAll = PersonGeneration(kind: .allowAll)
+
+    let rawValue: String
+  }
+}

FirebaseVertexAI/Sources/VertexAI.swift

@@ -104,12 +104,13 @@ public class VertexAI {
     )
   }
 
-  public func imagenModel(modelName: String, requestOptions: RequestOptions = RequestOptions())
-    -> ImagenModel {
+  public func imagenModel(modelName: String, safetySettings: ImagenSafetySettings? = nil,
+                          requestOptions: RequestOptions = RequestOptions()) -> ImagenModel {
     return ImagenModel(
       name: modelResourceName(modelName: modelName),
       projectID: projectID,
       apiKey: apiKey,
+      safetySettings: safetySettings,
       requestOptions: requestOptions,
       appCheck: appCheck,
       auth: auth

FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift

@@ -62,7 +62,11 @@ final class IntegrationTests: XCTestCase {
       systemInstruction: systemInstruction
     )
     imagenModel = vertex.imagenModel(
-      modelName: "imagen-3.0-fast-generate-001"
+      modelName: "imagen-3.0-fast-generate-001",
+      safetySettings: ImagenSafetySettings(
+        safetyFilterLevel: .blockLowAndAbove,
+        personGeneration: .blockAll
+      )
     )
 
     storage = Storage.storage()