parent
c17c80ff8b
commit
d6647a4551
@ -0,0 +1,82 @@
|
||||
//
|
||||
// File.swift
|
||||
//
|
||||
//
|
||||
// Created by Timothy Kautz on 1/30/23.
|
||||
//
|
||||
|
||||
import Foundation
|
||||
import CoreGraphics
|
||||
|
||||
@available(iOS 16.2, macOS 13.1, *)
|
||||
extension StableDiffusionPipeline {
|
||||
|
||||
/// Tyoe of processing that will be performed to generate an image
|
||||
public enum Mode {
|
||||
case textToImage
|
||||
case imageToImage
|
||||
// case inPainting
|
||||
}
|
||||
|
||||
/// Image generation configuration
|
||||
public struct SampleInput: Hashable {
|
||||
|
||||
/// Text prompt to guide sampling
|
||||
public var prompt: String
|
||||
/// Negative text prompt to guide sampling
|
||||
public var negativePrompt: String = ""
|
||||
/// Starting image for image2image or in-painting
|
||||
public var startingImage: CGImage? = nil
|
||||
//public var maskImage: CGImage? = nil
|
||||
public var strength: Float = 1.0
|
||||
/// Number of images to generate
|
||||
public var imageCount: Int = 1
|
||||
/// Number of inference steps to perform
|
||||
public var stepCount: Int = 50
|
||||
/// Random seed which to start generation
|
||||
public var seed: UInt32 = 0
|
||||
/// Controls the influence of the text prompt on sampling process (0=random images)
|
||||
public var guidanceScale: Float = 7.5
|
||||
/// Safety checks are only performed if `self.canSafetyCheck && !disableSafety`
|
||||
public var disableSafety: Bool = false
|
||||
/// The type of Scheduler to use.
|
||||
public var schedulerType: StableDiffusionScheduler = .pndmScheduler
|
||||
|
||||
/// Given the configuration, what mode will be used for generation
|
||||
public var mode: Mode {
|
||||
guard startingImage != nil else {
|
||||
return .textToImage
|
||||
}
|
||||
guard strength < 1.0 else {
|
||||
return .textToImage
|
||||
}
|
||||
return .imageToImage
|
||||
}
|
||||
|
||||
public init(
|
||||
prompt: String,
|
||||
negativePrompt: String = "",
|
||||
startingImage: CGImage? = nil,
|
||||
strength: Float = 1.0,
|
||||
imageCount: Int = 1,
|
||||
stepCount: Int = 50,
|
||||
seed: UInt32 = 0,
|
||||
guidanceScale: Float = 7.5,
|
||||
disableSafety: Bool = false,
|
||||
schedulerType: StableDiffusionScheduler = .pndmScheduler
|
||||
) {
|
||||
self.prompt = prompt
|
||||
self.negativePrompt = negativePrompt
|
||||
self.startingImage = startingImage
|
||||
self.strength = strength
|
||||
self.imageCount = imageCount
|
||||
self.stepCount = stepCount
|
||||
self.seed = seed
|
||||
self.guidanceScale = guidanceScale
|
||||
self.disableSafety = disableSafety
|
||||
self.schedulerType = schedulerType
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
Loading…
Reference in new issue