Skip to content

Commit 199aafd

Browse files
andrewheardandrewheard
committed
[Vertex AI] Refactor ImagenSafetySettings (#14307)
1 parent 58fd515 commit 199aafd

File tree

6 files changed

+76
-60
lines changed

6 files changed

+76
-60
lines changed

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -98,15 +98,15 @@ public final class ImagenModel {
9898
negativePrompt: generationConfig?.negativePrompt,
9999
aspectRatio: generationConfig?.aspectRatio?.rawValue,
100100
safetyFilterLevel: safetySettings?.safetyFilterLevel?.rawValue,
101-
personGeneration: safetySettings?.personGeneration?.rawValue,
101+
personGeneration: safetySettings?.personFilterLevel?.rawValue,
102102
outputOptions: generationConfig?.imageFormat.map {
103103
ImageGenerationOutputOptions(
104104
mimeType: $0.mimeType,
105105
compressionQuality: $0.compressionQuality
106106
)
107107
},
108108
addWatermark: generationConfig?.addWatermark,
109-
includeResponsibleAIFilterReason: safetySettings?.includeFilterReason ?? true
109+
includeResponsibleAIFilterReason: true
110110
)
111111
}
112112
}

‎FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenPersonFilterLevel.swift‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
// Copyright 2025 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
16+
public struct ImagenPersonFilterLevel: ProtoEnum {
17+
enum Kind: String {
18+
case blockAll = "dont_allow"
19+
case allowAdult = "allow_adult"
20+
case allowAll = "allow_all"
21+
}
22+
23+
public static let blockAll = ImagenPersonFilterLevel(kind: .blockAll)
24+
public static let allowAdult = ImagenPersonFilterLevel(kind: .allowAdult)
25+
public static let allowAll = ImagenPersonFilterLevel(kind: .allowAll)
26+
27+
let rawValue: String
28+
}

‎FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenSafetyFilterLevel.swift‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
// Copyright 2025 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
16+
public struct ImagenSafetyFilterLevel: ProtoEnum {
17+
enum Kind: String {
18+
case blockLowAndAbove = "block_low_and_above"
19+
case blockMediumAndAbove = "block_medium_and_above"
20+
case blockOnlyHigh = "block_only_high"
21+
case blockNone = "block_none"
22+
}
23+
24+
public static let blockLowAndAbove = ImagenSafetyFilterLevel(kind: .blockLowAndAbove)
25+
public static let blockMediumAndAbove = ImagenSafetyFilterLevel(kind: .blockMediumAndAbove)
26+
public static let blockOnlyHigh = ImagenSafetyFilterLevel(kind: .blockOnlyHigh)
27+
public static let blockNone = ImagenSafetyFilterLevel(kind: .blockNone)
28+
29+
let rawValue: String
30+
}

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

Lines changed: 5 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -16,50 +16,12 @@ import Foundation
1616

