[new-model] Add PAN2 text-to-video and image-to-video

Port of vllm-project/vllm-omni#8602. PAN2 is a 29B mixture-of-transformers
video DiT (separate video and text streams joined in one attention per
block) with a Qwen3.5 text encoder and the Wan2.2 VAE. One pipeline serves
text-to-video and first-frame image-to-video.

- Native PAN2 DiT; per-stream Q/K/V fused at load time, Ulysses SP with
  replicated text tokens through DistributedAttention.
- Native Qwen3.5 text encoder (Gated DeltaNet + gated attention) with a
  PyTorch chunked delta rule.
- Reuses the Wan2.2 VAE and the flow-match Euler scheduler.
- PAN2 stages: chat-template prompt crop at hidden_states[-3], first-frame
  conditioning, uniform sigmas with shift 7, fp32 latents. Adds a
  DenoisingStage.prepare_scheduler_model_output hook for the last one.
- Registry entry for the Modular Diffusers PAN2ModularPipeline manifest,
  pan2_t2v / pan2_i2v presets, example, docs, and local parity tests
  against a diffusers/transformers reference on the public tiny checkpoint.

docs/design/inference_schema_parity_inventory.yaml

@@ -283,6 +283,15 @@ surfaces:
           - fastvideo.configs.pipelines.hunyuan15.Hunyuan15T2V720PConfig
           - fastvideo.configs.pipelines.hyworld.HYWorldConfig
           - fastvideo.configs.pipelines.hyworld.Hunyuan15T2V480PConfig
+      text_system_prompt:
+        sources:
+          - fastvideo.configs.pipelines.pan2.PAN2PipelineConfig
+      text_encoder_skip_layers:
+        sources:
+          - fastvideo.configs.pipelines.pan2.PAN2PipelineConfig
+      text_max_sequence_length:
+        sources:
+          - fastvideo.configs.pipelines.pan2.PAN2PipelineConfig
       text_encoder_max_lengths:
         sources:
           - fastvideo.configs.pipelines.hunyuan15.Hunyuan15I2V480PStepDistilledConfig

docs/inference/pan2.md

@@ -0,0 +1,63 @@
+# PAN2 Text/Image to Video
+
+PAN2 is a 29B mixture-of-transformers video DiT with separate video and text streams, a Qwen3.5 text encoder and the
+Wan2.2 VAE. One pipeline serves text-to-video and image-to-video: pass `image_path` to condition on a first frame and
+leave it out for text only.
+
+## Model
+
+| Hub repo | Pipeline | Workloads | Example |
+|---|---|---|---|
+| `IFM/PAN2` | `PAN2ModularPipeline` | T2V, I2V | [`basic_pan2.py`](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_pan2.py) |
+
+PAN2 checkpoints use the Modular Diffusers layout (`modular_model_index.json`). Any repo id or local directory whose
+name contains `pan2` (or whose manifest declares `PAN2ModularPipeline`) resolves to this pipeline. The defaults come
+from the `pan2_t2v` and `pan2_i2v` presets:
+
+| Setting | Default |
+|---|---|
+| Resolution | 704x1248 (height and width must be multiples of 32) |
+| Frames | 313, 13 s at 24 fps (`num_frames - 1` must be a multiple of 4) |
+| Steps | 50, uniform sigmas from 1 to 0 with flow shift 7 |
+| Guidance | 3.0, classifier-free guidance with an empty negative prompt |
+
+PAN2 is trained on the structured captions its prompt rewriter produces, serialized as one-line JSON. Short raw
+prompts work, with weaker results.
+
+## Usage
+
+```bash
+python examples/inference/basic/basic_pan2.py
+```
+
+```python
+from fastvideo import VideoGenerator
+
+generator = VideoGenerator.from_pretrained(
+    "IFM/PAN2",
+    num_gpus=1,
+    dit_cpu_offload=False,
+    text_encoder_cpu_offload=True,
+)
+generator.generate_video(
+    "A red fox trots through fresh snow in a pine forest at dawn.",
+    image_path=None,  # or the path of a first frame for image-to-video
+    output_path="video_samples_pan2",
+)
+```
+
+For image-to-video, the first frame is scaled to cover the requested resolution and center cropped, then encoded by the
+VAE into the first latent frame of the conditioning input.
+
+## Parallelism
+
+With `num_gpus > 1`, FastVideo runs Ulysses sequence parallelism: the video tokens are sharded across ranks and the
+text tokens are replicated. The number of GPUs must divide the 32 attention heads. Sequence-parallel padding of the
+video tokens is trimmed before attention, so no attention mask is needed.
+
+## Differences from the vLLM-Omni integration
+
+- The optional prompt and video guardrails (`pan2-guardrail`) are not included.
+- The Cache-DiT `quality` policy and CFG-parallel execution are vLLM-Omni serving features and are not ported.
+- The text encoder is a FastVideo-native Qwen3.5 port. Its Gated DeltaNet layers use a PyTorch chunked delta rule
+  instead of the `fla` kernels.

docs/inference/support_matrix.md

@@ -74,6 +74,7 @@ column links a runnable script in `examples/inference/basic/` where one exists.
 | matrixgame | `FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers`<br>`FastVideo/Matrix-Game-2.0-GTA-Distilled-Diffusers`<br>`FastVideo/Matrix-Game-2.0-TempleRun-Distilled-Diffusers`<br>`FastVideo/Matrix-Game-2.0-Base-Diffusers`<br>`FastVideo/Matrix-Game-2.0-GTA-Diffusers`<br>`FastVideo/Matrix-Game-2.0-TempleRun-Diffusers`<br>`mignonjia/mg_longtuning_distilled_zelda`<br>`mignonjia/mg_sf_distilled_zelda_1k_steps`<br>`mignonjia/mg_sf_distilled_zelda`<br>`mignonjia/mg_causal_zelda`<br>`mignonjia/mg_bidirectional_zelda` | I2V | [basic_matrixgame2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_matrixgame2.py) |
 | matrixgame | `FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers` | I2V | [basic_matrixgame3.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_matrixgame3.py) |
 | minimax_h3 | `MiniMaxAI/MiniMax-H3` | T2V, I2V | [T2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_t2v.py)<br>[FL2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_fl2va.py)<br>[Ref2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_ref2va.py) |
+| pan2 | `IFM/PAN2` | T2V, I2V | [basic_pan2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_pan2.py) |
 | sd35 | `stabilityai/stable-diffusion-3.5-medium` | T2I | [basic_sd35_t2i.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_sd35_t2i.py) |
 | stable_audio | `FastVideo/stable-audio-open-1.0-Diffusers` | T2V | [basic_stable_audio.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_stable_audio.py) |
 | stable_audio | `FastVideo/stable-audio-open-small-Diffusers` | T2V | [basic_stable_audio_small.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_stable_audio_small.py) |

examples/inference/basic/basic_pan2.py

@@ -0,0 +1,40 @@
+import os
+
+from fastvideo import VideoGenerator
+
+# PAN2 is trained on the structured captions its prompt rewriter produces (serialized as one-line JSON). Plain
+# prompts like this one work, with weaker results.
+PROMPT = ("A red fox trots through fresh snow in a pine forest at dawn, its breath visible in the cold air. "
+          "The camera tracks it from the side at ground level as low golden light filters through the trees.")
+
+OUTPUT_PATH = "video_samples_pan2"
+MODEL_PATH = os.environ.get("PAN2_MODEL_PATH", "IFM/PAN2")
+
+# Set to a first frame (e.g. "assets/images/girl.png") for image-to-video; the same pipeline runs both modes.
+IMAGE_PATH = None
+
+
+def main():
+    # The 29B DiT runs in bf16; with more GPUs, num_gpus enables Ulysses sequence parallelism (the 32 attention
+    # heads must be divisible by it).
+    generator = VideoGenerator.from_pretrained(
+        MODEL_PATH,
+        num_gpus=int(os.environ.get("PAN2_NUM_GPUS", "1")),
+        use_fsdp_inference=False,
+        dit_cpu_offload=False,
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=True,
+    )
+
+    # Defaults come from the pan2 preset: 704x1248, 313 frames (13 s at 24 fps), 50 steps, guidance 3.0.
+    generator.generate_video(
+        PROMPT,
+        image_path=IMAGE_PATH,
+        output_path=OUTPUT_PATH,
+        save_video=True,
+    )
+
+
+if __name__ == "__main__":
+    main()

fastvideo/configs/models/dits/__init__.py

@@ -12,6 +12,7 @@ from fastvideo.configs.models.dits.longcat import LongCatVideoConfig
 from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig
 from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig
 from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config
