apple / ml-stable-diffusion

Stable Diffusion with Core ML on Apple Silicon
MIT License
16.95k stars 948 forks source link

Generation with ControlNet crashes on iPadOS #248

Open SaladDays831 opened 1 year ago

SaladDays831 commented 1 year ago

Generation with ControlNet crashes on iPadOS when using model converted with the script from the main branch The crash doesn't happen with a model converted with the same command, but using the script from this commit https://github.com/apple/ml-stable-diffusion/commit/b61c9aea05370d4bc06fce2dc00a002b21f13da5 (i.e. the Swift package is still on 'main', but for model conversion I used an older version)

Branch: main iPadOS 17

computeUnits = .cpuAndNeuralEngine
disableSafety: true
reduceMemory: true

SD + CN model used:

python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --convert-text-encoder --convert-vae-decoder --convert-vae-encoder --model-version "runwayml/stable-diffusion-v1-5" --unet-support-controlnet --quantize-nbits 6 --attention-implementation SPLIT_EINSUM_V2 --convert-controlnet "lllyasviel/control_v11p_sd15_canny" "lllyasviel/control_v11p_sd15_inpaint" --bundle-resources-for-swift-cli -o "/path/to/save"

The crash happens in this method of the ControlNet class:

execute(
        latents: [MLShapedArray<Float32>],
        timeStep: Int,
        hiddenStates: MLShapedArray<Float32>,
        images: [MLShapedArray<Float32>]
    ) throws -> [[String: MLShapedArray<Float32>]] { ... }
CleanShot 2023-08-30 at 14 01 34@2x

The crash happens when using both single/multiple CN models

atiorh commented 1 year ago

Hi @SaladDays831, thanks for the report! Could you please report (using the same exact same model bundle) whether the CLI on the Mac hit the same issue?

SaladDays831 commented 1 year ago

Hi @atiorh sorry for the delay I'm seeing a similar error when using the CLI:

Build complete! (0.16s)
Loading resources and creating pipeline
(Note: This can take a while the first time using these resources)
Sampling ...
CoreML/MLShapedArray.swift:557: Fatal error: MLMultiArray of data type Float16 is not supported.
atiorh commented 1 year ago

@SaladDays831 Could you please test latest main (c506322)?

SaladDays831 commented 1 year ago

Hi @atiorh I'm still running into the same crash on latest main

Thons commented 1 year ago

@SaladDays831 @atiorh This issue can be resolved by replacing the code with the following:

// add the parameter 'converting:'
outputs[n][k] = MLShapedArray<Float32>(converting: newValue)

If you use multiple ControlNets, you might encounter crashes at the line vDSP_vadd(inputPointer, 1, outputPointer, 1, outputPointer, 1, vDSP_Length(count)) in the else block. In that case, you can solve it this way:

if {
    outputs[n][k] = MLShapedArray<Float32>(converting: newValue)
}
else {
    if var outputArray = outputs[n][k] {
        let newValueShapedArray = MLShapedArray<Float32>(converting: newValue)
        let count = newValueShapedArray.count
        outputArray.withUnsafeMutableShapedBufferPointer { outputArr, _, _ in
            newValueShapedArray.withUnsafeShapedBufferPointer { inputArr, _, _ in
                if let inputAddress = inputArr.baseAddress,
                   let outputAddress = outputArr.baseAddress {
                    vDSP_vadd(inputAddress, 1, outputAddress, 1, outputAddress, 1, vDSP_Length(count))
                }
            }
        }
    }
}
SaladDays831 commented 1 year ago

@Thons thanks! I actually tried the solution you sent and ran into the same accelerate crash. Will try the fix tomorrow!

SaladDays831 commented 1 year ago

@atiorh the fix provided by @Thons works, I'm able to generate images using multiple ControlNet models, but this brings me to this https://github.com/apple/ml-stable-diffusion/issues/216 issue - only the first ControlNet model is being used

SaladDays831 commented 1 year ago

I tried bringing back the old CN logic (slow, but supports multiple CN models) while working around the MLShapedArray crash like so, and generating an image with inpaint + canny ControlNets

for n in 0..<results.count {
                let result = results.features(at: n)
                for k in result.featureNames {
                    let newValue = result.featureValue(for: k)!.multiArrayValue!
                    if modelIndex == 0 {
                        outputs[n][k] = MLShapedArray<Float32>(converting: newValue)
                    } else {
                        let shapedArray = MLShapedArray<Float32>(converting: newValue)
                        outputs[n][k]!.withUnsafeMutableShapedBufferPointer { pt, _, _ in
                            for (i, v) in shapedArray.scalars.enumerated() { pt[i] += v }
                        }
                    }
                }
            }

The generated image seems to use both ControlNets, but there's a problem with texturing (?) the subject. This wasn't an issue before b8b5e886df33f610c6986494f31b0206a55dc8d8

Thons commented 1 year ago

@SaladDays831 I'm sorry, it was my oversight that the second ControlNet wasn't functioning. My code shared with you did not crash but was ineffective. I also tried it and it does seem to have an issue. Logically speaking, that piece of code should be worked, more in-depth investigation is needed.