flux does not work on MPS devices

[bghira]

Describe the bug

    scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.

Reproduction

import torch
from diffusers import  FluxPipeline

pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, revision='refs/pr/1')
#pipe.enable_model_cpu_offload()
pipe.to(device='mps')

prompt = "A cat holding a sign that says hello world"
out = pipe(
    prompt=prompt, 
    guidance_scale=0., 
    height=768, 
    width=1360, 
    num_inference_steps=4, 
    max_sequence_length=256,
).images[0]
out.save("image.png")

it also doesn't work with cpu offload.

Logs

scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.

System Info

Git master

Who can help?

@sayakpaul

[bghira]

can't switch to fp32 😢

[bghira]

@pcuenca i've been looking at workarounds and there's really nothing, and this model is too big to run on CPU, it just never really completes the first step.

[sayakpaul]

No idea how to get around this problem :(

[bghira]

it also doesn't work on ROCm, as the dimensions of the operations overflow the ROCm kernel limits, so, it has to run layer-wise and takes about 2 minutes for one image

[bghira]

maybe @rromb or @pesser have some ideas

[pcuenca]

Great investigation @bghira! It's a bit surprising that it degrades so much with float32. Also unfortunate:

RuntimeError: "arange_mps" not implemented for 'BFloat16'

[Vargol]

Has anyone tried running just the arrange on the CPU for MPS if its supported and pushing the results back on the GPU, @bghira when you said flux didn't run on the CPU is that what you meant or were you referring to running the whole model ?

Nevermind had a proper look at the code now and can that's a load of rubbish :-)

[mgierschdev]

in my case:
RuntimeError: MPS backend out of memory (MPS allocated: 81.54 GB, other allocations: 384.00 KB, max allowed: 81.60 GB). Tried to allocate 72.00 MB on private pool. Use PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.0 to disable upper limit for memory allocations (may cause system failure).

any idea ? on how to run on a Mac M2

[bghira]

you cannot run it at all on any mac series but especially not m1 as it doesnt even have bf16 in hardware

[Vargol]

can't switch to fp32 😢

Try it with torch 2.3.1

[bghira]

i am on 2.4 for bf16 and its fixes

[Vargol]

If only there was a way you could have two python environments at once <joke> , but it shows that the noisey image from float32 isn't a fundamental issue with flux

[Vargol]

Interesting, if you monkey patch the rope function to move pos to the CPU you still get noise with 2.4.0 with both float32 and float64, the fp64 runs fine on CPU with 2.3.1.

import torch
from diffusers import  FluxPipeline
import diffusers

_flux_rope = diffusers.models.transformers.transformer_flux.rope
def new_flux_rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
    assert dim % 2 == 0, "The dimension must be even."

    if pos.device.type == "mps":
        print("I got called")
        return _flux_rope(pos.to("cpu"), dim, theta).to(device=pos.device)
    else:
        print("I should not be called")
        return _flux_rope(pos, dim, theta)

diffusers.models.transformers.transformer_flux.rope = new_flux_rope

pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", revision='refs/pr/1',  torch_dtype=torch.bfloat16).to("mps")

prompt = "A cat holding a sign that says hello world"
out = pipe(
     prompt=prompt,
     guidance_scale=0.,
     height=1024,
     width=1024,
     num_inference_steps=4,
     max_sequence_length=256,
).images[0]
out.save("flux_image.png")

[bghira]

:/ the trainer i'm using requires torch 2.4, and the model is basically useless to Apple users without torch 2.4 as that is also required to run it in lower precision levels