+from fastvideo.configs.models.dits.pan2 import PAN2VideoConfig
 from fastvideo.configs.models.dits.mmaudio import (
     MMAUDIO_44K_TRAINING_VARIANTS,
     MMAUDIO_TRAINING_VARIANTS,
@@ -37,7 +38,7 @@ __all__ = [
     "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "Kandinsky6ArchConfig",
     "Kandinsky6VideoAudioConfig", "Kandinsky6SRArchConfig", "Kandinsky6SRConfig", "MagiHumanVideoConfig",
     "StableAudioConfig", "GlmImageDiTConfig", "HeliosConfig", "LingBotWorld2CausalFastVideoConfig",
-    "LingBotWorldFastVideoConfig", "LingBotVideoConfig", "MiniMaxH3Config", "ZImageDiTConfig", "MMAudioArchConfig",
-    "MMAudioTransformerConfig", "MMAUDIO_VARIANT_ARCHITECTURES", "MMAUDIO_TRAINING_VARIANTS",
+    "LingBotWorldFastVideoConfig", "LingBotVideoConfig", "MiniMaxH3Config", "PAN2VideoConfig", "ZImageDiTConfig",
+    "MMAudioArchConfig", "MMAudioTransformerConfig", "MMAUDIO_VARIANT_ARCHITECTURES", "MMAUDIO_TRAINING_VARIANTS",
     "MMAUDIO_44K_TRAINING_VARIANTS", "get_mmaudio_transformer_config"
 ]

fastvideo/configs/models/dits/pan2.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+from fastvideo.platforms import AttentionBackendEnum
+
+
+def is_transformer_block(n: str, m) -> bool:
+    return "transformer_blocks" in n and str.isdigit(n.split(".")[-1])
+
+
+def is_refiner_block(n: str, m) -> bool:
+    return "context_refiner" in n and str.isdigit(n.split(".")[-1])
+
+
+@dataclass
+class PAN2ArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_transformer_block, is_refiner_block])
+    _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
+        AttentionBackendEnum.FLASH_ATTN,
+        AttentionBackendEnum.TORCH_SDPA,
+    )
+
+    # Diffusers PAN2Transformer3DModel keys -> native keys. The per-stream Q/K/V projections are fused on load.
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^x_embedder\.(.*)$": r"x_embedder.proj.\1",
+            r"^time_embedder\.linear_1\.(.*)$": r"time_embedder.mlp.fc_in.\1",
+            r"^time_embedder\.linear_2\.(.*)$": r"time_embedder.mlp.fc_out.\1",
+            r"^context_embedder\.0\.(.*)$": r"context_norm.\1",
+            r"^context_embedder\.1\.(.*)$": r"context_proj.\1",
+            r"^context_refiner\.(\d+)\.attn\.to_q\.(.*)$": (r"context_refiner.\1.attn.to_qkv.\2", 0, 3),
+            r"^context_refiner\.(\d+)\.attn\.to_k\.(.*)$": (r"context_refiner.\1.attn.to_qkv.\2", 1, 3),
+            r"^context_refiner\.(\d+)\.attn\.to_v\.(.*)$": (r"context_refiner.\1.attn.to_qkv.\2", 2, 3),
+            r"^context_refiner\.(\d+)\.attn\.to_out\.0\.(.*)$": r"context_refiner.\1.attn.to_out.\2",
+            r"^context_refiner\.(\d+)\.ff\.net\.0\.proj\.(.*)$": r"context_refiner.\1.ff.fc_in.\2",
+            r"^context_refiner\.(\d+)\.ff\.net\.2\.(.*)$": r"context_refiner.\1.ff.fc_out.\2",
+            r"^transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$": (r"transformer_blocks.\1.attn.to_qkv.\2", 0, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$": (r"transformer_blocks.\1.attn.to_qkv.\2", 1, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$": (r"transformer_blocks.\1.attn.to_qkv.\2", 2, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.add_q_proj\.(.*)$":
+            (r"transformer_blocks.\1.attn.add_qkv_proj.\2", 0, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.add_k_proj\.(.*)$":
+            (r"transformer_blocks.\1.attn.add_qkv_proj.\2", 1, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.add_v_proj\.(.*)$":
+            (r"transformer_blocks.\1.attn.add_qkv_proj.\2", 2, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$": r"transformer_blocks.\1.attn.to_out.\2",
+            r"^transformer_blocks\.(\d+)\.(ff|ff_context)\.net\.0\.proj\.(.*)$": r"transformer_blocks.\1.\2.fc_in.\3",
+            r"^transformer_blocks\.(\d+)\.(ff|ff_context)\.net\.2\.(.*)$": r"transformer_blocks.\1.\2.fc_out.\3",
+        })
+
+    reverse_param_names_mapping: dict = field(default_factory=lambda: {})
+
+    patch_size: tuple[int, ...] = (1, 2, 2)
+    in_channels: int = 97
+    out_channels: int = 48
+    num_attention_heads: int = 32
+    attention_head_dim: int = 128
+    num_layers: int = 44
+    num_refiner_layers: int = 4
+    mlp_ratio: float = 8.0
+    context_mlp_ratio: float = 4.0
+    refiner_mlp_ratio: float = 4.0
+    text_embed_dim: int = 4096
+    rope_axes_dim: tuple[int, ...] = (16, 56, 56)
+    rope_theta: float = 256.0
+    # Frame positions are rescaled to the frame rate the rotary embedding was trained at (16 fps of 24 fps).
+    rope_temporal_scale: float = 16 / 24
+    # PAN2 runs its text stream in the DiT dtype.
+    cast_prompt_embeds_to_dit_dtype: bool = True
+    exclude_lora_layers: list[str] = field(
+        default_factory=lambda: ["x_embedder", "time_embedder", "context_norm", "context_proj", "context_refiner"])
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.patch_size = tuple(self.patch_size)
+        self.rope_axes_dim = tuple(self.rope_axes_dim)
+        self.hidden_size: int = self.attention_head_dim * self.num_attention_heads
+        self.num_channels_latents: int = self.out_channels
+
+
+@dataclass
+class PAN2VideoConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=PAN2ArchConfig)
+
+    prefix: str = "PAN2"

fastvideo/configs/models/encoders/__init__.py

@@ -15,6 +15,7 @@ from fastvideo.configs.models.encoders.gemma import LTX2GemmaConfig
 from fastvideo.configs.models.encoders.mistral3 import Mistral3TextConfig
 from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import (MiniMaxH3Qwen3VLArchConfig, MiniMaxH3Qwen3VLConfig)
 from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig
+from fastvideo.configs.models.encoders.qwen3_5 import Qwen3_5TextArchConfig, Qwen3_5TextConfig
 from fastvideo.configs.models.encoders.lingbot_video import LingBotVideoQwen3VLTextConfig
 from fastvideo.configs.models.encoders.stable_audio_conditioner import (
     StableAudioConditionerArchConfig,
@@ -49,6 +50,8 @@ __all__ = [
     "StableAudioConditionerConfig",
     "T5GemmaEncoderConfig",
     "Qwen3TextConfig",
+    "Qwen3_5TextArchConfig",
+    "Qwen3_5TextConfig",
     "Mistral3TextConfig",
     "LingBotWorld2UMT5ArchConfig",
     "LingBotWorld2UMT5Config",

fastvideo/configs/models/encoders/qwen3_5.py

@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Qwen3.5 text encoder configuration (hybrid Gated DeltaNet + gated attention)."""
+from dataclasses import dataclass, field
+from typing import Any
+
+from fastvideo.configs.models.encoders.base import (
+    TextEncoderArchConfig,
+    TextEncoderConfig,
+)
+
+_LEGACY_LAYER_TYPES = {
+    "mamba": "linear_attention",
+    "conv": "linear_attention",
+    "attention": "full_attention",
+}
+_DEFAULT_ROPE_THETA = 10000.0
+_DEFAULT_PARTIAL_ROTARY_FACTOR = 0.25
+
+
+def _is_transformer_layer(n: str, m: Any) -> bool:
+    return "layers" in n and str.isdigit(n.split(".")[-1])
+
+
+def _is_embeddings(n: str, m: Any) -> bool:
+    return n.endswith("embed_tokens")
+
+
+def _is_final_norm(n: str, m: Any) -> bool:
+    return n == "norm"
+
+
+@dataclass
+class Qwen3_5TextArchConfig(TextEncoderArchConfig):
+    """Architecture config for the Qwen3.5 text model (``Qwen3_5TextModel``).
+
+    Layers alternate between Gated DeltaNet linear attention and gated full
+    attention according to ``layer_types``. The defaults describe the 9B text
+    tower; checkpoint ``config.json`` values override them at load time.
+    """
+
+    vocab_size: int = 248320
+    hidden_size: int = 4096
+    intermediate_size: int = 12288
+    num_hidden_layers: int = 32
+    num_attention_heads: int = 16
+    num_key_value_heads: int = 4
+    head_dim: int = 256
+    hidden_act: str = "silu"
+    max_position_embeddings: int = 32768
+    initializer_range: float = 0.02
+    rms_norm_eps: float = 1e-6
+    use_cache: bool = True
+    tie_word_embeddings: bool = False
+    attention_bias: bool = False
+    attention_dropout: float = 0.0
+    linear_conv_kernel_dim: int = 4
+    linear_key_head_dim: int = 128
+    linear_value_head_dim: int = 128
+    linear_num_key_heads: int = 16
+    linear_num_value_heads: int = 32
+    layer_types: list[str] | None = None
+    full_attention_interval: int = 4
+    # RoPE settings as written by Transformers 5 (``rope_parameters``) or by
+    # older configs (top-level ``rope_theta`` / ``partial_rotary_factor`` /
+    # ``rope_scaling``). Values inside ``rope_parameters`` take precedence.
+    rope_parameters: dict[str, Any] | None = None
+    rope_scaling: dict[str, Any] | None = None
+    rope_theta: float | None = None
+    partial_rotary_factor: float | None = None
+    pad_token_id: int | None = None
+    bos_token_id: int | None = None
+    eos_token_id: int | list[int] | None = None
+
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [_is_transformer_layer, _is_embeddings, _is_final_norm])
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        # Validate eagerly; the resolved views are recomputed on demand because
+        # ``update_model_arch`` re-runs this hook after replacing raw fields.
+        self.resolved_layer_types()
+        self.resolved_rope_parameters()
+
+    def resolved_layer_types(self) -> list[str]:
+        """Return one ``linear_attention`` / ``full_attention`` entry per layer."""
+        if self.layer_types is None:
+            interval = self.full_attention_interval
+            layer_types = [
+                "linear_attention" if (index + 1) % interval else "full_attention"
+                for index in range(self.num_hidden_layers)
+            ]
+        else:
+            layer_types = [_LEGACY_LAYER_TYPES.get(kind, kind) for kind in self.layer_types]
+        if len(layer_types) < self.num_hidden_layers:
+            raise ValueError(f"Qwen3.5 layer_types has {len(layer_types)} entries for "
+                             f"{self.num_hidden_layers} layers")
+        unknown = set(layer_types) - {"linear_attention", "full_attention"}
+        if unknown:
+            raise ValueError(f"Unsupported Qwen3.5 layer types: {sorted(unknown)}")
+        return layer_types[:self.num_hidden_layers]
+
+    def resolved_rope_parameters(self) -> dict[str, Any]:
+        """Return the effective text RoPE settings with Transformers' precedence."""
+        rope = dict(self.rope_scaling or self.rope_parameters or {})
+        rope.setdefault("rope_type", rope.get("type", "default"))
+        rope.setdefault("rope_theta", self.rope_theta if self.rope_theta is not None else _DEFAULT_ROPE_THETA)
+        rope.setdefault(
+            "partial_rotary_factor",
+            self.partial_rotary_factor if self.partial_rotary_factor is not None else _DEFAULT_PARTIAL_ROTARY_FACTOR,
+        )
+        return rope
+
+
+@dataclass
+class Qwen3_5TextConfig(TextEncoderConfig):
+    """Top-level config for the Qwen3.5 text encoder."""
+
+    arch_config: TextEncoderArchConfig = field(default_factory=Qwen3_5TextArchConfig)
+    prefix: str = "qwen3_5"

fastvideo/configs/pipelines/__init__.py

@@ -16,6 +16,7 @@ from fastvideo.configs.pipelines.lingbot_video import LingBotVideoT2VConfig
 from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig
 from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig
 from fastvideo.configs.pipelines.mmaudio import MMAudioV2AConfig
+from fastvideo.configs.pipelines.pan2 import PAN2PipelineConfig
 from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
 from fastvideo.registry import get_pipeline_config_cls_from_name
 from fastvideo.models.wan.pipeline_config import (LucyEditDevConfig, SelfForcingWanT2V480PConfig, WanI2V480PConfig,
@@ -29,5 +30,5 @@ __all__ = [
     "Kandinsky5T2VConfig", "Kandinsky5I2VConfig", "Kandinsky5DMDConfig", "Kandinsky6TI2VAConfig",
     "Kandinsky6SRPipelineConfig", "LingBotWorld2CausalFastI2V480PConfig", "LingBotWorldFastI2V480PConfig",
     "LingBotVideoT2VConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "MMAudioV2AConfig",
-    "get_pipeline_config_cls_from_name"
+    "PAN2PipelineConfig", "get_pipeline_config_cls_from_name"
 ]

fastvideo/configs/pipelines/pan2.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2 pipeline configuration: text-to-video, and image-to-video from a first frame."""
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
+from fastvideo.configs.models.dits import PAN2VideoConfig
+from fastvideo.configs.models.encoders import Qwen3_5TextConfig
+from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.models.wan.vae_config import WanVAEConfig
+
+PAN2_SYSTEM_PROMPT = "You are a helpful assistant. Generate the video with the following instructions."
+
+
+@dataclass
+class PAN2PipelineConfig(PipelineConfig):
+    """One configuration serves text-to-video and, given a first frame, image-to-video."""
+
+    dit_config: DiTConfig = field(default_factory=PAN2VideoConfig)
+    # The Wan2.2 48-channel VAE; its architecture and latent statistics come from vae/config.json at load time.
+    vae_config: VAEConfig = field(default_factory=WanVAEConfig)
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (Qwen3_5TextConfig(), ))
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", ))
+
+    dit_precision: str = "bf16"
+    vae_precision: str = "fp32"
+    flow_shift: float | None = 7.0
+    # PAN2 samples on fp32 latents; keep the Euler update out of autocast.
+    scheduler_step_in_fp32: bool = True
+
+    # Prompt encoding: PAN2 conditions on the tokens after the chat template's user header, taken from the text
+    # encoder's hidden states with its last two layers skipped (``hidden_states[-3]``).
+    text_system_prompt: str = PAN2_SYSTEM_PROMPT
+    text_encoder_skip_layers: int = 2
+    text_max_sequence_length: int = 1536
+
+    def __post_init__(self) -> None:
+        # The encoder is needed for the image-to-video first frame.
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True
+        self.text_encoder_configs[0].arch_config.output_hidden_states = True

fastvideo/models/dits/pan2.py

@@ -0,0 +1,397 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2 video transformer.
+
+A mixture-of-transformers DiT: the video and text streams keep separate weights and meet only in one joint
+attention per block, over ``[video; text]``. The text stream is modulated by a learned, timestep-free table.
+Under sequence parallelism the video tokens are sharded and the text tokens are replicated on every rank.
+"""
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from fastvideo.attention import DistributedAttention, LocalAttention
+from fastvideo.configs.models.dits.pan2 import PAN2VideoConfig
+from fastvideo.distributed.communication_op import (sequence_model_parallel_all_gather_with_unpad,
+                                                    sequence_model_parallel_shard)
+from fastvideo.distributed.parallel_state import get_sp_world_size
+from fastvideo.layers.layernorm import RMSNorm
+from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.layers.mlp import MLP
+from fastvideo.layers.rotary_embedding import get_1d_rotary_pos_embed
+from fastvideo.layers.visual_embedding import PatchEmbed, TimestepEmbedder
+from fastvideo.models.dits.base import BaseDiT
+from fastvideo.platforms import AttentionBackendEnum
+
+
+class PAN2SwiGLU(nn.Module):
+    """SwiGLU MLP of the text refiner; the fused input projection stores ``[value; gate]``."""
+
+    def __init__(self, dim: int, inner_dim: int, prefix: str = ""):
+        super().__init__()
+        self.fc_in = ReplicatedLinear(dim, 2 * inner_dim, bias=False, prefix=f"{prefix}.fc_in")
+        self.fc_out = ReplicatedLinear(inner_dim, dim, bias=False, prefix=f"{prefix}.fc_out")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x, _ = self.fc_in(x)
+        value, gate = x.chunk(2, dim=-1)
+        x, _ = self.fc_out(value * F.silu(gate))
+        return x
+
+
+class PAN2RefinerAttention(nn.Module):
+    """Text-only self-attention of the refiner; the text is replicated on every sequence-parallel rank."""
+
+    def __init__(self,
+                 dim: int,
+                 num_heads: int,
+                 head_dim: int,
+                 eps: float,
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+                 prefix: str = ""):
+        super().__init__()
+        self.num_heads = num_heads
+        self.to_qkv = ReplicatedLinear(dim, 3 * num_heads * head_dim, bias=False, prefix=f"{prefix}.to_qkv")
+        self.norm_q = RMSNorm(head_dim, eps=eps)
+        self.norm_k = RMSNorm(head_dim, eps=eps)
+        self.to_out = ReplicatedLinear(num_heads * head_dim, dim, bias=False, prefix=f"{prefix}.to_out")
+        self.attn = LocalAttention(num_heads=num_heads,
+                                   head_size=head_dim,
+                                   causal=False,
+                                   supported_attention_backends=supported_attention_backends)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        qkv, _ = self.to_qkv(hidden_states)
+        query, key, value = qkv.unflatten(-1, (3, self.num_heads, -1)).unbind(2)
+        query = self.norm_q(query).to(value.dtype)
+        key = self.norm_k(key).to(value.dtype)
+        hidden_states = self.attn(query, key, value)
+        hidden_states, _ = self.to_out(hidden_states.flatten(2))
+        return hidden_states
+
+
+class PAN2TokenRefinerBlock(nn.Module):
+
+    def __init__(self,
+                 dim: int,
+                 num_heads: int,
+                 head_dim: int,
+                 mlp_ratio: float,
+                 eps: float,
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+                 prefix: str = ""):
+        super().__init__()
+        self.norm1 = RMSNorm(dim, eps=eps)
+        self.attn = PAN2RefinerAttention(dim,
+                                         num_heads,
+                                         head_dim,
+                                         eps,
+                                         supported_attention_backends=supported_attention_backends,
+                                         prefix=f"{prefix}.attn")
+        self.norm2 = RMSNorm(dim, eps=eps)
+        self.ff = PAN2SwiGLU(dim, int(dim * mlp_ratio), prefix=f"{prefix}.ff")
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = hidden_states + self.attn(self.norm1(hidden_states))
+        hidden_states = hidden_states + self.ff(self.norm2(hidden_states))
+        return hidden_states
+
+
+class PAN2AdaLayerNormZero(nn.Module):
+    """Timestep modulation of the video stream: modulated attention input plus the gate/shift/scale for the rest."""
+
+    def __init__(self, dim: int, eps: float = 1e-6, prefix: str = ""):
+        super().__init__()
+        self.linear = ReplicatedLinear(dim, 6 * dim, bias=True, prefix=f"{prefix}.linear")
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+
+    def forward(self, x: torch.Tensor,
+                temb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        emb, _ = self.linear(F.silu(temb))
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=-1)
+        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
+
+
+class PAN2AdaLayerNormOut(nn.Module):
+
+    def __init__(self, dim: int, eps: float = 1e-6, prefix: str = ""):
+        super().__init__()
+        self.linear = ReplicatedLinear(dim, 2 * dim, bias=True, prefix=f"{prefix}.linear")
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+
+    def forward(self, x: torch.Tensor, temb: torch.Tensor) -> torch.Tensor:
+        emb, _ = self.linear(F.silu(temb).to(x.dtype))
+        scale, shift = emb.chunk(2, dim=-1)
+        return self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
+
+
+class PAN2Attention(nn.Module):
+    """Joint attention over ``[video; text]`` with per-stream projections and Q/K RMSNorm.
+
+    RoPE applies to the video tokens only; ``DistributedAttention`` applies it after its all-to-all, where each rank
+    holds the full (unpadded) video sequence for its heads.
+    """
+
+    def __init__(self,
+                 dim: int,
+                 num_heads: int,
+                 head_dim: int,
+                 eps: float,
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+                 prefix: str = ""):
+        super().__init__()
+        self.num_heads = num_heads
+        inner_dim = num_heads * head_dim
+
+        self.to_qkv = ReplicatedLinear(dim, 3 * inner_dim, bias=True, prefix=f"{prefix}.to_qkv")
+        self.norm_q = RMSNorm(head_dim, eps=eps)
+        self.norm_k = RMSNorm(head_dim, eps=eps)
+        self.to_out = ReplicatedLinear(inner_dim, dim, bias=True, prefix=f"{prefix}.to_out")
+
+        self.add_qkv_proj = ReplicatedLinear(dim, 3 * inner_dim, bias=True, prefix=f"{prefix}.add_qkv_proj")
+        self.norm_added_q = RMSNorm(head_dim, eps=eps)
+        self.norm_added_k = RMSNorm(head_dim, eps=eps)
+        self.to_add_out = ReplicatedLinear(inner_dim, dim, bias=True, prefix=f"{prefix}.to_add_out")
+
+        self.attn = DistributedAttention(num_heads=num_heads,
+                                         head_size=head_dim,
+                                         causal=False,
+                                         supported_attention_backends=supported_attention_backends,
+                                         prefix=f"{prefix}.attn")
+
+    def _project(self, projection: ReplicatedLinear,
+                 hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        qkv, _ = projection(hidden_states)
+        query, key, value = qkv.unflatten(-1, (3, self.num_heads, -1)).unbind(2)
+        return query, key, value
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        freqs_cis: tuple[torch.Tensor, torch.Tensor],
+        original_seq_len: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        query, key, value = self._project(self.to_qkv, hidden_states)
+        query = self.norm_q(query).to(value.dtype)
+        key = self.norm_k(key).to(value.dtype)
+
+        encoder_query, encoder_key, encoder_value = self._project(self.add_qkv_proj, encoder_hidden_states)
+        encoder_query = self.norm_added_q(encoder_query).to(encoder_value.dtype)
+        encoder_key = self.norm_added_k(encoder_key).to(encoder_value.dtype)
+
+        hidden_states, encoder_hidden_states = self.attn(query,
+                                                         key,
+                                                         value,
+                                                         original_seq_len,
+                                                         replicated_q=encoder_query,
+                                                         replicated_k=encoder_key,
+                                                         replicated_v=encoder_value,
+                                                         freqs_cis=freqs_cis)
+        hidden_states, _ = self.to_out(hidden_states.flatten(2))
+        encoder_hidden_states, _ = self.to_add_out(encoder_hidden_states.flatten(2))
+        return hidden_states, encoder_hidden_states
+
+
+class PAN2TransformerBlock(nn.Module):
+
+    def __init__(self,
+                 dim: int,
+                 num_heads: int,
+                 head_dim: int,
+                 mlp_ratio: float,
+                 context_mlp_ratio: float,
+                 eps: float = 1e-6,
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+                 prefix: str = ""):
+        super().__init__()
+        self.norm1 = PAN2AdaLayerNormZero(dim, eps=eps, prefix=f"{prefix}.norm1")
+        self.norm1_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.scale_shift_table_context = nn.Parameter(torch.zeros(6 * dim))
+
+        self.attn = PAN2Attention(dim,
+                                  num_heads,
+                                  head_dim,
+                                  eps,
+                                  supported_attention_backends=supported_attention_backends,
+                                  prefix=f"{prefix}.attn")
+
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.ff = MLP(dim, int(dim * mlp_ratio), act_type="gelu_pytorch_tanh", prefix=f"{prefix}.ff")
+        self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.ff_context = MLP(dim,
+                              int(dim * context_mlp_ratio),
+                              act_type="gelu_pytorch_tanh",
+                              prefix=f"{prefix}.ff_context")
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        freqs_cis: tuple[torch.Tensor, torch.Tensor],
+        original_seq_len: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, temb)
+        (
+            context_shift_msa,
+            context_scale_msa,
+            context_gate_msa,
+            context_shift_mlp,
+            context_scale_mlp,
+            context_gate_mlp,
+        ) = self.scale_shift_table_context.to(temb.dtype).chunk(6, dim=-1)
+        norm_encoder_hidden_states = (self.norm1_context(encoder_hidden_states) *
+                                      (1 + context_scale_msa) + context_shift_msa)
+
+        attn_output, context_attn_output = self.attn(norm_hidden_states, norm_encoder_hidden_states, freqs_cis,
+                                                     original_seq_len)
+        hidden_states = hidden_states + attn_output * gate_msa.unsqueeze(1)
+        encoder_hidden_states = encoder_hidden_states + context_attn_output * context_gate_msa
+
+        norm_hidden_states = self.norm2(hidden_states) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        norm_encoder_hidden_states = (self.norm2_context(encoder_hidden_states) *
+                                      (1 + context_scale_mlp) + context_shift_mlp)
+
+        hidden_states = hidden_states + self.ff(norm_hidden_states) * gate_mlp.unsqueeze(1)
+        encoder_hidden_states = encoder_hidden_states + self.ff_context(norm_encoder_hidden_states) * context_gate_mlp
+        return hidden_states, encoder_hidden_states
+
+
+class PAN2Transformer3DModel(BaseDiT):
+    """PAN2 DiT. The input stacks the noisy latents, the conditioning latents and a one-channel conditioning mask
+    along channels, so one model serves text-to-video and first-frame image-to-video."""
+
+    _fsdp_shard_conditions = PAN2VideoConfig()._fsdp_shard_conditions
+    _compile_conditions = PAN2VideoConfig()._compile_conditions
+    _supported_attention_backends = PAN2VideoConfig()._supported_attention_backends
+    param_names_mapping = PAN2VideoConfig().param_names_mapping
+    reverse_param_names_mapping = PAN2VideoConfig().reverse_param_names_mapping
+    lora_param_names_mapping = PAN2VideoConfig().lora_param_names_mapping
+
+    def __init__(self, config: PAN2VideoConfig, hf_config: dict[str, Any], **kwargs) -> None:
+        super().__init__(config=config, hf_config=hf_config)
+        inner_dim = config.num_attention_heads * config.attention_head_dim
+        self.hidden_size = inner_dim
+        self.num_attention_heads = config.num_attention_heads
+        self.num_channels_latents = config.num_channels_latents
+        self.out_channels = config.out_channels
+        self.patch_size = tuple(config.patch_size)
+        self.rope_axes_dim = tuple(config.rope_axes_dim)
+        if sum(self.rope_axes_dim) != config.attention_head_dim:
+            raise ValueError(f"rope_axes_dim {self.rope_axes_dim} must sum to the head dim {config.attention_head_dim}")
+        if config.num_attention_heads % get_sp_world_size() != 0:
+            raise ValueError(f"The number of attention heads ({config.num_attention_heads}) must be divisible by the "
+                             f"sequence parallel size ({get_sp_world_size()})")
+
+        self.x_embedder = PatchEmbed(self.patch_size,
+                                     config.in_channels,
+                                     inner_dim,
+                                     prefix=f"{config.prefix}.x_embedder")
+        self.time_embedder = TimestepEmbedder(inner_dim, act_layer="silu", frequency_embedding_size=256)
+
+        self.context_norm = RMSNorm(config.text_embed_dim, eps=1e-5)
+        self.context_proj = ReplicatedLinear(config.text_embed_dim,
+                                             inner_dim,
+                                             bias=True,
+                                             prefix=f"{config.prefix}.context_proj")
+        self.context_refiner = nn.ModuleList([
+            PAN2TokenRefinerBlock(inner_dim,
+                                  config.num_attention_heads,
+                                  config.attention_head_dim,
+                                  config.refiner_mlp_ratio,
+                                  eps=1e-5,
+                                  supported_attention_backends=self._supported_attention_backends,
+                                  prefix=f"{config.prefix}.context_refiner.{i}")
+            for i in range(config.num_refiner_layers)
+        ])
+        self.context_type_embedding = nn.Parameter(torch.zeros(inner_dim))
+
+        self.transformer_blocks = nn.ModuleList([
+            PAN2TransformerBlock(inner_dim,
+                                 config.num_attention_heads,
+                                 config.attention_head_dim,
+                                 config.mlp_ratio,
+                                 config.context_mlp_ratio,
+                                 supported_attention_backends=self._supported_attention_backends,
+                                 prefix=f"{config.prefix}.transformer_blocks.{i}") for i in range(config.num_layers)
+        ])
+
+        self.norm_out = PAN2AdaLayerNormOut(inner_dim, eps=1e-6, prefix=f"{config.prefix}.norm_out")
+        self.proj_out = ReplicatedLinear(inner_dim,
+                                         self.patch_size[0] * self.patch_size[1] * self.patch_size[2] *
+                                         self.out_channels,
+                                         bias=True,
+                                         prefix=f"{config.prefix}.proj_out")
+
+        self.gradient_checkpointing = False
+        self.__post_init__()
+
+    def _rotary_pos_embed(self, num_frames: int, height: int, width: int,
+                          device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
+        """Half-dim cos/sin of the 3-axis ``(t, h, w)`` RoPE over the full (unsharded) video token grid."""
+        axes = [torch.arange(size, device=device, dtype=torch.float32) for size in (num_frames, height, width)]
+        grid = torch.stack(torch.meshgrid(*axes, indexing="ij"), dim=0)
+        grid[0] = grid[0] * self.config.rope_temporal_scale
+        cos_sin = [
+            get_1d_rotary_pos_embed(dim, grid[i].reshape(-1), theta=self.config.rope_theta, use_real=False)
+            for i, dim in enumerate(self.rope_axes_dim)
+        ]
+        freqs_cos = torch.cat([cos for cos, _ in cos_sin], dim=1).float()
+        freqs_sin = torch.cat([sin for _, sin in cos_sin], dim=1).float()
+        return freqs_cos, freqs_sin
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                encoder_hidden_states: torch.Tensor | list[torch.Tensor],
+                timestep: torch.LongTensor,
+                encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None,
+                guidance=None,
+                **kwargs) -> torch.Tensor:
+        if not isinstance(encoder_hidden_states, torch.Tensor):
+            encoder_hidden_states = encoder_hidden_states[0]
+
+        batch_size, _, num_frames, height, width = hidden_states.shape
+        p_t, p_h, p_w = self.patch_size
+        post_patch_num_frames = num_frames // p_t
+        post_patch_height = height // p_h
+        post_patch_width = width // p_w
+
+        freqs_cis = self._rotary_pos_embed(post_patch_num_frames, post_patch_height, post_patch_width,
+                                           hidden_states.device)
+
+        hidden_states = self.x_embedder(hidden_states)
+        hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1)
+
+        if timestep.dim() == 0:
+            timestep = timestep.expand(batch_size)
+        temb = self.time_embedder(timestep).to(hidden_states.dtype)
+
+        encoder_hidden_states, _ = self.context_proj(self.context_norm(encoder_hidden_states))
+        for refiner_block in self.context_refiner:
+            encoder_hidden_states = refiner_block(encoder_hidden_states)
+        encoder_hidden_states = encoder_hidden_states + self.context_type_embedding
+
+        for block in self.transformer_blocks:
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                hidden_states, encoder_hidden_states = self._gradient_checkpointing_func(
+                    block, hidden_states, encoder_hidden_states, temb, freqs_cis, original_seq_len)
+            else:
+                hidden_states, encoder_hidden_states = block(hidden_states, encoder_hidden_states, temb, freqs_cis,
+                                                             original_seq_len)
+
+        hidden_states = self.norm_out(hidden_states, temb)
+        hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1)
+        hidden_states, _ = self.proj_out(hidden_states)
+
+        # proj_out emits channels before the patch offsets: (C, p_t, p_h, p_w).
+        hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, post_patch_height, post_patch_width,
+                                              -1, p_t, p_h, p_w)
+        hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7)
+        return hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+
+
+EntryClass = PAN2Transformer3DModel

fastvideo/models/encoders/qwen3_5.py

@@ -0,0 +1,574 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Qwen3.5 text encoder: hybrid Gated DeltaNet linear attention + gated full attention.
+
+The encoder runs replicated (no tensor or sequence parallelism) and
+inference-only, so the linear-attention layers use a pure-PyTorch chunked gated
+delta rule instead of the fused ``fla`` / ``causal_conv1d`` kernels. Full
+attention calls ``scaled_dot_product_attention`` directly with head-major
+``[B, H, L, D]`` inputs: that is the layout the reference text model feeds SDPA,
+and keeping it makes bf16 hidden states reproducible against that reference.
+"""
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from fastvideo.configs.models.encoders import BaseEncoderOutput
+from fastvideo.configs.models.encoders.qwen3_5 import Qwen3_5TextArchConfig, Qwen3_5TextConfig
+from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.layers.quantization import QuantizationConfig
+from fastvideo.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from fastvideo.logger import init_logger
+from fastvideo.models.encoders.base import TextEncoder
+from fastvideo.models.loader.weight_utils import default_weight_loader
+
+logger = init_logger(__name__)
+
+_GATED_DELTA_CHUNK_SIZE = 64
+# Checkpoints saved from the multimodal wrapper nest the text tower.
+_CHECKPOINT_PREFIXES = ("model.language_model.", "language_model.", "model.")
+_SKIPPED_CHECKPOINT_PREFIXES = ("lm_head.", "mtp.", "visual.", "model.visual.")
+
+
+class Qwen3_5RMSNorm(nn.Module):
+    """Zero-centered RMSNorm: the stored weight is an offset from one."""
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.zeros(hidden_size))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        output = x.float()
+        output = output * torch.rsqrt(output.pow(2).mean(-1, keepdim=True) + self.eps)
+        # The scale is applied in float32 before the downcast.
+        output = output * (1.0 + self.weight.float())
+        return output.type_as(x)
+
+
+class Qwen3_5RMSNormGated(nn.Module):
+    """Per-head RMSNorm followed by a SiLU output gate (plain, not zero-centered, weight)."""
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+
+    def forward(self, hidden_states: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
+        hidden_states = self.weight * hidden_states.to(input_dtype)
+        hidden_states = hidden_states * F.silu(gate.to(torch.float32))
+        return hidden_states.to(input_dtype)
+
+
+def _l2norm(x: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
+    return x * torch.rsqrt((x * x).sum(dim=-1, keepdim=True) + eps)
+
+
+def chunk_gated_delta_rule(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    chunk_size: int = _GATED_DELTA_CHUNK_SIZE,
+) -> torch.Tensor:
+    """Chunked gated delta rule over a full prefix with a zero initial state.
+
+    Inputs are ``[batch, seq, heads, dim]`` (``g`` / ``beta`` are
+    ``[batch, seq, heads]``); q/k must already be L2-normalized. The recurrence
+    runs in float32 and the output is returned in the query dtype.
+    """
+    output_dtype = query.dtype
+    query, key, value, beta, g = (x.transpose(1, 2).contiguous().to(torch.float32)
+                                  for x in (query, key, value, beta, g))
+
+    batch_size, num_heads, sequence_length, k_head_dim = key.shape
+    v_head_dim = value.shape[-1]
+    pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
+    query = F.pad(query, (0, 0, 0, pad_size))
+    key = F.pad(key, (0, 0, 0, pad_size))
+    value = F.pad(value, (0, 0, 0, pad_size))
+    beta = F.pad(beta, (0, pad_size))
+    g = F.pad(g, (0, pad_size))
+    total_sequence_length = sequence_length + pad_size
+    query = query * (1 / (query.shape[-1]**0.5))
+
+    v_beta = value * beta.unsqueeze(-1)
+    k_beta = key * beta.unsqueeze(-1)
+    query, key, value, k_beta, v_beta = (x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1])
+                                         for x in (query, key, value, k_beta, v_beta))
+    g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size)
+    upper_with_diagonal = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device),
+                                     diagonal=0)
+
+    # Intra-chunk decay and the (I + strictly-lower)^-1 WY transform, solved by
+    # forward substitution row by row.
+    g = g.cumsum(dim=-1)
+    decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
+    attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(upper_with_diagonal, 0)
+    for i in range(1, chunk_size):
+        row = attn[..., i, :i].clone()
+        sub = attn[..., :i, :i].clone()
+        attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)
+    attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
+    value = attn @ v_beta
+    k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
+
+    recurrent_state = torch.zeros(batch_size,
+                                  num_heads,
+                                  k_head_dim,
+                                  v_head_dim,
+                                  dtype=value.dtype,
+                                  device=value.device)
+    core_attn_out = torch.zeros_like(value)
+    for i in range(total_sequence_length // chunk_size):
+        q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i]
+        # decay_mask is already zero above the diagonal, which keeps this causal.
+        attn = q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]
+        v_prime = k_cumdecay[:, :, i] @ recurrent_state
+        v_new = v_i - v_prime
+        attn_inter = (q_i * g[:, :, i, :, None].exp()) @ recurrent_state
+        core_attn_out[:, :, i] = attn_inter + attn @ v_new
+        recurrent_state = (recurrent_state * g[:, :, i, -1, None, None].exp() +
+                           (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new)
+
+    core_attn_out = core_attn_out.reshape(core_attn_out.shape[0], core_attn_out.shape[1], -1, core_attn_out.shape[-1])
+    core_attn_out = core_attn_out[:, :, :sequence_length]
+    return core_attn_out.transpose(1, 2).contiguous().to(output_dtype)
+
+
+class Qwen3_5GatedDeltaNet(nn.Module):
+    """Gated DeltaNet token mixer used by the ``linear_attention`` layers."""
+
+    def __init__(
+        self,
+        config: Qwen3_5TextArchConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.num_v_heads = config.linear_num_value_heads
+        self.num_k_heads = config.linear_num_key_heads
+        if self.num_v_heads % self.num_k_heads:
+            raise ValueError(f"Qwen3.5 linear value heads ({self.num_v_heads}) must be a multiple of key heads "
+                             f"({self.num_k_heads})")
+        self.head_k_dim = config.linear_key_head_dim
+        self.head_v_dim = config.linear_value_head_dim
+        self.key_dim = self.head_k_dim * self.num_k_heads
+        self.value_dim = self.head_v_dim * self.num_v_heads
+        self.conv_kernel_size = config.linear_conv_kernel_dim
+        if config.hidden_act != "silu":
+            raise ValueError(f"Unsupported Qwen3.5 activation: {config.hidden_act}")
+
+        self.conv_dim = self.key_dim * 2 + self.value_dim
+        self.conv1d = nn.Conv1d(
+            in_channels=self.conv_dim,
+            out_channels=self.conv_dim,
+            bias=False,
+            kernel_size=self.conv_kernel_size,
+            groups=self.conv_dim,
+            padding=self.conv_kernel_size - 1,
+        )
+        self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads))
+        self.A_log = nn.Parameter(torch.zeros(self.num_v_heads))
+        self.norm = Qwen3_5RMSNormGated(self.head_v_dim, eps=config.rms_norm_eps)
+        self.out_proj = ReplicatedLinear(self.value_dim,
+                                         self.hidden_size,
+                                         bias=False,
+                                         quant_config=quant_config,
+                                         prefix=f"{prefix}.out_proj")
+        self.in_proj_qkv = ReplicatedLinear(self.hidden_size,
+                                            self.conv_dim,
+                                            bias=False,
+                                            quant_config=quant_config,
+                                            prefix=f"{prefix}.in_proj_qkv")
+        self.in_proj_z = ReplicatedLinear(self.hidden_size,
+                                          self.value_dim,
+                                          bias=False,
+                                          quant_config=quant_config,
+                                          prefix=f"{prefix}.in_proj_z")
+        self.in_proj_b = ReplicatedLinear(self.hidden_size,
+                                          self.num_v_heads,
+                                          bias=False,
+                                          quant_config=quant_config,
+                                          prefix=f"{prefix}.in_proj_b")
+        self.in_proj_a = ReplicatedLinear(self.hidden_size,
+                                          self.num_v_heads,
+                                          bias=False,
+                                          quant_config=quant_config,
+                                          prefix=f"{prefix}.in_proj_a")
+
+    def forward(self, hidden_states: torch.Tensor, padding_mask: torch.Tensor | None = None) -> torch.Tensor:
+        if padding_mask is not None:
+            # Zero padded tokens so they cannot leak into the recurrent state.
+            hidden_states = (hidden_states * padding_mask[:, :, None]).to(hidden_states.dtype)
+        batch_size, seq_len, _ = hidden_states.shape
+
+        mixed_qkv, _ = self.in_proj_qkv(hidden_states)
+        z, _ = self.in_proj_z(hidden_states)
+        b, _ = self.in_proj_b(hidden_states)
+        a, _ = self.in_proj_a(hidden_states)
+
+        # Causal depthwise conv: left padding of kernel-1, keep the first seq_len outputs.
+        mixed_qkv = mixed_qkv.transpose(1, 2)
+        mixed_qkv = F.silu(self.conv1d(mixed_qkv.to(self.conv1d.weight.dtype))[:, :, :seq_len]).to(mixed_qkv.dtype)
+        mixed_qkv = mixed_qkv.transpose(1, 2)
+        query, key, value = torch.split(mixed_qkv, [self.key_dim, self.key_dim, self.value_dim], dim=-1)
+        query = query.reshape(batch_size, seq_len, -1, self.head_k_dim)
+        key = key.reshape(batch_size, seq_len, -1, self.head_k_dim)
+        value = value.reshape(batch_size, seq_len, -1, self.head_v_dim)
+
+        beta = b.sigmoid()
+        # float32 keeps exp(A_log) finite for half-precision weights.
+        g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)
+        if self.num_v_heads > self.num_k_heads:
+            repeats = self.num_v_heads // self.num_k_heads
+            query = query.repeat_interleave(repeats, dim=2)
+            key = key.repeat_interleave(repeats, dim=2)
+
+        core_attn_out = chunk_gated_delta_rule(_l2norm(query), _l2norm(key), value, g=g, beta=beta)
+
+        core_attn_out = self.norm(core_attn_out.reshape(-1, self.head_v_dim), z.reshape(-1, self.head_v_dim))
+        core_attn_out = core_attn_out.reshape(batch_size, seq_len, -1)
+        output, _ = self.out_proj(core_attn_out)
+        return output
+
+
+class Qwen3_5RotaryEmbedding(nn.Module):
+    """Text-only partial RoPE.
+
+    Qwen3.5's interleaved M-RoPE collapses to standard RoPE when the temporal,
+    height and width positions are identical, which holds for text tokens.
+    """
+
+    def __init__(self, config: Qwen3_5TextArchConfig) -> None:
+        super().__init__()
+        rope = config.resolved_rope_parameters()
+        if rope["rope_type"] != "default":
+            raise NotImplementedError(f"Qwen3.5 rope_type {rope['rope_type']!r} is not supported")
+        rotary_dim = int(config.head_dim * rope["partial_rotary_factor"])
+        inv_freq = 1.0 / (rope["rope_theta"]**(
+            torch.arange(0, rotary_dim, 2, dtype=torch.float32, device="cpu") / rotary_dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    @torch.no_grad()
+    def forward(self, hidden_states: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        inv_freq = self.inv_freq.to(device=hidden_states.device, dtype=torch.float32)
+        device_type = hidden_states.device.type if hidden_states.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):
+            inv_freq = inv_freq[None, :, None].expand(position_ids.shape[0], -1, 1)
+            freqs = (inv_freq @ position_ids[:, None, :].float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos()
+            sin = emb.sin()
+        return cos.to(hidden_states.dtype), sin.to(hidden_states.dtype)
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def _repeat_kv(x: torch.Tensor, groups: int) -> torch.Tensor:
+    batch_size, num_kv_heads, seq_len, head_dim = x.shape
+    x = x[:, :, None].expand(batch_size, num_kv_heads, groups, seq_len, head_dim)
+    return x.reshape(batch_size, num_kv_heads * groups, seq_len, head_dim)
+
+
+def _apply_partial_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """Rotate the leading ``cos.shape[-1]`` channels of ``[B, H, L, D]`` and pass the rest through."""
+    rotary_dim = cos.shape[-1]
+    cos = cos.unsqueeze(1)
+    sin = sin.unsqueeze(1)
+    x_rot, x_pass = x[..., :rotary_dim], x[..., rotary_dim:]
+    x_rot = x_rot * cos + _rotate_half(x_rot) * sin
+    return torch.cat([x_rot, x_pass], dim=-1)
+
+
+class Qwen3_5Attention(nn.Module):
+    """Gated full attention with per-head QK-norm, partial RoPE and GQA."""
+
+    def __init__(
+        self,
+        config: Qwen3_5TextConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        arch = config.arch_config
+        self.num_heads = arch.num_attention_heads
+        self.num_kv_heads = arch.num_key_value_heads
+        if self.num_heads % self.num_kv_heads:
+            raise ValueError(f"Qwen3.5 attention heads ({self.num_heads}) must be a multiple of KV heads "
+                             f"({self.num_kv_heads})")
+        self.head_dim = arch.head_dim
+        self.scaling = self.head_dim**-0.5
+        # q_proj packs, per head, the query followed by its output gate.
+        self.q_proj = ReplicatedLinear(arch.hidden_size,
+                                       self.num_heads * self.head_dim * 2,
+                                       bias=arch.attention_bias,
+                                       quant_config=quant_config,
+                                       prefix=f"{prefix}.q_proj")
+        self.k_proj = ReplicatedLinear(arch.hidden_size,
+                                       self.num_kv_heads * self.head_dim,
+                                       bias=arch.attention_bias,
+                                       quant_config=quant_config,
+                                       prefix=f"{prefix}.k_proj")
+        self.v_proj = ReplicatedLinear(arch.hidden_size,
+                                       self.num_kv_heads * self.head_dim,
+                                       bias=arch.attention_bias,
+                                       quant_config=quant_config,
+                                       prefix=f"{prefix}.v_proj")
+        self.o_proj = ReplicatedLinear(self.num_heads * self.head_dim,
+                                       arch.hidden_size,
+                                       bias=arch.attention_bias,
+                                       quant_config=quant_config,
+                                       prefix=f"{prefix}.o_proj")
+        self.q_norm = Qwen3_5RMSNorm(self.head_dim, eps=arch.rms_norm_eps)
+        self.k_norm = Qwen3_5RMSNorm(self.head_dim, eps=arch.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        padding_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        batch_size, seq_len, _ = hidden_states.shape
+        packed, _ = self.q_proj(hidden_states)
+        query, gate = torch.chunk(packed.view(batch_size, seq_len, -1, self.head_dim * 2), 2, dim=-1)
+        gate = gate.reshape(batch_size, seq_len, -1)
+        key, _ = self.k_proj(hidden_states)
+        value, _ = self.v_proj(hidden_states)
+
+        query = self.q_norm(query).transpose(1, 2)
+        key = self.k_norm(key.view(batch_size, seq_len, -1, self.head_dim)).transpose(1, 2)
+        value = value.view(batch_size, seq_len, -1, self.head_dim).transpose(1, 2)
+        cos, sin = position_embeddings
+        query = _apply_partial_rotary(query, cos, sin)
+        key = _apply_partial_rotary(key, cos, sin)
+
+        groups = self.num_heads // self.num_kv_heads
+        # Native GQA is only used where SDPA keeps a fused kernel (no mask, head_dim <= 256).
+        native_gqa = padding_mask is None and self.head_dim <= 256
+        if groups > 1 and not native_gqa:
+            key = _repeat_kv(key, groups)
+            value = _repeat_kv(value, groups)
+        if padding_mask is None:
+            attn_output = F.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                dropout_p=0.0,
+                is_causal=seq_len > 1,
+                scale=self.scaling,
+                enable_gqa=groups > 1 and native_gqa,
+            )
+        else:
+            causal = torch.ones(seq_len, seq_len, device=query.device, dtype=torch.bool).tril()
+            mask = causal[None, None] & padding_mask.to(device=query.device, dtype=torch.bool)[:, None, None, :]
+            attn_output = F.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                attn_mask=mask,
+                dropout_p=0.0,
+                is_causal=False,
+                scale=self.scaling,
+            )
+
+        attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, -1)
+        attn_output = attn_output * torch.sigmoid(gate)
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class Qwen3_5MLP(nn.Module):
+
+    def __init__(
+        self,
+        config: Qwen3_5TextArchConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        if config.hidden_act != "silu":
+            raise ValueError(f"Unsupported Qwen3.5 activation: {config.hidden_act}")
+        self.gate_proj = ReplicatedLinear(config.hidden_size,
+                                          config.intermediate_size,
+                                          bias=False,
+                                          quant_config=quant_config,
+                                          prefix=f"{prefix}.gate_proj")
+        self.up_proj = ReplicatedLinear(config.hidden_size,
+                                        config.intermediate_size,
+                                        bias=False,
+                                        quant_config=quant_config,
+                                        prefix=f"{prefix}.up_proj")
+        self.down_proj = ReplicatedLinear(config.intermediate_size,
+                                          config.hidden_size,
+                                          bias=False,
+                                          quant_config=quant_config,
+                                          prefix=f"{prefix}.down_proj")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate, _ = self.gate_proj(x)
+        up, _ = self.up_proj(x)
+        output, _ = self.down_proj(F.silu(gate) * up)
+        return output
+
+
+class Qwen3_5DecoderLayer(nn.Module):
+
+    def __init__(
+        self,
+        config: Qwen3_5TextConfig,
+        layer_type: str,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        arch = config.arch_config
+        self.layer_type = layer_type
+        if layer_type == "linear_attention":
+            self.linear_attn = Qwen3_5GatedDeltaNet(arch, quant_config=quant_config, prefix=f"{prefix}.linear_attn")
+        else:
+            self.self_attn = Qwen3_5Attention(config, quant_config=quant_config, prefix=f"{prefix}.self_attn")
+        self.mlp = Qwen3_5MLP(arch, quant_config=quant_config, prefix=f"{prefix}.mlp")
+        self.input_layernorm = Qwen3_5RMSNorm(arch.hidden_size, eps=arch.rms_norm_eps)
+        self.post_attention_layernorm = Qwen3_5RMSNorm(arch.hidden_size, eps=arch.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        padding_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        if self.layer_type == "linear_attention":
+            hidden_states = self.linear_attn(hidden_states, padding_mask)
+        else:
+            hidden_states = self.self_attn(hidden_states, position_embeddings, padding_mask)
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        return residual + hidden_states
+
+
+class Qwen3_5TextModel(TextEncoder):
+    """Qwen3.5 text tower used as a diffusion text encoder (e.g. PAN2).
+
+    ``hidden_states`` follows the Transformers contract: entry 0 is the token
+    embeddings, entry ``i`` is the output of layer ``i - 1``, and the last entry
+    is the final-normed output (it replaces the raw last-layer output).
+    """
+
+    def __init__(self, config: Qwen3_5TextConfig) -> None:
+        super().__init__(config)
+        self.config = config
+        arch = config.arch_config
+        self.quant_config = getattr(config, "quant_config", None)
+        self.layer_types = arch.resolved_layer_types()
+
+        self.embed_tokens = VocabParallelEmbedding(
+            arch.vocab_size,
+            arch.hidden_size,
+            org_num_embeddings=arch.vocab_size,
+            quant_config=self.quant_config,
+        )
+        self.layers = nn.ModuleList([
+            Qwen3_5DecoderLayer(
+                config,
+                layer_type,
+                quant_config=self.quant_config,
+                prefix=f"{config.prefix}.layers.{index}",
+            ) for index, layer_type in enumerate(self.layer_types)
+        ])
+        self.norm = Qwen3_5RMSNorm(arch.hidden_size, eps=arch.rms_norm_eps)
+        self.rotary_emb = Qwen3_5RotaryEmbedding(arch)
+
+    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        **kwargs: Any,
+    ) -> BaseEncoderOutput:
+        output_hidden_states = (output_hidden_states
+                                if output_hidden_states is not None else self.config.arch_config.output_hidden_states)
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("Specify exactly one of input_ids or inputs_embeds")
+        if inputs_embeds is None:
+            assert input_ids is not None
+            if input_ids.ndim == 1:
+                input_ids = input_ids.unsqueeze(0)
+            inputs_embeds = self.get_input_embeddings(input_ids)
+        hidden_states = inputs_embeds
+        batch_size, seq_len = hidden_states.shape[:2]
+
+        if position_ids is None:
+            position_ids = torch.arange(seq_len, device=hidden_states.device).unsqueeze(0).expand(batch_size, -1)
+        elif position_ids.ndim == 1:
+            position_ids = position_ids.unsqueeze(0).expand(batch_size, -1)
+        elif position_ids.ndim != 2:
+            raise ValueError("Qwen3.5 text encoder expects text position_ids shaped [batch, seq]")
+
+        # An all-ones mask is equivalent to no mask; only real padding takes the masked path.
+        padding_mask = None
+        if attention_mask is not None:
+            if attention_mask.ndim != 2:
+                raise ValueError("Qwen3.5 text encoder expects a 2D [batch, seq] padding mask")
+            if not bool(attention_mask.to(torch.bool).all()):
+                padding_mask = attention_mask.to(hidden_states.device)
+
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+        all_hidden_states: tuple[torch.Tensor, ...] | None = () if output_hidden_states else None
+        for layer in self.layers:
+            if all_hidden_states is not None:
+                all_hidden_states += (hidden_states, )
+            hidden_states = layer(hidden_states, position_embeddings, padding_mask)
+
+        hidden_states = self.norm(hidden_states)
+        if all_hidden_states is not None:
+            all_hidden_states += (hidden_states, )
+
+        return BaseEncoderOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+            attention_mask=attention_mask,
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            if name.startswith(_SKIPPED_CHECKPOINT_PREFIXES) or "rotary_emb.inv_freq" in name:
+                continue
+            for checkpoint_prefix in _CHECKPOINT_PREFIXES:
+                if name.startswith(checkpoint_prefix):
+                    name = name[len(checkpoint_prefix):]
+                    break
+            if name not in params_dict:
+                logger.debug("Skipping unexpected Qwen3.5 checkpoint tensor %s", name)
+                continue
+            param = params_dict[name]
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
+
+
+EntryClass = Qwen3_5TextModel

fastvideo/models/registry.py

@@ -50,6 +50,7 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
     "Kandinsky5Transformer3DModel": ("dits", "kandinsky5", "Kandinsky5Transformer3DModel"),
     "Kandinsky6Transformer3DModel": ("dits", "kandinsky6", "Kandinsky6Transformer3DModel"),
     "Flux2Transformer2DModel": ("dits", "flux_2", "Flux2Transformer2DModel"),
+    "PAN2Transformer3DModel": ("dits", "pan2", "PAN2Transformer3DModel"),
 }
 
 _IMAGE_TO_VIDEO_DIT_MODELS = {
@@ -103,6 +104,7 @@ _TEXT_ENCODER_MODELS = {
     "LTX2GemmaTextEncoderModel": ("encoders", "gemma", "LTX2GemmaTextEncoderModel"),
     "Qwen3ForCausalLM": ("encoders", "qwen3", "Qwen3ForCausalLM"),
     "Mistral3ForConditionalGeneration": ("encoders", "mistral3", "Mistral3ForConditionalGeneration"),
+    "Qwen3_5TextModel": ("encoders", "qwen3_5", "Qwen3_5TextModel"),
 }
 
 _IMAGE_ENCODER_MODELS: dict[str, tuple] = {

fastvideo/pipelines/basic/pan2/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

fastvideo/pipelines/basic/pan2/pan2_pipeline.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2 video generation: text-to-video, and image-to-video from a first frame."""
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.basic.pan2.stages import (
+    PAN2ConditioningStage,
+    PAN2DenoisingStage,
+    PAN2InputValidationStage,
+    PAN2TextEncodingStage,
+    PAN2TimestepPreparationStage,
+)
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.stages import DecodingStage, LatentPreparationStage
+
+
+class PAN2Pipeline(ComposedPipelineBase):
+    """One stage chain serves both modes: a request with ``image_path`` (or a PIL ``pil_image``) runs image-to-video
+    from that first frame, and one without runs text-to-video."""
+
+    _required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "scheduler"]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        self.add_stage(stage_name="input_validation_stage", stage=PAN2InputValidationStage())
+        self.add_stage(stage_name="prompt_encoding_stage",
+                       stage=PAN2TextEncodingStage(
+                           text_encoders=[self.get_module("text_encoder")],
+                           tokenizers=[self.get_module("tokenizer")],
+                       ))
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=PAN2TimestepPreparationStage(scheduler=self.get_module("scheduler")))
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=LatentPreparationStage(scheduler=self.get_module("scheduler"),
+                                                    transformer=self.get_module("transformer")))
+        self.add_stage(stage_name="conditioning_stage", stage=PAN2ConditioningStage(vae=self.get_module("vae")))
+        self.add_stage(stage_name="denoising_stage",
+                       stage=PAN2DenoisingStage(transformer=self.get_module("transformer"),
+                                                scheduler=self.get_module("scheduler"),
+                                                pipeline=self))
+        self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"), pipeline=self))
+
+
+class PAN2ModularPipeline(PAN2Pipeline):
+    """The pipeline class PAN2 checkpoints declare in their ``modular_model_index.json``."""
+
+
+EntryClass = [PAN2Pipeline, PAN2ModularPipeline]

