Add safetySettings to imageGenerationParameters

firebase · Dec 9, 2024 · 24da495 · 24da495
1 parent 15d8704
commit 24da495
Showing 1 changed file with 11 additions and 5 deletions.
diff --git a/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenModel.swift b/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenModel.swift
@@ -54,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
+      )
     )
   }
 
@@ -63,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
       )
     )
   }
@@ -83,8 +88,9 @@ 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 {
     return ImageGenerationParameters(
       sampleCount: generationConfig?.numberOfImages ?? 1,