* 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
3 years 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