fastvideo/pipelines/basic/pan2/presets.py

@@ -0,0 +1,48 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Inference presets for PAN2 text-to-video and image-to-video."""
+
+from fastvideo.api.presets import InferencePreset, PresetStageSpec
+
+_DENOISE_STAGE = PresetStageSpec(
+    name="denoise",
+    kind="denoising",
+    description="PAN2 flow-matching denoising",
+    allowed_overrides=frozenset({"num_inference_steps", "guidance_scale"}),
+)
+
+# 704x1248, 313 frames (13 s at 24 fps), 50 steps, CFG 3. PAN2 has no default negative prompt; an empty one is
+# encoded as a single space.
+_SHARED_DEFAULTS = {
+    "height": 704,
+    "width": 1248,
+    "num_frames": 313,
+    "fps": 24,
+    "num_inference_steps": 50,
+    "guidance_scale": 3.0,
+    "negative_prompt": "",
+}
+
+PAN2_T2V = InferencePreset(
+    name="pan2_t2v",
+    version=1,
+    model_family="pan2",
+    description="PAN2 text-to-video",
+    workload_type="t2v",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults=dict(_SHARED_DEFAULTS),
+)
+
+# The same pipeline runs image-to-video when the request carries a first frame (image_path).
+PAN2_I2V = InferencePreset(
+    name="pan2_i2v",
+    version=1,
+    model_family="pan2",
+    description="PAN2 image-to-video from a first frame",
+    workload_type="i2v",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults=dict(_SHARED_DEFAULTS),
+)
+
+ALL_PRESETS = (PAN2_T2V, PAN2_I2V)
+
+__all__ = ["ALL_PRESETS", "PAN2_I2V", "PAN2_T2V"]

fastvideo/pipelines/basic/pan2/stages.py

@@ -0,0 +1,202 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2-specific pipeline stages: prompt encoding, first-frame conditioning, sigma schedule and fp32 sampling."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+import numpy as np
+import PIL.Image
+import torch
+import torchvision.transforms.functional as TF
+from torch.distributed.tensor import DTensor
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+from fastvideo.pipelines.stages.denoising import DenoisingStage
+from fastvideo.pipelines.stages.input_validation import InputValidationStage
+from fastvideo.pipelines.stages.text_encoding import TextEncodingStage
+from fastvideo.pipelines.stages.timestep_preparation import TimestepPreparationStage
+from fastvideo.pipelines.stages.validators import StageValidators as V
+from fastvideo.pipelines.stages.validators import VerificationResult
+from fastvideo.utils import PRECISION_TO_TYPE
+
+_USER_HEADER = "<|im_start|>user\n"
+
+
+class PAN2InputValidationStage(InputValidationStage):
+    """Rejects geometry the PAN2 VAE and patch embedding cannot represent."""
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        batch = super().forward(batch, fastvideo_args)
+        pipeline_config = fastvideo_args.pipeline_config
+        vae_arch = pipeline_config.vae_config.arch_config
+        spatial_multiple = vae_arch.spatial_compression_ratio * pipeline_config.dit_config.arch_config.patch_size[1]
+        if batch.height % spatial_multiple != 0 or batch.width % spatial_multiple != 0:
+            raise ValueError(f"PAN2 height and width must be divisible by {spatial_multiple}, "
+                             f"got {batch.height} and {batch.width}.")
+        if (batch.num_frames - 1) % vae_arch.temporal_compression_ratio != 0:
+            raise ValueError(f"PAN2 num_frames - 1 must be divisible by {vae_arch.temporal_compression_ratio}, "
+                             f"got {batch.num_frames}.")
+        if batch.pil_image is not None and not isinstance(batch.pil_image, PIL.Image.Image):
+            raise ValueError("PAN2 image-to-video expects the first frame as a PIL image or an image_path.")
+        return batch
+
+
+class PAN2TextEncodingStage(TextEncodingStage):
+    """Encodes a prompt with PAN2's chat template.
+
+    PAN2 conditions on the tokens after the ``<|im_start|>user\\n`` header (the prompt and the assistant header),
+    read from the text encoder's hidden states with its last ``text_encoder_skip_layers`` layers skipped.
+    """
+
+    @staticmethod
+    def tokenize(tokenizer: Any, prompt: str, system_prompt: str, max_length: int) -> tuple[torch.Tensor, int]:
+        """Token ids truncated to ``max_length`` tokens after the user header, and where that header ends."""
+        # An empty prompt is encoded as a single space, as in training.
+        messages = [{"role": "system", "content": system_prompt}, {"role": "user", "content": prompt or " "}]
+        input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=True,
+                                                  return_dict=True)["input_ids"]
+        if isinstance(input_ids, torch.Tensor):
+            input_ids = input_ids.tolist()
+        if input_ids and isinstance(input_ids[0], list):
+            input_ids = input_ids[0]
+        header_ids = tokenizer(_USER_HEADER, add_special_tokens=False)["input_ids"]
+        crop_start = next((i + len(header_ids)
+                           for i in range(len(input_ids)) if input_ids[i:i + len(header_ids)] == header_ids), None)
+        if crop_start is None:
+            raise ValueError(f"The PAN2 tokenizer's chat template has no {_USER_HEADER!r} header.")
+        return torch.tensor([input_ids[:crop_start + max_length]], dtype=torch.long), crop_start
+
+    @torch.no_grad()
+    def encode_text(
+        self,
+        text: str | list[str],
+        fastvideo_args: FastVideoArgs,
+        encoder_index: int | list[int] | None = None,
+        return_attention_mask: bool = False,
+        return_type: str = "list",
+        device: torch.device | str | None = None,
+        dtype: torch.dtype | None = None,
+        max_length: int | Sequence[int | None] | None = None,
+        truncation: bool | None = None,
+        padding: bool | str | None = None,
+    ):
+        pipeline_config = fastvideo_args.pipeline_config
+        texts = [text] if isinstance(text, str) else list(text)
+        if len(texts) != 1:
+            raise ValueError(f"PAN2 encodes one prompt per request, got {len(texts)}.")
+        if isinstance(max_length, Sequence):
+            max_length = max_length[0]
+        max_length = max_length or pipeline_config.text_max_sequence_length
+
+        tokenizer = self.tokenizers[0]
+        text_encoder = self.text_encoders[0]
+        target_device = torch.device(device) if device is not None else get_local_torch_device()
+        # A module parked on CPU by text_encoder_cpu_offload computes on the target device and goes back afterwards;
+        # FSDP-wrapped (DTensor) encoders move their own shards.
+        first_param = next(text_encoder.parameters(), None)
+        input_device = first_param.device if first_param is not None else target_device
+        moved_for_forward = False
+        if (first_param is not None and not isinstance(first_param, DTensor)
+                and input_device.type != target_device.type):
+            text_encoder = text_encoder.to(target_device)
+            input_device = target_device
+            moved_for_forward = True
+
+        input_ids, crop_start = self.tokenize(tokenizer, texts[0], pipeline_config.text_system_prompt, max_length)
+        with set_forward_context(current_timestep=0, attn_metadata=None):
+            outputs = text_encoder(input_ids=input_ids.to(input_device), output_hidden_states=True)
+        prompt_embeds = outputs.hidden_states[-(pipeline_config.text_encoder_skip_layers + 1)][:, crop_start:]
+        # fp32 is exact for the bf16 encoder output. The shared latent stage draws the initial noise in this dtype,
+        # and PAN2 samples on fp32 latents; the denoiser casts the embeddings to the DiT dtype.
+        prompt_embeds = prompt_embeds.to(device=target_device, dtype=dtype or torch.float32, copy=True)
+        attention_mask = torch.ones(prompt_embeds.shape[:2], dtype=torch.long, device=target_device)
+        if moved_for_forward and fastvideo_args.text_encoder_cpu_offload:
+            text_encoder.to("cpu")
+        return self.return_embeds([prompt_embeds], [attention_mask], return_type, return_attention_mask, [0])
+
+
+class PAN2TimestepPreparationStage(TimestepPreparationStage):
+    """PAN2 samples on uniformly spaced sigmas from 1 down to (excluding) 0; the scheduler applies the shift."""
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if batch.timesteps is None and batch.sigmas is None:
+            batch.sigmas = np.linspace(1.0, 0.0, batch.num_inference_steps + 1)[:-1].tolist()
+        return super().forward(batch, fastvideo_args)
+
+
+class PAN2ConditioningStage(PipelineStage):
+    """Builds the conditioning latents and mask that PAN2 stacks after the noisy latents along channels.
+
+    Text-to-video conditions on zeros. Image-to-video fills the first latent frame with the normalized VAE latents
+    of the first frame, scaled to cover the output and center cropped, and sets that frame's mask channel to one.
+    """
+
+    def __init__(self, vae) -> None:
+        super().__init__()
+        self.vae = vae
+
+    @staticmethod
+    def preprocess_image(image: PIL.Image.Image, height: int, width: int) -> torch.Tensor:
+        """``[1, 3, 1, H, W]`` pixels in ``[-1, 1]``: scale to cover the canvas, round the size, center crop."""
+        image = image.convert("RGB")
+        scale = max(width / image.width, height / image.height)
+        image = image.resize((round(image.width * scale), round(image.height * scale)), resample=PIL.Image.LANCZOS)
+        left, top = round((image.width - width) / 2), round((image.height - height) / 2)
+        image = image.crop((left, top, left + width, top + height))
+        return (TF.to_tensor(image) * 2.0 - 1.0)[None, :, None]
+
+    def _encode_first_frame(self, image: PIL.Image.Image, batch: ForwardBatch,
+                            fastvideo_args: FastVideoArgs) -> torch.Tensor:
+        device = get_local_torch_device()
+        vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
+        self.vae = self.vae.to(device)
+        pixels = self.preprocess_image(image, batch.height, batch.width).to(device=device, dtype=vae_dtype)
+        latents = self.vae.encode(pixels).mode().float()
+        shift = self.vae.shift_factor.to(latents.device, latents.dtype)
+        scale = self.vae.scaling_factor.to(latents.device, latents.dtype)
+        if fastvideo_args.vae_cpu_offload:
+            self.vae.to("cpu")
+        return (latents - shift) * scale
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        latents = batch.latents
+        batch_size, num_channels, num_frames, height, width = latents.shape
+        condition = torch.zeros(batch_size,
+                                num_channels + 1,
+                                num_frames,
+                                height,
+                                width,
+                                device=latents.device,
+                                dtype=torch.float32)
+        if batch.pil_image is not None:
+            image_latents = self._encode_first_frame(batch.pil_image, batch, fastvideo_args)
+            condition[:, :num_channels, :1] = image_latents.to(condition.device)
+            condition[:, num_channels:, :1] = 1.0
+        batch.image_latent = condition
+        return batch
+
+    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
+        result.add_check("height", batch.height, V.positive_int)
+        result.add_check("width", batch.width, V.positive_int)
+        return result
+
+    def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("image_latent", batch.image_latent, [V.is_tensor, V.with_dims(5)])
+        return result
+
+
+class PAN2DenoisingStage(DenoisingStage):
+    """Keeps the latents in fp32 across the loop; they are cast to the DiT dtype for each forward."""
+
+    def prepare_scheduler_model_output(self, noise_pred, state) -> torch.Tensor:
+        return noise_pred.float()

fastvideo/pipelines/stages/denoising.py

@@ -496,6 +496,7 @@ class DenoisingStage(PipelineStage):
                                 noise_pred_text,
                                 guidance_rescale=batch.guidance_rescale,
                             )
+                noise_pred = self.prepare_scheduler_model_output(noise_pred, state)
                 if scheduler_fp32:
                     # Diffusers-style: fp32 Euler update outside autocast avoids BF16 drift.
                     latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
@@ -629,6 +630,10 @@ class DenoisingStage(PipelineStage):
     def prepare_timestep(self, timestep, latent_model_input, batch, state) -> torch.Tensor:
         return timestep.repeat(latent_model_input.shape[0])
 
+    def prepare_scheduler_model_output(self, noise_pred, state) -> torch.Tensor:
+        """The model output handed to the scheduler step, which returns the next latents in its dtype."""
+        return noise_pred
+
     def finish_step(self, latents, state) -> torch.Tensor:
         return latents
 

fastvideo/registry.py

@@ -10,6 +10,7 @@ from __future__ import annotations
 
 import dataclasses
 import os
+import re
 from collections.abc import Callable
 from functools import lru_cache
 from typing import TYPE_CHECKING, Any
@@ -45,6 +46,7 @@ from fastvideo.configs.pipelines.flux_2 import (
 from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig
 from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig
 from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig
+from fastvideo.configs.pipelines.pan2 import PAN2PipelineConfig
 from fastvideo.configs.pipelines.mmaudio import (
     MMAudioLarge44kV2PipelineConfig,
     MMAudioLarge44kV2AConfig,
@@ -1295,6 +1297,22 @@ def _register_configs() -> None:
         pipeline_cls_name="MiniMaxH3ModularPipeline",
     )
 
+    # PAN2: one pipeline runs text-to-video, and image-to-video when the request carries a first frame.
+    def _is_pan2(name: str) -> bool:
+        name = name.lower()
+        return "pan2modularpipeline" in name or re.search(r"(^|[/_-])pan-?2([/_.-]|$)", name) is not None
+
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=PAN2PipelineConfig,
+        workload_types=(WorkloadType.T2V, WorkloadType.I2V),
+        hf_model_paths=["IFM/PAN2"],
+        model_detectors=[_is_pan2],
+        model_family="pan2",
+        default_preset="pan2_t2v",
+        pipeline_cls_name="PAN2ModularPipeline",
+    )
+
     # SD3.5
     register_configs(
         sampling_param_cls=None,
@@ -1487,6 +1505,8 @@ def _register_presets() -> None:
         ALL_PRESETS as MINIMAX_H3_PRESETS, )
     from fastvideo.pipelines.basic.mmaudio.presets import (
         ALL_PRESETS as MMAUDIO_PRESETS, )
+    from fastvideo.pipelines.basic.pan2.presets import (
+        ALL_PRESETS as PAN2_PRESETS, )
     from fastvideo.pipelines.basic.sd35.presets import (
         ALL_PRESETS as SD35_PRESETS, )
     from fastvideo.pipelines.basic.stable_audio.presets import (
@@ -1524,6 +1544,7 @@ def _register_presets() -> None:
         MATRIXGAME3_PRESETS,
         MINIMAX_H3_PRESETS,
         MMAUDIO_PRESETS,
+        PAN2_PRESETS,
         SD35_PRESETS,
         STABLE_AUDIO_PRESETS,
         TURBODIFFUSION_PRESETS,

mkdocs.yml

@@ -204,6 +204,7 @@ nav:
     - LTX-2.5: inference/ltx2_5.md
     - Kandinsky 6 T2IVA: inference/kandinsky6.md
     - Kandinsky 6 Video SR: inference/kandinsky6_sr.md
+    - PAN2: inference/pan2.md
     - CLI: inference/cli.md
     - Add Pipeline: inference/add_pipeline.md
     - Examples:

tests/local_tests/encoders/test_pan2_qwen3_5_parity.py

@@ -0,0 +1,339 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Qwen3.5 text-encoder parity for PAN2 against Transformers' ``Qwen3_5TextModel``.
+
+Coverage scope: both. The checkpoint test loads the public tiny PAN2 text
+encoder through FastVideo's production ``TextEncoderLoader`` and feeds both
+models the exact PAN2 chat-templated token ids. The random-config tests use a
+larger hybrid stack (GQA, partial RoPE, more value than key heads in the Gated
+DeltaNet, sequences that cross several 64-token delta-rule chunks) and also
+cover bf16 and padded batches.
+
+Everything runs on CPU. Set ``PAN2_TINY_MODEL_DIR`` to a local copy of
+``wuqing157/tiny-pan2-modular-pipe``; otherwise the HF cache is used and the
+checkpoint test skips when it is absent.
+"""
+
+from __future__ import annotations
+
+import copy
+import os
+import socket
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import torch
+from torch.testing import assert_close
+
+from fastvideo.configs.models.encoders.qwen3_5 import Qwen3_5TextConfig
+from fastvideo.forward_context import set_forward_context
+from fastvideo.models.loader.utils import set_default_torch_dtype
+
+PARITY_SCOPE = "both"
+TINY_REPO_ID = "wuqing157/tiny-pan2-modular-pipe"
+PAN2_SYSTEM_MESSAGE = "You are a helpful assistant. Generate the video with the following instructions."
+PAN2_TEXT_ENCODER_LAYERS_TO_SKIP = 2
+PAN2_PROMPTS = (
+    "A red fox crosses fresh snow at sunrise.",
+    "A slow dolly shot glides through a rain-soaked neon market at night. Steam rises from food stalls while "
+    "vendors in bright raincoats shout over the hum of generators. Reflections of pink, teal and amber signs ripple "
+    "across puddles as a cyclist weaves between crowds, her umbrella catching droplets that sparkle in the light. "
+    "In the background, a tram rattles past an old clock tower, its windows glowing warm yellow. The camera rises to "
+    "reveal rooftops crowded with antennas and laundry lines, then tilts toward a thunderstorm rolling in over the "
+    "harbor, lightning briefly silhouetting cranes and container ships. Cinematic lighting, shallow depth of field, "
+    "35mm film grain, realistic motion blur.",
+)
+FP32_TOLERANCE = {"atol": 1e-5, "rtol": 1e-5}
+BF16_TOLERANCE = {"atol": 2e-2, "rtol": 2e-2}
+
+
+def _hf_qwen3_5() -> tuple[Any, Any]:
+    module = pytest.importorskip("transformers.models.qwen3_5", reason="transformers without Qwen3.5 support")
+    return module.Qwen3_5TextConfig, module.Qwen3_5TextModel
+
+
+def _tiny_model_dir() -> Path | None:
+    local = os.getenv("PAN2_TINY_MODEL_DIR")
+    if local:
+        return Path(local)
+    try:
+        from huggingface_hub import snapshot_download
+
+        return Path(snapshot_download(TINY_REPO_ID, local_files_only=True))
+    except Exception:
+        return None
+
+
+def _free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(("127.0.0.1", 0))
+        return sock.getsockname()[1]
+
+
+@pytest.fixture(scope="module")
+def cpu_single_rank():
+    """Single-rank CPU runtime; the loader would otherwise target ``mps`` on CPU-only Linux hosts."""
+    import fastvideo.distributed.parallel_state as parallel_state
+    import fastvideo.models.loader.component_loader as component_loader
+    from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
+
+    with pytest.MonkeyPatch.context() as patch:
+        patch.setattr(parallel_state, "get_local_torch_device", lambda: torch.device("cpu"))
+        patch.setattr(component_loader, "get_local_torch_device", lambda: torch.device("cpu"))
+        for name, value in (("MASTER_ADDR", "127.0.0.1"), ("MASTER_PORT", str(_free_port())), ("RANK", "0"),
+                            ("WORLD_SIZE", "1"), ("LOCAL_RANK", "0")):
+            patch.setenv(name, value)
+        maybe_init_distributed_environment_and_model_parallel(1, 1)
+        yield
+        cleanup_dist_env_and_memory()
+
+
+def _loader_args(config: Qwen3_5TextConfig, precision: str = "fp32") -> SimpleNamespace:
+    return SimpleNamespace(
+        pipeline_config=SimpleNamespace(text_encoder_configs=(config, ), text_encoder_precisions=(precision, )),
+        text_encoder_cpu_offload=False,
+        override_text_encoder_quant=None,
+        override_text_encoder_safetensors=None,
+        pin_cpu_memory=False,
+        disable_offload_on_unified_memory=lambda device_id, offload_flag=None: False,
+    )
+
+
+def _pan2_input_ids(tokenizer: Any, prompt: str) -> tuple[torch.Tensor, int]:
+    """Tokenize exactly like PAN2's ``encode_prompt`` and return ids plus the user-header crop."""
+    messages = [{"role": "system", "content": PAN2_SYSTEM_MESSAGE}, {"role": "user", "content": prompt or " "}]
+    input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=True,
+                                              return_dict=True).input_ids
+    header_ids = tokenizer("<|im_start|>user\n", add_special_tokens=False).input_ids
+    crop_start = next(
+        (i + len(header_ids) for i in range(len(input_ids)) if input_ids[i:i + len(header_ids)] == header_ids),
+        None,
+    )
+    assert crop_start is not None, "PAN2 chat template has no user header"
+    return torch.tensor([input_ids]), crop_start
+
+
+def _run_both(reference: torch.nn.Module, native: torch.nn.Module, **inputs: Any) -> tuple[Any, Any]:
+    with torch.inference_mode():
+        expected = reference(**inputs, output_hidden_states=True)
+        with set_forward_context(current_timestep=0, attn_metadata=None):
+            actual = native(**inputs, output_hidden_states=True)
+    return expected, actual
+
+
+def _assert_hidden_states_close(
+    expected: tuple[torch.Tensor, ...],
+    actual: tuple[torch.Tensor, ...] | None,
+    label: str,
+    tolerance: dict[str, float],
+    valid: torch.Tensor | None = None,
+) -> float:
+    assert actual is not None, f"{label}: native encoder returned no hidden_states"
+    assert len(actual) == len(expected), f"{label}: {len(actual)} hidden states, reference has {len(expected)}"
+    worst = 0.0
+    for index, (reference_state, native_state) in enumerate(zip(expected, actual, strict=True)):
+        assert native_state.shape == reference_state.shape, f"{label} hidden_states[{index}] shape"
+        assert torch.isfinite(native_state).all(), f"{label} hidden_states[{index}] has non-finite values"
+        if valid is not None:
+            reference_state = reference_state[valid]
+            native_state = native_state[valid]
+        drift = (native_state.float() - reference_state.float()).abs().max().item()
+        worst = max(worst, drift)
+        print(f"{label} hidden_states[{index}]: max_abs={drift:.3e}", flush=True)
+        assert_close(native_state, reference_state, **tolerance, msg=lambda m: f"{label} hidden_states[{index}]: {m}")
+    return worst
+
+
+def test_pan2_qwen3_5_arch_config_matches_transformers_rope_and_layer_types() -> None:
+    """Config.json fields resolve to the same RoPE settings and layer schedule as Transformers."""
+    hf_config_cls, _ = _hf_qwen3_5()
+    cases = (
+        {},
+        {
+            "rope_theta": 1_000_000.0
+        },
+        {
+            "partial_rotary_factor": 0.5,
+            "rope_parameters": {
+                "rope_type": "default",
+                "rope_theta": 5_000_000.0
+            }
+        },
+        {
+            "partial_rotary_factor": 0.5,
+            "rope_parameters": {
+                "rope_type": "default",
+                "rope_theta": 5_000_000.0,
+                "partial_rotary_factor": 0.25
+            }
+        },
+        {
+            "num_hidden_layers": 6,
+            "full_attention_interval": 3
+        },
+        {
+            "num_hidden_layers": 3,
+            "layer_types": ["mamba", "attention", "linear_attention"]
+        },
+    )
+    for case in cases:
+        reference = hf_config_cls.from_dict(dict(case))
+        config = Qwen3_5TextConfig()
+        config.update_model_arch(dict(case))
+        rope = config.arch_config.resolved_rope_parameters()
+        for key in ("rope_type", "rope_theta", "partial_rotary_factor"):
+            assert rope[key] == reference.rope_parameters[key], f"{case}: {key}"
+        assert config.arch_config.resolved_layer_types() == list(reference.layer_types), case
+
+
+def test_pan2_qwen3_5_tiny_checkpoint_loader_parity(cpu_single_rank) -> None:
+    """Production-loader encoder matches every HF hidden state on real PAN2 prompts (fp32 CPU)."""
+    _, hf_model_cls = _hf_qwen3_5()
+    model_dir = _tiny_model_dir()
+    if model_dir is None or not (model_dir / "text_encoder" / "config.json").is_file():
+        pytest.skip(f"{TINY_REPO_ID} is not available locally; set PAN2_TINY_MODEL_DIR")
+    if not (model_dir / "tokenizer").is_dir():
+        pytest.skip(f"{TINY_REPO_ID} tokenizer is not available locally")
+    from transformers import AutoTokenizer
+
+    from fastvideo.models.encoders.qwen3_5 import Qwen3_5TextModel
+    from fastvideo.models.loader.component_loader import TextEncoderLoader
+
+    text_encoder_dir = model_dir / "text_encoder"
+    tokenizer = AutoTokenizer.from_pretrained(model_dir / "tokenizer", local_files_only=True)
+    reference, loading_info = hf_model_cls.from_pretrained(
+        text_encoder_dir,
+        local_files_only=True,
+        dtype=torch.float32,
+        output_loading_info=True,
+    )
+    load_errors = {
+        name: loading_info.get(name)
+        for name in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs") if loading_info.get(name)
+    }
+    assert not load_errors, f"Reference Qwen3.5 checkpoint did not load strictly: {load_errors}"
+    reference = reference.eval()
+
+    native = TextEncoderLoader().load(str(text_encoder_dir), _loader_args(Qwen3_5TextConfig()))
+    assert isinstance(native, Qwen3_5TextModel)
+    assert native.layer_types == list(reference.config.layer_types)
+    assert next(native.parameters()).dtype == torch.float32
+    assert_close(native.rotary_emb.inv_freq, reference.rotary_emb.inv_freq.float(), atol=0.0, rtol=0.0)
+
+    for prompt_index, prompt in enumerate(PAN2_PROMPTS):
+        input_ids, crop_start = _pan2_input_ids(tokenizer, prompt)
+        label = f"tiny prompt{prompt_index} (L={input_ids.shape[1]})"
+        expected, actual = _run_both(reference, native, input_ids=input_ids)
+        worst = _assert_hidden_states_close(expected.hidden_states, actual.hidden_states, label, FP32_TOLERANCE)
+        assert_close(actual.last_hidden_state, expected.last_hidden_state, **FP32_TOLERANCE)
+
+        skip = PAN2_TEXT_ENCODER_LAYERS_TO_SKIP + 1
+        reference_embeds = expected.hidden_states[-skip][:, crop_start:]
+        native_embeds = actual.hidden_states[-skip][:, crop_start:]
+        assert_close(native_embeds, reference_embeds, **FP32_TOLERANCE)
+        print(
+            f"{label} PAN2 prompt_embeds hidden_states[-{skip}][:, {crop_start}:] "
+            f"max_abs={(native_embeds - reference_embeds).abs().max().item():.3e} all_layers_max_abs={worst:.3e}",
+            flush=True,
+        )
+
+
+@pytest.fixture(scope="module")
+def random_pair(cpu_single_rank):
+    """A randomly initialized 8-layer hybrid reference and the native model loaded from its state dict."""
+    hf_config_cls, hf_model_cls = _hf_qwen3_5()
+    from fastvideo.models.encoders.qwen3_5 import Qwen3_5TextModel
+
+    torch.manual_seed(0)
+    hf_config = hf_config_cls(
+        vocab_size=1024,
+        hidden_size=128,
+        intermediate_size=256,
+        num_hidden_layers=8,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=32,
+        linear_conv_kernel_dim=4,
+        linear_key_head_dim=16,
+        linear_value_head_dim=24,
+        linear_num_key_heads=2,
+        linear_num_value_heads=4,
+        layer_types=[
+            "linear_attention",
+            "linear_attention",
+            "full_attention",
+            "linear_attention",
+            "full_attention",
+            "linear_attention",
+            "linear_attention",
+            "full_attention",
+        ],
+        rope_parameters={
+            "rope_type": "default",
+            "rope_theta": 1_000_000.0,
+            "partial_rotary_factor": 0.25
+        },
+        initializer_range=0.08,
+    )
+    reference = hf_model_cls(hf_config).eval()
+    with torch.no_grad():
+        # Move the zero-centered norms, gates and decay parameters off their
+        # init values so every term contributes.
+        for name, parameter in reference.named_parameters():
+            if name.endswith("norm.weight") or name.endswith(("A_log", "dt_bias")):
+                parameter.add_(0.2 * torch.randn_like(parameter))
+
+    config = Qwen3_5TextConfig()
+    config.update_model_arch(hf_config.to_dict())
+    with set_default_torch_dtype(torch.float32):
+        native = Qwen3_5TextModel(config).eval()
+    loaded = native.load_weights(reference.state_dict().items())
+    expected_names = {name for name, _ in native.named_parameters()}
+    assert loaded == expected_names, f"unloaded native parameters: {sorted(expected_names - loaded)}"
+    assert set(reference.state_dict()) == expected_names
+    return reference, native
+
+
+def test_pan2_qwen3_5_random_config_fp32_parity(random_pair) -> None:
+    """Multi-chunk delta rule (L=301), GQA and partial RoPE match HF in fp32."""
+    reference, native = random_pair
+    generator = torch.Generator().manual_seed(1)
+    input_ids = torch.randint(0, reference.config.vocab_size, (2, 301), generator=generator)
+    expected, actual = _run_both(reference, native, input_ids=input_ids)
+    worst = _assert_hidden_states_close(expected.hidden_states, actual.hidden_states, "random fp32", FP32_TOLERANCE)
+    print(f"random fp32 all_layers_max_abs={worst:.3e}", flush=True)
+    # Each layer must move the residual stream, or the comparison is vacuous.
+    for index in range(1, len(expected.hidden_states) - 1):
+        change = (expected.hidden_states[index] - expected.hidden_states[index - 1]).abs().max().item()
+        assert change > 0.1, f"layer {index - 1} barely changes the hidden state ({change:.3e})"
+
+
+def test_pan2_qwen3_5_random_config_padding_parity(random_pair) -> None:
+    """Right- and left-padded rows match HF on every non-padded position."""
+    reference, native = random_pair
+    generator = torch.Generator().manual_seed(2)
+    input_ids = torch.randint(0, reference.config.vocab_size, (2, 200), generator=generator)
+    attention_mask = torch.ones_like(input_ids)
+    attention_mask[0, 150:] = 0
+    attention_mask[1, :37] = 0
+    expected, actual = _run_both(reference, native, input_ids=input_ids, attention_mask=attention_mask)
+    worst = _assert_hidden_states_close(
+        expected.hidden_states,
+        actual.hidden_states,
+        "random padded",
+        FP32_TOLERANCE,
+        valid=attention_mask.bool(),
+    )
+    print(f"random padded all_layers_max_abs={worst:.3e}", flush=True)
+
+
+def test_pan2_qwen3_5_random_config_bf16_parity(random_pair) -> None:
+    """Half-precision dtype handling (fp32 norms, gates and delta rule) follows HF."""
+    reference, native = (copy.deepcopy(model).to(torch.bfloat16) for model in random_pair)
+    generator = torch.Generator().manual_seed(3)
+    input_ids = torch.randint(0, reference.config.vocab_size, (1, 173), generator=generator)
+    expected, actual = _run_both(reference, native, input_ids=input_ids)
+    assert actual.hidden_states is not None and actual.hidden_states[-1].dtype == torch.bfloat16
+    worst = _assert_hidden_states_close(expected.hidden_states, actual.hidden_states, "random bf16", BF16_TOLERANCE)
+    print(f"random bf16 all_layers_max_abs={worst:.3e}", flush=True)

