firebase-ios-sdk/FirebaseAI/Sources/TemplateImagenModel.swift

// Copyright 2025 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

/// A type that represents a remote image generation model (like Imagen), with the ability to
/// generate
/// images based on various input types.
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
public final class TemplateImagenModel: Sendable {
  let generativeAIService: GenerativeAIService
  let apiConfig: APIConfig

  init(generativeAIService: GenerativeAIService, apiConfig: APIConfig) {
    self.generativeAIService = generativeAIService
    self.apiConfig = apiConfig
  }

  /// Generates images from a prompt template and variables.
  ///
  /// - Parameters:
  ///   - template: The prompt template to use.
  ///   - variables: A dictionary of variables to substitute into the template.
  /// - Returns: The images generated by the model.
  /// - Throws: An error if the request failed.
  public func generateImages(template: String,
                             variables: [String: Any],
                             options: RequestOptions = RequestOptions()) async throws
    -> GenerateImagesResponse {
    let templateVariables = try variables.mapValues { try TemplateVariable(value: $0) }
    let projectID = generativeAIService.firebaseInfo.projectID
    let request = GenerateImagesRequest(
      template: template,
      variables: templateVariables,
      projectID: projectID,
      apiConfig: apiConfig,
      options: options
    )
    let response: GenerateImagesResponse = try await generativeAIService
      .loadRequest(request: request)
    return response
  }
}

// A placeholder for the response from an image generation request.
public struct GenerateImagesResponse: Decodable, @unchecked Sendable {
  public let images: [Data]

  enum CodingKeys: String, CodingKey {
    case predictions
  }

  struct Prediction: Decodable {
    let images: [String]
  }

  public init(from decoder: Decoder) throws {
    let container = try decoder.container(keyedBy: CodingKeys.self)
    let predictions = try container.decode([Prediction].self, forKey: .predictions)
    guard let firstPrediction = predictions.first else {
      images = []
      return
    }
    images = firstPrediction.images.compactMap { Data(base64Encoded: $0) }
  }
}