diff --git a/FirebaseVertexAI/Sources/Types/Internal/Imagen/ImageGenerationOutputOptions.swift b/FirebaseVertexAI/Sources/Types/Internal/Imagen/ImageGenerationOutputOptions.swift new file mode 100644 index 00000000000..6ac24776116 --- /dev/null +++ b/FirebaseVertexAI/Sources/Types/Internal/Imagen/ImageGenerationOutputOptions.swift @@ -0,0 +1,26 @@ +// 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, *) +struct ImageGenerationOutputOptions { + let mimeType: String + let compressionQuality: Int? +} + +// MARK: - Codable Conformance + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +extension ImageGenerationOutputOptions: Encodable {} diff --git a/FirebaseVertexAI/Sources/Types/Internal/Imagen/ImageGenerationParameters.swift b/FirebaseVertexAI/Sources/Types/Internal/Imagen/ImageGenerationParameters.swift new file mode 100644 index 00000000000..6e5380666cd --- /dev/null +++ b/FirebaseVertexAI/Sources/Types/Internal/Imagen/ImageGenerationParameters.swift @@ -0,0 +1,62 @@ +// 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. + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +struct ImageGenerationParameters { + let sampleCount: Int? + let storageURI: String? + let seed: Int32? + let negativePrompt: String? + let aspectRatio: String? + let safetyFilterLevel: String? + let personGeneration: String? + let outputOptions: ImageGenerationOutputOptions? + let addWatermark: Bool? + let includeResponsibleAIFilterReason: Bool? +} + +// MARK: - Codable Conformance + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +extension ImageGenerationParameters: Encodable { + enum CodingKeys: String, CodingKey { + case sampleCount + case storageURI = "storageUri" + case seed + case negativePrompt + case aspectRatio + case safetyFilterLevel = "safetySetting" + case personGeneration + case outputOptions + case addWatermark + case includeResponsibleAIFilterReason = "includeRaiReason" + } + + func encode(to encoder: any Encoder) throws { + 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) + try container.encodeIfPresent(personGeneration, forKey: .personGeneration) + try container.encodeIfPresent(outputOptions, forKey: .outputOptions) + try container.encodeIfPresent(addWatermark, forKey: .addWatermark) + try container.encodeIfPresent( + includeResponsibleAIFilterReason, + forKey: .includeResponsibleAIFilterReason + ) + } +} diff --git a/FirebaseVertexAI/Tests/Unit/Types/Imagen/ImageGenerationOutputOptionsTests.swift b/FirebaseVertexAI/Tests/Unit/Types/Imagen/ImageGenerationOutputOptionsTests.swift new file mode 100644 index 00000000000..9a5d93e9dc3 --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/Types/Imagen/ImageGenerationOutputOptionsTests.swift @@ -0,0 +1,58 @@ +// 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 XCTest + +@testable import FirebaseVertexAI + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +final class ImageGenerationOutputOptionsTests: XCTestCase { + let encoder = JSONEncoder() + + override func setUp() { + encoder.outputFormatting = [.sortedKeys, .prettyPrinted, .withoutEscapingSlashes] + } + + // MARK: - Encoding Tests + + func testEncodeOutputOptions_jpeg_defaultCompressionQuality() throws { + let mimeType = "image/jpeg" + let options = ImageGenerationOutputOptions(mimeType: mimeType, compressionQuality: nil) + + let jsonData = try encoder.encode(options) + + let json = try XCTUnwrap(String(data: jsonData, encoding: .utf8)) + XCTAssertEqual(json, """ + { + "mimeType" : "\(mimeType)" + } + """) + } + + func testEncodeOutputOptions_jpeg_customCompressionQuality() throws { + let mimeType = "image/jpeg" + let quality = 50 + let options = ImageGenerationOutputOptions(mimeType: mimeType, compressionQuality: quality) + + let jsonData = try encoder.encode(options) + + let json = try XCTUnwrap(String(data: jsonData, encoding: .utf8)) + XCTAssertEqual(json, """ + { + "compressionQuality" : \(quality), + "mimeType" : "\(mimeType)" + } + """) + } +} diff --git a/FirebaseVertexAI/Tests/Unit/Types/Imagen/ImageGenerationParametersTests.swift b/FirebaseVertexAI/Tests/Unit/Types/Imagen/ImageGenerationParametersTests.swift new file mode 100644 index 00000000000..f5ce3a3b81e --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/Types/Imagen/ImageGenerationParametersTests.swift @@ -0,0 +1,129 @@ +// 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 XCTest + +@testable import FirebaseVertexAI + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +final class ImageGenerationParametersTests: XCTestCase { + let encoder = JSONEncoder() + + override func setUp() { + encoder.outputFormatting = [.sortedKeys, .prettyPrinted, .withoutEscapingSlashes] + } + + // MARK: - Encoding Tests + + func testEncodeParameters_allSpecified() throws { + let sampleCount = 4 + let storageURI = "gs://bucket/folder" + let seed: Int32 = 1_076_107_968 + let negativePrompt = "test-negative-prompt" + let aspectRatio = "16:9" + let safetyFilterLevel = "block_low_and_above" + let personGeneration = "allow_adult" + let mimeType = "image/png" + let outputOptions = ImageGenerationOutputOptions(mimeType: mimeType, compressionQuality: nil) + let addWatermark = false + let includeRAIReason = true + let parameters = ImageGenerationParameters( + sampleCount: sampleCount, + storageURI: storageURI, + seed: seed, + negativePrompt: negativePrompt, + aspectRatio: aspectRatio, + safetyFilterLevel: safetyFilterLevel, + personGeneration: personGeneration, + outputOptions: outputOptions, + addWatermark: addWatermark, + includeResponsibleAIFilterReason: includeRAIReason + ) + + let jsonData = try encoder.encode(parameters) + + let json = try XCTUnwrap(String(data: jsonData, encoding: .utf8)) + XCTAssertEqual(json, """ + { + "addWatermark" : \(addWatermark), + "aspectRatio" : "\(aspectRatio)", + "includeRaiReason" : \(includeRAIReason), + "negativePrompt" : "\(negativePrompt)", + "outputOptions" : { + "mimeType" : "\(mimeType)" + }, + "personGeneration" : "\(personGeneration)", + "safetySetting" : "\(safetyFilterLevel)", + "sampleCount" : \(sampleCount), + "seed" : \(seed), + "storageUri" : "\(storageURI)" + } + """) + } + + func testEncodeParameters_someSpecified() throws { + let sampleCount = 2 + let aspectRatio = "3:4" + let safetyFilterLevel = "block_medium_and_above" + let addWatermark = true + let parameters = ImageGenerationParameters( + sampleCount: sampleCount, + storageURI: nil, + seed: nil, + negativePrompt: nil, + aspectRatio: aspectRatio, + safetyFilterLevel: safetyFilterLevel, + personGeneration: nil, + outputOptions: nil, + addWatermark: addWatermark, + includeResponsibleAIFilterReason: nil + ) + + let jsonData = try encoder.encode(parameters) + + let json = try XCTUnwrap(String(data: jsonData, encoding: .utf8)) + XCTAssertEqual(json, """ + { + "addWatermark" : \(addWatermark), + "aspectRatio" : "\(aspectRatio)", + "safetySetting" : "\(safetyFilterLevel)", + "sampleCount" : \(sampleCount) + } + """) + } + + func testEncodeParameters_noneSpecified() throws { + let parameters = ImageGenerationParameters( + sampleCount: nil, + storageURI: nil, + seed: nil, + negativePrompt: nil, + aspectRatio: nil, + safetyFilterLevel: nil, + personGeneration: nil, + outputOptions: nil, + addWatermark: nil, + includeResponsibleAIFilterReason: nil + ) + + let jsonData = try encoder.encode(parameters) + + let json = try XCTUnwrap(String(data: jsonData, encoding: .utf8)) + XCTAssertEqual(json, """ + { + + } + """) + } +}