tests/local_tests/pan2/PORT_STATUS.md

@@ -0,0 +1,71 @@
+# PAN2 Port Status
+
+## Summary
+- model_family: pan2
+- workload_types: T2V, I2V (one pipeline; I2V when the request carries a first frame)
+- official_ref: vllm-project/vllm-omni#8602 at `59dd8c9a` (the diffusers integration is private)
+- official_ref_dir: none; rebuilt in `tests/local_tests/pan2/reference.py`
+- hf_weights_path: `IFM/PAN2` (not public); tests use `wuqing157/tiny-pan2-modular-pipe@3ff3130b`
+- local_weights_dir: HF cache, or `PAN2_TINY_MODEL_DIR`
+- source_layout: Modular Diffusers (`modular_model_index.json`), diffusers component subfolders
+- local_tests_readme: `tests/local_tests/pan2/README.md`
+
+## Current Phase
+- phase: 11 (post-parity review and handoff)
+- status: in_progress (GPU and real-weight validation pending, see I001 and I002)
+- owner: orchestrator
+- last_updated: 2026-10-15
+
+## Component Matrix
+| Component | Type | Reuse/Port | Official Definition | Official Instantiation | FastVideo Target | Prototype | Conversion | Parity | Open Issues |
+|---|---|---|---|---|---|---|---|---|---|
+| transformer | DiT | port | vLLM-Omni `pan2_transformer.py` (diffusers `PAN2Transformer3DModel`) | `transformer/config.json` | `fastvideo/models/dits/pan2.py`, `fastvideo/configs/models/dits/pan2.py` | done | load-time key mapping (Q/K/V fused) | non_skip_pass | I002 |
+| text_encoder | encoder | port | transformers `Qwen3_5TextModel` | `text_encoder/config.json` | `fastvideo/models/encoders/qwen3_5.py`, `fastvideo/configs/models/encoders/qwen3_5.py` | done | passthrough | non_skip_pass | I002 |
+| tokenizer | tokenizer | reuse | `Qwen2Tokenizer` | `tokenizer/` | AutoTokenizer | n/a | passthrough | covered by pipeline parity | — |
+| vae | VAE | reuse | diffusers `AutoencoderKLWan` (Wan2.2, z_dim 48) | `vae/config.json` | `fastvideo/models/wan/vae.py` | done | passthrough | non_skip_pass | — |
+| scheduler | scheduler | reuse | diffusers `FlowMatchEulerDiscreteScheduler` | `scheduler/scheduler_config.json` (shift 7) | `fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py` | done | passthrough | non_skip_pass | — |
+
+## Conversion State
+- conversion_script: none; the transformer maps diffusers keys at load time (`PAN2ArchConfig.param_names_mapping`)
+- converted_weights_dir: n/a
+- source_layout: diffusers
+- strict_load_status: tiny checkpoint loads strictly (no missing or unexpected keys)
+- passthrough_components: text_encoder, tokenizer, vae, scheduler
+- retry_history: none
+
+## Parity Commands
+| Scope | Command | Last Result | Notes |
+|---|---|---|---|
+| DiT | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py` | 12 passed (CPU) | includes SP=2 vs one rank over gloo |
+| Text encoder | `pytest tests/local_tests/encoders/test_pan2_qwen3_5_parity.py` | 5 passed (CPU) | exact |
+| VAE | `pytest tests/local_tests/vaes/test_pan2_vae_parity.py` | passed (CPU) | exact |
+| Scheduler | `pytest tests/local_tests/schedulers/test_pan2_scheduler_parity.py` | passed (CPU) | exact |
+| Pipeline | `pytest tests/local_tests/pipelines/test_pan2_pipeline_parity.py` | 3 passed (CPU) | T2V and I2V latents exact |
+
+## Open Questions
+| ID | Question | Owner | Needed By Phase | Status | Resolution |
+|---|---|---|---|---|---|
+| Q001 | Which reference, with the diffusers integration private? | orchestrator | 2 | resolved | vLLM-Omni #8602, rebuilt from diffusers/transformers blocks in `reference.py` |
+| Q002 | Port the `pan2-guardrail` prompt/video guardrails? | orchestrator | 7 | resolved | Not ported: the package is not public and FastVideo has no guardrail stage. |
+| Q003 | Port Cache-DiT `quality` and CFG-parallel? | orchestrator | 7 | resolved | Not ported: vLLM-Omni serving features. Ulysses SP is supported through `DistributedAttention`. |
+
+## Issues And Blockers
+| ID | Phase | Component | Severity | Issue | Evidence | Owner | Status | Resolution |
+|---|---|---|---|---|---|---|---|---|
+| I001 | 9-10 | all | medium | No real-weight parity or SSIM regression | `IFM/PAN2` is not public | orchestrator | open | Add `fastvideo/tests/ssim/test_pan2_similarity.py` once weights are public. |
+| I002 | 9 | transformer, text_encoder | medium | No GPU run yet (bf16 + FlashAttention, NCCL SP, 9B encoder speed) | All tests ran on a CPU login node | orchestrator | open | Run the same tests on a GPU node. |
+| I003 | 7 | pipeline | low | The shared denoiser casts latents to the model-output dtype in the scheduler step; PAN2 samples on fp32 latents | pipeline parity fails (8e-2) without the fix | orchestrator | resolved | `DenoisingStage.prepare_scheduler_model_output` hook; `PAN2DenoisingStage` returns `noise_pred.float()`. |
+
+## Escape Hatches
+| ID | Phase | Decision Type | Question | Recommended Option | Status | Resolution |
+|---|---|---|---|---|---|---|
+
+## Decisions
+| Date | Decision | Rationale | Impact |
+|---|---|---|---|
+| 2026-10-15 | Native Qwen3.5 encoder with a PyTorch chunked gated delta rule | No `fla`/`causal_conv1d` dependency; encoder runs once per prompt | Slower than fused kernels for long prompts |
+| 2026-10-15 | Fuse the per-stream Q/K/V projections at load time | One GEMM per stream instead of three | Mapping in `PAN2ArchConfig.param_names_mapping` |
+| 2026-10-15 | Text embeddings kept in fp32 after encoding | The shared latent stage draws noise in the embedding dtype; PAN2 uses fp32 noise and latents | Denoiser casts embeddings to the DiT dtype (`cast_prompt_embeds_to_dit_dtype`) |
+
+## Handoff Notes
+- The remaining work is I001 and I002. Both need hardware or weights, not code changes.

tests/local_tests/pan2/README.md

@@ -0,0 +1,54 @@
+# PAN2 Local Tests
+
+Runbook for validating the PAN2 port (text-to-video and first-frame image-to-video). These tests run manually and are
+not part of the default CI suite.
+
+## Reference
+
+PAN2's diffusers integration (`PAN2Transformer3DModel`, `PAN2ModularPipeline`) is not public yet, and the full
+`IFM/PAN2` weights are not released. The reference is the vLLM-Omni port,
+[vllm-project/vllm-omni#8602](https://github.com/vllm-project/vllm-omni/pull/8602) (head `59dd8c9a`), which reports
+parity with the diffusers implementation. `reference.py` rebuilds it from diffusers and transformers building blocks:
+the DiT from `AdaLayerNormZero`, `AdaLayerNormContinuous`, `FeedForward` (gelu-approximate and swiglu), `Timesteps`,
+`TimestepEmbedding` and the diffusers RoPE helpers, using the checkpoint's module names so the checkpoint loads
+strictly. The pipeline is `reference_generate`, a step-by-step copy of the vLLM-Omni `PAN2Pipeline.forward`
+using `Qwen3_5TextModel`, `AutoencoderKLWan` and `FlowMatchEulerDiscreteScheduler`.
+
+The checkpoint is the public random-weight [`wuqing157/tiny-pan2-modular-pipe`](https://huggingface.co/wuqing157/tiny-pan2-modular-pipe)
+(revision `3ff3130b`), the same one vLLM-Omni tests against. It has two DiT blocks, a four-layer Qwen3.5 text model
+(three Gated DeltaNet layers and one full-attention layer) and a small Wan2.2-layout VAE.
+
+## Setup
+
+```bash
+huggingface-cli download wuqing157/tiny-pan2-modular-pipe --revision 3ff3130b4527f61aaab7bde01465c9884dc84cdd
+# or point the tests at a local copy:
+export PAN2_TINY_MODEL_DIR=/path/to/tiny-pan2-modular-pipe
+```
+
+The tests need `diffusers` with `AutoencoderKLWan` and `transformers` with `Qwen3_5TextModel` (5.x). They run on CPU
+when no GPU is present. On a GPU-less host, `import fastvideo` also needs the `fastvideo_kernel` stub described in
+`fastvideo/tests/AGENTS.md`.
+
+## Tests
+
+| Area | Command | Result (CPU, fp32 unless noted) |
+|---|---|---|
+| DiT: tiny checkpoint, strict load + forward | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py -k tiny_checkpoint` | max abs 6.0e-7 |
+| DiT: random weights, batch 1 and 2 | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py -k random_weights` | max abs 0 and 1.9e-5 (outputs ~16) |
+| DiT: Ulysses SP=2 vs one rank, padded tokens | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py -k sp2` | max abs 5.7e-6 (gloo on CPU; NCCL on 2 GPUs) |
+| Qwen3.5 text encoder | `pytest tests/local_tests/encoders/test_pan2_qwen3_5_parity.py` | all hidden states exact (fp32, bf16, padded batch) |
+| Wan2.2 VAE encode/decode + normalization | `pytest tests/local_tests/vaes/test_pan2_vae_parity.py` | exact; normalization 2.4e-7 |
+| Flow-match scheduler, PAN2 sigmas + shift | `pytest tests/local_tests/schedulers/test_pan2_scheduler_parity.py` | exact |
+| Pipeline, T2V and I2V (bf16 DiT) | `pytest tests/local_tests/pipelines/test_pan2_pipeline_parity.py` | final latents and video exact |
+
+The pipeline test builds the FastVideo pipeline through the registry and component loaders and compares the final
+latents (relative L2 < 2e-2, cosine > 0.9995) and the decoded video (PSNR > 40 dB) with `reference_generate`. To
+confirm the threshold catches real deviations, three injected changes were each run against it: conditioning on
+`hidden_states[-2]`, letting the latents fall to bf16 between steps, and shifting the first-frame crop by one pixel.
+Each moved the latents by about 8e-2 relative L2.
+
+## Not covered
+
+- Real-weight parity and an SSIM regression: `IFM/PAN2` is not public.
+- GPU runs: everything above ran on CPU. The bf16 GPU path, FlashAttention, and SP over NCCL still need a GPU node.

