chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,377 @@
|
||||
// For licensing see accompanying LICENSE.md file.
|
||||
// Copyright (C) 2022 Apple Inc. All Rights Reserved.
|
||||
|
||||
import ArgumentParser
|
||||
import CoreGraphics
|
||||
import CoreML
|
||||
import Foundation
|
||||
import StableDiffusion
|
||||
import UniformTypeIdentifiers
|
||||
import Cocoa
|
||||
import CoreImage
|
||||
import NaturalLanguage
|
||||
|
||||
@available(iOS 16.2, macOS 13.1, *)
|
||||
struct StableDiffusionSample: ParsableCommand {
|
||||
|
||||
static let configuration = CommandConfiguration(
|
||||
abstract: "Run stable diffusion to generate images guided by a text prompt",
|
||||
version: "0.1"
|
||||
)
|
||||
|
||||
@Argument(help: "Input string prompt")
|
||||
var prompt: String
|
||||
|
||||
@Option(help: "Input string negative prompt")
|
||||
var negativePrompt: String = ""
|
||||
|
||||
@Option(
|
||||
help: ArgumentHelp(
|
||||
"Path to stable diffusion resources.",
|
||||
discussion: "The resource directory should contain\n" +
|
||||
" - *compiled* models: {TextEncoder,Unet,VAEDecoder}.mlmodelc\n" +
|
||||
" - tokenizer info: vocab.json, merges.txt",
|
||||
valueName: "directory-path"
|
||||
)
|
||||
)
|
||||
var resourcePath: String = "./"
|
||||
|
||||
@Flag(name: .customLong("xl"), help: "The resources correspond to a Stable Diffusion XL model")
|
||||
var isXL: Bool = false
|
||||
|
||||
@Flag(name: .customLong("sd3"), help: "The resources correspond to a Stable Diffusion 3 model")
|
||||
var isSD3: Bool = false
|
||||
|
||||
@Option(help: "Path to starting image.")
|
||||
var image: String? = nil
|
||||
|
||||
@Option(help: "Strength for image2image.")
|
||||
var strength: Float = 0.5
|
||||
|
||||
@Option(help: "Number of images to sample / generate")
|
||||
var imageCount: Int = 1
|
||||
|
||||
@Option(help: "Number of diffusion steps to perform")
|
||||
var stepCount: Int = 50
|
||||
|
||||
@Option(
|
||||
help: ArgumentHelp(
|
||||
"How often to save samples at intermediate steps",
|
||||
discussion: "Set to 0 to only save the final sample"
|
||||
)
|
||||
)
|
||||
var saveEvery: Int = 0
|
||||
|
||||
@Option(help: "Output path")
|
||||
var outputPath: String = "./"
|
||||
|
||||
@Option(help: "Random seed")
|
||||
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
|
||||
|
||||
@Option(help: "Compute units to load model with {all,cpuOnly,cpuAndGPU,cpuAndNeuralEngine}")
|
||||
var computeUnits: ComputeUnits = .all
|
||||
|
||||
@Option(help: "Scheduler to use, one of {pndm, dpmpp}")
|
||||
var scheduler: SchedulerOption = .pndm
|
||||
|
||||
@Option(help: "Random number generator to use, one of {numpy, torch, nvidia}")
|
||||
var rng: RNGOption = .numpy
|
||||
|
||||
@Option(
|
||||
parsing: .upToNextOption,
|
||||
help: "ControlNet models used in image generation (enter file names in Resources/controlnet without extension)"
|
||||
)
|
||||
var controlnet: [String] = []
|
||||
|
||||
@Option(
|
||||
parsing: .upToNextOption,
|
||||
help: "image for each controlNet model (corresponding to the same order as --controlnet)"
|
||||
)
|
||||
var controlnetInputs: [String] = []
|
||||
|
||||
@Flag(help: "Disable safety checking")
|
||||
var disableSafety: Bool = false
|
||||
|
||||
@Flag(help: "Reduce memory usage")
|
||||
var reduceMemory: Bool = false
|
||||
|
||||
@Flag(help: "Use system multilingual NLContextualEmbedding as encoder model")
|
||||
var useMultilingualTextEncoder: Bool = false
|
||||
|
||||
@Option(help: "The natural language script for the multilingual contextual embedding")
|
||||
var script: Script = .latin
|
||||
|
||||
mutating func run() throws {
|
||||
guard FileManager.default.fileExists(atPath: resourcePath) else {
|
||||
throw RunError.resources("Resource path does not exist \(resourcePath)")
|
||||
}
|
||||
|
||||
let config = MLModelConfiguration()
|
||||
config.computeUnits = computeUnits.asMLComputeUnits
|
||||
let resourceURL = URL(filePath: resourcePath)
|
||||
|
||||
log("Loading resources and creating pipeline\n")
|
||||
log("(Note: This can take a while the first time using these resources)\n")
|
||||
let pipeline: StableDiffusionPipelineProtocol
|
||||
var scaleFactor: Float32 = 0.18215
|
||||
var shiftFactor: Float32 = 0.0
|
||||
var timestepShift: Float32 = 1.0
|
||||
if #available(macOS 14.0, iOS 17.0, *) {
|
||||
if isXL {
|
||||
scaleFactor = 0.13025
|
||||
if !controlnet.isEmpty {
|
||||
throw RunError.unsupported("ControlNet is not supported for Stable Diffusion XL")
|
||||
}
|
||||
if useMultilingualTextEncoder {
|
||||
throw RunError.unsupported("Multilingual text encoder is not yet supported for Stable Diffusion XL")
|
||||
}
|
||||
pipeline = try StableDiffusionXLPipeline(
|
||||
resourcesAt: resourceURL,
|
||||
configuration: config,
|
||||
reduceMemory: reduceMemory
|
||||
)
|
||||
} else if isSD3 {
|
||||
scaleFactor = 1.5305
|
||||
shiftFactor = 0.0609
|
||||
timestepShift = 3.0
|
||||
if !controlnet.isEmpty {
|
||||
throw RunError.unsupported("ControlNet is not supported for Stable Diffusion 3")
|
||||
}
|
||||
if useMultilingualTextEncoder {
|
||||
throw RunError.unsupported("Multilingual text encoder is not yet supported for Stable Diffusion 3")
|
||||
}
|
||||
pipeline = try StableDiffusion3Pipeline(
|
||||
resourcesAt: resourceURL,
|
||||
configuration: config,
|
||||
reduceMemory: reduceMemory
|
||||
)
|
||||
} else {
|
||||
pipeline = try StableDiffusionPipeline(
|
||||
resourcesAt: resourceURL,
|
||||
controlNet: controlnet,
|
||||
configuration: config,
|
||||
disableSafety: disableSafety,
|
||||
reduceMemory: reduceMemory,
|
||||
useMultilingualTextEncoder: useMultilingualTextEncoder,
|
||||
script: script
|
||||
)
|
||||
}
|
||||
} else {
|
||||
pipeline = try StableDiffusionPipeline(
|
||||
resourcesAt: resourceURL,
|
||||
controlNet: controlnet,
|
||||
configuration: config,
|
||||
disableSafety: disableSafety,
|
||||
reduceMemory: reduceMemory
|
||||
)
|
||||
}
|
||||
|
||||
try pipeline.loadResources()
|
||||
|
||||
let startingImage: CGImage?
|
||||
if let image {
|
||||
let imageURL = URL(filePath: image)
|
||||
do {
|
||||
startingImage = try convertImageToCGImage(imageURL: imageURL)
|
||||
} catch let error {
|
||||
throw RunError.resources("Starting image not found \(imageURL), error: \(error)")
|
||||
}
|
||||
|
||||
} else {
|
||||
startingImage = nil
|
||||
}
|
||||
|
||||
// convert image for ControlNet into CGImage when controlNet available
|
||||
let controlNetInputs: [CGImage]
|
||||
if !controlnet.isEmpty {
|
||||
controlNetInputs = try controlnetInputs.map { imagePath in
|
||||
let imageURL = URL(filePath: imagePath)
|
||||
do {
|
||||
return try convertImageToCGImage(imageURL: imageURL)
|
||||
} catch let error {
|
||||
throw RunError.resources("Image for ControlNet not found \(imageURL), error: \(error)")
|
||||
}
|
||||
}
|
||||
} else {
|
||||
controlNetInputs = []
|
||||
}
|
||||
|
||||
log("Sampling ...\n")
|
||||
let sampleTimer = SampleTimer()
|
||||
sampleTimer.start()
|
||||
|
||||
var pipelineConfig = StableDiffusionPipeline.Configuration(prompt: prompt)
|
||||
|
||||
pipelineConfig.negativePrompt = negativePrompt
|
||||
pipelineConfig.startingImage = startingImage
|
||||
pipelineConfig.strength = strength
|
||||
pipelineConfig.imageCount = imageCount
|
||||
pipelineConfig.stepCount = stepCount
|
||||
pipelineConfig.seed = seed
|
||||
pipelineConfig.controlNetInputs = controlNetInputs
|
||||
pipelineConfig.guidanceScale = guidanceScale
|
||||
pipelineConfig.schedulerType = scheduler.stableDiffusionScheduler
|
||||
pipelineConfig.rngType = rng.stableDiffusionRNG
|
||||
pipelineConfig.useDenoisedIntermediates = true
|
||||
pipelineConfig.encoderScaleFactor = scaleFactor
|
||||
pipelineConfig.decoderScaleFactor = scaleFactor
|
||||
pipelineConfig.decoderShiftFactor = shiftFactor
|
||||
pipelineConfig.schedulerTimestepShift = timestepShift
|
||||
|
||||
let images = try pipeline.generateImages(
|
||||
configuration: pipelineConfig) { progress in
|
||||
sampleTimer.stop()
|
||||
handleProgress(progress,sampleTimer)
|
||||
if progress.stepCount != progress.step {
|
||||
sampleTimer.start()
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
_ = try saveImages(images, logNames: true)
|
||||
}
|
||||
|
||||
func convertImageToCGImage(imageURL: URL) throws -> CGImage {
|
||||
let imageData = try Data(contentsOf: imageURL)
|
||||
guard
|
||||
let nsImage = NSImage(data: imageData),
|
||||
let loadedImage = nsImage.cgImage(forProposedRect: nil, context: nil, hints: nil)
|
||||
else {
|
||||
throw RunError.resources("Image not available \(resourcePath)")
|
||||
}
|
||||
return loadedImage
|
||||
}
|
||||
|
||||
func handleProgress(
|
||||
_ progress: StableDiffusionPipeline.Progress,
|
||||
_ sampleTimer: SampleTimer
|
||||
) {
|
||||
log("\u{1B}[1A\u{1B}[K")
|
||||
log("Step \(progress.step) of \(progress.stepCount) ")
|
||||
log(" [")
|
||||
log(String(format: "mean: %.2f, ", 1.0/sampleTimer.mean))
|
||||
log(String(format: "median: %.2f, ", 1.0/sampleTimer.median))
|
||||
log(String(format: "last %.2f", 1.0/sampleTimer.allSamples.last!))
|
||||
log("] step/sec")
|
||||
|
||||
if saveEvery > 0, progress.step % saveEvery == 0 {
|
||||
let saveCount = (try? saveImages(progress.currentImages, step: progress.step)) ?? 0
|
||||
log(" saved \(saveCount) image\(saveCount != 1 ? "s" : "")")
|
||||
}
|
||||
log("\n")
|
||||
}
|
||||
|
||||
func saveImages(
|
||||
_ images: [CGImage?],
|
||||
step: Int? = nil,
|
||||
logNames: Bool = false
|
||||
) throws -> Int {
|
||||
let url = URL(filePath: outputPath)
|
||||
var saved = 0
|
||||
for i in 0 ..< images.count {
|
||||
|
||||
guard let image = images[i] else {
|
||||
if logNames {
|
||||
log("Image \(i) failed safety check and was not saved")
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
let name = imageName(i, step: step)
|
||||
let fileURL = url.appending(path:name)
|
||||
|
||||
guard let dest = CGImageDestinationCreateWithURL(fileURL as CFURL, UTType.png.identifier as CFString, 1, nil) else {
|
||||
throw RunError.saving("Failed to create destination for \(fileURL)")
|
||||
}
|
||||
CGImageDestinationAddImage(dest, image, nil)
|
||||
if !CGImageDestinationFinalize(dest) {
|
||||
throw RunError.saving("Failed to save \(fileURL)")
|
||||
}
|
||||
if logNames {
|
||||
log("Saved \(name)\n")
|
||||
}
|
||||
saved += 1
|
||||
}
|
||||
return saved
|
||||
}
|
||||
|
||||
func imageName(_ sample: Int, step: Int? = nil) -> String {
|
||||
let fileCharLimit = 75
|
||||
var name = prompt.prefix(fileCharLimit).replacingOccurrences(of: " ", with: "_")
|
||||
if imageCount != 1 {
|
||||
name += ".\(sample)"
|
||||
}
|
||||
|
||||
if image != nil {
|
||||
name += ".str\(Int(strength * 100))"
|
||||
}
|
||||
|
||||
name += ".\(seed)"
|
||||
|
||||
if let step = step {
|
||||
name += ".\(step)"
|
||||
} else {
|
||||
name += ".final"
|
||||
}
|
||||
name += ".png"
|
||||
return name
|
||||
}
|
||||
|
||||
func log(_ str: String, term: String = "") {
|
||||
print(str, terminator: term)
|
||||
}
|
||||
}
|
||||
|
||||
enum RunError: Error {
|
||||
case resources(String)
|
||||
case saving(String)
|
||||
case unsupported(String)
|
||||
}
|
||||
|
||||
@available(iOS 16.2, macOS 13.1, *)
|
||||
enum ComputeUnits: String, ExpressibleByArgument, CaseIterable {
|
||||
case all, cpuAndGPU, cpuOnly, cpuAndNeuralEngine
|
||||
var asMLComputeUnits: MLComputeUnits {
|
||||
switch self {
|
||||
case .all: return .all
|
||||
case .cpuAndGPU: return .cpuAndGPU
|
||||
case .cpuOnly: return .cpuOnly
|
||||
case .cpuAndNeuralEngine: return .cpuAndNeuralEngine
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@available(iOS 16.2, macOS 13.1, *)
|
||||
enum SchedulerOption: String, ExpressibleByArgument {
|
||||
case pndm, dpmpp
|
||||
var stableDiffusionScheduler: StableDiffusionScheduler {
|
||||
switch self {
|
||||
case .pndm: return .pndmScheduler
|
||||
case .dpmpp: return .dpmSolverMultistepScheduler
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@available(iOS 16.2, macOS 13.1, *)
|
||||
enum RNGOption: String, ExpressibleByArgument {
|
||||
case numpy, torch, nvidia
|
||||
var stableDiffusionRNG: StableDiffusionRNG {
|
||||
switch self {
|
||||
case .numpy: return .numpyRNG
|
||||
case .torch: return .torchRNG
|
||||
case .nvidia: return .nvidiaRNG
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@available(iOS 16.2, macOS 13.1, *)
|
||||
extension Script: ExpressibleByArgument {}
|
||||
|
||||
if #available(iOS 16.2, macOS 13.1, *) {
|
||||
StableDiffusionSample.main()
|
||||
} else {
|
||||
print("Unsupported OS")
|
||||
}
|
||||
Reference in New Issue
Block a user