diff --git a/swift/StableDiffusionCLI/main.swift b/swift/StableDiffusionCLI/main.swift index 7343cf2..1b1c857 100644 --- a/swift/StableDiffusionCLI/main.swift +++ b/swift/StableDiffusionCLI/main.swift @@ -32,6 +32,12 @@ struct StableDiffusionSample: ParsableCommand { ) ) var resourcePath: String = "./" + + @Option(help: "Path to starting image.") + var image: String = "none" + + @Option(help: "Strength for image2image.") + var strength: Float = 0.5 @Option(help: "Number of images to sample / generate") var imageCount: Int = 1 @@ -51,7 +57,7 @@ struct StableDiffusionSample: ParsableCommand { var outputPath: String = "./" @Option(help: "Random seed") - var seed: UInt32 = 93 + var seed: UInt32 = UInt32.random(in: 0...UInt32.max) @Option(help: "Controls the influence of the text prompt on sampling process (0=random images)") var guidanceScale: Float = 7.5 @@ -84,6 +90,32 @@ struct StableDiffusionSample: ParsableCommand { disableSafety: disableSafety, reduceMemory: reduceMemory) try pipeline.loadResources() + + let startingImage: CGImage? + if image != "none" { + let imageURL = URL(filePath: image) +// if FileManager.default.fileExists(atPath: imageURL.path()) { +// throw RunError.resources("Starting image not found \(imageURL)") +// } + do { + let imageData = try Data(contentsOf: imageURL) + guard + let imgDataProvider = CGDataProvider(data: imageData as CFData), + let loadedImage = CGImage( + pngDataProviderSource: imgDataProvider, + decode: nil, shouldInterpolate: false, + intent: CGColorRenderingIntent.defaultIntent) + else { + throw RunError.resources("Starting Image not available \(resourcePath)") + } + startingImage = loadedImage + } catch let error { + throw RunError.resources("Starting image not found \(imageURL), error: \(error)") + } + + } else { + startingImage = nil + } log("Sampling ...\n") let sampleTimer = SampleTimer() @@ -92,6 +124,8 @@ struct StableDiffusionSample: ParsableCommand { let images = try pipeline.generateImages( prompt: prompt, negativePrompt: negativePrompt, + startingImage: startingImage, + strength: strength, imageCount: imageCount, stepCount: stepCount, seed: seed, @@ -168,6 +202,10 @@ struct StableDiffusionSample: ParsableCommand { if imageCount != 1 { name += ".\(sample)" } + + if image != "none" { + name += ".str\(Int(strength * 100))" + } name += ".\(seed)"