Possible Apple MPS issue - prompts array fails to generate any images

[FahimF]

Describe the bug

If you pass an array of prompts (a list of strings) rather than single prompt (a single string) to the pipeline, under Apple mps you never get an actual image. You get a lot of NSFW warnings and some images with noise but that's it.

If you run the image generator multiple times in a loop instead, it works.

Reproduction

Try the following code for image generation under an Apple mps device and check the output.

images = pipe(
        [prompt] * num_images,
        guidance_scale=7.5,
        latents = latents,
    )["sample"]

Logs

No response

System Info

Using the source from the GitHub repo main branch since the current release version was missing some stuff.

[anton-l]

Linking this to #355 to check once it's merged :)

[pcuenca]

Hi @FahimF!

I'm not getting corrupted images in this case, but a hard failure instead (my program aborts):

/AppleInternal/Library/BuildRoots/3535eba1-198f-11ed-aa94-e2c7c6032f02/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShaders/MPSCore/Types/MPSNDArray.mm:705: failed assertion `[MPSTemporaryNDArray initWithDevice:descriptor:] Error: product of dimension sizes > 2**31'

This appears to be a manifestation of this issue. I'm not sure if there's anything we can do to prevent this problem. For now, I'll document it as a known issue until we can find a solution.

I'm curious as to why it doesn't behave the same for you. Does inference complete all the iterations in your system? What version of PyTorch are you using, and how are you testing – are you using a Jupyter notebook or a script? Thanks a lot!

[FahimF]

Hi @pcuenca 😄 Ah yes, I forgot to mention that side-effect. Sorry! Been writing about these issues for the past few days and I guess I'm not able to keep up with where I wrote what. Basically, (at least in my testing) if I tried to do [prompt] * 2 I'd get the error you got. But if I do [prompt] * 3 I do get the images but they are either full of noise or black and I do get a NSFW warning.

I was busy doing several things at that point and so didn't try to replicate further to see if the same happens for all odd numbered batches vs. even numbered batches or if 2 was the only instance where I got the error.

If interested, I wrote about a bunch of issues that I ran into with mps in the following two articles. Not sure if any of that will help you but it does show the tweaks I'm making to the code to get it to work (like patching for another crash) etc. Just linking in case it helps:
https://write.farook.org/adventures-in-imagineering-mining-the-apple-silicon/
https://write.farook.org/adventures-in-imagineering-stable-diffusion-the-unstable-bits/

[FahimF]

@pcuenca Just tested again quickly, and it appears that the crash is only for batches of 2. Batches of 3, 4, and 5 work fine ... At least with my particular set up. If you need further information, please let me know and I'd be happy to do additional testing and or provide more info 😄

[pcuenca]

Awesome @FahimF, thanks a lot! I'll take a look and test a bit more.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[pcuenca]

Removing the stale label, as I think this is still being worked on by the PyTorch team. Not sure whether there's a solution or we are facing a fundamental size limitation in the MPS engine. I'll review the latest comments in the PyTorch issues.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.