forked from sagnik/Project_Velocity
101 lines
2.9 KiB
Swift
101 lines
2.9 KiB
Swift
import Foundation
|
|
import UIKit
|
|
import Alamofire
|
|
|
|
final class ComfyClient {
|
|
static let shared = ComfyClient()
|
|
|
|
private let endpoint = "http://192.168.x.x:8000/dream-weaver"
|
|
private let session: Session
|
|
|
|
private init(session: Session = .default) {
|
|
self.session = session
|
|
}
|
|
|
|
func generateImage(source: UIImage, prompt: String) async throws -> UIImage {
|
|
let resized = source.resizedSquare(to: 1024)
|
|
guard let imageData = resized.jpegData(compressionQuality: 0.9) else {
|
|
throw ComfyClientError.encodingFailed
|
|
}
|
|
|
|
let payload = DreamWeaverRequest(
|
|
imageBase64: imageData.base64EncodedString(),
|
|
prompt: prompt
|
|
)
|
|
|
|
let response = try await session.request(
|
|
endpoint,
|
|
method: .post,
|
|
parameters: payload,
|
|
encoder: JSONParameterEncoder.default,
|
|
headers: [.contentType("application/json")]
|
|
)
|
|
.validate(statusCode: 200..<300)
|
|
.serializingDecodable(DreamWeaverResponse.self)
|
|
.value
|
|
|
|
guard
|
|
let data = Data(base64Encoded: response.outputBase64),
|
|
let generated = UIImage(data: data)
|
|
else {
|
|
throw ComfyClientError.decodingFailed
|
|
}
|
|
|
|
return generated
|
|
}
|
|
}
|
|
|
|
private struct DreamWeaverRequest: Encodable {
|
|
let imageBase64: String
|
|
let prompt: String
|
|
}
|
|
|
|
private struct DreamWeaverResponse: Decodable {
|
|
let outputBase64: String
|
|
|
|
enum CodingKeys: String, CodingKey {
|
|
case outputBase64 = "output_base64"
|
|
case imageBase64 = "image_base64"
|
|
case image
|
|
}
|
|
|
|
init(from decoder: Decoder) throws {
|
|
let container = try decoder.container(keyedBy: CodingKeys.self)
|
|
if let preferred = try container.decodeIfPresent(String.self, forKey: .outputBase64) {
|
|
outputBase64 = preferred
|
|
return
|
|
}
|
|
if let legacy = try container.decodeIfPresent(String.self, forKey: .imageBase64) {
|
|
outputBase64 = legacy
|
|
return
|
|
}
|
|
outputBase64 = try container.decode(String.self, forKey: .image)
|
|
}
|
|
}
|
|
|
|
enum ComfyClientError: Error {
|
|
case encodingFailed
|
|
case decodingFailed
|
|
}
|
|
|
|
private extension UIImage {
|
|
func resizedSquare(to side: CGFloat) -> UIImage {
|
|
let format = UIGraphicsImageRendererFormat.default()
|
|
format.scale = 1
|
|
let renderer = UIGraphicsImageRenderer(size: CGSize(width: side, height: side), format: format)
|
|
|
|
return renderer.image { _ in
|
|
let aspect = size.width / size.height
|
|
let targetRect: CGRect
|
|
if aspect > 1 {
|
|
let width = side * aspect
|
|
targetRect = CGRect(x: (side - width) / 2, y: 0, width: width, height: side)
|
|
} else {
|
|
let height = side / aspect
|
|
targetRect = CGRect(x: 0, y: (side - height) / 2, width: side, height: height)
|
|
}
|
|
draw(in: targetRect)
|
|
}
|
|
}
|
|
}
|