Add safetySettings to imageGenerationParameters

andrewheard · andrewheard · commit 24da4958bd1c · 2024-12-09T18:14:55.000-05:00
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,

Original file line number	Diff line number	Diff line change
`@@ -54,7 +54,11 @@ public final class ImagenModel {`
`54`	`54`	`-> ImageGenerationResponse<ImagenInlineDataImage> {`
`55`	`55`	`return try await generateImages(`
`56`	`56`	`prompt: prompt,`
`57`		`- parameters: imageGenerationParameters(storageURI: nil, generationConfig: generationConfig)`
	`57`	`+ parameters: ImagenModel.imageGenerationParameters(`
	`58`	`+ storageURI: nil,`
	`59`	`+ generationConfig: generationConfig,`
	`60`	`+ safetySettings: safetySettings`
	`61`	`+ )`
`58`	`62`	`)`
`59`	`63`	`}`
`60`	`64`
`@@ -63,9 +67,10 @@ public final class ImagenModel {`
`63`	`67`	`-> ImageGenerationResponse<ImagenFileDataImage> {`
`64`	`68`	`return try await generateImages(`
`65`	`69`	`prompt: prompt,`
`66`		`- parameters: imageGenerationParameters(`
	`70`	`+ parameters: ImagenModel.imageGenerationParameters(`
`67`	`71`	`storageURI: storageURI,`
`68`		`- generationConfig: generationConfig`
	`72`	`+ generationConfig: generationConfig,`
	`73`	`+ safetySettings: safetySettings`
`69`	`74`	`)`
`70`	`75`	`)`
`71`	`76`	`}`
`@@ -83,8 +88,9 @@ public final class ImagenModel {`
`83`	`88`	`return try await generativeAIService.loadRequest(request: request)`
`84`	`89`	`}`
`85`	`90`
`86`		`- func imageGenerationParameters(storageURI: String?,`
`87`		`- generationConfig: ImagenGenerationConfig? = nil)`
	`91`	`+ static func imageGenerationParameters(storageURI: String?,`
	`92`	`+ generationConfig: ImagenGenerationConfig?,`
	`93`	`+ safetySettings: ImagenSafetySettings?)`
`88`	`94`	`-> ImageGenerationParameters {`
`89`	`95`	`return ImageGenerationParameters(`
`90`	`96`	`sampleCount: generationConfig?.numberOfImages ?? 1,`