Skip to content

Commit e773941

Browse files
andrewheardandrewheard
authored
[Vertex AI] Add ImageGenerationResponse for decoding PredictResponse (#14216)
1 parent 78fe33c commit e773941

File tree

9 files changed

+640
-0
lines changed

9 files changed

+640
-0
lines changed

‎FirebaseVertexAI/Sources/Types/Internal/Imagen/DecodableImagenImage.swift‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
// Copyright 2024 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+
protocol DecodableImagenImage: ImagenImage, Decodable {
17+
init(mimeType: String, bytesBase64Encoded: String?, gcsURI: String?)
18+
}
19+
20+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
21+
enum ImagenImageCodingKeys: String, CodingKey {
22+
case mimeType
23+
case bytesBase64Encoded
24+
case gcsURI = "gcsUri"
25+
}
26+
27+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
28+
extension DecodableImagenImage {
29+
init(from decoder: any Decoder) throws {
30+
let container = try decoder.container(keyedBy: ImagenImageCodingKeys.self)
31+
let mimeType = try container.decode(String.self, forKey: .mimeType)
32+
let bytesBase64Encoded = try container.decodeIfPresent(
33+
String.self,
34+
forKey: .bytesBase64Encoded
35+
)
36+
let gcsURI = try container.decodeIfPresent(String.self, forKey: .gcsURI)
37+
guard bytesBase64Encoded != nil || gcsURI != nil else {
38+
throw DecodingError.dataCorrupted(
39+
DecodingError.Context(
40+
codingPath: [ImagenImageCodingKeys.bytesBase64Encoded, ImagenImageCodingKeys.gcsURI],
41+
debugDescription: """
42+
Expected one of \(ImagenImageCodingKeys.bytesBase64Encoded.rawValue) or \
43+
\(ImagenImageCodingKeys.gcsURI.rawValue); both are nil.
44+
"""
45+
)
46+
)
47+
}
48+
guard bytesBase64Encoded == nil || gcsURI == nil else {
49+
throw DecodingError.dataCorrupted(
50+
DecodingError.Context(
51+
codingPath: [ImagenImageCodingKeys.bytesBase64Encoded, ImagenImageCodingKeys.gcsURI],
52+
debugDescription: """
53+
Expected one of \(ImagenImageCodingKeys.bytesBase64Encoded.rawValue) or \
54+
\(ImagenImageCodingKeys.gcsURI.rawValue); both are specified.
55+
"""
56+
)
57+
)
58+
}
59+
60+
self.init(mimeType: mimeType, bytesBase64Encoded: bytesBase64Encoded, gcsURI: gcsURI)
61+
}
62+
}

‎FirebaseVertexAI/Sources/Types/Internal/Imagen/ImageGenerationResponse.swift‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
// Copyright 2024 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+
import Foundation
16+
17+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
18+
struct ImageGenerationResponse {
19+
let images: [InternalImagenImage]
20+
let raiFilteredReason: String?
21+
}
22+
23+
// MARK: - Codable Conformances
24+
25+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
26+
extension ImageGenerationResponse: Decodable {
27+
enum CodingKeys: CodingKey {
28+
case predictions
29+
}
30+
31+
public init(from decoder: any Decoder) throws {
32+
let container = try decoder.container(keyedBy: CodingKeys.self)
33+
guard container.contains(.predictions) else {
34+
images = []
35+
raiFilteredReason = nil
36+
// TODO: Log warning if no predictions.
37+
return
38+
}
39+
var predictionsContainer = try container.nestedUnkeyedContainer(forKey: .predictions)
40+
41+
var images = [InternalImagenImage]()
42+
var raiFilteredReasons = [String]()
43+
while !predictionsContainer.isAtEnd {
44+
if let image = try? predictionsContainer.decode(InternalImagenImage.self) {
45+
images.append(image)
46+
} else if let filterReason = try? predictionsContainer.decode(RAIFilteredReason.self) {
47+
raiFilteredReasons.append(filterReason.raiFilteredReason)
48+
} else if let _ = try? predictionsContainer.decode(JSONObject.self) {
49+
// TODO: Log or throw unsupported prediction type
50+
} else {
51+
// This should never be thrown since JSONObject accepts any valid JSON.
52+
throw DecodingError.dataCorruptedError(
53+
in: predictionsContainer,
54+
debugDescription: "Failed to decode Prediction."
55+
)
56+
}
57+
}
58+
59+
self.images = images
60+
raiFilteredReason = raiFilteredReasons.first
61+
// TODO: Log if more than one RAI Filtered Reason; unexpected behaviour.
62+
}
63+
}

‎FirebaseVertexAI/Sources/Types/Internal/Imagen/InternalImagenImage.swift‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
// Copyright 2024 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+
struct InternalImagenImage {
17+
let mimeType: String
18+
let bytesBase64Encoded: String?
19+
let gcsURI: String?
20+
21+
init(mimeType: String, bytesBase64Encoded: String?, gcsURI: String?) {
22+
self.mimeType = mimeType
23+
self.bytesBase64Encoded = bytesBase64Encoded
24+
self.gcsURI = gcsURI
25+
}
26+
}
27+
28+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
29+
extension InternalImagenImage: DecodableImagenImage {}
30+
31+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
32+
extension InternalImagenImage: Equatable {}

‎FirebaseVertexAI/Sources/Types/Internal/Imagen/RAIFilteredReason.swift‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
// Copyright 2024 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+
struct RAIFilteredReason {
17+
let raiFilteredReason: String
18+
}
19+
20+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
21+
extension RAIFilteredReason: Decodable {
22+
enum CodingKeys: CodingKey {
23+
case raiFilteredReason
24+
}
25+
}

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

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
// Copyright 2024 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+
import Foundation
16+
17+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
18+
public protocol ImagenImage: ImagenImageRepresentable {
19+
var mimeType: String { get }
20+
var bytesBase64Encoded: String? { get }
21+
var gcsURI: String? { get }
22+
}

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

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
// Copyright 2024 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+
import Foundation
16+
17+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
18+
public protocol ImagenImageRepresentable {
19+
var imagenImage: any ImagenImage { get }
20+
}
21+
22+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
23+
public extension ImagenImage {
24+
var imagenImage: any ImagenImage {
25+
return self
26+
}
27+
}

0 commit comments

Comments
 (0)