Merge branch 'main' into feat/sd3-modular-pipeline

Author: AlanPonnachan

.ai/models.md

@@ -73,4 +73,14 @@ Consult the implementations in `src/diffusers/models/transformers/` if you need
 
 7. **Forgetting to update `_import_structure` and `_lazy_modules`.** The top-level `src/diffusers/__init__.py` has both -- missing either one causes partial import failures.
 
-8. **Hardcoded dtype in model forward.** Don't hardcode `torch.float32` or `torch.bfloat16` in the model's forward pass. Use the dtype of the input tensors or `self.dtype` so the model works with any precision.
+8. **Hardcoded dtype in model forward.** Don't hardcode `torch.float32` or `torch.bfloat16`, and don't cast activations by reading a weight's dtype (`self.linear.weight.dtype`) — the stored weight dtype isn't the compute dtype under gguf / quantized loading. Always derive the cast target from the input tensor's dtype or `self.dtype`.
+
+9. **`torch.float64` anywhere in the model.** MPS and several NPU backends don't support float64 -- ops will either error out or silently fall back. Reference repos commonly reach for float64 in RoPE frequency bases, timestep embeddings, sinusoidal position encodings, and similar "precision-sensitive" precompute code (`torch.arange(..., dtype=torch.float64)`, `.double()`, `torch.float64` literals). When porting a model, grep for `float64` / `double()` up front and resolve as follows:
+    - **Default: just use `torch.float32`.** For inference it is almost always sufficient -- the precision difference in RoPE angles, timestep embeddings, etc. is immaterial to image/video quality. Flip it and move on.
+    - **Only if float32 visibly degrades output, fall back to the device-gated pattern** we use in the repo:
+      ```python
+      is_mps = hidden_states.device.type == "mps"
+      is_npu = hidden_states.device.type == "npu"
+      freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64
+      ```
+      See `transformer_flux.py`, `transformer_flux2.py`, `transformer_wan.py`, `unet_2d_condition.py` for reference usages. Never leave an unconditional `torch.float64` in the model.

examples/research_projects/pytorch_xla/inference/flux/README.md

@@ -51,7 +51,42 @@ python flux_inference.py
 
 The script loads the text encoders onto the CPU and the Flux transformer and VAE models onto the TPU. The first time the script runs, the compilation time is longer, while the cache stores the compiled programs. On subsequent runs, compilation is much faster and the subsequent passes being the fastest. 
 
