Add MPS (Apple Silicon) support

[imperatormk]

Fixes #74

Summary

Full MPS (Apple Silicon) support for Mochi video generation - no more "Torch not compiled with CUDA enabled" errors or black/noisy outputs.

MPS Compatibility Fixes

The core issue was CUDA-specific code paths and precision problems on MPS:

• RoPE embeddings: Compute in float32 (bf16 loses precision on MPS, causes artifacts)
• CFG computation: Keep in float32 (critical for proper guidance)
• VAE decode: Disable bf16 autocast on MPS (causes corruption)
• Timestep embeddings: Float32 sigma tensors
• FinalLayer / patch embed: Cast to weight dtype before ops
• Attention: Added mps_flash mode, SDPA fallback works on MPS
• Device detection: CUDA → MPS → CPU fallback throughout codebase
• CUDA guards: Wrapped torch.cuda.* calls to not crash on MPS

Memory Optimizations

• NF4 quantization: ~20GB → ~5GB model size via --quantize flag
• Lazy loading: Models loaded on-demand
• T5 cleanup: Encoder freed after prompt encoding
• DiT cleanup: Freed before VAE decode on MPS

Optional Dependencies

For additional features, install these MPS-optimized packages:

# NF4 quantization support
pip install mps-bitsandbytes

# Flash attention for MPS
pip install mps-flash-attn

# Native Conv3D for MPS
pip install mps-conv3d

CLI Flags

• --quantize - NF4 4-bit quantization (requires mps-bitsandbytes)
• --attention-mode - override: flash, mps_flash, sdpa, sage

Scripts

• scripts/quantize_dit.py - pre-quantize Genmo DiT
• scripts/quantize_diffusers_transformer.py - pre-quantize Diffusers

Usage

# Basic MPS
python demos/cli.py --model_dir weights/ --prompt "A cat"

# With quantization (~5GB instead of ~20GB)
python demos/cli.py --model_dir weights/ --quantize --prompt "A cat"

# Pre-quantize weights
python scripts/quantize_dit.py weights/dit.safetensors --dtype bf16