tests/local_tests/pan2/reference.py

@@ -0,0 +1,299 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2 reference implementation for local parity tests.
+
+The diffusers ``PAN2Transformer3DModel`` and ``PAN2ModularPipeline`` are not public yet. The reference here is
+assembled from diffusers' and transformers' own building blocks (``AdaLayerNormZero``, ``AdaLayerNormContinuous``,
+``FeedForward`` with ``gelu-approximate``/``swiglu``, ``Timesteps``/``TimestepEmbedding``,
+``get_1d_rotary_pos_embed``/``apply_rotary_emb``, ``Qwen3_5TextModel``, ``AutoencoderKLWan``,
+``FlowMatchEulerDiscreteScheduler``) and follows the vLLM-Omni PAN2 port (vllm-project/vllm-omni#8602), which was
+validated against the diffusers implementation. Module names are the checkpoint's, so checkpoints load strictly.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import numpy as np
+import PIL.Image
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.attention import FeedForward
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed
+from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, RMSNorm
+
+TINY_REPO_ID = "wuqing157/tiny-pan2-modular-pipe"
+TINY_REVISION = "3ff3130b4527f61aaab7bde01465c9884dc84cdd"
+SYSTEM_PROMPT = "You are a helpful assistant. Generate the video with the following instructions."
+TEXT_ENCODER_LAYERS_TO_SKIP = 2
+MAX_SEQUENCE_LENGTH = 1536
+
+
+def tiny_model_dir() -> Path | None:
+    """``PAN2_TINY_MODEL_DIR``, else the pinned tiny checkpoint in the local HF cache."""
+    local = os.getenv("PAN2_TINY_MODEL_DIR")
+    if local:
+        return Path(local)
+    try:
+        from huggingface_hub import snapshot_download
+
+        return Path(snapshot_download(TINY_REPO_ID, revision=TINY_REVISION, local_files_only=True))
+    except Exception:
+        return None
+
+
+def _sdpa(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
+    # [B, S, H, D] -> [B, S, H*D]
+    out = F.scaled_dot_product_attention(query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2))
+    return out.transpose(1, 2).flatten(2)
+
+
+class _RefRefinerAttention(nn.Module):
+
+    def __init__(self, dim: int, heads: int, dim_head: int, eps: float):
+        super().__init__()
+        self.heads = heads
+        self.to_q = nn.Linear(dim, heads * dim_head, bias=False)
+        self.to_k = nn.Linear(dim, heads * dim_head, bias=False)
+        self.to_v = nn.Linear(dim, heads * dim_head, bias=False)
+        self.norm_q = RMSNorm(dim_head, eps=eps)
+        self.norm_k = RMSNorm(dim_head, eps=eps)
+        self.to_out = nn.ModuleList([nn.Linear(heads * dim_head, dim, bias=False)])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        query = self.norm_q(self.to_q(x).unflatten(-1, (self.heads, -1)))
+        key = self.norm_k(self.to_k(x).unflatten(-1, (self.heads, -1)))
+        value = self.to_v(x).unflatten(-1, (self.heads, -1))
+        return self.to_out[0](_sdpa(query, key, value))
+
+
+class _RefRefinerBlock(nn.Module):
+
+    def __init__(self, dim: int, heads: int, dim_head: int, mlp_ratio: float, eps: float):
+        super().__init__()
+        self.norm1 = RMSNorm(dim, eps=eps)
+        self.attn = _RefRefinerAttention(dim, heads, dim_head, eps)
+        self.norm2 = RMSNorm(dim, eps=eps)
+        self.ff = FeedForward(dim, inner_dim=int(dim * mlp_ratio), activation_fn="swiglu", bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.norm1(x))
+        return x + self.ff(self.norm2(x))
+
+
+class _RefJointAttention(nn.Module):
+
+    def __init__(self, dim: int, heads: int, dim_head: int, eps: float):
+        super().__init__()
+        self.heads = heads
+        self.to_q = nn.Linear(dim, heads * dim_head)
+        self.to_k = nn.Linear(dim, heads * dim_head)
+        self.to_v = nn.Linear(dim, heads * dim_head)
+        self.norm_q = RMSNorm(dim_head, eps=eps)
+        self.norm_k = RMSNorm(dim_head, eps=eps)
+        self.to_out = nn.ModuleList([nn.Linear(heads * dim_head, dim)])
+        self.add_q_proj = nn.Linear(dim, heads * dim_head)
+        self.add_k_proj = nn.Linear(dim, heads * dim_head)
+        self.add_v_proj = nn.Linear(dim, heads * dim_head)
+        self.norm_added_q = RMSNorm(dim_head, eps=eps)
+        self.norm_added_k = RMSNorm(dim_head, eps=eps)
+        self.to_add_out = nn.Linear(heads * dim_head, dim)
+
+    def forward(self, x: torch.Tensor, context: torch.Tensor,
+                rope: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
+        query = self.norm_q(self.to_q(x).unflatten(-1, (self.heads, -1)))
+        key = self.norm_k(self.to_k(x).unflatten(-1, (self.heads, -1)))
+        value = self.to_v(x).unflatten(-1, (self.heads, -1))
+        # Interleaved (GPT-J) rotary embedding on the video tokens only.
+        query = apply_rotary_emb(query, rope, use_real=True, use_real_unbind_dim=-1, sequence_dim=1)
+        key = apply_rotary_emb(key, rope, use_real=True, use_real_unbind_dim=-1, sequence_dim=1)
+
+        context_query = self.norm_added_q(self.add_q_proj(context).unflatten(-1, (self.heads, -1)))
+        context_key = self.norm_added_k(self.add_k_proj(context).unflatten(-1, (self.heads, -1)))
+        context_value = self.add_v_proj(context).unflatten(-1, (self.heads, -1))
+
+        out = _sdpa(torch.cat([query, context_query], dim=1), torch.cat([key, context_key], dim=1),
+                    torch.cat([value, context_value], dim=1))
+        out, context_out = out.split([x.shape[1], context.shape[1]], dim=1)
+        return self.to_out[0](out), self.to_add_out(context_out)
+
+
+class _RefBlock(nn.Module):
+
+    def __init__(self, dim: int, heads: int, dim_head: int, mlp_ratio: float, context_mlp_ratio: float):
+        super().__init__()
+        eps = 1e-6
+        self.norm1 = AdaLayerNormZero(dim)
+        self.norm1_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.scale_shift_table_context = nn.Parameter(torch.zeros(6 * dim))
+        self.attn = _RefJointAttention(dim, heads, dim_head, eps)
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.ff = FeedForward(dim, inner_dim=int(dim * mlp_ratio), activation_fn="gelu-approximate")
+        self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.ff_context = FeedForward(dim, inner_dim=int(dim * context_mlp_ratio), activation_fn="gelu-approximate")
+
+    def forward(self, x, context, temb, rope):
+        norm_x, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(x, emb=temb)
+        c_shift_msa, c_scale_msa, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = (
+            self.scale_shift_table_context.chunk(6, dim=-1))
+        norm_context = self.norm1_context(context) * (1 + c_scale_msa) + c_shift_msa
+        attn_out, context_attn_out = self.attn(norm_x, norm_context, rope)
+        x = x + attn_out * gate_msa.unsqueeze(1)
+        context = context + context_attn_out * c_gate_msa
+        norm_x = self.norm2(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        norm_context = self.norm2_context(context) * (1 + c_scale_mlp) + c_shift_mlp
+        x = x + self.ff(norm_x) * gate_mlp.unsqueeze(1)
+        context = context + self.ff_context(norm_context) * c_gate_mlp
+        return x, context
+
+
+class ReferencePAN2Transformer(nn.Module):
+
+    def __init__(self, patch_size=(1, 2, 2), in_channels=97, out_channels=48, num_attention_heads=32,
+                 attention_head_dim=128, num_layers=44, num_refiner_layers=4, mlp_ratio=8.0, context_mlp_ratio=4.0,
+                 refiner_mlp_ratio=4.0, text_embed_dim=4096, rope_axes_dim=(16, 56, 56), rope_theta=256.0,
+                 rope_temporal_scale=16 / 24, **_):
+        super().__init__()
+        self.patch_size = tuple(patch_size)
+        self.rope_axes_dim = tuple(rope_axes_dim)
+        self.rope_theta = rope_theta
+        self.rope_temporal_scale = rope_temporal_scale
+        inner_dim = num_attention_heads * attention_head_dim
+        self.x_embedder = nn.Conv3d(in_channels, inner_dim, kernel_size=self.patch_size, stride=self.patch_size)
+        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
+        self.time_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=inner_dim)
+        self.context_embedder = nn.Sequential(RMSNorm(text_embed_dim, eps=1e-5), nn.Linear(text_embed_dim, inner_dim))
+        self.context_refiner = nn.ModuleList([
+            _RefRefinerBlock(inner_dim, num_attention_heads, attention_head_dim, refiner_mlp_ratio, eps=1e-5)
+            for _ in range(num_refiner_layers)
+        ])
+        self.context_type_embedding = nn.Parameter(torch.zeros(inner_dim))
+        self.transformer_blocks = nn.ModuleList([
+            _RefBlock(inner_dim, num_attention_heads, attention_head_dim, mlp_ratio, context_mlp_ratio)
+            for _ in range(num_layers)
+        ])
+        self.norm_out = AdaLayerNormContinuous(inner_dim, inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out = nn.Linear(inner_dim, self.patch_size[0] * self.patch_size[1] * self.patch_size[2] * out_channels)
+
+    def _rope(self, num_frames: int, height: int, width: int, device) -> tuple[torch.Tensor, torch.Tensor]:
+        axes = [torch.arange(size, device=device, dtype=torch.float32) for size in (num_frames, height, width)]
+        grid = torch.stack(torch.meshgrid(*axes, indexing="ij"), dim=0)
+        grid[0] = grid[0] * self.rope_temporal_scale
+        cos, sin = zip(*[
+            get_1d_rotary_pos_embed(dim, grid[i].reshape(-1), self.rope_theta, use_real=True,
+                                    repeat_interleave_real=True) for i, dim in enumerate(self.rope_axes_dim)
+        ])
+        return torch.cat(cos, dim=1).float(), torch.cat(sin, dim=1).float()
+
+    def forward(self, hidden_states, timestep, encoder_hidden_states):
+        batch_size, _, num_frames, height, width = hidden_states.shape
+        p_t, p_h, p_w = self.patch_size
+        t, h, w = num_frames // p_t, height // p_h, width // p_w
+        rope = self._rope(t, h, w, hidden_states.device)
+        x = self.x_embedder(hidden_states).flatten(2).transpose(1, 2)
+        temb = self.time_embedder(self.time_proj(timestep).to(x.dtype))
+        context = self.context_embedder(encoder_hidden_states)
+        for block in self.context_refiner:
+            context = block(context)
+        context = context + self.context_type_embedding
+        for block in self.transformer_blocks:
+            x, context = block(x, context, temb, rope)
+        x = self.proj_out(self.norm_out(x, temb))
+        x = x.reshape(batch_size, t, h, w, -1, p_t, p_h, p_w).permute(0, 4, 1, 5, 2, 6, 3, 7)
+        return x.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# Reference pipeline (vLLM-Omni ``PAN2Pipeline.forward``)
+# --------------------------------------------------------------------------------------------------------------------
+
+
+def reference_encode_prompt(tokenizer, text_encoder, prompt: str, device) -> torch.Tensor:
+    messages = [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": prompt or " "}]
+    input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=True,
+                                              return_dict=True).input_ids
+    header_ids = tokenizer("<|im_start|>user\n", add_special_tokens=False).input_ids
+    crop_start = next(i + len(header_ids) for i in range(len(input_ids))
+                      if input_ids[i:i + len(header_ids)] == header_ids)
+    input_ids = torch.tensor([input_ids[:crop_start + MAX_SEQUENCE_LENGTH]], device=device)
+    hidden_states = text_encoder(input_ids=input_ids, output_hidden_states=True).hidden_states
+    return hidden_states[-(TEXT_ENCODER_LAYERS_TO_SKIP + 1)][:, crop_start:]
+
+
+def reference_condition_latents(vae, latents: torch.Tensor, image: PIL.Image.Image | None, height: int,
+                                width: int) -> torch.Tensor:
+    batch_size, num_channels, num_frames, latent_height, latent_width = latents.shape
+    condition = latents.new_zeros(batch_size, num_channels + 1, num_frames, latent_height, latent_width)
+    if image is None:
+        return condition
+    from diffusers.video_processor import VideoProcessor
+
+    image = image.convert("RGB")
+    scale = max(width / image.width, height / image.height)
+    image = image.resize((round(image.width * scale), round(image.height * scale)), resample=PIL.Image.LANCZOS)
+    left, top = round((image.width - width) / 2), round((image.height - height) / 2)
+    image = image.crop((left, top, left + width, top + height))
+    pixels = VideoProcessor(vae_scale_factor=16).preprocess(image, height=height, width=width)
+    pixels = pixels.unsqueeze(2).to(device=latents.device, dtype=vae.dtype)
+    image_latents = vae.encode(pixels).latent_dist.mode()
+    mean = torch.tensor(vae.config.latents_mean).view(1, -1, 1, 1, 1).to(image_latents)
+    std = torch.tensor(vae.config.latents_std).view(1, -1, 1, 1, 1).to(image_latents)
+    condition[:, :num_channels, :1] = (image_latents - mean) / std
+    condition[:, num_channels:, :1] = 1.0
+    return condition
+
+
+@torch.no_grad()
+def reference_generate(components: dict, *, prompt: str, negative_prompt: str, image: PIL.Image.Image | None,
+                       height: int, width: int, num_frames: int, num_inference_steps: int, guidance_scale: float,
+                       seed: int, dit_dtype: torch.dtype, device="cpu") -> tuple[torch.Tensor, torch.Tensor]:
+    """Final latents and the decoded video in ``[0, 1]``, following the vLLM-Omni pipeline step by step."""
+    from diffusers.utils.torch_utils import randn_tensor
+
+    tokenizer, text_encoder = components["tokenizer"], components["text_encoder"]
+    transformer, vae, scheduler = components["transformer"], components["vae"], components["scheduler"]
+    do_cfg = guidance_scale > 1.0
+    prompt_embeds = reference_encode_prompt(tokenizer, text_encoder, prompt, device).to(dit_dtype)
+    negative_embeds = (reference_encode_prompt(tokenizer, text_encoder, negative_prompt, device).to(dit_dtype)
+                       if do_cfg else None)
+
+    shape = (1, vae.config.z_dim, (num_frames - 1) // vae.config.scale_factor_temporal + 1,
+             height // vae.config.scale_factor_spatial, width // vae.config.scale_factor_spatial)
+    generator = [torch.Generator("cpu").manual_seed(seed)]
+    latents = randn_tensor(shape, generator=generator, device=torch.device(device), dtype=torch.float32)
+    condition = reference_condition_latents(vae, latents, image, height, width).to(dit_dtype)
+
+    scheduler.set_timesteps(sigmas=np.linspace(1.0, 0.0, num_inference_steps + 1)[:-1], device=device)
+    for t in scheduler.timesteps:
+        model_input = torch.cat([latents.to(dit_dtype), condition], dim=1)
+        timestep = t.expand(model_input.shape[0])
+        noise_pred = transformer(model_input, timestep, prompt_embeds)
+        if do_cfg:
+            noise_uncond = transformer(model_input, timestep, negative_embeds)
+            noise_pred = noise_uncond + guidance_scale * (noise_pred - noise_uncond)
+        latents = scheduler.step(noise_pred.float(), t, latents, return_dict=False)[0]
+
+    mean = torch.tensor(vae.config.latents_mean).view(1, -1, 1, 1, 1).to(latents)
+    std = torch.tensor(vae.config.latents_std).view(1, -1, 1, 1, 1).to(latents)
+    video = vae.decode((latents * std + mean).to(vae.dtype), return_dict=False)[0]
+    return latents, (video.float() / 2 + 0.5).clamp(0, 1)
+
+
+def load_reference_components(model_dir: Path, dit_dtype: torch.dtype) -> dict:
+    import json
+
+    from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler
+    from safetensors.torch import load_file
+    from transformers import AutoTokenizer, Qwen3_5TextModel
+
+    transformer_dir = model_dir / "transformer"
+    transformer = ReferencePAN2Transformer(**json.loads((transformer_dir / "config.json").read_text()))
+    transformer.load_state_dict(load_file(str(transformer_dir / "diffusion_pytorch_model.safetensors")), strict=True)
+    return {
+        "tokenizer": AutoTokenizer.from_pretrained(model_dir / "tokenizer"),
+        "text_encoder": Qwen3_5TextModel.from_pretrained(model_dir / "text_encoder", torch_dtype=torch.float32).eval(),
+        "transformer": transformer.to(dit_dtype).eval(),
+        "vae": AutoencoderKLWan.from_pretrained(model_dir / "vae", torch_dtype=torch.float32).eval(),
+        "scheduler": FlowMatchEulerDiscreteScheduler.from_pretrained(model_dir / "scheduler"),
+    }

tests/local_tests/pipelines/test_pan2_pipeline_parity.py

@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2 end-to-end pipeline parity on the public tiny checkpoint.
+
+The FastVideo pipeline (built through the real registry and component loaders) is compared against
+``tests/local_tests/pan2/reference.py::reference_generate``, which follows the vLLM-Omni ``PAN2Pipeline.forward``
+step by step with transformers/diffusers components: prompt template and crop, fp32 noise, first-frame
+conditioning, the uniform sigma schedule with shift 7, sequential CFG, fp32 Euler steps and the VAE decode.
+
+Both sides run the DiT in bf16 (FastVideo's denoiser always feeds it bf16), so the comparison allows bf16 rounding.
+Runs on CPU when no GPU is present. Set ``PAN2_TINY_MODEL_DIR`` to a local copy of the tiny checkpoint.
+"""
+
+from __future__ import annotations
+
+import os
+import socket
+
+import PIL.Image
+import pytest
+import torch
+
+os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
+
+from tests.local_tests.pan2.reference import load_reference_components, reference_generate, tiny_model_dir  # noqa: E402
+
+HEIGHT, WIDTH, NUM_FRAMES = 64, 96, 9
+NUM_INFERENCE_STEPS = 3
+# A large scale makes the guided steps visible through the tiny random-weight VAE.
+GUIDANCE_SCALE = 30.0
+SEED = 42
+PROMPT = "a cat walks on the grass"
+NEGATIVE_PROMPT = "blurry"
+# Not the output aspect ratio, so the cover-resize and center crop are exercised.
+CONDITION_IMAGE = PIL.Image.radial_gradient("L").resize((120, 64)).convert("RGB")
+
+
+def _free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(("127.0.0.1", 0))
+        return sock.getsockname()[1]
+
+
+@pytest.fixture(scope="module")
+def pan2_runtime():
+    model_dir = tiny_model_dir()
+    if model_dir is None or not (model_dir / "modular_model_index.json").is_file():
+        pytest.skip("wuqing157/tiny-pan2-modular-pipe is not available locally; set PAN2_TINY_MODEL_DIR")
+    from fastvideo.configs.pipelines.pan2 import PAN2PipelineConfig
+    from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
+    import fastvideo.distributed.parallel_state as parallel_state
+    from fastvideo.fastvideo_args import FastVideoArgs
+    from fastvideo.pipelines import build_pipeline
+
+    monkeypatch = pytest.MonkeyPatch()
+    if not torch.cuda.is_available():
+        # On a GPU-less Linux host fastvideo resolves the local device to mps; run on CPU instead. Swapping the code
+        # object reaches every module that imported the function by name.
+        def _cpu_device() -> torch.device:
+            return torch.device("cpu")
+
+        monkeypatch.setattr(parallel_state.get_local_torch_device, "__code__", _cpu_device.__code__)
+    for name, value in (("MASTER_ADDR", "127.0.0.1"), ("MASTER_PORT", str(_free_port())), ("RANK", "0"),
+                        ("WORLD_SIZE", "1"), ("LOCAL_RANK", "0")):
+        monkeypatch.setenv(name, os.environ.get(name, value))
+    maybe_init_distributed_environment_and_model_parallel(1, 1)
+
+    pipeline_config = PAN2PipelineConfig(dit_precision="bf16",
+                                         vae_precision="fp32",
+                                         text_encoder_precisions=("fp32", ))
+    args = FastVideoArgs(model_path=str(model_dir),
+                         num_gpus=1,
+                         pipeline_config=pipeline_config,
+                         dit_cpu_offload=False,
+                         text_encoder_cpu_offload=False,
+                         vae_cpu_offload=False,
+                         use_fsdp_inference=False)
+    pipeline = build_pipeline(args)
+    device = parallel_state.get_local_torch_device()
+    try:
+        yield pipeline, args, model_dir, device
+    finally:
+        del pipeline
+        cleanup_dist_env_and_memory()
+        monkeypatch.undo()
+
+
+def _fastvideo_generate(pipeline, args, image, guidance_scale: float = GUIDANCE_SCALE):
+    from fastvideo.pipelines import ForwardBatch
+
+    batch = ForwardBatch(data_type="video",
+                         prompt=PROMPT,
+                         negative_prompt=NEGATIVE_PROMPT,
+                         seed=SEED,
+                         height=HEIGHT,
+                         width=WIDTH,
+                         num_frames=NUM_FRAMES,
+                         num_inference_steps=NUM_INFERENCE_STEPS,
+                         guidance_scale=guidance_scale,
+                         pil_image=image)
+    with torch.inference_mode():
+        batch = pipeline.forward(batch, args)
+    return batch.latents.float().cpu(), batch.output.float().cpu()
+
+
+def _psnr(a: torch.Tensor, b: torch.Tensor) -> float:
+    mse = (a - b).pow(2).mean().clamp_min(1e-12)
+    return float(10 * torch.log10(1.0 / mse))
+
+
+@pytest.mark.parametrize("image", [None, CONDITION_IMAGE], ids=["t2v", "i2v"])
+def test_pan2_pipeline_matches_reference(pan2_runtime, image):
+    pipeline, args, model_dir, device = pan2_runtime
+    latents, video = _fastvideo_generate(pipeline, args, image)
+
+    components = load_reference_components(model_dir, torch.bfloat16)
+    components = {
+        name: module.to(device) if isinstance(module, torch.nn.Module) else module
+        for name, module in components.items()
+    }
+    ref_latents, ref_video = reference_generate(components,
+                                                prompt=PROMPT,
+                                                negative_prompt=NEGATIVE_PROMPT,
+                                                image=image,
+                                                height=HEIGHT,
+                                                width=WIDTH,
+                                                num_frames=NUM_FRAMES,
+                                                num_inference_steps=NUM_INFERENCE_STEPS,
+                                                guidance_scale=GUIDANCE_SCALE,
+                                                seed=SEED,
+                                                dit_dtype=torch.bfloat16,
+                                                device=device)
+    ref_latents, ref_video = ref_latents.float().cpu(), ref_video.float().cpu()
+
+    assert latents.dtype == torch.float32
+    assert latents.shape == ref_latents.shape == (1, 48, 3, HEIGHT // 16, WIDTH // 16)
+    assert video.shape == ref_video.shape == (1, 3, NUM_FRAMES, HEIGHT, WIDTH)
+    latent_rel = float((latents - ref_latents).norm() / ref_latents.norm())
+    cosine = float(torch.nn.functional.cosine_similarity(latents.flatten(), ref_latents.flatten(), dim=0))
+    psnr = _psnr(video, ref_video)
+    print(f"latents rel_l2={latent_rel:.3e} cosine={cosine:.6f}; video psnr={psnr:.1f} dB")
+    assert latent_rel < 2e-2
+    assert cosine > 0.9995
+    assert psnr > 40.0
+
+
+def test_pan2_image_conditioning_and_guidance_change_every_frame(pan2_runtime):
+    pipeline, args, _, _ = pan2_runtime
+    _, t2v = _fastvideo_generate(pipeline, args, None)
+    _, i2v = _fastvideo_generate(pipeline, args, CONDITION_IMAGE)
+    _, unguided = _fastvideo_generate(pipeline, args, None, guidance_scale=1.0)
+
+    def per_frame_max_diff(a, b):
+        return (a - b).abs().amax(dim=(0, 1, 3, 4))
+
+    assert per_frame_max_diff(t2v, i2v).max() > 1e-3
+    assert per_frame_max_diff(t2v, unguided).min() > 1e-3

tests/local_tests/schedulers/test_pan2_scheduler_parity.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2 flow-matching Euler schedule on FastVideo's FlowMatchEulerDiscreteScheduler.
+
+Coverage scope: both. PAN2 loads Diffusers' ``FlowMatchEulerDiscreteScheduler``
+from ``<checkpoint>/scheduler`` (shift 7.0, no dynamic shifting), optionally
+overrides the shift with ``flow_shift``, calls
+``set_timesteps(sigmas=np.linspace(1.0, 0.0, N + 1)[:-1])`` and steps float32
+latents. The FastVideo side is loaded by the production ``SchedulerLoader`` and
+its schedule is built by the stock ``TimestepPreparationStage`` from
+``batch.sigmas``. Timesteps, sigmas, and every Euler update must be bit-exact.
+
+The default asset is the tiny PAN2 checkpoint (config only), so this runs on a
+CPU-only host::
+
+    pytest tests/local_tests/schedulers/test_pan2_scheduler_parity.py -v -s
+
+Set ``PAN2_MODEL_PATH`` to a local PAN2 root containing ``scheduler/`` to use
+another checkpoint.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+from torch.testing import assert_close
+
+try:
+    from diffusers import FlowMatchEulerDiscreteScheduler as DiffusersFlowMatchEulerDiscreteScheduler
+except ImportError:  # pragma: no cover - environment dependent
+    pytest.skip("diffusers with FlowMatchEulerDiscreteScheduler is required", allow_module_level=True)
+
+from fastvideo.models.loader.component_loader import SchedulerLoader
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler as FastVideoFlowMatchEulerDiscreteScheduler, )
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.timestep_preparation import TimestepPreparationStage
+
+PARITY_SCOPE = "both"
+TINY_PAN2_REPO = "wuqing157/tiny-pan2-modular-pipe"
+TINY_PAN2_REVISION = "3ff3130b4527f61aaab7bde01465c9884dc84cdd"
+MODEL_PATH_ENV = "PAN2_MODEL_PATH"
+LATENT_SHAPE = (1, 48, 3, 4, 6)
+
+
+def _pan2_root(subfolder: str) -> Path:
+    """Resolve a PAN2 root that contains ``subfolder``; skip when unavailable."""
+    override = os.getenv(MODEL_PATH_ENV)
+    if override:
+        root = Path(override)
+    else:
+        try:
+            from huggingface_hub import snapshot_download
+        except ImportError:
+            pytest.skip("huggingface_hub is required to resolve the tiny PAN2 checkpoint")
+        root = None
+        last_error: Exception | None = None
+        for local_files_only in (True, False):
+            try:
+                root = Path(
+                    snapshot_download(
+                        TINY_PAN2_REPO,
+                        revision=TINY_PAN2_REVISION,
+                        allow_patterns=[f"{subfolder}/*"],
+                        local_files_only=local_files_only,
+                    ))
+                break
+            except Exception as exc:  # noqa: BLE001 - offline/missing cache must skip.
+                last_error = exc
+        if root is None:
+            pytest.skip(f"Tiny PAN2 checkpoint unavailable: {last_error}")
+    if not (root / subfolder).is_dir():
+        pytest.skip(f"PAN2 {subfolder}/ missing under {root}")
+    return root
+
+
+@pytest.fixture(scope="module")
+def pan2_root() -> Path:
+    root = _pan2_root("scheduler")
+    if not (root / "scheduler" / "scheduler_config.json").is_file():
+        pytest.skip(f"PAN2 scheduler config missing under {root}")
+    return root
+
+
+@pytest.fixture
+def cpu_runtime(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Pin FastVideo's device helper to CPU; it reports ``mps`` on CPU-only Linux."""
+    monkeypatch.setattr("fastvideo.pipelines.stages.timestep_preparation.get_local_torch_device",
+                        lambda: torch.device("cpu"))
+
+
+def _pan2_sigmas(num_steps: int) -> np.ndarray:
+    """PAN2's pre-shift grid: uniform from 1 down to, but excluding, 0."""
+    return np.linspace(1.0, 0.0, num_steps + 1)[:-1]
+
+
+def _reference(root: Path, flow_shift: float | None = None) -> DiffusersFlowMatchEulerDiscreteScheduler:
+    """Mirror the upstream PAN2 pipeline's scheduler construction."""
+    scheduler = DiffusersFlowMatchEulerDiscreteScheduler.from_pretrained(str(root),
+                                                                         subfolder="scheduler",
+                                                                         local_files_only=True)
+    if flow_shift is not None:
+        scheduler._shift = flow_shift
+    return scheduler
+
+
+def _fastvideo(root: Path, flow_shift: float | None = None) -> FastVideoFlowMatchEulerDiscreteScheduler:
+    """Load through the production SchedulerLoader, which applies ``flow_shift``."""
+    fastvideo_args = SimpleNamespace(pipeline_config=SimpleNamespace(flow_shift=flow_shift))
+    scheduler = SchedulerLoader().load(str(root / "scheduler"), fastvideo_args)
+    assert type(scheduler) is FastVideoFlowMatchEulerDiscreteScheduler
+    return scheduler
+
+
+def _prepare_fastvideo_timesteps(scheduler: FastVideoFlowMatchEulerDiscreteScheduler,
+                                 num_steps: int) -> ForwardBatch:
+    """Drive the stock TimestepPreparationStage with PAN2's custom sigmas."""
+    batch = ForwardBatch(data_type="video", num_inference_steps=num_steps, sigmas=_pan2_sigmas(num_steps).tolist())
+    stage = TimestepPreparationStage(scheduler)
+    assert stage.verify_input(batch, SimpleNamespace()).is_valid()
+    batch = stage.forward(batch, SimpleNamespace())
+    assert batch.timesteps is scheduler.timesteps
+    return batch
+
+
+def test_pan2_scheduler_loader_keeps_checkpoint_config(pan2_root: Path) -> None:
+    """Without ``flow_shift`` the loader keeps every released config value, including shift 7.0."""
+    raw = json.loads((pan2_root / "scheduler" / "scheduler_config.json").read_text())
+    assert raw["_class_name"] == "FlowMatchEulerDiscreteScheduler"
+    scheduler = _fastvideo(pan2_root)
+    for key, value in raw.items():
+        if not key.startswith("_"):
+            assert scheduler.config[key] == value, key
+    assert scheduler.shift == raw["shift"] == 7.0
+    assert not scheduler.config.use_dynamic_shifting
+
+
+@pytest.mark.parametrize("num_steps", [1, 4, 8, 30, 50])
+def test_pan2_scheduler_schedule_matches_diffusers(num_steps: int, pan2_root: Path, cpu_runtime: None) -> None:
+    """Timesteps and sigmas from ``batch.sigmas`` are bit-exact for short and long schedules."""
+    reference = _reference(pan2_root)
+    reference.set_timesteps(sigmas=_pan2_sigmas(num_steps), device="cpu")
+    fastvideo = _fastvideo(pan2_root)
+    batch = _prepare_fastvideo_timesteps(fastvideo, num_steps)
+
+    print(f"[pan2-scheduler] N={num_steps} timesteps[:4]={reference.timesteps[:4].tolist()} "
+          f"sigmas[-3:]={reference.sigmas[-3:].tolist()}")
+    assert len(batch.timesteps) == num_steps
+    assert_close(batch.timesteps, reference.timesteps, atol=0, rtol=0)
+    assert_close(fastvideo.sigmas, reference.sigmas, atol=0, rtol=0)
+    assert fastvideo.sigmas[0].item() == 1.0 and fastvideo.sigmas[-1].item() == 0.0
+
+
+@pytest.mark.parametrize("flow_shift", [None, 5.0], ids=["checkpoint_shift", "flow_shift_override"])
+def test_pan2_scheduler_steps_match_diffusers(flow_shift: float | None, pan2_root: Path, cpu_runtime: None) -> None:
+    """Five float32 Euler updates are bit-exact, with and without a shift override."""
+    num_steps = 5
+    reference = _reference(pan2_root, flow_shift)
+    reference.set_timesteps(sigmas=_pan2_sigmas(num_steps), device="cpu")
+    fastvideo = _fastvideo(pan2_root, flow_shift)
+    _prepare_fastvideo_timesteps(fastvideo, num_steps)
+    assert fastvideo.shift == reference.shift == (flow_shift or 7.0)
+    assert_close(fastvideo.timesteps, reference.timesteps, atol=0, rtol=0)
+
+    generator = torch.Generator().manual_seed(2026)
+    reference_latents = torch.randn(LATENT_SHAPE, generator=generator)
+    fastvideo_latents = reference_latents.clone()
+    max_diff = 0.0
+    for index in range(num_steps):
+        noise_pred = torch.randn(LATENT_SHAPE, generator=generator)
+        reference_latents = reference.step(noise_pred, reference.timesteps[index], reference_latents).prev_sample
+        fastvideo_latents = fastvideo.step(noise_pred, fastvideo.timesteps[index], fastvideo_latents).prev_sample
+        assert fastvideo_latents.dtype == torch.float32
+        max_diff = max(max_diff, (fastvideo_latents - reference_latents).abs().max().item())
+        assert_close(fastvideo_latents, reference_latents, atol=0, rtol=0)
+    print(f"[pan2-scheduler] flow_shift={flow_shift} steps={num_steps} max_abs={max_diff:.3e}")
+    assert fastvideo.step_index == reference.step_index == num_steps
+
+
+def test_pan2_scheduler_step_returns_model_output_dtype(pan2_root: Path, cpu_runtime: None) -> None:
+    """Both implementations cast ``prev_sample`` to the model output dtype.
+
+    PAN2 keeps float32 latents by stepping with ``noise_pred.float()``; a bf16
+    noise prediction would silently drop the latents to bf16 on either side.
+    """
+    reference = _reference(pan2_root)
+    reference.set_timesteps(sigmas=_pan2_sigmas(2), device="cpu")
+    fastvideo = _fastvideo(pan2_root)
+    _prepare_fastvideo_timesteps(fastvideo, 2)
+    latents = torch.randn(LATENT_SHAPE, generator=torch.Generator().manual_seed(5))
+    noise_pred = torch.randn(LATENT_SHAPE, generator=torch.Generator().manual_seed(6)).to(torch.bfloat16)
+    expected = reference.step(noise_pred, reference.timesteps[0], latents).prev_sample
+    actual = fastvideo.step(noise_pred, fastvideo.timesteps[0], latents).prev_sample
+    assert expected.dtype == actual.dtype == torch.bfloat16
+    assert_close(actual, expected, atol=0, rtol=0)

tests/local_tests/transformers/test_pan2_transformer_parity.py

@@ -0,0 +1,311 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2 transformer weight-schema and forward parity.
+
+The reference is ``tests/local_tests/pan2/reference.py``: diffusers building blocks arranged as the (not yet public)
+diffusers ``PAN2Transformer3DModel``, following the vLLM-Omni port. Its module names are the checkpoint's, so the
+public tiny checkpoint loads into it strictly.
+
+Set ``PAN2_TINY_MODEL_DIR`` to a local copy of ``wuqing157/tiny-pan2-modular-pipe``; otherwise the HF cache is
+used and the checkpoint test skips when it is absent.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import socket
+import subprocess
+import sys
+from dataclasses import fields
+from pathlib import Path
+
+import pytest
+import torch
+import torch.nn as nn
+from torch.testing import assert_close
+
+os.environ.setdefault("DISABLE_SP", "1")
+os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
+
+from fastvideo.forward_context import set_forward_context  # noqa: E402
+from fastvideo.models.loader.fsdp_load import load_model_from_full_model_state_dict  # noqa: E402
+from fastvideo.models.loader.utils import get_param_names_mapping  # noqa: E402
+from tests.local_tests.pan2.reference import TINY_REPO_ID, ReferencePAN2Transformer, tiny_model_dir  # noqa: E402
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+SP_WORLD_SIZE = 2
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# Helpers
+# --------------------------------------------------------------------------------------------------------------------
+
+
+def _free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(("127.0.0.1", 0))
+        return sock.getsockname()[1]
+
+
+def _small_kwargs() -> dict:
+    """Larger than the tiny checkpoint: several heads, asymmetric RoPE axes, more blocks, distinct MLP ratios."""
+    return {
+        "patch_size": (1, 2, 2),
+        "in_channels": 9,
+        "out_channels": 4,
+        "num_attention_heads": 4,
+        "attention_head_dim": 32,
+        "num_layers": 3,
+        "num_refiner_layers": 2,
+        "mlp_ratio": 3.0,
+        "context_mlp_ratio": 2.0,
+        "refiner_mlp_ratio": 1.5,
+        "text_embed_dim": 48,
+        "rope_axes_dim": (8, 12, 12),
+        "rope_theta": 256.0,
+        "rope_temporal_scale": 16 / 24,
+    }
+
+
+@pytest.fixture
+def single_rank_dist():
+    from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
+
+    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+    os.environ.setdefault("MASTER_PORT", str(_free_port()))
+    os.environ.setdefault("RANK", "0")
+    os.environ.setdefault("WORLD_SIZE", "1")
+    os.environ.setdefault("LOCAL_RANK", "0")
+    maybe_init_distributed_environment_and_model_parallel(1, 1)
+    yield
+    cleanup_dist_env_and_memory()
+
+
+def _build_native(arch_kwargs: dict):
+    from fastvideo.configs.models.dits.pan2 import PAN2VideoConfig
+    from fastvideo.models.dits.pan2 import PAN2Transformer3DModel
+
+    config = PAN2VideoConfig()
+    config.update_model_arch(arch_kwargs)
+    return PAN2Transformer3DModel(config=config, hf_config=arch_kwargs).float().eval()
+
+
+def _load_into_native(native, weights) -> None:
+    incompatible = load_model_from_full_model_state_dict(
+        native,
+        iter(weights),
+        device=torch.device("cpu"),
+        param_dtype=torch.float32,
+        strict=True,
+        param_names_mapping=get_param_names_mapping(native.param_names_mapping),
+        training_mode=False,
+    )
+    assert incompatible.missing_keys == []
+    assert incompatible.unexpected_keys == []
+
+
+def _randomize(module: nn.Module, seed: int) -> None:
+    generator = torch.Generator().manual_seed(seed)
+    with torch.no_grad():
+        for param in module.parameters():
+            param.copy_(torch.randn(param.shape, generator=generator) * 0.2)
+
+
+def _make_inputs(in_channels: int, text_dim: int, batch_size: int = 2, seed: int = 0) -> dict[str, torch.Tensor]:
+    generator = torch.Generator().manual_seed(seed)
+    return {
+        # 9 frames at 64x96 in latent space: 3 latent frames of 4x6.
+        "hidden_states": torch.randn(batch_size, in_channels, 3, 4, 6, generator=generator),
+        "timestep": torch.tensor([700.0, 125.5][:batch_size]),
+        "encoder_hidden_states": torch.randn(batch_size, 13, text_dim, generator=generator),
+    }
+
+
+def _compare(native, reference, inputs: dict[str, torch.Tensor], scope: str) -> None:
+    with torch.inference_mode():
+        expected = reference(**inputs)
+        with set_forward_context(current_timestep=0, attn_metadata=None):
+            actual = native(**inputs)
+    assert actual.shape == expected.shape
+    diff = (actual - expected).abs()
+    print(f"{scope}: shape={tuple(actual.shape)} ref_abs_mean={expected.abs().mean().item():.3e} "
+          f"diff_max={diff.max().item():.3e} diff_mean={diff.mean().item():.3e}")
+    assert torch.isfinite(expected).all() and expected.abs().mean() > 1e-3
+    assert_close(actual, expected, atol=2e-5, rtol=2e-5)
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# Tests
+# --------------------------------------------------------------------------------------------------------------------
+
+
+def test_pan2_transformer_registry_resolves_native_class():
+    from fastvideo.models.dits.pan2 import PAN2Transformer3DModel
+    from fastvideo.models.registry import ModelRegistry
+
+    model_cls, architecture = ModelRegistry.resolve_model_cls("PAN2Transformer3DModel")
+    assert model_cls is PAN2Transformer3DModel
+    assert architecture == "PAN2Transformer3DModel"
+
+
+@pytest.mark.parametrize(("module_name", "expected"), [
+    ("transformer_blocks.0", True),
+    ("transformer_blocks.43", True),
+    ("context_refiner.3", True),
+    ("transformer_blocks.0.attn.to_qkv", False),
+    ("transformer_blocks.0.ff.fc_in", False),
+])
+def test_pan2_fsdp_shard_conditions_match_only_blocks(module_name: str, expected: bool):
+    from fastvideo.configs.models.dits.pan2 import PAN2ArchConfig
+
+    conditions = PAN2ArchConfig()._fsdp_shard_conditions
+    assert any(condition(module_name, None) for condition in conditions) is expected
+
+
+def test_pan2_default_arch_config_is_the_29b_model():
+    from fastvideo.configs.models.dits.pan2 import PAN2ArchConfig
+
+    config = PAN2ArchConfig()
+    assert config.hidden_size == 4096
+    assert config.num_channels_latents == 48
+    assert config.in_channels == 2 * 48 + 1
+    assert sum(config.rope_axes_dim) == config.attention_head_dim
+
+
+def test_pan2_tiny_checkpoint_config_fields_are_arch_fields():
+    model_dir = tiny_model_dir()
+    if model_dir is None or not (model_dir / "transformer" / "config.json").is_file():
+        pytest.skip(f"{TINY_REPO_ID} is not available locally; set PAN2_TINY_MODEL_DIR")
+    from fastvideo.configs.models.dits.pan2 import PAN2ArchConfig
+
+    checkpoint_config = json.loads((model_dir / "transformer" / "config.json").read_text())
+    assert checkpoint_config.pop("_class_name") == "PAN2Transformer3DModel"
+    checkpoint_config.pop("_diffusers_version", None)
+    assert set(checkpoint_config) <= {field.name for field in fields(PAN2ArchConfig)}
+
+
+def test_pan2_tiny_checkpoint_strict_load_and_forward_parity(single_rank_dist):
+    model_dir = tiny_model_dir()
+    if model_dir is None or not (model_dir / "transformer").is_dir():
+        pytest.skip(f"{TINY_REPO_ID} is not available locally; set PAN2_TINY_MODEL_DIR")
+    from safetensors.torch import load_file
+
+    transformer_dir = model_dir / "transformer"
+    arch_kwargs = json.loads((transformer_dir / "config.json").read_text())
+    state_dict = load_file(str(transformer_dir / "diffusion_pytorch_model.safetensors"))
+
+    reference = ReferencePAN2Transformer(**arch_kwargs).float().eval()
+    reference.load_state_dict(state_dict, strict=True)
+    native = _build_native(arch_kwargs)
+    _load_into_native(native, state_dict.items())
+
+    inputs = _make_inputs(arch_kwargs["in_channels"], arch_kwargs["text_embed_dim"])
+    _compare(native, reference, inputs, "tiny checkpoint")
+
+
+@pytest.mark.parametrize("batch_size", [1, 2])
+def test_pan2_random_weights_forward_parity(single_rank_dist, batch_size: int):
+    kwargs = _small_kwargs()
+    reference = ReferencePAN2Transformer(**kwargs).float().eval()
+    # Zero-initialized tables (text modulation, type embedding) must be exercised too.
+    _randomize(reference, seed=1234)
+    native = _build_native(kwargs)
+    _load_into_native(native, reference.state_dict().items())
+
+    inputs = _make_inputs(kwargs["in_channels"], kwargs["text_embed_dim"], batch_size=batch_size, seed=batch_size)
+    _compare(native, reference, inputs, f"random weights, batch {batch_size}")
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# Sequence parallel: SP=2 must match one rank (video tokens are padded to a multiple of the SP size)
+# --------------------------------------------------------------------------------------------------------------------
+
+
+def _run_sp_worker(mode: str, output_path: Path) -> None:
+    from fastvideo.configs.models.dits.pan2 import PAN2VideoConfig
+    from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
+    from fastvideo.models.dits.pan2 import PAN2Transformer3DModel
+    import torch.distributed as dist
+
+    sp_size = 1 if mode == "single" else SP_WORLD_SIZE
+    if torch.cuda.is_available():
+        device = torch.device(f"cuda:{int(os.environ.get('LOCAL_RANK', '0'))}")
+        torch.cuda.set_device(device)
+        torch.backends.cuda.matmul.allow_tf32 = False
+        torch.backends.cudnn.allow_tf32 = False
+    else:
+        # GPU-less host: gloo ranks on CPU. fastvideo resolves the local device to mps there; swapping the code
+        # object reaches every module that imported get_local_torch_device by name.
+        import fastvideo.distributed.parallel_state as parallel_state
+
+        def _cpu_device() -> torch.device:
+            return torch.device("cpu")
+
+        parallel_state.get_local_torch_device.__code__ = _cpu_device.__code__
+        device = torch.device("cpu")
+    try:
+        maybe_init_distributed_environment_and_model_parallel(1, sp_size)
+        kwargs = _small_kwargs()
+        reference = ReferencePAN2Transformer(**kwargs).float()
+        _randomize(reference, seed=1234)
+        config = PAN2VideoConfig()
+        config.update_model_arch(kwargs)
+        native = PAN2Transformer3DModel(config=config, hf_config=kwargs).float().eval()
+        _load_into_native(native, reference.state_dict().items())
+        native = native.to(device)
+        # 5x2x6 latents patchify to 15 video tokens, which SP=2 pads to 16.
+        inputs = _make_inputs(kwargs["in_channels"], kwargs["text_embed_dim"])
+        inputs["hidden_states"] = torch.randn(2, kwargs["in_channels"], 5, 2, 6, generator=torch.Generator().manual_seed(5))
+        inputs = {name: value.to(device) for name, value in inputs.items()}
+        with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None):
+            output = native(**inputs)
+        if int(os.environ.get("RANK", "0")) == 0:
+            torch.save({"output": output.detach().cpu()}, output_path)
+        dist.barrier()
+    finally:
+        cleanup_dist_env_and_memory()
+
+
+def _run_torchrun(mode: str, nproc_per_node: int, output_path: Path) -> None:
+    command = [
+        sys.executable, "-m", "torch.distributed.run", "--nnodes", "1", "--nproc_per_node",
+        str(nproc_per_node), "--master_port",
+        str(_free_port()),
+        str(Path(__file__).resolve()), "--pan2-sp-worker", "--mode", mode, "--output",
+        str(output_path)
+    ]
+    environment = os.environ.copy()
+    environment["PYTHONPATH"] = os.pathsep.join(filter(None, [str(REPO_ROOT), environment.get("PYTHONPATH")]))
+    environment["DISABLE_SP"] = "0"
+    environment["FASTVIDEO_ATTENTION_BACKEND"] = "TORCH_SDPA"
+    process = subprocess.run(command, capture_output=True, text=True, env=environment)
+    if process.returncode != 0:
+        raise RuntimeError(f"{mode} worker failed with code {process.returncode}\n"
+                           f"STDOUT:\n{process.stdout}\nSTDERR:\n{process.stderr}")
+
+
+def test_pan2_sp2_matches_single_rank(tmp_path: Path):
+    """Ulysses SP=2 on an odd video token count (padded to the SP size) must match one rank. Runs on CPU with gloo
+    when no GPU is present."""
+    if torch.cuda.is_available() and torch.cuda.device_count() < SP_WORLD_SIZE:
+        pytest.skip(f"PAN2 SP parity requires {SP_WORLD_SIZE} CUDA devices.")
+    single_path, sp_path = tmp_path / "single.pt", tmp_path / "sp2.pt"
+    _run_torchrun("single", 1, single_path)
+    _run_torchrun("sp", SP_WORLD_SIZE, sp_path)
+    single = torch.load(single_path, map_location="cpu", weights_only=True)["output"]
+    sp = torch.load(sp_path, map_location="cpu", weights_only=True)["output"]
+    diff = (sp - single).abs()
+    print(f"sp2-vs-single diff_max={diff.max().item():.3e}")
+    assert_close(sp, single, atol=1e-5, rtol=1e-5)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--pan2-sp-worker", action="store_true")
+    parser.add_argument("--mode", choices=["single", "sp"], required=True)
+    parser.add_argument("--output", type=str, required=True)
+    args = parser.parse_args()
+    if not args.pan2_sp_worker:
+        raise SystemExit("This module is intended to be run by pytest.")
+    _run_sp_worker(args.mode, Path(args.output))

tests/local_tests/vaes/test_pan2_vae_parity.py

@@ -0,0 +1,387 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2 reuse of FastVideo's Wan2.2 (48-channel) VAE.
+
+Coverage scope: both. PAN2 ships a Diffusers ``AutoencoderKLWan`` in the
+Wan2.2 layout (z_dim 48, residual blocks, 2x2 patchify, 16x spatial and 4x
+temporal compression). The FastVideo side loads it through ``VAELoader``, the
+path a pipeline uses, and is compared on CPU fp32 against
+``diffusers.AutoencoderKLWan.from_pretrained``:
+
+* posterior mode/logvar for a single first frame (PAN2's I2V conditioning),
+* posterior mode/logvar for a 9-frame clip (4x temporal downsampling + cache),
+* decode of 3 latent frames (first-chunk handling + temporal upsampling cache).
+
+The normalization tests use a copy of the checkpoint with non-trivial
+``latents_mean``/``latents_std`` (the tiny checkpoint ships 0/1, which would
+hide a missing or doubled normalization) and pin where it happens: the raw
+``encode``/``decode`` never normalize, ``PAN2ConditioningStage`` applies
+PAN2's ``(z - mean) / std`` once, and ``DecodingStage`` applies
+``z * std + mean`` once.
+
+The default asset is the random-weight tiny PAN2 checkpoint, so this runs on a
+CPU-only host::
+
+    pytest tests/local_tests/vaes/test_pan2_vae_parity.py -v -s
+
+Set ``PAN2_MODEL_PATH`` to a local PAN2 root containing ``vae/`` to use other
+weights.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import load_file
+from torch.testing import assert_close
+
+try:
+    from diffusers import AutoencoderKLWan as DiffusersAutoencoderKLWan
+except ImportError:  # pragma: no cover - environment dependent
+    pytest.skip("diffusers with AutoencoderKLWan is required", allow_module_level=True)
+
+from fastvideo.models.loader.component_loader import VAELoader
+from fastvideo.models.wan.vae import AutoencoderKLWan as FastVideoAutoencoderKLWan
+from fastvideo.models.wan.vae_config import WanVAEConfig
+
+PARITY_SCOPE = "both"
+TINY_PAN2_REPO = "wuqing157/tiny-pan2-modular-pipe"
+TINY_PAN2_REVISION = "3ff3130b4527f61aaab7bde01465c9884dc84cdd"
+MODEL_PATH_ENV = "PAN2_MODEL_PATH"
+
+# H and W must be multiples of 32 (16x spatial compression times the 2x2
+# patchify inside the VAE).
+HEIGHT, WIDTH = 64, 96
+VIDEO_FRAMES = 9
+LATENT_FRAMES = (VIDEO_FRAMES - 1) // 4 + 1
+ATOL = RTOL = 1e-5
+
+# Structural fields of Wan-AI/Wan2.2-TI2V-5B-Diffusers vae/config.json. The
+# statistics do not affect parameter shapes, so placeholders keep this inline.
+WAN22_TI2V_VAE_CONFIG = {
+    "attn_scales": [],
+    "base_dim": 160,
+    "clip_output": False,
+    "decoder_base_dim": 256,
+    "dim_mult": [1, 2, 4, 4],
+    "dropout": 0.0,
+    "in_channels": 12,
+    "is_residual": True,
+    "latents_mean": [0.0] * 48,
+    "latents_std": [1.0] * 48,
+    "num_res_blocks": 2,
+    "out_channels": 12,
+    "patch_size": 2,
+    "scale_factor_spatial": 16,
+    "scale_factor_temporal": 4,
+    "temperal_downsample": [False, True, True],
+    "z_dim": 48,
+}
+
+
+def _pan2_root(subfolder: str) -> Path:
+    """Resolve a PAN2 root that contains ``subfolder``; skip when unavailable."""
+    override = os.getenv(MODEL_PATH_ENV)
+    if override:
+        root = Path(override)
+    else:
+        try:
+            from huggingface_hub import snapshot_download
+        except ImportError:
+            pytest.skip("huggingface_hub is required to resolve the tiny PAN2 checkpoint")
+        root = None
+        last_error: Exception | None = None
+        for local_files_only in (True, False):
+            try:
+                root = Path(
+                    snapshot_download(
+                        TINY_PAN2_REPO,
+                        revision=TINY_PAN2_REVISION,
+                        allow_patterns=[f"{subfolder}/*"],
+                        local_files_only=local_files_only,
+                    ))
+                break
+            except Exception as exc:  # noqa: BLE001 - offline/missing cache must skip.
+                last_error = exc
+        if root is None:
+            pytest.skip(f"Tiny PAN2 checkpoint unavailable: {last_error}")
+    if not (root / subfolder).is_dir():
+        pytest.skip(f"PAN2 {subfolder}/ missing under {root}")
+    return root
+
+
+@pytest.fixture(scope="module")
+def vae_dir() -> Path:
+    path = _pan2_root("vae") / "vae"
+    for name in ("config.json", "diffusion_pytorch_model.safetensors"):
+        if not (path / name).is_file():
+            pytest.skip(f"PAN2 VAE asset missing: {path / name}")
+    return path
+
+
+@pytest.fixture
+def cpu_runtime(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Pin FastVideo's device helper to CPU; it reports ``mps`` on CPU-only Linux."""
+    for module in (
+            "fastvideo.models.loader.component_loader",
+            "fastvideo.pipelines.stages.decoding",
+            "fastvideo.pipelines.basic.pan2.stages",
+    ):
+        monkeypatch.setattr(f"{module}.get_local_torch_device", lambda: torch.device("cpu"))
+
+
+def _pan2_vae_config() -> WanVAEConfig:
+    """The VAE config a PAN2 pipeline config would carry (arch filled by the loader)."""
+    config = WanVAEConfig()
+    config.load_encoder = True
+    config.load_decoder = True
+    return config
+
+
+def _load_fastvideo(vae_dir: Path) -> FastVideoAutoencoderKLWan:
+    """Load through the production VAELoader, exactly as a pipeline does."""
+    fastvideo_args = SimpleNamespace(
+        model_paths={},
+        pipeline_config=SimpleNamespace(vae_config=_pan2_vae_config(), vae_precision="fp32"),
+        vae_cpu_offload=False,
+    )
+    vae = VAELoader().load(str(vae_dir), fastvideo_args)
+    assert isinstance(vae, FastVideoAutoencoderKLWan)
+    assert fastvideo_args.model_paths["vae"] == str(vae_dir)
+    assert all(param.dtype == torch.float32 and param.device.type == "cpu" for param in vae.parameters())
+    return vae
+
+
+def _load_diffusers(vae_dir: Path) -> DiffusersAutoencoderKLWan:
+    model, info = DiffusersAutoencoderKLWan.from_pretrained(
+        str(vae_dir),
+        torch_dtype=torch.float32,
+        local_files_only=True,
+        output_loading_info=True,
+    )
+    problems = {key: value for key, value in info.items() if value}
+    assert not problems, f"Diffusers reference did not load strictly: {problems}"
+    return model.eval()
+
+
+def _pixels(num_frames: int, seed: int) -> torch.Tensor:
+    generator = torch.Generator().manual_seed(seed)
+    return torch.rand((1, 3, num_frames, HEIGHT, WIDTH), generator=generator) * 2.0 - 1.0
+
+
+def _latents(seed: int) -> torch.Tensor:
+    generator = torch.Generator().manual_seed(seed)
+    return torch.randn((1, 48, LATENT_FRAMES, HEIGHT // 16, WIDTH // 16), generator=generator)
+
+
+def _report(name: str, actual: torch.Tensor, expected: torch.Tensor, clamp_bound: float | None = None) -> None:
+    assert actual.shape == expected.shape, f"{name}: {tuple(actual.shape)} != {tuple(expected.shape)}"
+    diff = (actual.float() - expected.float()).abs()
+    message = (f"[pan2-vae] {name}: shape={tuple(actual.shape)} max_abs={diff.max().item():.3e} "
+               f"mean_abs={diff.mean().item():.3e} ref_absmax={expected.abs().max().item():.3e}")
+    if clamp_bound is not None:
+        # Clamped pixels compare trivially; report how many there are.
+        message += f" clamped={(expected.abs() >= clamp_bound).float().mean().item():.2%}"
+    print(message)
+
+
+def _nontrivial_stats() -> tuple[list[float], list[float]]:
+    mean = [round(0.05 * (index - 24), 4) for index in range(48)]
+    std = [round(0.4 + 0.03 * index, 4) for index in range(48)]
+    return mean, std
+
+
+@pytest.fixture
+def stats_vae_dir(vae_dir: Path, tmp_path: Path) -> Path:
+    """Checkpoint copy with non-trivial latent statistics (weights symlinked)."""
+    config = json.loads((vae_dir / "config.json").read_text())
+    config["latents_mean"], config["latents_std"] = _nontrivial_stats()
+    target = tmp_path / "vae"
+    target.mkdir()
+    (target / "config.json").write_text(json.dumps(config))
+    (target / "diffusion_pytorch_model.safetensors").symlink_to(
+        (vae_dir / "diffusion_pytorch_model.safetensors").resolve())
+    return target
+
+
+def _diffusers_stats(model: DiffusersAutoencoderKLWan, like: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    mean = torch.tensor(model.config.latents_mean).view(1, -1, 1, 1, 1).to(like)
+    std = torch.tensor(model.config.latents_std).view(1, -1, 1, 1, 1).to(like)
+    return mean, std
+
+
+def test_pan2_vae_production_loader_matches_checkpoint(vae_dir: Path, cpu_runtime: None) -> None:
+    """The loader fills the tiny Wan2.2 arch from config.json and covers every weight.
+
+    ``VAELoader`` loads ``AutoencoderKLWan`` with ``strict=False``, so key
+    coverage is checked here instead of trusting the loader.
+    """
+    raw = json.loads((vae_dir / "config.json").read_text())
+    vae = _load_fastvideo(vae_dir)
+    arch = vae.config.arch_config
+    for key in ("base_dim", "decoder_base_dim", "z_dim", "num_res_blocks", "in_channels", "out_channels",
+                "is_residual", "patch_size", "scale_factor_spatial", "scale_factor_temporal", "dropout"):
+        assert getattr(arch, key) == raw[key], key
+    for key in ("dim_mult", "attn_scales", "temperal_downsample", "latents_mean", "latents_std"):
+        assert list(getattr(arch, key)) == list(raw[key]), key
+    assert vae.config.use_feature_cache and not vae.config.use_tiling
+    assert arch.spatial_compression_ratio == 16 and arch.temporal_compression_ratio == 4
+    assert tuple(vae.scaling_factor.shape) == tuple(vae.shift_factor.shape) == (1, 48, 1, 1, 1)
+    assert_close(vae.scaling_factor.flatten(), 1.0 / torch.tensor(raw["latents_std"]))
+    assert_close(vae.shift_factor.flatten(), torch.tensor(raw["latents_mean"]))
+
+    checkpoint = load_file(str(vae_dir / "diffusion_pytorch_model.safetensors"))
+    state = vae.state_dict()
+    assert set(state) == set(checkpoint), (sorted(set(checkpoint) - set(state)), sorted(set(state) - set(checkpoint)))
+    for key, tensor in checkpoint.items():
+        assert torch.equal(state[key], tensor.float()), key
+
+
+@pytest.mark.parametrize("variant", ["wan22_ti2v_5b", "tiny_pan2"])
+def test_pan2_vae_config_shapes_match_diffusers(variant: str, request: pytest.FixtureRequest) -> None:
+    """``WanVAEConfig`` expresses both the real Wan2.2 and the tiny PAN2 VAE configs.
+
+    Both models are built on the meta device, so the full-size config costs no
+    memory; parameter names and shapes must match one-to-one.
+    """
+    if variant == "tiny_pan2":
+        raw = json.loads((request.getfixturevalue("vae_dir") / "config.json").read_text())
+        config = {key: value for key, value in raw.items() if not key.startswith("_")}
+    else:
+        config = WAN22_TI2V_VAE_CONFIG
+    fastvideo_config = _pan2_vae_config()
+    fastvideo_config.update_model_arch(config)
+    with torch.device("meta"):
+        fastvideo = FastVideoAutoencoderKLWan(fastvideo_config)
+        reference = DiffusersAutoencoderKLWan.from_config(config)
+    fastvideo_shapes = {key: tuple(value.shape) for key, value in fastvideo.state_dict().items()}
+    reference_shapes = {key: tuple(value.shape) for key, value in reference.state_dict().items()}
+    assert fastvideo_shapes == reference_shapes
+
+
+def test_pan2_vae_wan22_topology_random_weight_parity() -> None:
+    """Numerically cover the Wan2.2 paths the tiny config skips.
+
+    The tiny checkpoint uses ``dim_mult`` [1, 1, 1, 1] and no separate decoder
+    width, so channel-changing shortcuts and resamplers never run. This narrow
+    random-weight model keeps the real topology (``dim_mult`` [1, 2, 4, 4], two
+    res blocks, ``decoder_base_dim`` > ``base_dim``).
+    """
+    config = dict(WAN22_TI2V_VAE_CONFIG, base_dim=8, decoder_base_dim=16)
+    torch.manual_seed(0)
+    reference = DiffusersAutoencoderKLWan.from_config(config).eval()
+    fastvideo_config = _pan2_vae_config()
+    fastvideo_config.update_model_arch(config)
+    fastvideo = FastVideoAutoencoderKLWan(fastvideo_config).eval()
+    fastvideo.load_state_dict(reference.state_dict(), strict=True)
+
+    video = _pixels(VIDEO_FRAMES, seed=11)
+    latents = _latents(seed=12)
+    with torch.inference_mode():
+        expected_mode = reference.encode(video).latent_dist.mode()
+        actual_mode = fastvideo.encode(video).mode()
+        expected_video = reference.decode(latents).sample
+        actual_video = fastvideo.decode(latents)
+    _report("wan22-topology encode mode", actual_mode, expected_mode)
+    _report("wan22-topology decode", actual_video, expected_video, clamp_bound=1.0)
+    assert_close(actual_mode, expected_mode, atol=ATOL, rtol=RTOL)
+    assert_close(actual_video, expected_video, atol=ATOL, rtol=RTOL)
+
+
+@pytest.mark.parametrize("num_frames", [1, VIDEO_FRAMES], ids=["first_frame", "video_9f"])
+def test_pan2_vae_encode_matches_diffusers(num_frames: int, vae_dir: Path, cpu_runtime: None) -> None:
+    """Posterior mode (PAN2 uses ``latent_dist.mode()``) and logvar match Diffusers.
+
+    FastVideo's ``encode`` returns the ``DiagonalGaussianDistribution`` itself,
+    so the port calls ``vae.encode(x).mode()`` instead of ``.latent_dist.mode()``.
+    """
+    reference = _load_diffusers(vae_dir)
+    fastvideo = _load_fastvideo(vae_dir)
+    pixels = _pixels(num_frames, seed=100 + num_frames)
+    with torch.inference_mode():
+        expected = reference.encode(pixels).latent_dist
+        actual = fastvideo.encode(pixels)
+    assert tuple(actual.mode().shape) == (1, 48, (num_frames - 1) // 4 + 1, HEIGHT // 16, WIDTH // 16)
+    _report(f"encode mode ({num_frames}f)", actual.mode(), expected.mode())
+    _report(f"encode logvar ({num_frames}f)", actual.logvar, expected.logvar)
+    assert_close(actual.mode(), expected.mode(), atol=ATOL, rtol=RTOL)
+    assert_close(actual.logvar, expected.logvar, atol=ATOL, rtol=RTOL)
+
+
+def test_pan2_vae_decode_matches_diffusers(vae_dir: Path, cpu_runtime: None) -> None:
+    """Decode three latent frames so first-chunk and temporal cache paths run."""
+    reference = _load_diffusers(vae_dir)
+    fastvideo = _load_fastvideo(vae_dir)
+    latents = _latents(seed=7)
+    with torch.inference_mode():
+        expected = reference.decode(latents).sample
+        actual = fastvideo.decode(latents)
+    assert tuple(actual.shape) == (1, 3, VIDEO_FRAMES, HEIGHT, WIDTH)
+    _report("decode", actual, expected, clamp_bound=1.0)
+    assert_close(actual, expected, atol=ATOL, rtol=RTOL)
+
+
+def test_pan2_conditioning_stage_matches_reference(stats_vae_dir: Path, cpu_runtime: None) -> None:
+    """``PAN2ConditioningStage`` builds the reference I2V condition: cover-resize and center crop the first frame,
+    encode it, apply ``(z - mean) / std`` once, and set the mask channel on the first latent frame only.
+
+    The raw ``encode`` is unnormalized, so the stage is the only place the statistics apply.
+    """
+    import PIL.Image
+
+    from fastvideo.pipelines.basic.pan2.stages import PAN2ConditioningStage
+    from tests.local_tests.pan2.reference import reference_condition_latents
+
+    reference = _load_diffusers(stats_vae_dir)
+    fastvideo = _load_fastvideo(stats_vae_dir)
+    # Not the output aspect ratio, so the resize and crop are exercised.
+    image = PIL.Image.radial_gradient("L").resize((120, 64)).convert("RGB")
+    latents = torch.zeros(1, 48, LATENT_FRAMES, HEIGHT // 16, WIDTH // 16)
+    with torch.inference_mode():
+        pixels = PAN2ConditioningStage.preprocess_image(image, HEIGHT, WIDTH)
+        raw = reference.encode(pixels).latent_dist.mode()
+        assert_close(fastvideo.encode(pixels).mode(), raw, atol=ATOL, rtol=RTOL)
+        expected = reference_condition_latents(reference, latents, image, HEIGHT, WIDTH)
+
+        batch = SimpleNamespace(latents=latents, pil_image=image, height=HEIGHT, width=WIDTH, image_latent=None)
+        stage_args = SimpleNamespace(pipeline_config=SimpleNamespace(vae_precision="fp32"), vae_cpu_offload=False)
+        actual = PAN2ConditioningStage(fastvideo).forward(batch, stage_args).image_latent
+    _report("PAN2 I2V condition", actual, expected)
+    assert actual.shape == (1, 49, LATENT_FRAMES, HEIGHT // 16, WIDTH // 16)
+    assert_close(actual, expected, atol=ATOL, rtol=RTOL)
+    assert torch.equal(actual[:, 48, 0], torch.ones_like(actual[:, 48, 0]))
+    assert not actual[:, :, 1:].any()
+    assert not torch.allclose(raw, expected[:, :48, :1], atol=1e-2), "statistics too close to identity"
+
+
+def test_pan2_decoding_stage_denormalizes_once(stats_vae_dir: Path, cpu_runtime: None) -> None:
+    """``DecodingStage`` performs PAN2's ``z * std + mean`` itself.
+
+    The PAN2 denoising loop must hand normalized latents to ``DecodingStage``.
+    Denormalizing beforehand (as the upstream pipeline does inline) would apply
+    the statistics twice, which the negative check below detects.
+    """
+    from fastvideo.pipelines.stages.decoding import DecodingStage
+
+    reference = _load_diffusers(stats_vae_dir)
+    fastvideo = _load_fastvideo(stats_vae_dir)
+    latents = _latents(seed=31)
+    # vae_tiling=True is the PipelineConfig default; the feature-cache decode ignores tiling, so it must not matter.
+    decode_args = SimpleNamespace(
+        pipeline_config=SimpleNamespace(vae_precision="fp32", vae_decode_precision="fp32", vae_tiling=True),
+        disable_autocast=False,
+    )
+    stage = DecodingStage(fastvideo)
+    with torch.inference_mode():
+        mean, std = _diffusers_stats(reference, latents)
+        expected = reference.decode(latents * std + mean).sample
+        expected = (expected * 0.5 + 0.5).clamp(0, 1)
+        actual = stage.decode(latents, decode_args)
+        doubled = stage.decode(latents * std + mean, decode_args)
+    _report("DecodingStage output", actual, expected)
+    assert_close(actual, expected, atol=ATOL, rtol=RTOL)
+    assert (doubled - expected).abs().max().item() > 1e-3, "double denormalization went undetected"