* Refactor code to split it up a bit more into specific classes & rename consistently * Add extra supported parameters to generate in Pipeline * Update UI in TextoImageView to handle all values to be supported (but some are not usable currently)pull/2/head
Fahim Farook
1 year ago
parent
d7d0039969
commit
c79312fa8e
@ -0,0 +1,46 @@
|
||||
//
|
||||
// PreviewView.swift
|
||||
// Diffusion
|
||||
//
|
||||
// Created by Fahim Farook on 15/12/2022.
|
||||
//
|
||||
|
||||
import SwiftUI
|
||||
|
||||
struct PreviewView: View {
|
||||
var image: Binding<CGImage?>
|
||||
var state: Binding<GenerationState>
|
||||
|
||||
var body: some View {
|
||||
switch state.wrappedValue {
|
||||
case .startup: return AnyView(Image("placeholder").resizable())
|
||||
case .running(let progress):
|
||||
guard let progress = progress, progress.stepCount > 0 else {
|
||||
// The first time it takes a little bit before generation starts
|
||||
return AnyView(ProgressView())
|
||||
}
|
||||
let step = Int(progress.step) + 1
|
||||
let fraction = Double(step) / Double(progress.stepCount)
|
||||
let label = "Step \(step) of \(progress.stepCount)"
|
||||
return AnyView(ProgressView(label, value: fraction, total: 1).padding())
|
||||
case .idle(let lastPrompt):
|
||||
guard let theImage = image.wrappedValue else {
|
||||
return AnyView(Image(systemName: "exclamationmark.triangle").resizable())
|
||||
}
|
||||
|
||||
let imageView = Image(theImage, scale: 1, label: Text("generated"))
|
||||
return AnyView(
|
||||
VStack {
|
||||
imageView.resizable().clipShape(RoundedRectangle(cornerRadius: 20))
|
||||
ShareLink(item: imageView, preview: SharePreview(lastPrompt, image: imageView))
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
struct PreviewView_Previews: PreviewProvider {
|
||||
static var previews: some View {
|
||||
PreviewView(image: .constant(nil), state: .constant(.startup))
|
||||
}
|
||||
}
|
||||
@ -1,104 +0,0 @@
|
||||
//
|
||||
// TextToImage.swift
|
||||
// Diffusion
|
||||
//
|
||||
// Created by Pedro Cuenca on December 2022.
|
||||
// See LICENSE at https://github.com/huggingface/swift-coreml-diffusers/LICENSE
|
||||
//
|
||||
|
||||
import SwiftUI
|
||||
import Combine
|
||||
import StableDiffusion
|
||||
|
||||
// TODO: bind to UI controls
|
||||
let scheduler = StableDiffusionScheduler.dpmpp
|
||||
let steps = 25
|
||||
let seed: UInt32? = nil
|
||||
|
||||
func generate(pipeline: Pipeline?, prompt: String) async -> CGImage? {
|
||||
guard let pipeline = pipeline else { return nil }
|
||||
return try? pipeline.generate(prompt: prompt, scheduler: scheduler, numInferenceSteps: steps, seed: seed)
|
||||
}
|
||||
|
||||
enum GenerationState {
|
||||
case startup
|
||||
case running(StableDiffusionProgress?)
|
||||
case idle(String)
|
||||
}
|
||||
|
||||
struct ImageWithPlaceholder: View {
|
||||
var image: Binding<CGImage?>
|
||||
var state: Binding<GenerationState>
|
||||
|
||||
var body: some View {
|
||||
switch state.wrappedValue {
|
||||
case .startup: return AnyView(Image("placeholder").resizable())
|
||||
case .running(let progress):
|
||||
guard let progress = progress, progress.stepCount > 0 else {
|
||||
// The first time it takes a little bit before generation starts
|
||||
return AnyView(ProgressView())
|
||||
}
|
||||
let step = Int(progress.step) + 1
|
||||
let fraction = Double(step) / Double(progress.stepCount)
|
||||
let label = "Step \(step) of \(progress.stepCount)"
|
||||
return AnyView(ProgressView(label, value: fraction, total: 1).padding())
|
||||
case .idle(let lastPrompt):
|
||||
guard let theImage = image.wrappedValue else {
|
||||
return AnyView(Image(systemName: "exclamationmark.triangle").resizable())
|
||||
}
|
||||
|
||||
let imageView = Image(theImage, scale: 1, label: Text("generated"))
|
||||
return AnyView(
|
||||
VStack {
|
||||
imageView.resizable().clipShape(RoundedRectangle(cornerRadius: 20))
|
||||
ShareLink(item: imageView, preview: SharePreview(lastPrompt, image: imageView))
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct TextToImage: View {
|
||||
@EnvironmentObject var context: DiffusionGlobals
|
||||
|
||||
@State private var prompt = "Labrador in the style of Vermeer"
|
||||
@State private var image: CGImage? = nil
|
||||
@State private var state: GenerationState = .startup
|
||||
|
||||
@State private var progressSubscriber: Cancellable?
|
||||
|
||||
func submit() {
|
||||
if case .running = state { return }
|
||||
Task {
|
||||
state = .running(nil)
|
||||
image = await generate(pipeline: context.pipeline, prompt: prompt)
|
||||
state = .idle(prompt)
|
||||
}
|
||||
}
|
||||
|
||||
var body: some View {
|
||||
VStack {
|
||||
HStack {
|
||||
TextField("Prompt", text: $prompt)
|
||||
.textFieldStyle(.roundedBorder)
|
||||
.onSubmit {
|
||||
submit()
|
||||
}
|
||||
Button("Generate") {
|
||||
submit()
|
||||
}
|
||||
.padding()
|
||||
.buttonStyle(.borderedProminent)
|
||||
}
|
||||
ImageWithPlaceholder(image: $image, state: $state)
|
||||
.scaledToFit()
|
||||
Spacer()
|
||||
}
|
||||
.padding()
|
||||
.onAppear {
|
||||
progressSubscriber = context.pipeline!.progressPublisher.sink { progress in
|
||||
guard let progress = progress else { return }
|
||||
state = .running(progress)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,111 @@
|
||||
//
|
||||
// TextToImageView.swift
|
||||
// Diffusion
|
||||
//
|
||||
// Created by Pedro Cuenca on December 2022.
|
||||
// See LICENSE at https://github.com/huggingface/swift-coreml-diffusers/LICENSE
|
||||
//
|
||||
|
||||
import SwiftUI
|
||||
import Combine
|
||||
import StableDiffusion
|
||||
|
||||
enum GenerationState {
|
||||
case startup
|
||||
case running(StableDiffusionProgress?)
|
||||
case idle(String)
|
||||
}
|
||||
|
||||
struct TextToImageView: View {
|
||||
@EnvironmentObject var context: DiffusionGlobals
|
||||
|
||||
@State private var image: CGImage? = nil
|
||||
@State private var state: GenerationState = .startup
|
||||
@State private var prompt = "Labrador in the style of Vermeer"
|
||||
@State private var scheduler = StableDiffusionScheduler.dpmpp
|
||||
@State private var width = 512.0
|
||||
@State private var height = 512.0
|
||||
@State private var steps = 25.0
|
||||
@State private var numImages = 1.0
|
||||
@State private var seed: UInt32? = nil
|
||||
@State private var safetyOn: Bool = true
|
||||
|
||||
@State private var progressSubscriber: Cancellable?
|
||||
|
||||
func submit() {
|
||||
if case .running = state { return }
|
||||
Task {
|
||||
state = .running(nil)
|
||||
image = await generate(pipeline: context.pipeline, prompt: prompt)
|
||||
state = .idle(prompt)
|
||||
}
|
||||
}
|
||||
|
||||
var body: some View {
|
||||
VStack(alignment: .leading) {
|
||||
HStack {
|
||||
TextField("Prompt", text: $prompt)
|
||||
.textFieldStyle(.roundedBorder)
|
||||
.onSubmit {
|
||||
submit()
|
||||
}
|
||||
Button("Generate") {
|
||||
submit()
|
||||
}
|
||||
.padding()
|
||||
.buttonStyle(.borderedProminent)
|
||||
}
|
||||
Spacer()
|
||||
HStack(alignment: .top) {
|
||||
VStack(alignment: .leading) {
|
||||
Group {
|
||||
Text("Image Width")
|
||||
Slider(value: $width, in: 64...2048, step: 8, label: {},
|
||||
minimumValueLabel: {Text("64")},
|
||||
maximumValueLabel: {Text("2048")})
|
||||
Text("Image Height")
|
||||
Slider(value: $height, in: 64...2048, step: 8, label: {},
|
||||
minimumValueLabel: {Text("64")},
|
||||
maximumValueLabel: {Text("2048")})
|
||||
}
|
||||
Text("Number of Inference Steps")
|
||||
Slider(value: $steps, in: 1...300, step: 1, label: {},
|
||||
minimumValueLabel: {Text("1")},
|
||||
maximumValueLabel: {Text("300")})
|
||||
Text("Number of Images")
|
||||
Slider(value: $numImages, in: 1...8, step: 1, label: {},
|
||||
minimumValueLabel: {Text("1")},
|
||||
maximumValueLabel: {Text("8")})
|
||||
Text("Safety")
|
||||
Toggle("", isOn: $safetyOn)
|
||||
Text("Seed Check On?")
|
||||
TextField("", value: $seed, format: .number)
|
||||
}
|
||||
Spacer()
|
||||
VStack {
|
||||
PreviewView(image: $image, state: $state)
|
||||
.scaledToFit()
|
||||
}
|
||||
}
|
||||
Spacer()
|
||||
}
|
||||
.padding()
|
||||
.onAppear {
|
||||
progressSubscriber = context.pipeline?.progressPublisher.sink { progress in
|
||||
guard let progress = progress else { return }
|
||||
state = .running(progress)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func generate(pipeline: Pipeline?, prompt: String) async -> CGImage? {
|
||||
guard let pipeline = pipeline else { return nil }
|
||||
return try? pipeline.generate(prompt: prompt, scheduler: scheduler, numInferenceSteps: Int(steps), safetyOn: safetyOn, seed: seed)
|
||||
}
|
||||
}
|
||||
|
||||
struct TextToImageView_Previews: PreviewProvider {
|
||||
static var previews: some View {
|
||||
TextToImageView().environmentObject(DiffusionGlobals())
|
||||
}
|
||||
}
|
||||
Loading…
Reference in new issue