1717
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
1818
public struct ImagenSafetySettings {
19-
let safetyFilterLevel: SafetyFilterLevel?
20-
let includeFilterReason: Bool?
21-
let personGeneration: PersonGeneration?
19+
let safetyFilterLevel: ImagenSafetyFilterLevel?
20+
let personFilterLevel: ImagenPersonFilterLevel?
2221

23-
public init(safetyFilterLevel: SafetyFilterLevel? = nil, includeFilterReason: Bool? = nil,
24-
personGeneration: PersonGeneration? = nil) {
22+
public init(safetyFilterLevel: ImagenSafetyFilterLevel? = nil,
23+
personFilterLevel: ImagenPersonFilterLevel? = nil) {
2524
self.safetyFilterLevel = safetyFilterLevel
26-
self.includeFilterReason = includeFilterReason
27-
self.personGeneration = personGeneration
28-
}
29-
}
30-
31-
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
32-
public extension ImagenSafetySettings {
33-
struct SafetyFilterLevel: ProtoEnum {
34-
enum Kind: String {
35-
case blockLowAndAbove = "block_low_and_above"
36-
case blockMediumAndAbove = "block_medium_and_above"
37-
case blockOnlyHigh = "block_only_high"
38-
case blockNone = "block_none"
39-
}
40-
41-
public static let blockLowAndAbove = SafetyFilterLevel(kind: .blockLowAndAbove)
42-
public static let blockMediumAndAbove = SafetyFilterLevel(kind: .blockMediumAndAbove)
43-
public static let blockOnlyHigh = SafetyFilterLevel(kind: .blockOnlyHigh)
44-
public static let blockNone = SafetyFilterLevel(kind: .blockNone)
45-
46-
let rawValue: String
47-
}
48-
}
49-
50-
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
51-
public extension ImagenSafetySettings {
52-
struct PersonGeneration: ProtoEnum {
53-
enum Kind: String {
54-
case blockAll = "dont_allow"
55-
case allowAdult = "allow_adult"
56-
case allowAll = "allow_all"
57-
}
58-
59-
public static let blockAll = PersonGeneration(kind: .blockAll)
60-
public static let allowAdult = PersonGeneration(kind: .allowAdult)
61-
public static let allowAll = PersonGeneration(kind: .allowAll)
62-
63-
let rawValue: String
25+
self.personFilterLevel = personFilterLevel
6426
}
6527
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ final class IntegrationTests: XCTestCase {
6565
modelName: "imagen-3.0-fast-generate-001",
6666
safetySettings: ImagenSafetySettings(
6767
safetyFilterLevel: .blockLowAndAbove,
68-
personGeneration: .blockAll
68+
personFilterLevel: .blockAll
6969
)
7070
)
7171

‎FirebaseVertexAI/Tests/Unit/Types/Imagen/ImageGenerationParametersTests.swift‎

Lines changed: 10 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -111,24 +111,22 @@ final class ImageGenerationParametersTests: XCTestCase {
111111
}
112112

113113
func testDefaultParameters_includeSafetySettings() throws {
114-
let safetyFilterLevel = ImagenSafetySettings.SafetyFilterLevel.blockOnlyHigh
115-
let personGeneration = ImagenSafetySettings.PersonGeneration.allowAll
116-
let includeFilterReason = false
114+
let safetyFilterLevel = ImagenSafetyFilterLevel.blockOnlyHigh
115+
let personFilterLevel = ImagenPersonFilterLevel.allowAll
117116
let safetySettings = ImagenSafetySettings(
118117
safetyFilterLevel: safetyFilterLevel,
119-
includeFilterReason: includeFilterReason,
120-
personGeneration: personGeneration
118+
personFilterLevel: personFilterLevel
121119
)
122120
let expectedParameters = ImageGenerationParameters(
123121
sampleCount: 1,
124122
storageURI: nil,
125123
negativePrompt: nil,
126124
aspectRatio: nil,
127125
safetyFilterLevel: safetyFilterLevel.rawValue,
128-
personGeneration: personGeneration.rawValue,
126+
personGeneration: personFilterLevel.rawValue,
129127
outputOptions: nil,
130128
addWatermark: nil,
131-
includeResponsibleAIFilterReason: includeFilterReason
129+
includeResponsibleAIFilterReason: true
132130
)
133131

134132
let parameters = ImagenModel.imageGenerationParameters(
@@ -156,27 +154,25 @@ final class ImageGenerationParametersTests: XCTestCase {
156154
imageFormat: imageFormat,
157155
addWatermark: addWatermark
158156
)
159-
let safetyFilterLevel = ImagenSafetySettings.SafetyFilterLevel.blockNone
160-
let personGeneration = ImagenSafetySettings.PersonGeneration.blockAll
161-
let includeFilterReason = false
157+
let safetyFilterLevel = ImagenSafetyFilterLevel.blockNone
158+
let personFilterLevel = ImagenPersonFilterLevel.blockAll
162159
let safetySettings = ImagenSafetySettings(
163160
safetyFilterLevel: safetyFilterLevel,
164-
includeFilterReason: includeFilterReason,
165-
personGeneration: personGeneration
161+
personFilterLevel: personFilterLevel
166162
)
167163
let expectedParameters = ImageGenerationParameters(
168164
sampleCount: sampleCount,
169165
storageURI: storageURI,
170166
negativePrompt: negativePrompt,
171167
aspectRatio: aspectRatio.rawValue,
172168
safetyFilterLevel: safetyFilterLevel.rawValue,
173-
personGeneration: personGeneration.rawValue,
169+
personGeneration: personFilterLevel.rawValue,
174170
outputOptions: ImageGenerationOutputOptions(
175171
mimeType: imageFormat.mimeType,
176172
compressionQuality: imageFormat.compressionQuality
177173
),
178174
addWatermark: addWatermark,
179-
includeResponsibleAIFilterReason: includeFilterReason
175+
includeResponsibleAIFilterReason: true
180176
)
181177

182178
let parameters = ImagenModel.imageGenerationParameters(

0 commit comments

Comments
 (0)