apple / ml-stable-diffusion

Stable Diffusion with Core ML on Apple Silicon
MIT License
16.66k stars 920 forks source link

Add a way to preview an images interim status #85

Open martinlexow opened 1 year ago

martinlexow commented 1 year ago

It would be great if we could get the data of each intermediate step as an image.

This way we could build a preview in our UIs like this:

catAI

This would increase user experience especially on iPhones and iPads, where computing times are usually longer than on the Mac.

ynagatomo commented 1 year ago

As an example, in my sample iOS app, I display the intermediate images step by step.

icpryde commented 1 year ago

Yes I’d love this

martinlexow commented 1 year ago

@ynagatomo Thanks so much for your example code! Could you point us to the right direction in which place you have implemented the preview method? I wasn’t able to find it.

ynagatomo commented 1 year ago

Hi. In the ImageGenerator.swift,

  1. set the progress handler when calling generateImages() method; let cgImages = try sdPipeline.generateImages(prompt: param.prompt, negativePrompt: param.negativePrompt, imageCount: param.imageCount, stepCount: param.stepCount, seed: UInt32(param.seed), guidanceScale: param.guidanceScale, disableSafety: param.disableSafety, progressHandler: self.progressHandler) // <-- progress handler
  2. The apple/ml-stable-diffusion library provides the immediate images via the progress handler's StableDiffusionPipeline.Progress parameter; nonisolated func progressHandler(progress: StableDiffusionPipeline.Progress) -> Bool { debugLog("IG: Progress: step / stepCount = (progress.step) / (progress.stepCount)")

    if ProcessInfo.processInfo.isiOSAppOnMac {
        let generatedImages = GeneratedImages(parameter: GenerationParameter(prompt: progress.prompt,
                                             negativePrompt: "unknow", // progress does not provide this now
                                             guidanceScale: 0.0, // progress does not provide this now
                                             seed: 0,
                                             stepCount: progress.stepCount,
                                             imageCount: progress.currentImages.count,
                                             disableSafety: progress.isSafetyEnabled),
                                             images: progress.currentImages.compactMap {
            if let cgImage = $0 {
                return UIImage(cgImage: cgImage)
            } else {
                return nil
            }
        })
    
        DispatchQueue.main.async {
            self.setGeneratedImages(generatedImages)
            self.setProgressStep(step: progress.step, stepCount: progress.stepCount)
        }
    } else {
        DispatchQueue.main.async {
            self.setProgressStep(step: progress.step, stepCount: progress.stepCount)
        }
    }
    
    return true // continue

    }

In my sample code, I convert the immediage images (CGImage) to UIImage and display them in SwiftUI views.

ynagatomo commented 1 year ago

Getting the intermediate image takes time, so the above code, as the simplest case, doesn't do it on iPhone or iPad, but does it on macOS. Ex. using an iPad Pro 2020/A12Z, for generating a image (512x512 px) with 20 steps;