feat: Apple Silicon (MPS) compatibility — fix PR #170 quality & memory issues - #816
feat: Apple Silicon (MPS) compatibility — fix PR #170 quality & memory issues#816donghao1393 wants to merge 6 commits into
Conversation
- memory.py: auto device detection (MPS→CUDA→CPU), MPS memory via psutil, unified _empty_cache() - hunyuan_video_packed.py: torch.backends.cuda guard, dynamic cu_seqlens device - utils.py: MPS print_free_mem, PyAV mp4 save fallback - dit_common.py: LayerNorm/RMSNorm fp32 internal compute to prevent bf16 accumulation drift on MPS - demo_gradio.py: global HF cache, adaptive VAE (direct on ≥96GB, chunked otherwise for 128GB Mac safety) - requirements.txt: relaxed version pins for MPS compatibility All CUDA code paths preserved when CUDA is available. Backward compatible — no breaking changes for existing CUDA users. Fixes PR lllyasviel#170: blown-out colours, blocky noise, OOM on Apple Silicon. © 2026 董昊.
…lur on 128 GB Macs
…undary blur on 128 GB Macs" This reverts commit 0f931bd.
On ≥128 GB Macs (free >= 82 GB): offload DiT to CPU before VAE decode, use direct vae_decode() instead of chunked — eliminates chunk-boundary soft-blending entirely, removing ghosting on fast limb motion. DiT is reloaded automatically at the next section's sampling start (device check). Offload/reload overhead ~2s/section. Lower-memory Macs (≤96 GB) keep the chunked path unchanged.
MPS SDPA backend loses bf16 softmax precision beyond ~3068 tokens. Chunk Q into 3068-token windows with full K/V (mathematically exact). Also fixes varlen path that raised NotImplementedError on MPS. CUDA paths unchanged.
MPS bf16 SDPA loses ~2x softmax precision vs fp32 on large sequences (14k+ tokens), manifesting as edge smearing / stutter on fast limb motion. Convert Q/K/V to fp32 before attention, cast output back. Overhead: ~1ms dtype conversion per attention call.
Update: fp32 chunked SDPA — fast-motion ghosting resolvedThe latest commit (77ead07) upgrades the MPS attention path to fp32 compute inside the chunked SDPA kernel. MPS bf16 softmax loses ~2× precision at FramePack's 14k+ token scale, causing edge smearing and frame-to-frame stutter on rapid limb motion (dancing, hand-waving). Converting Q/K/V to fp32 before attention eliminates this — the output is cast back to the input dtype for downstream compatibility. ResultsBefore (bf16 SDPA): hands and feet appeared as smeared blobs during fast motion — "ghost fist" effect. After (fp32 SDPA): limbs are clearly resolved frame-to-frame. The dancer's hands and feet maintain structural integrity even at peak velocity. A slight softness remains on the fastest-moving edges — this is residual bf16 accumulation across ~900 DiT block passes (60 blocks × 15 steps), below the detection threshold of any individual operator and attributable to MPS hardware-level bf16 accumulation — same category as the Demo
260720_205630_332_8244_10.mp41 s, 15 steps — crisp, zero ghosting
260720_202648_046_8466_10.mp41 s, 15 steps — limbs resolved, slight softness on fastest edges All 16 MPS compatibility fixes (device detection, fp32 norms, chunked+fp32 SDPA, adaptive VAE, PyAV mp4, global HF cache) are now integrated and tested on M4 Max 128 GB. Ready for review. |
bf16/fp16 softmax loses numerical precision on sequences > ~3068 tokens — MPSGraph softMax does not internally accumulate in fp32 for half-precision inputs. At 14k tokens the relative error vs fp32 grows to ~2.4%, producing visible edge smearing in video DiT. Split Q into 3068-token chunks when dtype is bf16 or fp16 and qSize exceeds threshold. K/V are kept full-size for the complete softmax denominator — row-independence makes this exact. Test Plan: python -c "import torch; q=torch.randn(1,24,14000,128,device='mps',dtype=torch.bfloat16); k=torch.randn(1,24,14000,128,device='mps',dtype=torch.bfloat16); v=torch.randn(1,24,14000,128,device='mps',dtype=torch.bfloat16); out=torch.nn.functional.scaled_dot_product_attention(q,k,v); print(out.shape)" Discovered via FramePack image-to-video generation (lllyasviel/FramePack#816). Co-authored with AI assistant.
1 participant
Summary
Adds native Apple Silicon (MPS) support to FramePack, resolving the three quality/memory thresholds documented in #170 by @brandon929: blown-out colours, blocky noise, and OOM.
Changes (6 files, +144/-32)
diffusers_helper/memory.pypsutilwithsysctlfallback. Unified_empty_cache()helper replaces scatteredtorch.cuda.empty_cache()calls.diffusers_helper/models/hunyuan_video_packed.pytorch.backends.cudachecks (no-op on Mac). Dynamic device inget_cu_seqlens— usestext_mask.deviceinstead of hardcoded"cuda".diffusers_helper/utils.pyprint_free_mem()with MPS fallback.save_bcthw_as_mp4()uses PyAV whentorchvision >=0.28removedwrite_video.diffusers_helper/dit_common.pydemo_gradio.py~/.cache/huggingfaceinstead of a localhf_download/directory (avoids re-downloading 45 GB if models are already cached).requirements.txtBackward Compatibility
All CUDA code paths are preserved behind
is_available()guards. Existing CUDA users are unaffected — this is purely additive.Tested
Demo — FramePack on Apple Silicon (M4 Max)
*Input image — a fighter jet soaring thr
Gradio WebUI running natively on Mac — no CUDA required. High-VRAM mode detected (128 GB unified memory), all models on MPS.
260719_172543_926_3463_37.mp4
Generated video: 5 seconds, 15 steps, 480×832 resolution. Output is temporally coherent with crisp detail — no blown-out colours, no blocky noise, no structural collapse (issues documented in #170, resolved by fp32 LayerNorm computation and adaptive VAE chunking).
Prompt: Launches a missile, turns right and banks away, cinematic aerial combat (I don't know why the missile not launched, maybe the model did not spot out the missile.
Quick start on Mac:
Models auto-download to ~/.cache/huggingface on first run. 10–15 steps recommended for MPS (optimal quality/speed balance).
Additional Context
This work implements the Mac compatibility layer discussed in #170 — the
blown-out coloursat the "640" resolution bucket are resolved by the fp32 norm computation. The adaptive VAE chunking prevents OOM on machines with ≤128 GB unified memory.© 2026 董昊.