Appple MPS support for Diffusion! - Error: MLIR pass manager failed

[lukebelz]

Ran on MBP m4 Max 64gb unified memory with macOS 15.4.

This is running the default diffusion script. I git cloned and installed the latest diffusion to get the AutoencoderKLWan and WanPipeline support. I believe the issue stems from WanPipeline's use of bfloat16, which MPS does not support. Using float16 or float32 fails with this error.

python3 wan2-1.py
Loading checkpoint shards: 100%|██████████████████████████████████████████████████████████| 5/5 [00:05<00:00, 1.07s/it]
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████| 12/12 [00:35<00:00, 2.96s/it]
Loading pipeline components...: 100%|█████████████████████████████████████████████████████| 5/5 [00:41<00:00, 8.21s/it]
/AppleInternal/Library/BuildRoots/799aba5a-efcd-11ef-b353-122ba06eff56/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShadersGraph/mpsgraph/MetalPerformanceShadersGraph/Core/Files/MPSGraphExecutable.mm:3561: failed assertion `Error: MLIR pass manager failed'
[1] 31767 abort python3 wan2-1.py
/opt/homebrew/anaconda3/envs/ml/lib/python3.12/multiprocessing/resource_tracker.py:255: UserWarning: resource_tracker: There appear to be 1 leaked semaphore objects to clean up at shutdown
warnings.warn('resource_tracker: There appear to be %d '

Here is the script in case it helps. It does work if I choose cpu instead of mps:

import torch
from diffusers.utils import export_to_video
from diffusers import AutoencoderKLWan, WanPipeline
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler

model_id = "Wan-AI/Wan2.1-T2V-14B-Diffusers"
vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
flow_shift = 5.0 # 5.0 for 720P, 3.0 for 480P
scheduler = UniPCMultistepScheduler(
prediction_type='flow_prediction',
use_flow_sigmas=True,
num_train_timesteps=1000,
flow_shift=flow_shift
)

pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.float16)
pipe.scheduler = scheduler

pipe.to("mps")

prompt = "A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon. The kitchen is cozy, with sunlight streaming through the window."
negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"

output = pipe(
prompt=prompt,
negative_prompt=negative_prompt,
height=720,
width=1280,
num_frames=81,
guidance_scale=5.0,
).frames[0]

export_to_video(output, "output.mp4", fps=16)