Add MPS-compatible component dtypes

by mzx - opened 1 day ago

←

Add minimal MPS dtype handling while preserving the existing CUDA, CPU, and default-device behavior.

Load SDXL-Turbo in float32 by default on MPS.
Load LTX-Video in bfloat16 by default on MPS.
Allow the encoder to be cast before device transfer, and use float32 for the pipeline-created encoder on MPS. This converts the released V-JEPA2 checkpoint's float64 buffers before moving them to MPS.
Continue to honor explicitly supplied generator dtypes.

Local validation:

Existing test suite: 21 passed.
End-to-end NevoPipeline run on a 96 GB M3 Ultra with all components automatically loaded and resident together.
512x512 SDXL image generation and 25-frame, 8-step LTX-Video generation completed successfully.
Runtime dtypes: SDXL float32, LTX bfloat16, V-JEPA2 float32.
Observed MPS driver memory after generation: 52.4 GiB.

Environment: macOS 26.5.1 arm64, PyTorch 2.13.0, Diffusers 0.39.0, Transformers 5.13.0. The integration run used PYTORCH_ENABLE_MPS_FALLBACK=1.

Upload images, audio, and videos by dragging in the text input, pasting, or clicking here.

Tap or paste here to upload images

Ready to merge

This branch is ready to get merged automatically.

· Sign up or log in to comment