diff --git a/Diffusion-macOS/ControlsView.swift b/Diffusion-macOS/ControlsView.swift index 3213dd5..c3a967e 100644 --- a/Diffusion-macOS/ControlsView.swift +++ b/Diffusion-macOS/ControlsView.swift @@ -54,6 +54,7 @@ struct ControlsView: View { @State private var model = Settings.shared.currentModel.modelVersion @State private var disclosedModel = true @State private var disclosedPrompt = true + @State private var disclosedGuidance = false @State private var disclosedSteps = false @State private var disclosedSeed = false @@ -134,6 +135,17 @@ struct ControlsView: View { Divider() + DisclosureGroup(isExpanded: $disclosedGuidance) { + CompactSlider(value: $generation.guidanceScale, in: 0...20, step: 0.5) { + Text("Guidance Scale") + Spacer() + Text(generation.guidanceScale.formatted("%.1f")) + }.padding(.leading, 10) + } label: { + Label("Guidance Scale", systemImage: "scalemass").foregroundColor(.secondary) + } + Divider() + DisclosureGroup(isExpanded: $disclosedSteps) { CompactSlider(value: $generation.steps, in: 0...150, step: 5) { Text("Steps") diff --git a/Diffusion/Pipeline/Pipeline.swift b/Diffusion/Pipeline/Pipeline.swift index 570dc9f..1226ac6 100644 --- a/Diffusion/Pipeline/Pipeline.swift +++ b/Diffusion/Pipeline/Pipeline.swift @@ -29,7 +29,15 @@ class Pipeline { self.pipeline = pipeline } - func generate(prompt: String, negativePrompt: String = "", scheduler: StableDiffusionScheduler, numInferenceSteps stepCount: Int = 50, seed: UInt32? = nil) throws -> (CGImage, TimeInterval) { + func generate( + prompt: String, + negativePrompt: String = "", + scheduler: StableDiffusionScheduler, + numInferenceSteps stepCount: Int = 50, + seed: UInt32? = nil, + guidanceScale: Float = 7.5, + disableSafety: Bool = false + ) throws -> (CGImage, TimeInterval) { let beginDate = Date() print("Generating...") let theSeed = seed ?? UInt32.random(in: 0.. (CGImage, TimeInterval)? { guard let pipeline = pipeline else { return nil } let seed = self.seed >= 0 ? UInt32(self.seed) : nil - return try? pipeline.generate(prompt: positivePrompt, negativePrompt: negativePrompt, scheduler: scheduler, numInferenceSteps: Int(steps), seed: seed) + return try? pipeline.generate( + prompt: positivePrompt, + negativePrompt: negativePrompt, + scheduler: scheduler, + numInferenceSteps: Int(steps), + seed: seed, + guidanceScale: Float(guidanceScale), + disableSafety: disableSafety + ) } } diff --git a/Diffusion/Utils.swift b/Diffusion/Utils.swift index cd0db8e..480c265 100644 --- a/Diffusion/Utils.swift +++ b/Diffusion/Utils.swift @@ -9,3 +9,9 @@ import Foundation extension String: Error {} + +extension Double { + func formatted(_ format: String) -> String { + return String(format: "\(format)", self) + } +}