-On a Trillium v6e-4, you should expect ~6 sec / 4 images or 1.5 sec / image (as devices run generation in parallel):
+On a Trillium v6e-4, you should expect ~6 sec / 4 images or 1.5 sec / image (as devices run generation in parallel).
+
+> **Note:** `flux_inference.py` uses `xmp.spawn` (one process per chip) and requires the full model to fit on a single chip. If you run into OOM errors (e.g., on v5e with 16GB HBM per chip), use the SPMD version instead — see below.
+
+### SPMD version (for v5e-8 and similar)
+
+On TPU configurations where a single chip cannot hold the full FLUX transformer (~16GB in bf16), use `flux_inference_spmd.py`. This script uses PyTorch/XLA SPMD to shard the transformer across multiple chips using a `(data, model)` mesh — 4-way model parallel so each chip holds ~4GB of weights, with the remaining chips for data parallelism.
+
+```bash
+python flux_inference_spmd.py --schnell
+```
+
+Key differences from `flux_inference.py`:
+- **Single-process SPMD** instead of multi-process `xmp.spawn` — the XLA compiler handles all collective communication transparently.
+- **Transformer weights are sharded** across the `"model"` mesh axis using `xs.mark_sharding`.
+- **VAE lives on CPU**, moved to XLA only for decode (then moved back), since the transformer stays on device throughout.
+- **Text encoding** runs on CPU before loading the transformer.
+
+On a v5litepod-8 (v5e, 8 chips, 16GB HBM each) with FLUX.1-schnell, expect ~1.76 sec/image at steady state (after compilation):
+
+```
+2026-04-15 02:24:30 [info     ] SPMD mesh: (2, 4), axes: ('data', 'model'), devices: 8
+2026-04-15 02:24:30 [info     ] encoding prompt on CPU...
+2026-04-15 02:26:20 [info     ] loading VAE on CPU...
+2026-04-15 02:26:20 [info     ] loading flux transformer from black-forest-labs/FLUX.1-schnell
+2026-04-15 02:27:22 [info     ] starting compilation run...
+2026-04-15 02:52:55 [info     ] compilation took 1533.4575625509997 sec.
+2026-04-15 02:52:56 [info     ] starting inference run...
+2026-04-15 02:56:11 [info     ] inference time: 195.74092420299985
+2026-04-15 02:56:13 [info     ] inference time: 1.7625778899996476
+2026-04-15 02:56:13 [info     ] avg. inference over 2 iterations took 98.75175104649975 sec.
+```
+
+The first inference iteration includes VAE compilation (~195s). The second iteration shows the true steady-state speed (~1.76s).
+
+### v6e-4 results (original `flux_inference.py`)
 
 ```bash
 WARNING:root:libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU.

examples/research_projects/pytorch_xla/inference/flux/flux_inference_spmd.py

@@ -0,0 +1,193 @@
+"""FLUX inference on TPU using PyTorch/XLA SPMD.
+
+Uses SPMD to shard the transformer across multiple TPU chips, enabling
+inference on devices where the model doesn't fit on a single chip (e.g., v5e).
+The VAE is loaded on CPU at startup, moved to XLA for decode, then moved back.
+"""
+
+from argparse import ArgumentParser
+from pathlib import Path
+from time import perf_counter
+
+import numpy as np
+import structlog
+import torch
+import torch_xla.core.xla_model as xm
+import torch_xla.debug.metrics as met
+import torch_xla.debug.profiler as xp
+import torch_xla.distributed.spmd as xs
+import torch_xla.runtime as xr
+from torch_xla.experimental.custom_kernel import FlashAttention
+
+from diffusers import AutoencoderKL, FluxPipeline
+
+
+cache_path = Path("/tmp/data/compiler_cache_eXp")
+cache_path.mkdir(parents=True, exist_ok=True)
+xr.initialize_cache(str(cache_path), readonly=False)
+xr.use_spmd()
+
+logger = structlog.get_logger()
+metrics_filepath = "/tmp/metrics_report.txt"
+VAE_SCALE_FACTOR = 8
+
+
+def _vae_decode(latents, vae, height, width, device):
+    """Move VAE to XLA, decode latents, move VAE back to CPU."""
+    vae.to(device)
+    latents = FluxPipeline._unpack_latents(latents, height, width, VAE_SCALE_FACTOR)
+    latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor
+    with torch.no_grad():
+        image = vae.decode(latents, return_dict=False)[0]
+    vae.to("cpu")
+    return image
+
+
+def main(args):
+    # --- SPMD mesh: 4-way model parallel to fit transformer + VAE on v5e chips ---
+    num_devices = xr.global_runtime_device_count()
+    if num_devices >= 4:
+        mesh = xs.Mesh(np.arange(num_devices), (num_devices // 4, 4), ("data", "model"))
+    else:
+        NotImplementedError
+    xs.set_global_mesh(mesh)
+    logger.info(f"SPMD mesh: {mesh.mesh_shape}, axes: {mesh.axis_names}, devices: {num_devices}")
+
+    # --- Profiler ---
+    profile_path = Path("/tmp/data/profiler_out_eXp")
+    profile_path.mkdir(parents=True, exist_ok=True)
+    profiler_port = 9012
+    profile_duration = args.profile_duration
+    if args.profile:
+        logger.info(f"starting profiler on port {profiler_port}")
+        _ = xp.start_server(profiler_port)
+
+    device = xm.xla_device()
+
+    # --- Checkpoint ---
+    if args.schnell:
+        ckpt_id = "black-forest-labs/FLUX.1-schnell"
+    else:
+        ckpt_id = "black-forest-labs/FLUX.1-dev"
+
+    # --- Text encoding (CPU) ---
+    prompt = "photograph of an electronics chip in the shape of a race car with trillium written on its side"
+    logger.info("encoding prompt on CPU...")
+    text_pipe = FluxPipeline.from_pretrained(ckpt_id, transformer=None, vae=None, torch_dtype=torch.bfloat16).to("cpu")
+    with torch.no_grad():
+        prompt_embeds, pooled_prompt_embeds, _ = text_pipe.encode_prompt(
+            prompt=prompt, prompt_2=None, max_sequence_length=512
+        )
+    image_processor = text_pipe.image_processor
+    del text_pipe
+
+    # --- Load VAE on CPU (moved to XLA only for decode) ---
+    logger.info("loading VAE on CPU...")
+    vae = AutoencoderKL.from_pretrained(ckpt_id, subfolder="vae", torch_dtype=torch.bfloat16)
+
+    # --- Load transformer and shard ---
+    logger.info(f"loading flux transformer from {ckpt_id}")
+    flux_pipe = FluxPipeline.from_pretrained(
+        ckpt_id,
+        text_encoder=None,
+        tokenizer=None,
+        text_encoder_2=None,
+        tokenizer_2=None,
+        vae=None,
+        torch_dtype=torch.bfloat16,
+    ).to(device)
+
+    for name, param in flux_pipe.transformer.named_parameters():
+        if param.dim() >= 2:
+            spec = [None] * param.dim()
+            largest_dim = max(range(param.dim()), key=lambda d: param.shape[d])
+            spec[largest_dim] = "model"
+            xs.mark_sharding(param, mesh, tuple(spec))
+
+    flux_pipe.transformer.enable_xla_flash_attention(partition_spec=("data", None, None, None), is_flux=True)
+    FlashAttention.DEFAULT_BLOCK_SIZES = {
+        "block_q": 1536,
+        "block_k_major": 1536,
+        "block_k": 1536,
+        "block_b": 1536,
+        "block_q_major_dkv": 1536,
+        "block_k_major_dkv": 1536,
+        "block_q_dkv": 1536,
+        "block_k_dkv": 1536,
+        "block_q_dq": 1536,
+        "block_k_dq": 1536,
+        "block_k_major_dq": 1536,
+    }
+
+    width = args.width
+    height = args.height
+    guidance = args.guidance
+    n_steps = 4 if args.schnell else 28
+
+    prompt_embeds = prompt_embeds.to(device)
+    pooled_prompt_embeds = pooled_prompt_embeds.to(device)
+    xs.mark_sharding(prompt_embeds, mesh, ("data", None, None))
+    xs.mark_sharding(pooled_prompt_embeds, mesh, ("data", None))
+
+    # --- Compilation run ---
+    logger.info("starting compilation run...")
+    ts = perf_counter()
+    latents = flux_pipe(
+        prompt_embeds=prompt_embeds,
+        pooled_prompt_embeds=pooled_prompt_embeds,
+        num_inference_steps=28,
+        guidance_scale=guidance,
+        height=height,
+        width=width,
+        output_type="latent",
+    ).images
+    image = _vae_decode(latents, vae, height, width, device)
+    image = image_processor.postprocess(image)[0]
+    logger.info(f"compilation took {perf_counter() - ts} sec.")
+    image.save("/tmp/compile_out.png")
+
+    # --- Inference loop ---
+    seed = 4096 if args.seed is None else args.seed
+    xm.set_rng_state(seed=seed, device=device)
+    times = []
+    logger.info("starting inference run...")
+    for _ in range(args.itters):
+        ts = perf_counter()
+
+        if args.profile:
+            xp.trace_detached(f"localhost:{profiler_port}", str(profile_path), duration_ms=profile_duration)
+        latents = flux_pipe(
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            num_inference_steps=n_steps,
+            guidance_scale=guidance,
+            height=height,
+            width=width,
+            output_type="latent",
+        ).images
+        image = _vae_decode(latents, vae, height, width, device)
+        image = image_processor.postprocess(image)[0]
+        inference_time = perf_counter() - ts
+        logger.info(f"inference time: {inference_time}")
+        times.append(inference_time)
+
+    logger.info(f"avg. inference over {args.itters} iterations took {sum(times) / len(times)} sec.")
+    image.save("/tmp/inference_out.png")
+    metrics_report = met.metrics_report()
+    with open(metrics_filepath, "w+") as fout:
+        fout.write(metrics_report)
+    logger.info(f"saved metric information as {metrics_filepath}")
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("--schnell", action="store_true", help="run flux schnell instead of dev")
+    parser.add_argument("--width", type=int, default=1024, help="width of the image to generate")
+    parser.add_argument("--height", type=int, default=1024, help="height of the image to generate")
+    parser.add_argument("--guidance", type=float, default=3.5, help="guidance strength for dev")
+    parser.add_argument("--seed", type=int, default=None, help="seed for inference")
+    parser.add_argument("--profile", action="store_true", help="enable profiling")
+    parser.add_argument("--profile-duration", type=int, default=10000, help="duration for profiling in msec.")
+    parser.add_argument("--itters", type=int, default=15, help="items to run inference and get avg time in sec.")
+    args = parser.parse_args()
+    main(args)

scripts/convert_longcat_audio_dit_to_diffusers.py

@@ -131,6 +131,7 @@ def convert_longcat_audio_dit(
         cross_attn_norm=config.get("dit_cross_attn_norm", False),
         eps=config.get("dit_eps", 1e-6),
         use_latent_condition=config.get("dit_use_latent_condition", True),
+        ff_mult=config.get("dit_ff_mult", 4),
     )
     transformer.load_state_dict(transformer_state_dict, strict=True)
     transformer = transformer.to(dtype=torch_dtype)

setup.py

@@ -124,7 +124,7 @@
     "pytest-xdist",
     "python>=3.10.0",
     "ruff==0.9.10",
-    "safetensors>=0.3.1",
+    "safetensors>=0.8.0-rc.0",
     "sentencepiece>=0.1.91,!=0.1.92",
     "GitPython<3.1.19",
     "scipy",

src/diffusers/__init__.py

@@ -458,6 +458,8 @@
             "HeliosPyramidDistilledAutoBlocks",
             "HeliosPyramidDistilledModularPipeline",
             "HeliosPyramidModularPipeline",
+            "HunyuanVideo15AutoBlocks",
+            "HunyuanVideo15ModularPipeline",
             "LTXAutoBlocks",
             "LTXModularPipeline",
             "QwenImageAutoBlocks",
@@ -1246,6 +1248,8 @@
             HeliosPyramidDistilledAutoBlocks,
             HeliosPyramidDistilledModularPipeline,
             HeliosPyramidModularPipeline,
+            HunyuanVideo15AutoBlocks,
+            HunyuanVideo15ModularPipeline,
             LTXAutoBlocks,
             LTXModularPipeline,
             QwenImageAutoBlocks,

src/diffusers/dependency_versions_table.py

@@ -31,7 +31,7 @@
     "pytest-xdist": "pytest-xdist",
     "python": "python>=3.10.0",
     "ruff": "ruff==0.9.10",
-    "safetensors": "safetensors>=0.3.1",
+    "safetensors": "safetensors>=0.8.0-rc.0",
     "sentencepiece": "sentencepiece>=0.1.91,!=0.1.92",
     "GitPython": "GitPython<3.1.19",
     "scipy": "scipy",

src/diffusers/models/transformers/transformer_longcat_audio_dit.py

@@ -475,6 +475,7 @@ def __init__(
         cross_attn_norm: bool = False,
         eps: float = 1e-6,
         use_latent_condition: bool = True,
+        ff_mult: float = 4.0,
     ):
         super().__init__()
         dim = dit_dim
@@ -498,7 +499,7 @@ def __init__(
                     cross_attn_norm=cross_attn_norm,
                     adaln_type=adaln_type,
                     adaln_use_text_cond=adaln_use_text_cond,
-                    ff_mult=4.0,
+                    ff_mult=ff_mult,
                 )
                 for _ in range(dit_depth)
             ]

src/diffusers/modular_pipelines/__init__.py

@@ -89,6 +89,10 @@
         "QwenImageLayeredModularPipeline",
         "QwenImageLayeredAutoBlocks",
     ]
+    _import_structure["hunyuan_video1_5"] = [
+        "HunyuanVideo15AutoBlocks",
+        "HunyuanVideo15ModularPipeline",
+    ]
     _import_structure["ltx"] = [
         "LTXAutoBlocks",
         "LTXModularPipeline",
@@ -124,6 +128,10 @@
             HeliosPyramidDistilledModularPipeline,
             HeliosPyramidModularPipeline,
         )
+        from .hunyuan_video1_5 import (
+            HunyuanVideo15AutoBlocks,
+            HunyuanVideo15ModularPipeline,
+        )
         from .ltx import LTXAutoBlocks, LTXModularPipeline
         from .modular_pipeline import (
             AutoPipelineBlocks,

src/diffusers/modular_pipelines/hunyuan_video1_5/__init__.py

@@ -0,0 +1,49 @@
+from typing import TYPE_CHECKING
+
+from ...utils import (
+    DIFFUSERS_SLOW_IMPORT,
+    OptionalDependencyNotAvailable,
+    _LazyModule,
+    get_objects_from_module,
+    is_torch_available,
+    is_transformers_available,
+)
+
+
+_dummy_objects = {}
+_import_structure = {}
+
+try:
+    if not (is_transformers_available() and is_torch_available()):
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    from ...utils import dummy_torch_and_transformers_objects  # noqa F403
+
+    _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
+else:
+    _import_structure["modular_blocks_hunyuan_video1_5"] = [
+        "HunyuanVideo15AutoBlocks",
+    ]
+    _import_structure["modular_pipeline"] = ["HunyuanVideo15ModularPipeline"]
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+    try:
+        if not (is_transformers_available() and is_torch_available()):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
+    else:
+        from .modular_blocks_hunyuan_video1_5 import HunyuanVideo15AutoBlocks
+        from .modular_pipeline import HunyuanVideo15ModularPipeline
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+        module_spec=__spec__,
+    )
+
+    for name, value in _dummy_objects.items():
+        setattr(sys.modules[__name__], name, value)