Apple Silicon / MPS support: small set of patches makes PuLID-Flux run on M3

[UncleSaurus]

We got PuLID-Flux v0.9.1 running end-to-end on an Apple M3 Ultra (PyTorch 2.4, MPS backend, no CUDA) by applying a small set of mechanical patches to the upstream code. Inference timings on our machine: 22 s at 512×512, 98 s at 1024×1024 (20 steps, bf16). Output quality matches the Replicate-hosted zsxkib/flux-pulid reference on the same prompts and face references.

I'd like to upstream the strict-MPS-required fixes (and possibly a couple of quality-of-life ones) if you're open to PRs. Filing this issue first to confirm the direction.

What's actually broken on MPS today

1. flux/math.py::rope() uses float64 — MPS doesn't support fp64

TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework
doesn't support float64. Please use float32 instead.

This fires on the first denoise step on any Apple machine, so PuLID-Flux is completely blocked there.

Patch: branch the dtype on pos.device.type == "mps" so MPS uses fp32; CPU and CUDA keep the upstream fp64 precision.

def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
    assert dim % 2 == 0
    scale_dtype = torch.float32 if pos.device.type == "mps" else torch.float64
    scale = torch.arange(0, dim, 2, dtype=scale_dtype, device=pos.device) / dim
    # ... rest unchanged

2. Unconditional torch.cuda.empty_cache() / torch.cuda.manual_seed_all() calls

pulid/pipeline_flux.py, pulid/pipeline.py, pulid/pipeline_v1_1.py, and pulid/utils.py all call torch.cuda.empty_cache() (and manual_seed_all) unconditionally. On MPS-only machines these raise / silently no-op depending on PyTorch version, and the strategic free points they represent end up dead code.

Patch: guard with torch.cuda.is_available() and optionally also call the torch.mps.empty_cache() counterpart so the memory-management intent is preserved on Apple. We hit a real OOM on a 512 GB M3 Ultra the first time, because the upstream calls were silently no-op'd on our backend and MPS allocator grew monotonically across 20 denoise steps.

3. @torch.inference_mode() is load-bearing for MPS

Already present on upstream's FluxGenerator.generate_image — but worth a note in the README / docs that running PuLID-Flux without inference_mode (e.g. by inlining the body for a one-shot script) holds the autograd graph for all 19 double + 38 single = 57 Flux attention blocks across all 20 denoise steps and OOM's. We learned this the hard way 🙂.

Quality-of-life: skip downloads when local files exist

pulid/pipeline_flux.py calls snapshot_download('DIAMONIK7777/antelopev2', …) and hf_hub_download('guozinan/PuLID', …) unconditionally in __init__ / load_pretrain. Even when the caller passes pretrain_path=, the load_pretrain still hits HF Hub first.

Patch: wrap both with os.path.exists() checks so a fully-local setup doesn't hit the network. Backward-compatible (the download branch is preserved when files aren't there).

Configurability (optional, separate PR)

pulid/pipeline_flux.py hardcodes root='.' for FaceAnalysis and 'models/antelopev2/glintr100.onnx' for the InsightFace model. Forces callers to os.chdir() to a workdir with the right layout — annoying under threaded callers because chdir is process-wide.

Reading os.environ.get('PULID_MODELS_ROOT', '.') would let callers point at an absolute root without changing cwd. Default '.' preserves upstream behavior.

Validation

• Setup: Apple M3 Ultra, 512 GB unified RAM, PyTorch 2.4, MPS backend.
• Models: FLUX.1-dev (bf16), pulid_flux_v0.9.1.safetensors, DIAMONIK7777/antelopev2.
• MPS memory safety: setting PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.6 at process start as a safety net; never approached during 1024×1024 runs.
• Results: validated on a Fantasy Elf, Hybrid Cat, and Anthro Fox character (avatar→selfie with different scenes). Face identity + form decoration + scene composition all preserved across 9 test selfies. Visually indistinguishable from the same prompts via Replicate.

Happy to send PRs for items #1 and #2 first (smallest blast radius, fixes the actual MPS blockers), and a separate PR for the QoL bits if there's interest. Let me know what fits your direction.