[ci]: capture performance worker logs on failure

Co-authored-by: Mrinaal Dogra <mdogra@ucsd.edu>
Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>

.buildkite/pipeline.yml

@@ -1,6 +1,8 @@
 env:
   IMAGE_VERSION: "py3.12-latest"
   BUILDKITE_CLEAN_CHECKOUT: true
+  # Buildkite only launches Modal; remote jobs initialize their own submodules.
+  BUILDKITE_GIT_SUBMODULES: false
 
 notify:
   - github_commit_status:

.gitignore

@@ -6,6 +6,7 @@ results/
 wandb/
 *.ipynb
 *.jpg
+!examples/dataset/lingbotworld2/image.jpg
 *.safetensors
 *.mp4
 *.png
@@ -34,6 +35,7 @@ env
 *.log
 weights/
 logs/
+/Z-Image/
 official_weights/
 converted_weights/
 

docs/design/inference_schema_parity_inventory.yaml

@@ -333,15 +333,25 @@ surfaces:
       use_distill:
         sources: [fastvideo.configs.pipelines.longcat.LongCatT2V480PConfig, fastvideo.configs.pipelines.longcat.LongCatT2V704PConfig]
       scheduler_arch:
-        sources: [fastvideo.configs.pipelines.sd35.SD35Config]
+        sources:
+          - fastvideo.configs.pipelines.sd35.SD35Config
+          - fastvideo.configs.pipelines.zimage.ZImagePipelineConfig
       text_encoder_archs:
-        sources: [fastvideo.configs.pipelines.sd35.SD35Config]
+        sources:
+          - fastvideo.configs.pipelines.sd35.SD35Config
+          - fastvideo.configs.pipelines.zimage.ZImagePipelineConfig
       tokenizer_archs:
-        sources: [fastvideo.configs.pipelines.sd35.SD35Config]
+        sources:
+          - fastvideo.configs.pipelines.sd35.SD35Config
+          - fastvideo.configs.pipelines.zimage.ZImagePipelineConfig
       transformer_arch:
-        sources: [fastvideo.configs.pipelines.sd35.SD35Config]
+        sources:
+          - fastvideo.configs.pipelines.sd35.SD35Config
+          - fastvideo.configs.pipelines.zimage.ZImagePipelineConfig
       vae_arch:
-        sources: [fastvideo.configs.pipelines.sd35.SD35Config]
+        sources:
+          - fastvideo.configs.pipelines.sd35.SD35Config
+          - fastvideo.configs.pipelines.zimage.ZImagePipelineConfig
       expand_timesteps:
         sources:
           - fastvideo.configs.pipelines.wan.FastWan2_2_TI2V_5B_Config
@@ -421,6 +431,8 @@ surfaces:
       frame_receptive_field: "MagiHuman internal data-proxy receptive-field setting."
       image_conditioning: "MagiHuman preset variant marker for reference-image conditioning."
       ref_audio_offset: "MagiHuman internal data-proxy audio alignment offset."
+      scheduler_sigma_min: "Z-Image scheduler parity invariant; not part of the public typed inference API."
+      scheduler_use_reference_discrete_timesteps: "Z-Image scheduler parity invariant; not part of the public typed inference API."
       sr_local_attn_layers: "MagiHuman SR internal sparse-attention layer selection."
       text_offset: "MagiHuman internal data-proxy text alignment offset."
       vae_stride: "MagiHuman internal VAE/data-proxy stride setting."
@@ -438,6 +450,7 @@ surfaces:
       grid_sizes: request.inputs.grid_sizes
       pose: request.inputs.pose
       c2ws_plucker_emb: request.inputs.c2ws_plucker_emb
+      action_path: request.inputs.action_path
       refine_from: request.inputs.refine_from
       stage1_video: request.inputs.stage1_video
       prompt: request.prompt
@@ -447,6 +460,7 @@ surfaces:
       output_video_name: request.output.output_video_name
       num_videos_per_prompt: request.sampling.num_videos_per_prompt
       seed: request.sampling.seed
+      max_sequence_length: request.sampling.max_sequence_length
       num_frames: request.sampling.num_frames
       height: request.sampling.height
       width: request.sampling.width
@@ -457,6 +471,8 @@ surfaces:
       num_inference_steps_sr: request.sampling.num_inference_steps_sr
       guidance_scale: request.sampling.guidance_scale
       guidance_scale_2: request.sampling.guidance_scale_2
+      cfg_normalization: request.sampling.cfg_normalization
+      cfg_truncation: request.sampling.cfg_truncation
       guidance_rescale: request.sampling.guidance_rescale
       use_embedded_guidance: request.sampling.use_embedded_guidance
       true_cfg_scale: request.sampling.true_cfg_scale
@@ -508,7 +524,6 @@ surfaces:
     internal_only:
       data_type: "Derived from the request shape and not a public input."
       latents: "Pre-generated diffusion latents supplied by parity/debug harnesses; not a public input."
-      max_sequence_length: "Model-specific text-encoder sequence cap; not part of the public typed inference API."
 
   sampling_param_extensions: {}
 

examples/dataset/lingbotworld2/README.md

@@ -0,0 +1,17 @@
+# LingBot World 2 Example Dataset
+
+These files were copied unchanged from the LingBot World 2 repository for the
+FastVideo causal-fast inference example.
+
+- Repository: `https://github.com/Robbyant/lingbot-world-v2.git`
+- Source commit: `94f43115de8d4a4f9f282126528c300a0b232c5f`
+- Source directory: `examples/03`
+
+## Files
+
+- `image.jpg`: source image for image-to-video generation. SHA-256:
+  `6ee3dacfef32cfef504dd698adb8a660cf15f686535c52fed4903fef27c0edd0`
+- `poses.npy`: camera-to-world trajectory matrices. SHA-256:
+  `bd0a23a696e184b0b43e7767eb432bfe644690560fe327fa96961affc941c404`
+- `intrinsics.npy`: camera intrinsic parameters. SHA-256:
+  `821fca6cf957ae8fbb1181307f02479efb1705e04c9e05734cd02fb43462e082`

examples/inference/basic/basic_lingbotworld2_causal_fast.py

@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Run LingBot World 2 14B causal-fast I2V generation with FastVideo."""
+
+import os
+from pathlib import Path
+
+from fastvideo import VideoGenerator
+
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+DATASET_DIR = REPO_ROOT / "examples" / "dataset" / "lingbotworld2"
+OUTPUT_PATH = REPO_ROOT / "outputs" / "lingbotworld2_causal_fast.mp4"
+
+
+def main() -> None:
+    """Load the native FastVideo LingBot World 2 causal-fast pipeline and generate one video."""
+    generator = VideoGenerator.from_pretrained(
+        os.environ["LINGBOTWORLD2_MODEL_PATH"],
+        num_gpus=8,
+        sp_size=8,
+        hsdp_shard_dim=8,
+        use_fsdp_inference=True,
+        dit_layerwise_offload=False,
+        dit_cpu_offload=False,
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=False,
+        pin_cpu_memory=True,
+        override_pipeline_cls_name="LingBotWorld2CausalFastPipeline",
+    )
+
+    try:
+        generator.generate_video(
+            "A serene lakeside scene with a lone tree standing in calm water, surrounded by distant snow-capped mountains under a bright blue sky with drifting white clouds; gentle ripples reflect the tree and sky, creating a tranquil, meditative atmosphere.",
+            image_path=str(DATASET_DIR / "image.jpg"),
+            action_path=str(DATASET_DIR),
+            output_path=str(OUTPUT_PATH),
+            save_video=True,
+            height=480,
+            width=832,
+            num_frames=65,
+            num_inference_steps=4,
+            guidance_scale=1.0,
+            negative_prompt="",
+            fps=16,
+            seed=42,
+        )
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_zimage.py

@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Run Z-Image-Turbo text-to-image generation through FastVideo.
+
+User story:
+    "I want the official Z-Image-Turbo defaults and a deterministic PNG from
+    a local or Hugging Face checkpoint."
+"""
+
+import argparse
+from pathlib import Path
+
+from fastvideo import VideoGenerator
+from fastvideo.api import (
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    OutputConfig,
+    ParallelismConfig,
+    PipelineSelection,
+    SamplingConfig,
+)
+
+
+DEFAULT_PROMPT = (
+    "Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. "
+    "Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. "
+    "Neon lightning-bolt lamp (⚡️), bright yellow glow, above extended left palm. Soft-lit outdoor night background, "
+    "silhouetted tiered pagoda (西安大雁塔), blurred colorful distant lights."
+)
+DEFAULT_REVISION = "f332072aa78be7aecdf3ee76d5c247082da564a6"
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run Z-Image-Turbo text-to-image generation.")
+    parser.add_argument("--model-path", default="Tongyi-MAI/Z-Image-Turbo")
+    parser.add_argument("--revision", default=DEFAULT_REVISION)
+    parser.add_argument("--output", default="outputs/zimage/zimage_turbo.png")
+    parser.add_argument("--prompt", default=DEFAULT_PROMPT)
+    parser.add_argument("--negative-prompt", default="")
+    parser.add_argument("--height", type=int, default=1024)
+    parser.add_argument("--width", type=int, default=1024)
+    parser.add_argument("--steps", type=int, default=8)
+    parser.add_argument("--guidance-scale", type=float, default=0.0)
+    parser.add_argument("--max-sequence-length", type=int, default=512)
+    parser.add_argument("--cfg-normalization", action=argparse.BooleanOptionalAction, default=False)
+    parser.add_argument("--cfg-truncation", type=float, default=1.0)
+    parser.add_argument("--seed", type=int, default=42)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+
+    generator = VideoGenerator.from_config(
+        GeneratorConfig(
+            model_path=args.model_path,
+            revision=args.revision,
+            engine=EngineConfig(
+                num_gpus=1,
+                parallelism=ParallelismConfig(tp_size=1, sp_size=1),
+                use_fsdp_inference=False,
+            ),
+            # The model registry selects the native zimage_turbo preset.
+            pipeline=PipelineSelection(workload_type="t2i"),
+        ))
+    try:
+        generator.generate(
+            GenerationRequest(
+                prompt=args.prompt,
+                negative_prompt=args.negative_prompt,
+                sampling=SamplingConfig(
+                    height=args.height,
+                    width=args.width,
+                    num_frames=1,
+                    fps=1,
+                    num_inference_steps=args.steps,
+                    guidance_scale=args.guidance_scale,
+                    max_sequence_length=args.max_sequence_length,
+                    cfg_normalization=args.cfg_normalization,
+                    cfg_truncation=args.cfg_truncation,
+                    seed=args.seed,
+                ),
+                output=OutputConfig(
+                    output_path=str(output),
+                    save_video=True,
+                    return_frames=False,
+                ),
+            ))
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

fastvideo/api/sampling_param.py

@@ -51,8 +51,9 @@ class SamplingParam:
     gt_latents: Any | None = None  # Ground truth latents [B, 16, T, H, W]
     conditioning_mask: Any | None = None  # Mask [B, 1, T, H, W]
 
-    # Camera control inputs (LingBotWorld)
+    # Camera control inputs (LingBotWorld and LingBotWorld2)
     c2ws_plucker_emb: Any | None = None  # Plucker embedding: [B, C, F_lat, H_lat, W_lat]
+    action_path: str | None = None  # Directory containing poses.npy and intrinsics.npy
 
     # Refine inputs (LongCat 480p->720p upscaling)
     # Path-based refine (load stage1 video from disk, e.g. MP4)
@@ -90,6 +91,11 @@ class SamplingParam:
     num_inference_steps_sr: int = 50
     guidance_scale: float = 1.0
     guidance_scale_2: float | None = None
+    # Z-Image CFG controls. ``cfg_normalization=True`` caps the guided
+    # prediction norm at the positive-prediction norm; ``cfg_truncation``
+    # disables CFG above the normalized-noise threshold.
+    cfg_normalization: bool = False
+    cfg_truncation: float | None = 1.0
     # Embedded guidance (FLUX): do not treat ``guidance_scale > 1`` as classic CFG.
     use_embedded_guidance: bool = False
     # Diffusers-style true CFG for FLUX when > 1 (requires negative prompt encoding).
@@ -323,6 +329,18 @@ class SamplingParam:
             default=SamplingParam.guidance_scale,
             help="Classifier-free guidance scale",
         )
+        parser.add_argument(
+            "--cfg-normalization",
+            action=StoreBoolean,
+            default=SamplingParam.cfg_normalization,
+            help="Cap Z-Image CFG prediction norm to the positive-prediction norm",
+        )
+        parser.add_argument(
+            "--cfg-truncation",
+            type=float,
+            default=SamplingParam.cfg_truncation,
+            help="Disable Z-Image CFG above this normalized-noise threshold",
+        )
         parser.add_argument(
             "--guidance-rescale",
             type=float,

fastvideo/api/schema.py

@@ -130,6 +130,7 @@ class InputConfig:
     keyboard_cond: Any | None = None
     grid_sizes: Any | None = None
     c2ws_plucker_emb: Any | None = None
+    action_path: str | None = None
     refine_from: str | None = None
     stage1_video: Any | None = None
 
@@ -138,6 +139,7 @@ class InputConfig:
 class SamplingConfig:
     num_videos_per_prompt: int = 1
     seed: int = 1024
+    max_sequence_length: int | None = None
     num_frames: int = 125
     height: int = 720
     width: int = 1280
@@ -148,6 +150,8 @@ class SamplingConfig:
     num_inference_steps_sr: int = 50
     guidance_scale: float = 1.0
     guidance_scale_2: float | None = None
+    cfg_normalization: bool = False
+    cfg_truncation: float | None = 1.0
     guidance_rescale: float = 0.0
     true_cfg_scale: float | None = None
     use_embedded_guidance: bool | None = None

fastvideo/configs/models/dits/__init__.py

@@ -12,12 +12,14 @@ from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig
 from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig
 from fastvideo.configs.models.dits.stable_audio import StableAudioConfig
 from fastvideo.configs.models.dits.wanvideo import WanVideoConfig
+from fastvideo.configs.models.dits.zimage import ZImageDiTConfig
 from fastvideo.configs.models.dits.hyworld import HYWorldConfig
 from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig
+from fastvideo.configs.models.dits.lingbotworld2 import LingBotWorld2CausalFastVideoConfig
 
 __all__ = [
     "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "DreamXWorldConfig",
     "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "FluxDiTConfig", "Flux2Config",
     "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig",
-    "StableAudioConfig", "GlmImageDiTConfig"
+    "StableAudioConfig", "GlmImageDiTConfig", "LingBotWorld2CausalFastVideoConfig", "ZImageDiTConfig"
 ]

fastvideo/configs/models/dits/lingbotworld2.py

@@ -0,0 +1,55 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+def is_blocks(n: str, m) -> bool:
+    return "blocks" in n and str.isdigit(n.split(".")[-1])
+
+
+@dataclass
+class LingBotWorld2CausalFastArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
+    param_names_mapping: dict = field(default_factory=dict)
+    reverse_param_names_mapping: dict = field(default_factory=dict)
+    lora_param_names_mapping: dict = field(default_factory=dict)
+
+    model_type: str = "i2v"
+    patch_size: tuple[int, int, int] = (1, 2, 2)
+    text_len: int = 512
+    in_dim: int = 36
+    dim: int = 5120
+    ffn_dim: int = 13824
+    freq_dim: int = 256
+    text_dim: int = 4096
+    out_dim: int = 16
+    num_heads: int = 40
+    num_layers: int = 40
+    qk_norm: bool = True
+    cross_attn_norm: bool = True
+    eps: float = 1e-6
+
+    local_attn_size: int = 18
+    sink_size: int = 6
+    chunk_size: int = 4
+    sample_shift: float = 10.0
+    num_train_timesteps: int = 1000
+    timesteps_index: tuple[int, int, int, int] = (0, 250, 500, 750)
+    max_area: int = 480 * 832
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.hidden_size = self.dim
+        self.num_attention_heads = self.num_heads
+        self.attention_head_dim = self.dim // self.num_heads
+        self.in_channels = self.in_dim
+        self.out_channels = self.out_dim
+        self.num_channels_latents = self.out_dim
+
+
+@dataclass
+class LingBotWorld2CausalFastVideoConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=LingBotWorld2CausalFastArchConfig)
+
+    prefix: str = "Wan"

fastvideo/configs/models/dits/zimage.py

@@ -0,0 +1,60 @@
+# 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_zimage_block(name: str, module) -> bool:
+    parts = name.split(".")
+    return len(parts) >= 2 and parts[-2] in {"noise_refiner", "context_refiner", "layers"} and parts[-1].isdigit()
+
+
+@dataclass
+class ZImageDiTArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_zimage_block])
+    _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (AttentionBackendEnum.TORCH_SDPA, )
+
+    all_patch_size: tuple[int, ...] = (2, )
+    all_f_patch_size: tuple[int, ...] = (1, )
+    in_channels: int = 16
+    dim: int = 3840
+    n_layers: int = 30
+    n_refiner_layers: int = 2
+    n_heads: int = 30
+    n_kv_heads: int = 30
+    norm_eps: float = 1e-5
+    qk_norm: bool = True
+    cap_feat_dim: int = 2560
+    rope_theta: float = 256.0
+    t_scale: float = 1000.0
+    axes_dims: tuple[int, ...] = (32, 48, 48)
+    axes_lens: tuple[int, ...] = (1536, 512, 512)
+
+    adaln_embed_dim: int = 256
+    frequency_embedding_size: int = 256
+    timestep_mid_size: int = 1024
+    max_period: int = 10000
+    seq_multi_of: int = 32
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        if len(self.all_patch_size) != len(self.all_f_patch_size):
+            raise ValueError("all_patch_size and all_f_patch_size must have equal length")
+        if self.dim % self.n_heads:
+            raise ValueError("dim must be divisible by n_heads")
+        if self.dim // self.n_heads != sum(self.axes_dims):
+            raise ValueError("attention head dimension must equal sum(axes_dims)")
+        if len(self.axes_dims) != len(self.axes_lens) or any(dim % 2 for dim in self.axes_dims):
+            raise ValueError("RoPE axes require matching lengths and even dimensions")
+
+        self.hidden_size = self.dim
+        self.num_attention_heads = self.n_heads
+        self.num_channels_latents = self.in_channels
+        self.out_channels = self.in_channels
+
+
+@dataclass
+class ZImageDiTConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=ZImageDiTArchConfig)
+    prefix: str = "ZImage"

fastvideo/configs/models/encoders/__init__.py

@@ -2,6 +2,7 @@ from fastvideo.configs.models.encoders.base import (BaseEncoderOutput, EncoderCo
                                                     TextEncoderConfig)
 from fastvideo.configs.models.encoders.clip import (CLIPTextConfig, CLIPVisionConfig, WAN2_1ControlCLIPVisionConfig)
 from fastvideo.configs.models.encoders.llama import LlamaConfig
+from fastvideo.configs.models.encoders.lingbotworld2_t5 import LingBotWorld2UMT5ArchConfig, LingBotWorld2UMT5Config
 from fastvideo.configs.models.encoders.t5 import T5Config, T5LargeConfig
 from fastvideo.configs.models.encoders.qwen2_5 import Qwen2_5_VLConfig
 from fastvideo.configs.models.encoders.siglip import SiglipVisionConfig
@@ -17,5 +18,6 @@ __all__ = [
     "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig", "BaseEncoderOutput", "CLIPTextConfig",
     "CLIPVisionConfig", "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig", "Qwen2_5_VLConfig",
     "Reason1ArchConfig", "Reason1Config", "LTX2GemmaConfig", "SiglipVisionConfig", "StableAudioConditionerArchConfig",
-    "StableAudioConditionerConfig", "T5GemmaEncoderConfig", "Qwen3TextConfig", "Mistral3TextConfig"
+    "StableAudioConditionerConfig", "T5GemmaEncoderConfig", "Qwen3TextConfig", "Mistral3TextConfig",
+    "LingBotWorld2UMT5ArchConfig", "LingBotWorld2UMT5Config"
 ]

fastvideo/configs/models/encoders/base.py

@@ -83,6 +83,7 @@ class TextEncoderConfig(EncoderConfig):
     arch_config: ArchConfig = field(default_factory=TextEncoderArchConfig)
     is_chat_model: bool = False
     treat_empty_as_dot: bool = False
+    chat_template_enable_thinking: bool = field(default=False, kw_only=True)
 
 
 @dataclass

fastvideo/configs/models/encoders/lingbotworld2_t5.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.encoders.base import (
+    TextEncoderArchConfig,
+    TextEncoderConfig,
+)
+
+
+@dataclass
+class LingBotWorld2UMT5ArchConfig(TextEncoderArchConfig):
+    architectures: list[str] = field(default_factory=lambda: ["LingBotWorld2T5EncoderModel"])
+    vocab_size: int = 256384
+    dim: int = 4096
+    dim_attn: int = 4096
+    dim_ffn: int = 10240
+    num_heads: int = 64
+    num_layers: int = 24
+    num_buckets: int = 32
+    text_len: int = 512
+    hidden_size: int = 4096
+    dropout: float = 0.1
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        self.tokenizer_kwargs = {
+            "padding": "max_length",
+            "truncation": True,
+            "max_length": self.text_len,
+            "add_special_tokens": True,
+            "return_attention_mask": True,
+            "return_tensors": "pt",
+        }
+
+
+@dataclass
+class LingBotWorld2UMT5Config(TextEncoderConfig):
+    arch_config: TextEncoderArchConfig = field(default_factory=LingBotWorld2UMT5ArchConfig)
+
+    prefix: str = "text_encoder"

fastvideo/configs/pipelines/__init__.py

@@ -7,6 +7,7 @@ from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyua
 from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
 from fastvideo.configs.pipelines.hyworld import HYWorldConfig
 from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5I2VConfig, Kandinsky5T2VConfig
+from fastvideo.configs.pipelines.lingbotworld2 import LingBotWorld2CausalFastI2V480PConfig
 from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig
 from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig
 from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
@@ -19,5 +20,6 @@ __all__ = [
     "Hunyuan15T2V720PConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig",
     "SelfForcingWanT2V480PConfig", "LucyEditDevConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig",
     "DreamXWorld5BCamPipelineConfig", "DreamXWorld5BARPipelineConfig", "HYWorldConfig", "Kandinsky5T2VConfig",
-    "Kandinsky5I2VConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
+    "Kandinsky5I2VConfig", "LingBotWorld2CausalFastI2V480PConfig", "MatrixGame2I2V480PConfig",
+    "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
 ]

fastvideo/configs/pipelines/lingbotworld2.py

@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: Apache-2.0
+import html
+from dataclasses import dataclass, field
+
+import ftfy
+import torch
+
+from fastvideo.configs.models import DiTConfig
+from fastvideo.configs.models.dits.lingbotworld2 import LingBotWorld2CausalFastVideoConfig
+from fastvideo.configs.models.encoders import BaseEncoderOutput, LingBotWorld2UMT5Config
+from fastvideo.configs.models.vaes import WanVAEConfig
+from fastvideo.configs.pipelines.wan import Wan2_2_I2V_A14B_Config
+
+
+def lingbotworld2_whitespace_preprocess(prompt: str) -> str:
+    text = ftfy.fix_text(prompt)
+    text = html.unescape(html.unescape(text))
+    return " ".join(text.strip().split())
+
+
+def lingbotworld2_t5_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
+    assert outputs.last_hidden_state is not None
+    return outputs.last_hidden_state
+
+
+@dataclass
+class LingBotWorld2CausalFastI2V480PConfig(Wan2_2_I2V_A14B_Config):
+    dit_config: DiTConfig = field(default_factory=LingBotWorld2CausalFastVideoConfig)
+    vae_config: WanVAEConfig = field(default_factory=WanVAEConfig)
+    text_encoder_configs: tuple = field(default_factory=lambda: (LingBotWorld2UMT5Config(), ))
+    preprocess_text_funcs: tuple = field(default_factory=lambda: (lingbotworld2_whitespace_preprocess, ))
+    postprocess_text_funcs: tuple = field(default_factory=lambda: (lingbotworld2_t5_postprocess_text, ))
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", ))
+    flow_shift: float | None = 10.0
+    boundary_ratio: float | None = 0.947
+    vae_precision: str = "fp32"
+    vae_decode_precision: str | None = "fp32"
+    vae_tiling: bool = False
+    vae_sp: bool = False
+    dit_precision: str = "bf16"
+    is_causal: bool = False
+
+    def __post_init__(self) -> None:
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True

fastvideo/configs/pipelines/zimage.py

@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import torch
+
+from fastvideo.configs.models import EncoderConfig
+from fastvideo.configs.models.dits.zimage import ZImageDiTConfig
+from fastvideo.configs.models.encoders import BaseEncoderOutput
+from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig
+from fastvideo.configs.models.vaes.autoencoder_kl import AutoencoderKLVAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig, preprocess_text
+
+
+def _zimage_text_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
+    if outputs.hidden_states is None:
+        raise RuntimeError("Z-Image requires Qwen3 hidden states")
+    return outputs.hidden_states[-2]
+
+
+@dataclass
+class ZImagePipelineConfig(PipelineConfig):
+    """Configuration for the native Z-Image text-to-image pipeline."""
+
+    scheduler_arch: str = "FlowMatchEulerDiscreteScheduler"
+    transformer_arch: str = "ZImageTransformer2DModel"
+    vae_arch: str = "AutoencoderKL"
+    text_encoder_archs: tuple[str, ...] = ("Qwen3Model", )
+    tokenizer_archs: tuple[str, ...] = ("Qwen2Tokenizer", )
+
+    dit_config: ZImageDiTConfig = field(default_factory=ZImageDiTConfig)
+    vae_config: AutoencoderKLVAEConfig = field(default_factory=AutoencoderKLVAEConfig)
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
+        default_factory=lambda: (Qwen3TextConfig(chat_template_enable_thinking=True), ))
+    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(default_factory=lambda: (preprocess_text, ))
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
+                                  ...] = field(default_factory=lambda: (_zimage_text_postprocess, ))
+
+    dit_precision: str = "bf16"
+    vae_precision: str = "fp32"
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", ))
+    vae_tiling: bool = False
+    vae_sp: bool = False
+
+    embedded_cfg_scale: float = 0.0
+    flow_shift: float | None = 3.0
+    scheduler_step_in_fp32: bool = True
+    scheduler_sigma_min: float = 0.0
+    scheduler_use_reference_discrete_timesteps: bool = True

fastvideo/models/dits/lingbotworld2/__init__.py

@@ -0,0 +1,5 @@
+from .causal_fast import LingBotWorld2CausalFastTransformer3DModel
+
+__all__ = ["LingBotWorld2CausalFastTransformer3DModel"]
+
+EntryClass = LingBotWorld2CausalFastTransformer3DModel

fastvideo/models/dits/lingbotworld2/cam_utils.py

@@ -0,0 +1,204 @@
+# SPDX-License-Identifier: Apache-2.0
+import numpy as np
+import os
+import torch
+from scipy.interpolate import interp1d
+from scipy.spatial.transform import Rotation, Slerp
+
+
+# --- Official Code (Leave Unchanged) ---
+ 
+def interpolate_camera_poses(
+    src_indices: np.ndarray, 
+    src_rot_mat: np.ndarray, 
+    src_trans_vec: np.ndarray, 
+    tgt_indices: np.ndarray,
+) -> torch.Tensor:
+    # interpolate translation
+    interp_func_trans = interp1d(
+        src_indices, 
+        src_trans_vec, 
+        axis=0, 
+        kind='linear', 
+        bounds_error=False,
+        fill_value="extrapolate",
+    )
+    interpolated_trans_vec = interp_func_trans(tgt_indices)
+
+    # interpolate rotation
+    src_quat_vec = Rotation.from_matrix(src_rot_mat)
+    # ensure there is no sudden change in qw
+    quats = src_quat_vec.as_quat().copy()  # [N, 4]
+    for i in range(1, len(quats)):
+        if np.dot(quats[i], quats[i-1]) < 0:
+            quats[i] = -quats[i]
+    src_quat_vec = Rotation.from_quat(quats)
+    slerp_func_rot = Slerp(src_indices, src_quat_vec)
+    interpolated_rot_quat = slerp_func_rot(tgt_indices)
+    interpolated_rot_mat = interpolated_rot_quat.as_matrix()
+
+    poses = np.zeros((len(tgt_indices), 4, 4))
+    poses[:, :3, :3] = interpolated_rot_mat
+    poses[:, :3, 3] = interpolated_trans_vec
+    poses[:, 3, 3] = 1.0
+    return torch.from_numpy(poses).float()
+
+
+def SE3_inverse(T: torch.Tensor) -> torch.Tensor:
+    Rot = T[:, :3, :3] # [B,3,3]
+    trans = T[:, :3, 3:] # [B,3,1]
+    R_inv = Rot.transpose(-1, -2)
+    t_inv = -torch.bmm(R_inv, trans)
+    T_inv = torch.eye(4, device=T.device, dtype=T.dtype)[None, :, :].repeat(T.shape[0], 1, 1)
+    T_inv[:, :3, :3] = R_inv
+    T_inv[:, :3, 3:] = t_inv
+    return T_inv
+
+
+def compute_relative_poses(
+    c2ws_mat: torch.Tensor, 
+    framewise: bool = False, 
+    normalize_trans: bool = True, 
+) -> torch.Tensor:
+    ref_w2cs = SE3_inverse(c2ws_mat[0:1])
+    relative_poses = torch.matmul(ref_w2cs, c2ws_mat)
+    # ensure identity matrix for 1st frame
+    relative_poses[0] = torch.eye(4, device=c2ws_mat.device, dtype=c2ws_mat.dtype)
+    if framewise:
+        # compute pose between i and i+1
+        relative_poses_framewise = torch.bmm(SE3_inverse(relative_poses[:-1]), relative_poses[1:])
+        relative_poses[1:] = relative_poses_framewise
+    if normalize_trans: # note refer to camctrl2: "we scale the coordinate inputs to roughly 1 standard deviation to simplify model learning."
+        translations = relative_poses[:, :3, 3] # [f, 3]
+        max_norm = torch.norm(translations, dim=-1).max()
+        # only normlaize when moving
+        if max_norm > 0:
+            relative_poses[:, :3, 3] = translations / max_norm
+    return relative_poses
+
+
+@torch.no_grad()
+def create_meshgrid(n_frames: int, height: int, width: int, bias: float = 0.5, device='cuda', dtype=torch.float32) -> torch.Tensor:
+    x_range = torch.arange(width, device=device, dtype=dtype)
+    y_range = torch.arange(height, device=device, dtype=dtype)
+    grid_y, grid_x = torch.meshgrid(y_range, x_range, indexing='ij')
+    grid_xy = torch.stack([grid_x, grid_y], dim=-1).view([-1, 2]) + bias # [h*w, 2]
+    grid_xy = grid_xy[None, ...].repeat(n_frames, 1, 1) # [f, h*w, 2]
+    return grid_xy
+
+
+def get_plucker_embeddings(
+    c2ws_mat: torch.Tensor,
+    Ks: torch.Tensor,
+    height: int,
+    width: int,
+):
+    n_frames = c2ws_mat.shape[0]
+    device = c2ws_mat.device
+    dtype = c2ws_mat.dtype
+
+    grid_y, grid_x = torch.meshgrid(
+        torch.arange(height, device=device, dtype=dtype) + 0.5,
+        torch.arange(width, device=device, dtype=dtype) + 0.5,
+        indexing='ij',
+    )
+    x_flat = grid_x.reshape(-1)
+    y_flat = grid_y.reshape(-1)
+
+    fx, fy, cx, cy = Ks[0, 0], Ks[0, 1], Ks[0, 2], Ks[0, 3]
+    dirs = torch.stack([
+        (x_flat - cx) / fx,
+        (y_flat - cy) / fy,
+        torch.ones_like(x_flat),
+    ], dim=-1)
+    dirs = dirs / dirs.norm(dim=-1, keepdim=True)
+
+    rays_d = (c2ws_mat[:, :3, :3] @ dirs.T).transpose(1, 2)
+    rays_o = c2ws_mat[:, :3, 3].unsqueeze(1).expand_as(rays_d)
+    plucker_embeddings = torch.cat([rays_o, rays_d], dim=-1)
+    return plucker_embeddings.view(n_frames, height, width, 6)
+
+
+def get_Ks_transformed(
+    Ks: torch.Tensor,
+    height_org: int,
+    width_org: int,
+    height_resize: int,
+    width_resize: int,
+    height_final: int,
+    width_final: int,
+):
+    fx, fy, cx, cy = Ks.chunk(4, dim=-1) # [f, 1]
+
+    scale_x = width_resize / width_org
+    scale_y = height_resize / height_org
+
+    fx_resize = fx * scale_x
+    fy_resize = fy * scale_y
+    cx_resize = cx * scale_x
+    cy_resize = cy * scale_y
+
+    crop_offset_x = (width_resize - width_final) / 2
+    crop_offset_y = (height_resize - height_final) / 2
+
+    cx_final = cx_resize - crop_offset_x
+    cy_final = cy_resize - crop_offset_y
+    
+    Ks_transformed = torch.zeros_like(Ks)
+    Ks_transformed[:, 0:1] = fx_resize
+    Ks_transformed[:, 1:2] = fy_resize
+    Ks_transformed[:, 2:3] = cx_final
+    Ks_transformed[:, 3:4] = cy_final
+
+    return Ks_transformed
+
+
+# --- Custom ---
+
+def prepare_camera_embedding(
+    action_path: str,
+    num_frames: int,
+    height: int,
+    width: int,
+    spatial_scale: int = 8,
+) -> tuple[torch.Tensor, int]:
+    c2ws = np.load(os.path.join(action_path, "poses.npy"))
+    len_c2ws = ((len(c2ws) - 1) // 4) * 4 + 1
+    num_frames = min(num_frames, len_c2ws)
+    c2ws = c2ws[:num_frames]
+
+    Ks = torch.from_numpy(
+        np.load(os.path.join(action_path, "intrinsics.npy"))
+    ).float()
+    Ks = get_Ks_transformed(
+        Ks,
+        height_org=480,
+        width_org=832,
+        height_resize=height,
+        width_resize=width,
+        height_final=height,
+        width_final=width,
+    )
+    Ks = Ks[0]  # use first frame
+
+    len_c2ws = len(c2ws)
+    num_latent_frames = (len_c2ws - 1) // 4 + 1
+    c2ws_infer = interpolate_camera_poses(
+        src_indices=np.linspace(0, len_c2ws - 1, len_c2ws),
+        src_rot_mat=c2ws[:, :3, :3],
+        src_trans_vec=c2ws[:, :3, 3],
+        tgt_indices=np.linspace(0, len_c2ws - 1, num_latent_frames),
+    )
+    c2ws_infer = compute_relative_poses(c2ws_infer, framewise=True)
+    Ks = Ks.repeat(num_latent_frames, 1)
+    plucker = get_plucker_embeddings(c2ws_infer, Ks, height, width)  # [F, H, W, 6]
+
+    # reshpae
+    latent_height = height // spatial_scale
+    latent_width = width // spatial_scale
+    plucker = plucker.view(num_latent_frames, latent_height, spatial_scale, latent_width, spatial_scale, 6)
+    plucker = plucker.permute(0, 1, 3, 5, 2, 4).contiguous()
+    plucker = plucker.view(num_latent_frames, latent_height, latent_width, 6 * spatial_scale * spatial_scale)
+    c2ws_plucker_emb = plucker.permute(3, 0, 1, 2).contiguous().unsqueeze(0)
+
+    return c2ws_plucker_emb, num_frames

fastvideo/models/dits/lingbotworld2/causal_fast.py

@@ -0,0 +1,776 @@
+# SPDX-License-Identifier: Apache-2.0
+"""LingBot World 2 causal-fast DiT implemented inside FastVideo."""
+
+import math
+import warnings
+from typing import Any
+
+from einops import rearrange
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from fastvideo.configs.models.dits.lingbotworld2 import (
+    LingBotWorld2CausalFastVideoConfig,
+)
+from fastvideo.distributed.communication_op import (
+    sequence_model_parallel_all_gather,
+    sequence_model_parallel_all_to_all_4D,
+)
+from fastvideo.distributed.parallel_state import get_sp_parallel_rank, get_sp_world_size
+from fastvideo.models.dits.base import BaseDiT
+from fastvideo.platforms import AttentionBackendEnum
+
+try:
+    import flash_attn_interface
+
+    FLASH_ATTN_3_AVAILABLE = True
+except ModuleNotFoundError:
+    FLASH_ATTN_3_AVAILABLE = False
+
+try:
+    import flash_attn
+
+    FLASH_ATTN_2_AVAILABLE = True
+except ModuleNotFoundError:
+    FLASH_ATTN_2_AVAILABLE = False
+
+
+def is_blocks(n: str, m) -> bool:
+    return "blocks" in n and str.isdigit(n.split(".")[-1])
+
+
+def sinusoidal_embedding_1d(dim: int, position: torch.Tensor) -> torch.Tensor:
+    """Build Wan/LingBot World 2 sinusoidal timestep embeddings."""
+    assert dim % 2 == 0
+    half = dim // 2
+    position = position.type(torch.float64)
+    sinusoid = torch.outer(
+        position,
+        torch.pow(10000, -torch.arange(half, device=position.device).to(position).div(half)),
+    )
+    return torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
+
+
+@torch.amp.autocast("cuda", enabled=False)
+def rope_params(max_seq_len: int, dim: int, theta: int = 10000) -> torch.Tensor:
+    """Return complex RoPE frequencies used by the released LingBot World 2 model."""
+    assert dim % 2 == 0
+    freqs = torch.outer(
+        torch.arange(max_seq_len),
+        1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float64).div(dim)),
+    )
+    return torch.polar(torch.ones_like(freqs), freqs)
+
+
+def flash_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q_lens: torch.Tensor | None = None,
+    k_lens: torch.Tensor | None = None,
+    dropout_p: float = 0.0,
+    softmax_scale: float | None = None,
+    q_scale: float | None = None,
+    causal: bool = False,
+    window_size: tuple[int, int] = (-1, -1),
+    deterministic: bool = False,
+    dtype: torch.dtype = torch.bfloat16,
+    version: int | None = None,
+) -> torch.Tensor:
+    """Run LingBot World 2-compatible FlashAttention on packed varlen inputs."""
+    half_dtypes = (torch.float16, torch.bfloat16)
+    assert dtype in half_dtypes
+    assert q.device.type == "cuda" and q.size(-1) <= 256
+
+    b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype
+
+    def half(x: torch.Tensor) -> torch.Tensor:
+        return x if x.dtype in half_dtypes else x.to(dtype)
+
+    if q_lens is None:
+        q = half(q.flatten(0, 1))
+        q_lens = torch.tensor([lq] * b, dtype=torch.int32, device=q.device)
+    else:
+        q = half(torch.cat([u[:v] for u, v in zip(q, q_lens, strict=True)]))
+
+    if k_lens is None:
+        k = half(k.flatten(0, 1))
+        v = half(v.flatten(0, 1))
+        k_lens = torch.tensor([lk] * b, dtype=torch.int32, device=k.device)
+    else:
+        k = half(torch.cat([u[:v] for u, v in zip(k, k_lens, strict=True)]))
+        v = half(torch.cat([u[:v] for u, v in zip(v, k_lens, strict=True)]))
+
+    q = q.to(v.dtype)
+    k = k.to(v.dtype)
+    if q_scale is not None:
+        q = q * q_scale
+
+    if version == 3 and not FLASH_ATTN_3_AVAILABLE:
+        warnings.warn("FlashAttention 3 is not available; using FlashAttention 2.")
+
+    cu_q = torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
+        0, dtype=torch.int32
+    ).to(q.device, non_blocking=True)
+    cu_k = torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
+        0, dtype=torch.int32
+    ).to(k.device, non_blocking=True)
+
+    if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE:
+        x = flash_attn_interface.flash_attn_varlen_func(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=cu_q,
+            cu_seqlens_k=cu_k,
+            seqused_q=None,
+            seqused_k=None,
+            max_seqlen_q=lq,
+            max_seqlen_k=lk,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            deterministic=deterministic,
+        ).unflatten(0, (b, lq))
+    else:
+        assert FLASH_ATTN_2_AVAILABLE
+        x = flash_attn.flash_attn_varlen_func(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=cu_q,
+            cu_seqlens_k=cu_k,
+            max_seqlen_q=lq,
+            max_seqlen_k=lk,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size=window_size,
+            deterministic=deterministic,
+        ).unflatten(0, (b, lq))
+    return x.type(out_dtype)
+
+
+def attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q_lens: torch.Tensor | None = None,
+    k_lens: torch.Tensor | None = None,
+    dropout_p: float = 0.0,
+    softmax_scale: float | None = None,
+    q_scale: float | None = None,
+    causal: bool = False,
+    window_size: tuple[int, int] = (-1, -1),
+    deterministic: bool = False,
+    dtype: torch.dtype = torch.bfloat16,
+    fa_version: int | None = None,
+) -> torch.Tensor:
+    """Dispatch LingBot World 2 attention to FlashAttention when available."""
+    if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE:
+        return flash_attention(
+            q=q,
+            k=k,
+            v=v,
+            q_lens=q_lens,
+            k_lens=k_lens,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            q_scale=q_scale,
+            causal=causal,
+            window_size=window_size,
+            deterministic=deterministic,
+            dtype=dtype,
+            version=fa_version,
+        )
+
+    if q_lens is not None or k_lens is not None:
+        warnings.warn("Padding masks are disabled without FlashAttention.")
+    q = q.transpose(1, 2).to(dtype)
+    k = k.transpose(1, 2).to(dtype)
+    v = v.transpose(1, 2).to(dtype)
+    out = F.scaled_dot_product_attention(
+        q, k, v, attn_mask=None, is_causal=causal, dropout_p=dropout_p
+    )
+    return out.transpose(1, 2).contiguous()
+
+
+@torch.amp.autocast("cuda", enabled=False)
+def causal_rope_apply(
+    x: torch.Tensor,
+    grid_sizes: torch.Tensor,
+    freqs: torch.Tensor,
+    start_frame: int = 0,
+) -> torch.Tensor:
+    """Apply LingBot World 2 causal RoPE with the current chunk frame offset."""
+    n, c = x.size(2), x.size(3) // 2
+    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+    output = []
+    for i, (f, h, w) in enumerate(grid_sizes.tolist()):
+        seq_len = f * h * w
+        x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(seq_len, n, -1, 2))
+        freqs_i = torch.cat(
+            [
+                freqs[0][start_frame : start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1),
+                freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+                freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1),
+            ],
+            dim=-1,
+        ).reshape(seq_len, 1, -1)
+        x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
+        x_i = torch.cat([x_i, x[i, seq_len:]])
+        output.append(x_i)
+    return torch.stack(output).type_as(x)
+
+
+class WanRMSNorm(nn.Module):
+    """RMSNorm used by Wan/LingBot World 2 attention projections."""
+
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Normalize the last dimension in fp32 and restore input dtype."""
+        return self._norm(x.float()).type_as(x) * self.weight
+
+    def _norm(self, x: torch.Tensor) -> torch.Tensor:
+        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+
+
+class WanLayerNorm(nn.LayerNorm):
+    """LayerNorm variant that computes in fp32 and returns the input dtype."""
+
+    def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine: bool = False):
+        super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply layer norm in fp32 for LingBot World 2 numerical parity."""
+        return super().forward(x.float()).type_as(x)
+
+
+class CausalWanSelfAttention(nn.Module):
+    """LingBot World 2 causal self-attention with rolling KV cache."""
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        local_attn_size: int = -1,
+        sink_size: int = 0,
+        qk_norm: bool = True,
+        eps: float = 1e-6,
+    ):
+        super().__init__()
+        assert dim % num_heads == 0
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.local_attn_size = local_attn_size
+        self.sink_size = sink_size
+        self.qk_norm = qk_norm
+        self.eps = eps
+        self.q = nn.Linear(dim, dim)
+        self.k = nn.Linear(dim, dim)
+        self.v = nn.Linear(dim, dim)
+        self.o = nn.Linear(dim, dim)
+        self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        seq_lens: torch.Tensor,
+        grid_sizes: torch.Tensor,
+        freqs: torch.Tensor,
+        kv_cache: dict,
+        current_start: int = 0,
+        max_attention_size: int = 1_000_000,
+        frame_seqlen: int | None = None,
+        seq_lens_int: int | None = None,
+    ) -> torch.Tensor:
+        """Project QKV, update the rolling cache, and attend to its active window."""
+        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+        q = self.norm_q(self.q(x)).view(b, s, n, d)
+        k = self.norm_k(self.k(x)).view(b, s, n, d)
+        v = self.v(x).view(b, s, n, d)
+
+        if frame_seqlen is None:
+            frame_seqlen = math.prod(grid_sizes[0][1:]).item()
+        current_start_frame = current_start // frame_seqlen
+        if seq_lens_int is None:
+            seq_lens_int = int(seq_lens[0].item() if seq_lens.dim() > 0 else seq_lens.item())
+
+        sp_size = get_sp_world_size()
+        if sp_size > 1:
+            q = sequence_model_parallel_all_to_all_4D(q, scatter_dim=2, gather_dim=1)
+            k = sequence_model_parallel_all_to_all_4D(k, scatter_dim=2, gather_dim=1)
+            v = sequence_model_parallel_all_to_all_4D(v, scatter_dim=2, gather_dim=1)
+            padded_seq_len = s * sp_size
+            roped_query = causal_rope_apply(q, grid_sizes, freqs, start_frame=current_start_frame).type_as(v)
+            roped_key = causal_rope_apply(k, grid_sizes, freqs, start_frame=current_start_frame).type_as(v)
+            roped_query = roped_query[:, :seq_lens_int]
+            roped_key = roped_key[:, :seq_lens_int]
+            v = v[:, :seq_lens_int]
+            num_new_tokens = seq_lens_int
+        else:
+            padded_seq_len = s
+            roped_query = causal_rope_apply(q, grid_sizes, freqs, start_frame=current_start_frame).type_as(v)
+            roped_key = causal_rope_apply(k, grid_sizes, freqs, start_frame=current_start_frame).type_as(v)
+            num_new_tokens = roped_query.shape[1]
+
+        current_end = current_start + num_new_tokens
+        sink_tokens = self.sink_size * frame_seqlen
+        kv_cache_size = kv_cache["k"].shape[1]
+
+        if self.local_attn_size == -1:
+            local_end_index = current_start + num_new_tokens
+            local_start_index = current_start
+            kv_cache["k"][:, local_start_index:local_end_index] = roped_key
+            kv_cache["v"][:, local_start_index:local_end_index] = v
+        elif (current_end > kv_cache["global_end_index"].item()) and (
+            num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size
+        ):
+            num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size
+            num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens
+            kv_cache["k"][:, sink_tokens : sink_tokens + num_rolled_tokens] = kv_cache["k"][
+                :, sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens
+            ].clone()
+            kv_cache["v"][:, sink_tokens : sink_tokens + num_rolled_tokens] = kv_cache["v"][
+                :, sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens
+            ].clone()
+            local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item() - num_evicted_tokens
+            local_start_index = local_end_index - num_new_tokens
+            kv_cache["k"][:, local_start_index:local_end_index] = roped_key
+            kv_cache["v"][:, local_start_index:local_end_index] = v
+        else:
+            local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item()
+            local_start_index = local_end_index - num_new_tokens
+            kv_cache["k"][:, local_start_index:local_end_index] = roped_key
+            kv_cache["v"][:, local_start_index:local_end_index] = v
+
+        k_cache = kv_cache["k"][:, max(0, local_end_index - max_attention_size) : local_end_index]
+        v_cache = kv_cache["v"][:, max(0, local_end_index - max_attention_size) : local_end_index]
+        x = attention(roped_query, k_cache, v_cache)
+        kv_cache["global_end_index"].fill_(current_end)
+        kv_cache["local_end_index"].fill_(local_end_index)
+
+        if sp_size > 1:
+            sp_pad = padded_seq_len - seq_lens_int
+            if sp_pad > 0:
+                x = torch.cat([x, x.new_zeros(b, sp_pad, x.size(2), d)], dim=1)
+            x = sequence_model_parallel_all_to_all_4D(x, scatter_dim=1, gather_dim=2)
+        return self.o(x.flatten(2))
+
+class WanCrossAttention(CausalWanSelfAttention):
+    """LingBot World 2 cross-attention with reusable text K/V cache."""
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        context: torch.Tensor,
+        context_lens: torch.Tensor | None,
+        crossattn_cache: dict | None = None,
+        cross_attn_first_call: bool | None = None,
+    ) -> torch.Tensor:
+        """Attend hidden states to text context, populating cache on first use."""
+        b, n, d = x.size(0), self.num_heads, self.head_dim
+        q = self.norm_q(self.q(x)).view(b, -1, n, d)
+        if crossattn_cache is not None:
+            is_first = crossattn_cache["is_init"].item() == 0 if cross_attn_first_call is None else cross_attn_first_call
+            if is_first:
+                crossattn_cache["is_init"].fill_(1)
+                k = self.norm_k(self.k(context)).view(b, -1, n, d)
+                v = self.v(context).view(b, -1, n, d)
+                crossattn_cache["k"].copy_(k)
+                crossattn_cache["v"].copy_(v)
+            else:
+                k = crossattn_cache["k"]
+                v = crossattn_cache["v"]
+        else:
+            k = self.norm_k(self.k(context)).view(b, -1, n, d)
+            v = self.v(context).view(b, -1, n, d)
+        x = flash_attention(q, k, v, k_lens=context_lens)
+        return self.o(x.flatten(2))
+
+
+class CausalWanAttentionBlock(nn.Module):
+    """One LingBot World 2 causal transformer block including camera injection."""
+
+    def __init__(
+        self,
+        dim: int,
+        ffn_dim: int,
+        num_heads: int,
+        local_attn_size: int = -1,
+        sink_size: int = 0,
+        qk_norm: bool = True,
+        cross_attn_norm: bool = False,
+        eps: float = 1e-6,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.num_heads = num_heads
+        self.local_attn_size = local_attn_size
+        self.qk_norm = qk_norm
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+        self.norm1 = WanLayerNorm(dim, eps)
+        self.self_attn = CausalWanSelfAttention(dim, num_heads, local_attn_size, sink_size, qk_norm, eps)
+        self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
+        self.cross_attn = WanCrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps)
+        self.norm2 = WanLayerNorm(dim, eps)
+        self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim))
+        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+        self.cam_injector_layer1 = nn.Linear(dim, dim)
+        self.cam_injector_layer2 = nn.Linear(dim, dim)
+        self.cam_scale_layer = nn.Linear(dim, dim)
+        self.cam_shift_layer = nn.Linear(dim, dim)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        e: torch.Tensor,
+        seq_lens: torch.Tensor,
+        grid_sizes: torch.Tensor,
+        freqs: torch.Tensor,
+        context: torch.Tensor,
+        context_lens: torch.Tensor | None,
+        dit_cond_dict: dict[str, Any] | None = None,
+        kv_cache: dict | None = None,
+        crossattn_cache: dict | None = None,
+        current_start: int = 0,
+        max_attention_size: int = 1_000_000,
+        frame_seqlen: int | None = None,
+        cross_attn_first_call: bool | None = None,
+        seq_lens_int: int | None = None,
+    ) -> torch.Tensor:
+        """Apply self-attention, camera modulation, cross-attention, and FFN."""
+        assert kv_cache is not None
+        assert e.dtype == torch.float32
+        with torch.amp.autocast("cuda", dtype=torch.float32):
+            e = (self.modulation.unsqueeze(0) + e).chunk(6, dim=2)
+
+        y = self.self_attn(
+            self.norm1(x).float() * (1 + e[1].squeeze(2)) + e[0].squeeze(2),
+            seq_lens,
+            grid_sizes,
+            freqs,
+            kv_cache,
+            current_start,
+            max_attention_size,
+            frame_seqlen=frame_seqlen,
+            seq_lens_int=seq_lens_int,
+        )
+        with torch.amp.autocast("cuda", dtype=torch.float32):
+            x = x + y * e[2].squeeze(2)
+
+        if dit_cond_dict is not None and "c2ws_plucker_emb" in dit_cond_dict:
+            c2ws_plucker_emb = dit_cond_dict["c2ws_plucker_emb"]
+            c2ws_hidden_states = self.cam_injector_layer2(
+                F.silu(self.cam_injector_layer1(c2ws_plucker_emb))
+            )
+            c2ws_hidden_states = c2ws_hidden_states + c2ws_plucker_emb
+            x = (1.0 + self.cam_scale_layer(c2ws_hidden_states)) * x + self.cam_shift_layer(c2ws_hidden_states)
+
+        x = x + self.cross_attn(
+            self.norm3(x),
+            context,
+            context_lens,
+            crossattn_cache=crossattn_cache,
+            cross_attn_first_call=cross_attn_first_call,
+        )
+        y = self.ffn(self.norm2(x).float() * (1 + e[4].squeeze(2)) + e[3].squeeze(2))
+        with torch.amp.autocast("cuda", dtype=torch.float32):
+            x = x + y * e[5].squeeze(2)
+        return x
+
+
+class CausalHead(nn.Module):
+    """Output projection head for LingBot World 2 causal-fast DiT."""
+
+    def __init__(self, dim: int, out_dim: int, patch_size: tuple[int, int, int], eps: float = 1e-6):
+        super().__init__()
+        self.dim = dim
+        self.out_dim = out_dim
+        self.patch_size = patch_size
+        self.eps = eps
+        self.norm = WanLayerNorm(dim, eps)
+        self.head = nn.Linear(dim, math.prod(patch_size) * out_dim)
+        self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+
+    def forward(self, x: torch.Tensor, e: torch.Tensor) -> torch.Tensor:
+        """Normalize, modulate, and project hidden states to latent patches."""
+        assert e.dtype == torch.float32
+        with torch.amp.autocast("cuda", dtype=torch.float32):
+            e = (self.modulation.unsqueeze(0) + e.unsqueeze(2)).chunk(2, dim=2)
+            x = self.head(self.norm(x) * (1 + e[1].squeeze(2)) + e[0].squeeze(2))
+        return x
+
+
+class LingBotWorld2CausalFastTransformer3DModel(BaseDiT):
+    """Released LingBot World 2 14B causal-fast model with native FastVideo loading."""
+
+    _fsdp_shard_conditions = [is_blocks]
+    _compile_conditions: list = []
+    _supported_attention_backends = (
+        AttentionBackendEnum.FLASH_ATTN,
+        AttentionBackendEnum.TORCH_SDPA,
+    )
+    param_names_mapping: dict = {}
+    reverse_param_names_mapping: dict = {}
+    lora_param_names_mapping: dict = {}
+
+    def __init__(self, config: LingBotWorld2CausalFastVideoConfig, hf_config: dict[str, Any]) -> None:
+        super().__init__(config=config, hf_config=hf_config)
+        self.model_type = config.model_type
+        self.patch_size = tuple(config.patch_size)
+        self.text_len = config.text_len
+        self.in_dim = config.in_dim
+        self.dim = config.dim
+        self.hidden_size = config.dim
+        self.ffn_dim = config.ffn_dim
+        self.freq_dim = config.freq_dim
+        self.text_dim = config.text_dim
+        self.out_dim = config.out_dim
+        self.out_channels = config.out_dim
+        self.num_heads = config.num_heads
+        self.num_attention_heads = config.num_heads
+        self.attention_head_dim = config.dim // config.num_heads
+        self.num_layers = config.num_layers
+        self.local_attn_size = config.local_attn_size
+        self.sink_size = config.sink_size
+        self.qk_norm = config.qk_norm
+        self.cross_attn_norm = config.cross_attn_norm
+        self.eps = config.eps
+        self.num_channels_latents = config.out_dim
+
+        control_dim = 6
+        self.patch_embedding = nn.Conv3d(self.in_dim, self.dim, kernel_size=self.patch_size, stride=self.patch_size)
+        self.patch_embedding_wancamctrl = nn.Linear(
+            control_dim * 64 * self.patch_size[0] * self.patch_size[1] * self.patch_size[2],
+            self.dim,
+        )
+        self.c2ws_hidden_states_layer1 = nn.Linear(self.dim, self.dim)
+        self.c2ws_hidden_states_layer2 = nn.Linear(self.dim, self.dim)
+        self.text_embedding = nn.Sequential(
+            nn.Linear(self.text_dim, self.dim),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(self.dim, self.dim),
+        )
+        self.time_embedding = nn.Sequential(
+            nn.Linear(self.freq_dim, self.dim),
+            nn.SiLU(),
+            nn.Linear(self.dim, self.dim),
+        )
+        self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(self.dim, self.dim * 6))
+        self.blocks = nn.ModuleList(
+            [
+                CausalWanAttentionBlock(
+                    self.dim,
+                    self.ffn_dim,
+                    self.num_heads,
+                    self.local_attn_size,
+                    self.sink_size,
+                    self.qk_norm,
+                    self.cross_attn_norm,
+                    self.eps,
+                )
+                for _ in range(self.num_layers)
+            ]
+        )
+        self.head = CausalHead(self.dim, self.out_dim, self.patch_size, self.eps)
+        self.freqs: torch.Tensor | None = None
+        self.init_weights()
+        self.__post_init__()
+
+    def _get_freqs(self, device: torch.device) -> torch.Tensor:
+        """Materialize the non-persistent RoPE frequency table outside meta init."""
+        if self.freqs is None or self.freqs.is_meta or self.freqs.device != device:
+            d = self.dim // self.num_heads
+            self.freqs = torch.cat(
+                [
+                    rope_params(1024, d - 4 * (d // 6)),
+                    rope_params(1024, 2 * (d // 6)),
+                    rope_params(1024, 2 * (d // 6)),
+                ],
+                dim=1,
+            ).to(device)
+        return self.freqs
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor | list[torch.Tensor] | None = None,
+        encoder_hidden_states: torch.Tensor | list[torch.Tensor] | None = None,
+        timestep: torch.Tensor | None = None,
+        encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None,
+        guidance=None,
+        *,
+        x: list[torch.Tensor] | None = None,
+        t: torch.Tensor | None = None,
+        context: list[torch.Tensor] | torch.Tensor | None = None,
+        seq_len: int | None = None,
+        y: list[torch.Tensor] | None = None,
+        dit_cond_dict: dict[str, Any] | None = None,
+        kv_cache: list[dict] | None = None,
+        crossattn_cache: list[dict] | None = None,
+        current_start: int = 0,
+        max_attention_size: int = 1_000_000,
+        frame_seqlen: int | None = None,
+        cross_attn_first_call: bool | None = None,
+        **kwargs,
+    ) -> list[torch.Tensor]:
+        """Run one cached causal-fast DiT forward using the released LingBot World 2 ABI."""
+        del encoder_hidden_states_image, guidance, kwargs
+        if x is None:
+            assert isinstance(hidden_states, torch.Tensor)
+            x = [hidden_states[0]]
+        if t is None:
+            assert timestep is not None
+            t = timestep
+        if context is None:
+            context = encoder_hidden_states
+        if isinstance(context, torch.Tensor):
+            context = [u for u in context]
+        assert context is not None
+        assert seq_len is not None
+        assert kv_cache is not None
+        assert crossattn_cache is not None
+        if self.model_type == "i2v":
+            assert y is not None
+
+        device = self.patch_embedding.weight.device
+        freqs = self._get_freqs(device)
+        if y is not None:
+            x = [torch.cat([u, v], dim=0) for u, v in zip(x, y, strict=True)]
+
+        x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+        grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long, device=u.device) for u in x])
+        x = [u.flatten(2).transpose(1, 2) for u in x]
+        seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long, device=device)
+        assert seq_lens.max() <= seq_len
+        x = torch.cat(x)
+        seq_lens_int = int(seq_lens[0].item())
+        sp_size = get_sp_world_size()
+        sp_rank = get_sp_parallel_rank()
+        padded_seq_len = ((seq_lens_int + sp_size - 1) // sp_size) * sp_size
+        sp_pad_len = padded_seq_len - seq_lens_int
+        if sp_pad_len > 0:
+            x = torch.cat([x, x.new_zeros(x.size(0), sp_pad_len, x.size(2))], dim=1)
+
+        if t.dim() == 1:
+            t = t.expand(t.size(0), padded_seq_len)
+        with torch.amp.autocast("cuda", dtype=torch.float32):
+            bt = t.size(0)
+            t = t.flatten()
+            e = self.time_embedding(
+                sinusoidal_embedding_1d(self.freq_dim, t).unflatten(0, (bt, padded_seq_len)).float()
+            )
+            e0 = self.time_projection(e).unflatten(2, (6, self.dim))
+
+        context_lens = None
+        context = self.text_embedding(
+            torch.stack([torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context])
+        )
+
+        if dit_cond_dict is not None and "c2ws_plucker_emb" in dit_cond_dict:
+            c2ws_plucker_emb = dit_cond_dict["c2ws_plucker_emb"]
+            c2ws_plucker_emb = [
+                rearrange(
+                    i,
+                    "1 c (f c1) (h c2) (w c3) -> 1 (f h w) (c c1 c2 c3)",
+                    c1=self.patch_size[0],
+                    c2=self.patch_size[1],
+                    c3=self.patch_size[2],
+                )
+                for i in c2ws_plucker_emb
+            ]
+            c2ws_plucker_emb = torch.cat(c2ws_plucker_emb, dim=1)
+            c2ws_plucker_emb = self.patch_embedding_wancamctrl(c2ws_plucker_emb)
+            c2ws_hidden_states = self.c2ws_hidden_states_layer2(
+                F.silu(self.c2ws_hidden_states_layer1(c2ws_plucker_emb))
+            )
+            c2ws_plucker_emb = c2ws_plucker_emb + c2ws_hidden_states
+            cam_len = c2ws_plucker_emb.size(1)
+            if cam_len < padded_seq_len:
+                c2ws_plucker_emb = torch.cat(
+                    [
+                        c2ws_plucker_emb,
+                        c2ws_plucker_emb.new_zeros(
+                            c2ws_plucker_emb.size(0),
+                            padded_seq_len - cam_len,
+                            c2ws_plucker_emb.size(2),
+                        ),
+                    ],
+                    dim=1,
+                )
+            elif cam_len > padded_seq_len:
+                c2ws_plucker_emb = c2ws_plucker_emb[:, :padded_seq_len, :]
+            if sp_size > 1:
+                c2ws_plucker_emb = torch.chunk(c2ws_plucker_emb, sp_size, dim=1)[sp_rank]
+            dit_cond_dict = dict(dit_cond_dict)
+            dit_cond_dict["c2ws_plucker_emb"] = c2ws_plucker_emb
+
+        if sp_size > 1:
+            x = torch.chunk(x, sp_size, dim=1)[sp_rank]
+            e = torch.chunk(e, sp_size, dim=1)[sp_rank]
+            e0 = torch.chunk(e0, sp_size, dim=1)[sp_rank]
+
+        for block_index, block in enumerate(self.blocks):
+            x = block(
+                x,
+                e=e0,
+                seq_lens=seq_lens,
+                grid_sizes=grid_sizes,
+                freqs=freqs,
+                context=context,
+                context_lens=context_lens,
+                dit_cond_dict=dit_cond_dict,
+                kv_cache=kv_cache[block_index],
+                crossattn_cache=crossattn_cache[block_index],
+                current_start=current_start,
+                max_attention_size=max_attention_size,
+                frame_seqlen=frame_seqlen,
+                cross_attn_first_call=cross_attn_first_call,
+                seq_lens_int=seq_lens_int,
+            )
+
+        x = self.head(x, e)
+        if sp_size > 1:
+            x = sequence_model_parallel_all_gather(x, dim=1)
+        return [u.float() for u in self.unpatchify(x, grid_sizes)]
+
+    def unpatchify(self, x: torch.Tensor, grid_sizes: torch.Tensor) -> list[torch.Tensor]:
+        """Reconstruct latent videos from flattened patch tokens."""
+        c = self.out_dim
+        out = []
+        for u, v in zip(x, grid_sizes.tolist(), strict=True):
+            u = u[: math.prod(v)].view(*v, *self.patch_size, c)
+            u = torch.einsum("fhwpqrc->cfphqwr", u)
+            u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size, strict=True)])
+            out.append(u)
+        return out
+
+    def init_weights(self) -> None:
+        """Initialize modules for non-meta construction; checkpoint load overwrites them."""
+        if self.patch_embedding.weight.is_meta:
+            return
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+        nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
+        for m in self.text_embedding.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=0.02)
+        for m in self.time_embedding.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=0.02)
+        nn.init.zeros_(self.head.head.weight)
+
+
+EntryClass = LingBotWorld2CausalFastTransformer3DModel

fastvideo/models/dits/zimage.py

@@ -0,0 +1,570 @@
+# SPDX-License-Identifier: Apache-2.0
+"""FastVideo-native Z-Image transformer.
+
+Z-Image attends over padded variable-length image/text streams and requires a
+key-padding mask. FastVideo's distributed attention wrappers do not expose that
+mask contract yet, so this implementation uses torch SDPA and is SP=1 only.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils.rnn import pad_sequence
+
+from fastvideo.configs.models.dits.zimage import ZImageDiTConfig
+from fastvideo.distributed.parallel_state import get_sp_world_size, model_parallel_is_initialized
+from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.models.dits.base import BaseDiT
+from fastvideo.platforms import AttentionBackendEnum
+
+
+def _linear(layer: ReplicatedLinear, x: torch.Tensor) -> torch.Tensor:
+    return layer(x)[0]
+
+
+def _prepare_attention_mask(attention_mask: torch.Tensor | None, dtype: torch.dtype) -> torch.Tensor | None:
+    if attention_mask is None:
+        return None
+    if attention_mask.ndim == 2:
+        attention_mask = attention_mask[:, None, None, :]
+    if attention_mask.dtype == torch.bool:
+        additive_mask = torch.zeros_like(attention_mask, dtype=dtype)
+        additive_mask.masked_fill_(~attention_mask, float("-inf"))
+        return additive_mask
+    return attention_mask
+
+
+class TimestepEmbedder(nn.Module):
+
+    def __init__(
+        self,
+        out_size: int,
+        mid_size: int,
+        frequency_embedding_size: int,
+        max_period: int,
+    ) -> None:
+        super().__init__()
+        self.mlp = nn.ModuleList([
+            ReplicatedLinear(frequency_embedding_size, mid_size, bias=True),
+            nn.SiLU(),
+            ReplicatedLinear(mid_size, out_size, bias=True),
+        ])
+        self.frequency_embedding_size = frequency_embedding_size
+        self.max_period = max_period
+
+    @staticmethod
+    def timestep_embedding(t: torch.Tensor, dim: int, max_period: int) -> torch.Tensor:
+        with torch.amp.autocast("cuda", enabled=False):
+            half = dim // 2
+            freqs = torch.exp(
+                -math.log(max_period) * torch.arange(half, dtype=torch.float32, device=t.device) / half)
+            args = t[:, None].float() * freqs[None]
+            embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+            if dim % 2:
+                embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+            return embedding
+
+    def forward(self, t: torch.Tensor) -> torch.Tensor:
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size, self.max_period)
+        weight_dtype = self.mlp[0].weight.dtype
+        if weight_dtype.is_floating_point:
+            t_freq = t_freq.to(weight_dtype)
+        return _linear(self.mlp[2], self.mlp[1](_linear(self.mlp[0], t_freq)))
+
+
+class RMSNorm(nn.Module):
+
+    def __init__(self, dim: int, eps: float = 1e-5) -> None:
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        output = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return output * self.weight
+
+
+class FeedForward(nn.Module):
+
+    def __init__(self, dim: int, hidden_dim: int) -> None:
+        super().__init__()
+        self.w1 = ReplicatedLinear(dim, hidden_dim, bias=False)
+        self.w2 = ReplicatedLinear(hidden_dim, dim, bias=False)
+        self.w3 = ReplicatedLinear(dim, hidden_dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return _linear(self.w2, F.silu(_linear(self.w1, x)) * _linear(self.w3, x))
+
+
+def apply_rotary_emb(x_in: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+    with torch.amp.autocast("cuda", enabled=False):
+        x = torch.view_as_complex(x_in.float().reshape(*x_in.shape[:-1], -1, 2))
+        x_out = torch.view_as_real(x * freqs_cis.unsqueeze(2)).flatten(3)
+        return x_out.type_as(x_in)
+
+
+class ZImageAttention(nn.Module):
+
+    def __init__(self, dim: int, n_heads: int, n_kv_heads: int, qk_norm: bool = True, eps: float = 1e-5) -> None:
+        super().__init__()
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.head_dim = dim // n_heads
+
+        self.to_q = ReplicatedLinear(dim, n_heads * self.head_dim, bias=False)
+        self.to_k = ReplicatedLinear(dim, n_kv_heads * self.head_dim, bias=False)
+        self.to_v = ReplicatedLinear(dim, n_kv_heads * self.head_dim, bias=False)
+        self.to_out = nn.ModuleList([ReplicatedLinear(n_heads * self.head_dim, dim, bias=False)])
+
+        self.norm_q = RMSNorm(self.head_dim, eps=eps) if qk_norm else None
+        self.norm_k = RMSNorm(self.head_dim, eps=eps) if qk_norm else None
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        freqs_cis: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        query = _linear(self.to_q, hidden_states).unflatten(-1, (self.n_heads, -1))
+        key = _linear(self.to_k, hidden_states).unflatten(-1, (self.n_kv_heads, -1))
+        value = _linear(self.to_v, hidden_states).unflatten(-1, (self.n_kv_heads, -1))
+
+        if self.norm_q is not None:
+            query = self.norm_q(query)
+        if self.norm_k is not None:
+            key = self.norm_k(key)
+        if freqs_cis is not None:
+            query = apply_rotary_emb(query, freqs_cis)
+            key = apply_rotary_emb(key, freqs_cis)
+
+        mask = _prepare_attention_mask(attention_mask, query.dtype)
+        hidden_states = F.scaled_dot_product_attention(
+            query.transpose(1, 2),
+            key.transpose(1, 2),
+            value.transpose(1, 2),
+            attn_mask=mask,
+            dropout_p=0.0,
+            is_causal=False,
+        ).transpose(1, 2).contiguous()
+        return _linear(self.to_out[0], hidden_states.flatten(2, 3).to(query.dtype))
+
+
+class ZImageTransformerBlock(nn.Module):
+
+    def __init__(
+        self,
+        layer_id: int,
+        dim: int,
+        n_heads: int,
+        n_kv_heads: int,
+        norm_eps: float,
+        qk_norm: bool,
+        adaln_embed_dim: int,
+        modulation: bool = True,
+    ) -> None:
+        super().__init__()
+        self.dim = dim
+        self.head_dim = dim // n_heads
+        self.layer_id = layer_id
+        self.modulation = modulation
+
+        self.attention = ZImageAttention(dim, n_heads, n_kv_heads, qk_norm, norm_eps)
+        self.feed_forward = FeedForward(dim=dim, hidden_dim=int(dim / 3 * 8))
+        self.attention_norm1 = RMSNorm(dim, eps=norm_eps)
+        self.ffn_norm1 = RMSNorm(dim, eps=norm_eps)
+        self.attention_norm2 = RMSNorm(dim, eps=norm_eps)
+        self.ffn_norm2 = RMSNorm(dim, eps=norm_eps)
+
+        if modulation:
+            self.adaLN_modulation = nn.ModuleList(
+                [ReplicatedLinear(min(dim, adaln_embed_dim), 4 * dim, bias=True)])
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attn_mask: torch.Tensor,
+        freqs_cis: torch.Tensor,
+        adaln_input: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if self.modulation:
+            assert adaln_input is not None
+            scale_msa, gate_msa, scale_mlp, gate_mlp = _linear(
+                self.adaLN_modulation[0], adaln_input).unsqueeze(1).chunk(4, dim=2)
+            gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh()
+            scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp
+
+            attn_out = self.attention(
+                self.attention_norm1(x) * scale_msa,
+                attention_mask=attn_mask,
+                freqs_cis=freqs_cis,
+            )
+            x = x + gate_msa * self.attention_norm2(attn_out)
+            x = x + gate_mlp * self.ffn_norm2(self.feed_forward(self.ffn_norm1(x) * scale_mlp))
+        else:
+            attn_out = self.attention(
+                self.attention_norm1(x),
+                attention_mask=attn_mask,
+                freqs_cis=freqs_cis,
+            )
+            x = x + self.attention_norm2(attn_out)
+            x = x + self.ffn_norm2(self.feed_forward(self.ffn_norm1(x)))
+        return x
+
+
+class FinalLayer(nn.Module):
+
+    def __init__(self, hidden_size: int, out_channels: int, adaln_embed_dim: int) -> None:
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = ReplicatedLinear(hidden_size, out_channels, bias=True)
+        self.adaLN_modulation = nn.ModuleList([
+            nn.SiLU(),
+            ReplicatedLinear(min(hidden_size, adaln_embed_dim), hidden_size, bias=True),
+        ])
+
+    def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
+        scale = 1.0 + _linear(self.adaLN_modulation[1], self.adaLN_modulation[0](c))
+        return _linear(self.linear, self.norm_final(x) * scale.unsqueeze(1))
+
+
+class RopeEmbedder:
+
+    def __init__(self, theta: float, axes_dims: tuple[int, ...], axes_lens: tuple[int, ...]) -> None:
+        if len(axes_dims) != len(axes_lens):
+            raise ValueError("RoPE axes require matching dimensions and lengths")
+        self.theta = theta
+        self.axes_dims = axes_dims
+        self.axes_lens = axes_lens
+        self.freqs_cis: list[torch.Tensor] | None = None
+
+    @staticmethod
+    def precompute_freqs_cis(dim: tuple[int, ...], end: tuple[int, ...], theta: float) -> list[torch.Tensor]:
+        with torch.device("cpu"):
+            freqs_cis = []
+            for axis_dim, axis_end in zip(dim, end):
+                freqs = 1.0 / (theta**(torch.arange(0, axis_dim, 2, dtype=torch.float64) / axis_dim))
+                timestep = torch.arange(axis_end, dtype=torch.float64)
+                angles = torch.outer(timestep, freqs).float()
+                freqs_cis.append(torch.polar(torch.ones_like(angles), angles).to(torch.complex64))
+            return freqs_cis
+
+    def __call__(self, ids: torch.Tensor) -> torch.Tensor:
+        if ids.ndim != 2 or ids.shape[-1] != len(self.axes_dims):
+            raise ValueError("RoPE ids must have shape [sequence, number_of_axes]")
+        if self.freqs_cis is None:
+            self.freqs_cis = [
+                freqs.to(ids.device)
+                for freqs in self.precompute_freqs_cis(self.axes_dims, self.axes_lens, self.theta)
+            ]
+        elif self.freqs_cis[0].device != ids.device:
+            self.freqs_cis = [freqs.to(ids.device) for freqs in self.freqs_cis]
+        return torch.cat([self.freqs_cis[i][ids[:, i]] for i in range(len(self.axes_dims))], dim=-1)
+
+
+class ZImageTransformer2DModel(BaseDiT):
+    _default_config = ZImageDiTConfig()
+    _fsdp_shard_conditions = _default_config.arch_config._fsdp_shard_conditions
+    _compile_conditions = _default_config.arch_config._compile_conditions
+    _supported_attention_backends = (AttentionBackendEnum.TORCH_SDPA, )
+    param_names_mapping = _default_config.arch_config.param_names_mapping
+    reverse_param_names_mapping = _default_config.arch_config.reverse_param_names_mapping
+
+    def __init__(self, config: ZImageDiTConfig, hf_config: dict[str, Any]) -> None:
+        super().__init__(config=config, hf_config=hf_config)
+        arch = config.arch_config
+
+        self.in_channels = arch.in_channels
+        self.out_channels = arch.in_channels
+        self.all_patch_size = tuple(arch.all_patch_size)
+        self.all_f_patch_size = tuple(arch.all_f_patch_size)
+        self.dim = arch.dim
+        self.n_heads = arch.n_heads
+        self.rope_theta = arch.rope_theta
+        self.t_scale = arch.t_scale
+        self.seq_multi_of = arch.seq_multi_of
+
+        self.hidden_size = arch.dim
+        self.num_attention_heads = arch.n_heads
+        self.num_channels_latents = arch.in_channels
+
+        self.all_x_embedder = nn.ModuleDict({
+            f"{patch_size}-{f_patch_size}": ReplicatedLinear(
+                f_patch_size * patch_size * patch_size * arch.in_channels, arch.dim, bias=True)
+            for patch_size, f_patch_size in zip(self.all_patch_size, self.all_f_patch_size)
+        })
+        self.all_final_layer = nn.ModuleDict({
+            f"{patch_size}-{f_patch_size}": FinalLayer(
+                arch.dim,
+                patch_size * patch_size * f_patch_size * self.out_channels,
+                arch.adaln_embed_dim,
+            )
+            for patch_size, f_patch_size in zip(self.all_patch_size, self.all_f_patch_size)
+        })
+
+        block_kwargs = {
+            "dim": arch.dim,
+            "n_heads": arch.n_heads,
+            "n_kv_heads": arch.n_kv_heads,
+            "norm_eps": arch.norm_eps,
+            "qk_norm": arch.qk_norm,
+            "adaln_embed_dim": arch.adaln_embed_dim,
+        }
+        self.noise_refiner = nn.ModuleList([
+            ZImageTransformerBlock(1000 + layer_id, modulation=True, **block_kwargs)
+            for layer_id in range(arch.n_refiner_layers)
+        ])
+        self.context_refiner = nn.ModuleList([
+            ZImageTransformerBlock(layer_id, modulation=False, **block_kwargs)
+            for layer_id in range(arch.n_refiner_layers)
+        ])
+        self.t_embedder = TimestepEmbedder(
+            min(arch.dim, arch.adaln_embed_dim),
+            mid_size=arch.timestep_mid_size,
+            frequency_embedding_size=arch.frequency_embedding_size,
+            max_period=arch.max_period,
+        )
+        self.cap_embedder = nn.ModuleList([
+            RMSNorm(arch.cap_feat_dim, eps=arch.norm_eps),
+            ReplicatedLinear(arch.cap_feat_dim, arch.dim, bias=True),
+        ])
+        self.x_pad_token = nn.Parameter(torch.empty((1, arch.dim)))
+        self.cap_pad_token = nn.Parameter(torch.empty((1, arch.dim)))
+        self.layers = nn.ModuleList([
+            ZImageTransformerBlock(layer_id, modulation=True, **block_kwargs) for layer_id in range(arch.n_layers)
+        ])
+        self.axes_dims = tuple(arch.axes_dims)
+        self.axes_lens = tuple(arch.axes_lens)
+        self.rope_embedder = RopeEmbedder(arch.rope_theta, self.axes_dims, self.axes_lens)
+        self.__post_init__()
+
+    def unpatchify(
+        self,
+        x: list[torch.Tensor],
+        size: list[tuple[int, int, int]],
+        patch_size: int,
+        f_patch_size: int,
+    ) -> list[torch.Tensor]:
+        patch_height = patch_width = patch_size
+        patch_frames = f_patch_size
+        if len(x) != len(size):
+            raise ValueError("output batch and original sizes must have equal length")
+        for i, (frames, height, width) in enumerate(size):
+            original_length = (frames // patch_frames) * (height // patch_height) * (width // patch_width)
+            x[i] = (x[i][:original_length].view(
+                frames // patch_frames,
+                height // patch_height,
+                width // patch_width,
+                patch_frames,
+                patch_height,
+                patch_width,
+                self.out_channels,
+            ).permute(6, 0, 3, 1, 4, 2, 5).reshape(self.out_channels, frames, height, width))
+        return x
+
+    @staticmethod
+    def create_coordinate_grid(
+        size: tuple[int, int, int],
+        start: tuple[int, int, int] | None = None,
+        device: torch.device | None = None,
+    ) -> torch.Tensor:
+        start = start or (0, ) * len(size)
+        axes = [
+            torch.arange(axis_start, axis_start + span, dtype=torch.int32, device=device)
+            for axis_start, span in zip(start, size)
+        ]
+        return torch.stack(torch.meshgrid(axes, indexing="ij"), dim=-1)
+
+    def patchify_and_embed(
+        self,
+        all_image: list[torch.Tensor],
+        all_cap_feats: list[torch.Tensor],
+        patch_size: int,
+        f_patch_size: int,
+    ) -> tuple[
+        list[torch.Tensor],
+        list[torch.Tensor],
+        list[tuple[int, int, int]],
+        list[torch.Tensor],
+        list[torch.Tensor],
+        list[torch.Tensor],
+        list[torch.Tensor],
+    ]:
+        patch_height = patch_width = patch_size
+        patch_frames = f_patch_size
+        device = all_image[0].device
+
+        image_out = []
+        image_sizes = []
+        image_pos_ids = []
+        image_pad_masks = []
+        cap_pos_ids = []
+        cap_pad_masks = []
+        cap_feats_out = []
+
+        for image, cap_feat in zip(all_image, all_cap_feats):
+            cap_length = len(cap_feat)
+            cap_padding = (-cap_length) % self.seq_multi_of
+            cap_pos_ids.append(
+                self.create_coordinate_grid(
+                    (cap_length + cap_padding, 1, 1),
+                    start=(1, 0, 0),
+                    device=device,
+                ).flatten(0, 2))
+            cap_pad_masks.append(
+                torch.cat([
+                    torch.zeros(cap_length, dtype=torch.bool, device=device),
+                    torch.ones(cap_padding, dtype=torch.bool, device=device),
+                ]) if cap_padding else torch.zeros(cap_length, dtype=torch.bool, device=device))
+            cap_feats_out.append(
+                torch.cat([cap_feat, cap_feat[-1:].repeat(cap_padding, 1)]) if cap_padding else cap_feat)
+
+            channels, frames, height, width = image.size()
+            image_sizes.append((frames, height, width))
+            frame_tokens, height_tokens, width_tokens = (
+                frames // patch_frames,
+                height // patch_height,
+                width // patch_width,
+            )
+            image = image.view(
+                channels,
+                frame_tokens,
+                patch_frames,
+                height_tokens,
+                patch_height,
+                width_tokens,
+                patch_width,
+            )
+            image = image.permute(1, 3, 5, 2, 4, 6, 0).reshape(
+                frame_tokens * height_tokens * width_tokens,
+                patch_frames * patch_height * patch_width * channels,
+            )
+
+            image_length = len(image)
+            image_padding = (-image_length) % self.seq_multi_of
+            original_pos_ids = self.create_coordinate_grid(
+                (frame_tokens, height_tokens, width_tokens),
+                start=(cap_length + cap_padding + 1, 0, 0),
+                device=device,
+            ).flatten(0, 2)
+            if image_padding:
+                padding_pos_ids = self.create_coordinate_grid((1, 1, 1), device=device).flatten(0, 2).repeat(
+                    image_padding, 1)
+                image_pos_ids.append(torch.cat([original_pos_ids, padding_pos_ids]))
+            else:
+                image_pos_ids.append(original_pos_ids)
+            image_pad_masks.append(
+                torch.cat([
+                    torch.zeros(image_length, dtype=torch.bool, device=device),
+                    torch.ones(image_padding, dtype=torch.bool, device=device),
+                ]) if image_padding else torch.zeros(image_length, dtype=torch.bool, device=device))
+            image_out.append(
+                torch.cat([image, image[-1:].repeat(image_padding, 1)]) if image_padding else image)
+
+        return (
+            image_out,
+            cap_feats_out,
+            image_sizes,
+            image_pos_ids,
+            cap_pos_ids,
+            image_pad_masks,
+            cap_pad_masks,
+        )
+
+    @staticmethod
+    def _attention_mask(lengths: list[int], device: torch.device) -> torch.Tensor:
+        mask = torch.zeros((len(lengths), max(lengths)), dtype=torch.bool, device=device)
+        for i, length in enumerate(lengths):
+            mask[i, :length] = True
+        return mask
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor | list[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,
+        patch_size: int = 2,
+        f_patch_size: int = 1,
+        **kwargs,
+    ) -> tuple[list[torch.Tensor], dict]:
+        del encoder_hidden_states_image, guidance, kwargs
+        if model_parallel_is_initialized() and get_sp_world_size() != 1:
+            raise NotImplementedError(
+                "Z-Image masked SDPA does not support sequence parallelism; run with sp_size=1")
+        if patch_size not in self.all_patch_size or f_patch_size not in self.all_f_patch_size:
+            raise ValueError(f"unsupported patch sizes: spatial={patch_size}, temporal={f_patch_size}")
+        if isinstance(hidden_states, torch.Tensor):
+            hidden_states = list(hidden_states.unbind(0))
+        if isinstance(encoder_hidden_states, torch.Tensor):
+            encoder_hidden_states = list(encoder_hidden_states.unbind(0))
+
+        device = hidden_states[0].device
+        timestep_embedding = self.t_embedder(timestep * self.t_scale)
+        (
+            hidden_states,
+            encoder_hidden_states,
+            image_sizes,
+            image_pos_ids,
+            cap_pos_ids,
+            image_inner_pad_masks,
+            cap_inner_pad_masks,
+        ) = self.patchify_and_embed(hidden_states, encoder_hidden_states, patch_size, f_patch_size)
+
+        image_lengths = [len(item) for item in hidden_states]
+        if not all(length % self.seq_multi_of == 0 for length in image_lengths):
+            raise ValueError("padded image sequence lengths must be aligned")
+        hidden_states = torch.cat(hidden_states)
+        hidden_states = _linear(self.all_x_embedder[f"{patch_size}-{f_patch_size}"], hidden_states)
+        adaln_input = timestep_embedding.type_as(hidden_states)
+        hidden_states[torch.cat(image_inner_pad_masks)] = self.x_pad_token
+        hidden_states = list(hidden_states.split(image_lengths))
+        image_freqs_cis = list(
+            self.rope_embedder(torch.cat(image_pos_ids)).split([len(item) for item in image_pos_ids]))
+        hidden_states = pad_sequence(hidden_states, batch_first=True, padding_value=0.0)
+        image_freqs_cis = pad_sequence(image_freqs_cis, batch_first=True, padding_value=0.0)
+        image_freqs_cis = image_freqs_cis[:, :hidden_states.shape[1]]
+        image_attn_mask = self._attention_mask(image_lengths, device)
+        for layer in self.noise_refiner:
+            hidden_states = layer(hidden_states, image_attn_mask, image_freqs_cis, adaln_input)
+
+        cap_lengths = [len(item) for item in encoder_hidden_states]
+        if not all(length % self.seq_multi_of == 0 for length in cap_lengths):
+            raise ValueError("padded caption sequence lengths must be aligned")
+        encoder_hidden_states = torch.cat(encoder_hidden_states)
+        encoder_hidden_states = _linear(self.cap_embedder[1], self.cap_embedder[0](encoder_hidden_states))
+        encoder_hidden_states[torch.cat(cap_inner_pad_masks)] = self.cap_pad_token
+        encoder_hidden_states = list(encoder_hidden_states.split(cap_lengths))
+        cap_freqs_cis = list(self.rope_embedder(torch.cat(cap_pos_ids)).split([len(item) for item in cap_pos_ids]))
+        encoder_hidden_states = pad_sequence(encoder_hidden_states, batch_first=True, padding_value=0.0)
+        cap_freqs_cis = pad_sequence(cap_freqs_cis, batch_first=True, padding_value=0.0)
+        cap_freqs_cis = cap_freqs_cis[:, :encoder_hidden_states.shape[1]]
+        cap_attn_mask = self._attention_mask(cap_lengths, device)
+        for layer in self.context_refiner:
+            encoder_hidden_states = layer(encoder_hidden_states, cap_attn_mask, cap_freqs_cis)
+
+        unified = []
+        unified_freqs_cis = []
+        for i, (image_length, cap_length) in enumerate(zip(image_lengths, cap_lengths)):
+            unified.append(
+                torch.cat([hidden_states[i][:image_length], encoder_hidden_states[i][:cap_length]]))
+            unified_freqs_cis.append(
+                torch.cat([image_freqs_cis[i][:image_length], cap_freqs_cis[i][:cap_length]]))
+        unified_lengths = [image_length + cap_length for image_length, cap_length in zip(image_lengths, cap_lengths)]
+        unified = pad_sequence(unified, batch_first=True, padding_value=0.0)
+        unified_freqs_cis = pad_sequence(unified_freqs_cis, batch_first=True, padding_value=0.0)
+        unified_attn_mask = self._attention_mask(unified_lengths, device)
+        for layer in self.layers:
+            unified = layer(unified, unified_attn_mask, unified_freqs_cis, adaln_input)
+
+        unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"](unified, adaln_input)
+        outputs = self.unpatchify(list(unified.unbind(0)), image_sizes, patch_size, f_patch_size)
+        return outputs, {}
+
+
+EntryClass = ZImageTransformer2DModel

fastvideo/models/encoders/lingbotworld2_t5.py

@@ -0,0 +1,269 @@
+# SPDX-License-Identifier: Apache-2.0
+"""LingBot World 2 UMT5 encoder with the released checkpoint's module names."""
+
+from collections.abc import Iterable
+import html
+import math
+import string
+
+import ftfy
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from fastvideo.configs.models.encoders import BaseEncoderOutput
+from fastvideo.configs.models.encoders.lingbotworld2_t5 import LingBotWorld2UMT5Config
+from fastvideo.models.encoders.base import TextEncoder
+from fastvideo.models.loader.weight_utils import default_weight_loader
+
+
+def basic_clean(text: str) -> str:
+    """Apply LingBot World 2's ftfy/html cleanup before tokenization."""
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+
+def whitespace_clean(text: str) -> str:
+    """Collapse all whitespace runs to single spaces."""
+    return " ".join(text.split())
+
+
+def canonicalize(text: str, keep_punctuation_exact_string: str | None = None) -> str:
+    """Normalize prompts with LingBot World 2's optional punctuation handling."""
+    text = text.replace("_", " ")
+    if keep_punctuation_exact_string:
+        text = keep_punctuation_exact_string.join(
+            part.translate(str.maketrans("", "", string.punctuation))
+            for part in text.split(keep_punctuation_exact_string)
+        )
+    else:
+        text = text.translate(str.maketrans("", "", string.punctuation))
+    return " ".join(text.lower().split())
+
+
+def lingbotworld2_whitespace_preprocess(prompt: str) -> str:
+    """Match the LingBot World 2 source tokenizer's `clean='whitespace'` behavior."""
+    return whitespace_clean(basic_clean(prompt))
+
+
+class GELU(nn.Module):
+    """T5 gated GELU approximation used by the source checkpoint."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply the tanh GELU approximation."""
+        return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))
+
+
+class T5LayerNorm(nn.Module):
+    """T5 RMS-style layer norm with source-compatible parameter name."""
+
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Normalize in fp32 and apply the learned scale."""
+        x = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps)
+        if self.weight.dtype in (torch.float16, torch.bfloat16):
+            x = x.type_as(self.weight)
+        return self.weight * x
+
+
+class T5Attention(nn.Module):
+    """LingBot World 2 source T5 attention block."""
+
+    def __init__(self, dim: int, dim_attn: int, num_heads: int, dropout: float = 0.1):
+        super().__init__()
+        assert dim_attn % num_heads == 0
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.num_heads = num_heads
+        self.head_dim = dim_attn // num_heads
+        self.q = nn.Linear(dim, dim_attn, bias=False)
+        self.k = nn.Linear(dim, dim_attn, bias=False)
+        self.v = nn.Linear(dim, dim_attn, bias=False)
+        self.o = nn.Linear(dim_attn, dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        context: torch.Tensor | None = None,
+        mask: torch.Tensor | None = None,
+        pos_bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Project QKV, add relative bias/mask, and return attended states."""
+        context = x if context is None else context
+        b, n, c = x.size(0), self.num_heads, self.head_dim
+        q = self.q(x).view(b, -1, n, c)
+        k = self.k(context).view(b, -1, n, c)
+        v = self.v(context).view(b, -1, n, c)
+        attn_bias = x.new_zeros(b, n, q.size(1), k.size(1))
+        if pos_bias is not None:
+            attn_bias += pos_bias
+        if mask is not None:
+            assert mask.ndim in (2, 3)
+            mask = mask.view(b, 1, 1, -1) if mask.ndim == 2 else mask.unsqueeze(1)
+            attn_bias.masked_fill_(mask == 0, torch.finfo(x.dtype).min)
+        attn = torch.einsum("binc,bjnc->bnij", q, k) + attn_bias
+        attn = F.softmax(attn.float(), dim=-1).type_as(attn)
+        x = torch.einsum("bnij,bjnc->binc", attn, v)
+        return self.dropout(self.o(x.reshape(b, -1, n * c)))
+
+
+class T5FeedForward(nn.Module):
+    """LingBot World 2 source T5 gated feed-forward block."""
+
+    def __init__(self, dim: int, dim_ffn: int, dropout: float = 0.1):
+        super().__init__()
+        self.dim = dim
+        self.dim_ffn = dim_ffn
+        self.gate = nn.Sequential(nn.Linear(dim, dim_ffn, bias=False), GELU())
+        self.fc1 = nn.Linear(dim, dim_ffn, bias=False)
+        self.fc2 = nn.Linear(dim_ffn, dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply the gated feed-forward projection."""
+        x = self.fc1(x) * self.gate(x)
+        x = self.dropout(x)
+        x = self.fc2(x)
+        return self.dropout(x)
+
+
+class T5RelativeEmbedding(nn.Module):
+    """Per-block relative position embedding used by LingBot World 2 UMT5."""
+
+    def __init__(self, num_buckets: int, num_heads: int, bidirectional: bool, max_dist: int = 128):
+        super().__init__()
+        self.num_buckets = num_buckets
+        self.num_heads = num_heads
+        self.bidirectional = bidirectional
+        self.max_dist = max_dist
+        self.embedding = nn.Embedding(num_buckets, num_heads)
+
+    def forward(self, lq: int, lk: int) -> torch.Tensor:
+        """Build a relative-position bias tensor for attention logits."""
+        device = self.embedding.weight.device
+        rel_pos = torch.arange(lk, device=device).unsqueeze(0) - torch.arange(lq, device=device).unsqueeze(1)
+        rel_pos = self._relative_position_bucket(rel_pos)
+        rel_pos_embeds = self.embedding(rel_pos)
+        return rel_pos_embeds.permute(2, 0, 1).unsqueeze(0).contiguous()
+
+    def _relative_position_bucket(self, rel_pos: torch.Tensor) -> torch.Tensor:
+        """Map token offsets to T5 relative-position buckets."""
+        if self.bidirectional:
+            num_buckets = self.num_buckets // 2
+            rel_buckets = (rel_pos > 0).long() * num_buckets
+            rel_pos = torch.abs(rel_pos)
+        else:
+            num_buckets = self.num_buckets
+            rel_buckets = 0
+            rel_pos = -torch.min(rel_pos, torch.zeros_like(rel_pos))
+        max_exact = num_buckets // 2
+        rel_pos_large = max_exact + (
+            torch.log(rel_pos.float() / max_exact) / math.log(self.max_dist / max_exact) * (num_buckets - max_exact)
+        ).long()
+        rel_pos_large = torch.min(rel_pos_large, torch.full_like(rel_pos_large, num_buckets - 1))
+        rel_buckets += torch.where(rel_pos < max_exact, rel_pos, rel_pos_large)
+        return rel_buckets
+
+
+class T5SelfAttention(nn.Module):
+    """One source-compatible UMT5 encoder block."""
+
+    def __init__(
+        self,
+        dim: int,
+        dim_attn: int,
+        dim_ffn: int,
+        num_heads: int,
+        num_buckets: int,
+        shared_pos: bool = False,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        self.norm1 = T5LayerNorm(dim)
+        self.attn = T5Attention(dim, dim_attn, num_heads, dropout)
+        self.norm2 = T5LayerNorm(dim)
+        self.ffn = T5FeedForward(dim, dim_ffn, dropout)
+        self.pos_embedding = None if shared_pos else T5RelativeEmbedding(num_buckets, num_heads, bidirectional=True)
+
+    def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None, pos_bias: torch.Tensor | None = None) -> torch.Tensor:
+        """Run self-attention and feed-forward residual updates."""
+        e = pos_bias if self.pos_embedding is None else self.pos_embedding(x.size(1), x.size(1))
+        x = self._fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e))
+        return self._fp16_clamp(x + self.ffn(self.norm2(x)))
+
+    @staticmethod
+    def _fp16_clamp(x: torch.Tensor) -> torch.Tensor:
+        if x.dtype == torch.float16 and torch.isinf(x).any():
+            clamp = torch.finfo(x.dtype).max - 1000
+            x = torch.clamp(x, min=-clamp, max=clamp)
+        return x
+
+
+class LingBotWorld2T5EncoderModel(TextEncoder):
+    """FastVideo-native LingBot World 2 UMT5 encoder for the released `.pth` weights."""
+
+    fall_back_to_pt_during_load = True
+    allow_patterns_overrides = ["*.pt"]
+
+    def __init__(self, config: LingBotWorld2UMT5Config, prefix: str = ""):
+        super().__init__(config)
+        del prefix
+        arch = config.arch_config
+        self.token_embedding = nn.Embedding(arch.vocab_size, arch.dim)
+        self.dropout = nn.Dropout(arch.dropout)
+        self.blocks = nn.ModuleList(
+            [
+                T5SelfAttention(
+                    arch.dim,
+                    arch.dim_attn,
+                    arch.dim_ffn,
+                    arch.num_heads,
+                    arch.num_buckets,
+                    shared_pos=False,
+                    dropout=arch.dropout,
+                )
+                for _ in range(arch.num_layers)
+            ]
+        )
+        self.norm = T5LayerNorm(arch.dim)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | 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,
+    ) -> BaseEncoderOutput:
+        """Encode token ids and return source-compatible hidden states."""
+        del position_ids, inputs_embeds, output_hidden_states, kwargs
+        assert input_ids is not None
+        x = self.dropout(self.token_embedding(input_ids))
+        for block in self.blocks:
+            x = block(x, attention_mask)
+        x = self.dropout(self.norm(x))
+        return BaseEncoderOutput(last_hidden_state=x, attention_mask=attention_mask)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load source `.pth` weights whose names already match this module."""
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            if name not in params_dict:
+                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 = LingBotWorld2T5EncoderModel

fastvideo/models/encoders/qwen3.py

@@ -327,17 +327,15 @@ class Qwen3ForCausalLM(TextEncoder):
         dtype: torch.dtype,
         device: torch.device,
     ) -> nn.Module:
-        from transformers import AutoModelForCausalLM
+        from transformers import AutoModel
 
-        if device.type == "cpu" and torch.cuda.is_available():
-            from fastvideo.distributed import get_local_torch_device
-
-            device = get_local_torch_device()
-
-        return AutoModelForCausalLM.from_pretrained(
+        # FastVideo uses Qwen3 only as a text encoder. Loading the body avoids
+        # materializing an unused LM head and full-vocabulary logits, including
+        # for checkpoints whose metadata names Qwen3ForCausalLM.
+        return AutoModel.from_pretrained(
             model_path,
             local_files_only=True,
-            torch_dtype=dtype,
+            dtype=dtype,
             low_cpu_mem_usage=True,
         ).eval().to(device)
 
@@ -368,9 +366,18 @@ class Qwen3ForCausalLM(TextEncoder):
         residual = None
 
         if position_ids is None:
+            # Expand to [batch_size, seq_len]: the rotary layer flattens
+            # positions to ``num_tokens`` and reshapes q/k to
+            # ``(num_tokens, -1, head_dim)``. A bare [1, seq_len] only matches
+            # ``num_tokens`` when batch_size == 1; for batched inputs it folds
+            # the batch dim into the head dim and misaligns RoPE. Expanding to
+            # ``batch_size * seq_len`` tokens keeps the layout correct.
             position_ids = torch.arange(
-                0, hidden_states.shape[1], device=hidden_states.device
-            ).unsqueeze(0)
+                0,
+                hidden_states.shape[1],
+                device=hidden_states.device,
+                dtype=torch.long,
+            ).unsqueeze(0).expand(hidden_states.shape[0], -1)
 
         all_hidden_states: tuple[Any, ...] | None = (
             () if output_hidden_states else None
@@ -405,6 +412,20 @@ class Qwen3ForCausalLM(TextEncoder):
     ) -> set[str]:
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
+        stacked_params_mapping = self.config.arch_config.stacked_params_mapping
+        # A fused destination is initialized either by one already-fused tensor
+        # or after every split source projection has been loaded. Include
+        # auxiliary quantization parameters (for example scale_weight) rather
+        # than limiting completeness checks to weight/bias tensors.
+        expected_stacked_shards = {
+            (name, shard_id)
+            for name in params_dict
+            for param_name, _, shard_id in stacked_params_mapping
+            if param_name in name
+        }
+        fused_param_names = {name for name, _ in expected_stacked_shards}
+        loaded_stacked_shards: set[tuple[str, str | int]] = set()
+        loaded_fused_params: set[str] = set()
 
         for name, loaded_weight in weights:
             if name.startswith("model."):
@@ -423,37 +444,66 @@ class Qwen3ForCausalLM(TextEncoder):
                     continue
                 name = kv_scale_name
 
-            for (
-                param_name,
-                weight_name,
-                shard_id,
-            ) in self.config.arch_config.stacked_params_mapping:
+            matched_stacked_param = False
+            for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:
                     continue
-                name = name.replace(weight_name, param_name)
+                matched_stacked_param = True
+                target_name = name.replace(weight_name, param_name)
 
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
+                if target_name.endswith(".bias") and target_name not in params_dict:
+                    break
 
-                if name not in params_dict:
-                    continue
+                if target_name not in params_dict:
+                    break
 
-                param = params_dict[name]
+                param = params_dict[target_name]
                 weight_loader = param.weight_loader
                 weight_loader(param, loaded_weight, shard_id)
+                loaded_params.add(target_name)
+                loaded_stacked_shards.add((target_name, shard_id))
                 break
+
+            if matched_stacked_param:
+                continue
+
+            if name.endswith(".bias") and name not in params_dict:
+                continue
+
+            if name not in params_dict:
+                continue
+
+            param = params_dict[name]
+            if name in fused_param_names and name.endswith(".scale_weight"):
+                # Merged scale loaders interpret a missing shard id as shard 0.
+                # An exact fused key is already a complete vector, so copy it
+                # atomically and retain the default loader's shape validation.
+                default_weight_loader(param, loaded_weight)
             else:
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-
-                if name not in params_dict:
-                    continue
-
-                param = params_dict[name]
                 weight_loader = getattr(param, "weight_loader", default_weight_loader)
                 weight_loader(param, loaded_weight)
-
             loaded_params.add(name)
+            if name in fused_param_names:
+                loaded_fused_params.add(name)
+
+        required_split_shards = {
+            (name, shard_id)
+            for name, shard_id in expected_stacked_shards
+            if name not in loaded_fused_params
+        }
+        missing_stacked_shards = required_split_shards - loaded_stacked_shards
+        if missing_stacked_shards:
+            formatted_missing = ", ".join(
+                f"{name}[{shard_id}]"
+                for name, shard_id in sorted(
+                    missing_stacked_shards,
+                    key=lambda item: (item[0], str(item[1])),
+                )
+            )
+            raise ValueError(
+                "Missing required stacked checkpoint shards: "
+                f"{formatted_missing}"
+            )
 
         return loaded_params
 

fastvideo/models/loader/component_loader.py

@@ -391,18 +391,22 @@ class TextEncoderLoader(ComponentLoader):
             model_config.quant_config = quant_cls()
 
         with set_default_torch_dtype(PRECISION_TO_TYPE[dtype]):
-            with target_device:
-                architectures = getattr(model_config, "architectures", [])
-                model_cls, _ = ModelRegistry.resolve_model_cls(architectures)
-                if getattr(model_cls, "supports_hf_from_pretrained", False):
-                    model = model_cls.from_pretrained_local(  # type: ignore[attr-defined]
-                        model_path,
-                        model_config,  # type: ignore[arg-type]
-                        dtype=PRECISION_TO_TYPE[dtype],
-                        device=target_device,
-                    )
-                    return model.eval()
+            architectures = getattr(model_config, "architectures", [])
+            model_cls, _ = ModelRegistry.resolve_model_cls(architectures)
+            if getattr(model_cls, "supports_hf_from_pretrained", False):
+                model = model_cls.from_pretrained_local(  # type: ignore[attr-defined]
+                    model_path,
+                    model_config,  # type: ignore[arg-type]
+                    dtype=PRECISION_TO_TYPE[dtype],
+                    device=target_device,
+                )
+                # HF passthrough encoders return before FastVideo's FSDP
+                # wrapping path, so the text stage needs their placement to
+                # put token tensors on the same device.
+                model._fastvideo_input_device = target_device
+                return model.eval()
 
+            with target_device:
                 model = model_cls(model_config)  # type: ignore
 
             weights_to_load = {name for name, _ in model.named_parameters()}
@@ -820,6 +824,24 @@ class VAELoader(ComponentLoader):
                 vae.load_state_dict(sd, strict=False)
                 return vae.eval()
 
+            if class_name == "LingBotWorld2WanVAE":
+                dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
+                config.pop("_class_name", None)
+                vae_config = fastvideo_args.pipeline_config.vae_config
+                vae_config.update_model_arch(config)
+                vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+                weight_path = os.path.join(model_path, "Wan2.1_VAE.pth")
+                if not os.path.exists(weight_path):
+                    raise FileNotFoundError(
+                        f"Missing LingBot World 2 VAE weights: {weight_path}"
+                    )
+                vae = vae_cls(
+                    vae_config,
+                    checkpoint_path=weight_path,
+                    dtype=dtype,
+                ).to(target_device)
+                return vae.eval()
+
             # LTX-2 uses CausalVideoAutoencoder with nested "vae" config
             if class_name == "CausalVideoAutoencoder" and "vae" in config:
                 vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)

fastvideo/models/loader/utils.py

@@ -18,8 +18,10 @@ def set_default_torch_dtype(dtype: torch.dtype):
     """Sets the default torch dtype to the given dtype."""
     old_dtype = torch.get_default_dtype()
     torch.set_default_dtype(dtype)
-    yield
-    torch.set_default_dtype(old_dtype)
+    try:
+        yield
+    finally:
+        torch.set_default_dtype(old_dtype)
 
 
 def get_param_names_mapping(

fastvideo/models/registry.py

@@ -42,6 +42,11 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
     "LTX2Transformer3DModel": ("dits", "ltx2", "LTX2Transformer3DModel"),
     "SD3Transformer2DModel": ("dits", "sd3", "SD3Transformer2DModel"),
     "LingBotWorldTransformer3DModel": ("dits", "lingbotworld", "LingBotWorldTransformer3DModel"),
+    "LingBotWorld2CausalFastTransformer3DModel": (
+        "dits",
+        "lingbotworld2",
+        "LingBotWorld2CausalFastTransformer3DModel",
+    ),
     "Gen3CTransformer3DModel": ("dits", "gen3c", "Gen3CTransformer3DModel"),
     "Kandinsky5Transformer3DModel": ("dits", "kandinsky5", "Kandinsky5Transformer3DModel"),
     "Flux2Transformer2DModel": ("dits", "flux_2", "Flux2Transformer2DModel"),
@@ -53,6 +58,11 @@ _IMAGE_TO_VIDEO_DIT_MODELS = {
     "DreamXWorldTransformer3DModel": ("dits", "dreamx_world", "DreamXWorldTransformer3DModel"),
     "DreamXWorldARTransformer3DModel": ("dits", "dreamx_world_ar", "DreamXWorldARTransformer3DModel"),
     "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
+    "LingBotWorld2CausalFastTransformer3DModel": (
+        "dits",
+        "lingbotworld2",
+        "LingBotWorld2CausalFastTransformer3DModel",
+    ),
     "MatrixGame2WanModel": ("dits", "matrixgame2", "MatrixGame2WanModel"),
     "CausalMatrixGame2WanModel": ("dits", "matrixgame2", "CausalMatrixGame2WanModel"),
     # Legacy aliases for older HF model_index.json files
@@ -64,6 +74,7 @@ _IMAGE_TO_VIDEO_DIT_MODELS = {
 # Text-to-image DiT models (2D image generation)
 _TEXT_TO_IMAGE_DIT_MODELS = {
     "GlmImageTransformer2DModel": ("dits", "glm_image", "GlmImageTransformer2DModel"),
+    "ZImageTransformer2DModel": ("dits", "zimage", "ZImageTransformer2DModel"),
 }
 
 _TEXT_ENCODER_MODELS = {
@@ -72,12 +83,16 @@ _TEXT_ENCODER_MODELS = {
     ("encoders", "clip", "CLIPTextModelWithProjection"),
     "LlamaModel": ("encoders", "llama", "LlamaModel"),
     "UMT5EncoderModel": ("encoders", "t5", "UMT5EncoderModel"),
+    "LingBotWorld2T5EncoderModel": ("encoders", "lingbotworld2_t5", "LingBotWorld2T5EncoderModel"),
     "T5EncoderModel": ("encoders", "t5_hf", "T5EncoderModel"),
     "BertModel": ("encoders", "clip", "CLIPTextModel"),
     "Qwen2_5_VLTextModel": ("encoders", "qwen2_5", "Qwen2_5_VLTextModel"),
     "Reason1TextEncoder": ("encoders", "reason1", "Reason1TextEncoder"),
     "Qwen2_5_VLForConditionalGeneration":
     ("encoders", "reason1", "Reason1TextEncoder"),
+    # Z-Image-Turbo's text_encoder/config.json declares architecture
+    # "Qwen3Model"; route it to the shared Qwen3 encoder (added for Flux2 Klein).
+    "Qwen3Model": ("encoders", "qwen3", "Qwen3ForCausalLM"),
     "LTX2GemmaTextEncoderModel": ("encoders", "gemma", "LTX2GemmaTextEncoderModel"),
     "Qwen3ForCausalLM": ("encoders", "qwen3", "Qwen3ForCausalLM"),
     "Mistral3ForConditionalGeneration":
@@ -98,6 +113,7 @@ _VAE_MODELS = {
     "AutoencoderKLHYWorld": ("vaes", "hyworldvae", "AutoencoderKLHYWorld"),
     "AutoencoderKLHunyuanVideo15": ("vaes", "hunyuan15vae", "AutoencoderKLHunyuanVideo15"),
     "AutoencoderKLWan": ("vaes", "wanvae", "AutoencoderKLWan"),
+    "LingBotWorld2WanVAE": ("vaes", "lingbotworld2_wanvae", "LingBotWorld2WanVAE"),
     "AutoencoderKL": ("vaes", "autoencoder_kl", "AutoencoderKL"),
     "AutoencoderKLGen3CTokenizer":
     ("vaes", "gen3c_tokenizer_vae", "AutoencoderKLGen3CTokenizer"),

fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py

@@ -96,6 +96,14 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
             The minimum sigma value for the noise schedule.
         sigma_data (`float`, *optional*):
             The sigma data value for scaling.
+        use_reference_discrete_timesteps (`bool`, defaults to False):
+            Some reference schedulers (e.g. Z-Image) construct the timestep
+            schedule by linspacing `num_inference_steps + 1` points from
+            `t_max` to `t_min` and dropping the terminal point. Default
+            (`False`) preserves the original `np.linspace(t_max, t_min,
+            num_inference_steps)` (float64) behaviour used by every existing
+            model. Enable this flag only when matching a reference scheduler
+            that expects the +1 + drop-terminal construction.
     """
 
     _compatibles: list[Any] = []
@@ -122,6 +130,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
         sigma_max: float | None = None,
         sigma_min: float | None = None,
         sigma_data: float | None = None,
+        use_reference_discrete_timesteps: bool = False,
     ):
         if sum([
                 self.config.use_beta_sigmas, self.config.use_exponential_sigmas,
@@ -155,7 +164,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
 
         self.sigmas = sigmas.to(
             "cpu")  # to avoid too much CPU/GPU communication
-        self.sigma_min = self.sigmas[-1].item()
+        self.sigma_min = sigma_min if sigma_min is not None else self.sigmas[-1].item()
         self.sigma_max = self.sigmas[0].item()
 
         BaseScheduler.__init__(self)
@@ -350,7 +359,19 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
             if timesteps_array is None:
                 t_max = self._sigma_to_t(self.sigma_max)
                 t_min = self._sigma_to_t(self.sigma_min)
-                timesteps_array = np.linspace(t_max, t_min, num_inference_steps)
+                if self.config.use_reference_discrete_timesteps:
+                    # Some reference schedulers (for example Z-Image) build a
+                    # float64 num_steps+1 linspace and drop the terminal point.
+                    timesteps_array = np.linspace(
+                        t_max,
+                        t_min,
+                        num_inference_steps + 1,
+                    )[:-1]
+                else:
+                    # Preserve the original numpy default (float64) here —
+                    # casting to float32 silently shifts rounded timestep
+                    # values for every existing model that uses this branch.
+                    timesteps_array = np.linspace(t_max, t_min, num_inference_steps)
             sigmas_array = timesteps_array / self.config.num_train_timesteps
         else:
             sigmas_array = np.array(sigmas).astype(np.float32)

fastvideo/models/vaes/lingbotworld2_wanvae.py

@@ -0,0 +1,722 @@
+import logging
+from types import SimpleNamespace
+
+import torch
+import torch.cuda.amp as amp
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+__all__ = [
+    'Wan2_1_VAE',
+    'LingBotWorld2WanVAE',
+]
+
+CACHE_T = 2
+
+
+class CausalConv3d(nn.Conv3d):
+    """
+    Causal 3d convolusion.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._padding = (self.padding[2], self.padding[2], self.padding[1],
+                         self.padding[1], 2 * self.padding[0], 0)
+        self.padding = (0, 0, 0)
+
+    def forward(self, x, cache_x=None):
+        padding = list(self._padding)
+        if cache_x is not None and self._padding[4] > 0:
+            cache_x = cache_x.to(x.device)
+            x = torch.cat([cache_x, x], dim=2)
+            padding[4] -= cache_x.shape[2]
+        x = F.pad(x, padding)
+
+        return super().forward(x)
+
+
+class RMS_norm(nn.Module):
+
+    def __init__(self, dim, channel_first=True, images=True, bias=False):
+        super().__init__()
+        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+
+        self.channel_first = channel_first
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(shape))
+        self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.
+
+    def forward(self, x):
+        return F.normalize(
+            x, dim=(1 if self.channel_first else
+                    -1)) * self.scale * self.gamma + self.bias
+
+
+class Upsample(nn.Upsample):
+
+    def forward(self, x):
+        """
+        Fix bfloat16 support for nearest neighbor interpolation.
+        """
+        return super().forward(x.float()).type_as(x)
+
+
+class Resample(nn.Module):
+
+    def __init__(self, dim, mode):
+        assert mode in ('none', 'upsample2d', 'upsample3d', 'downsample2d',
+                        'downsample3d')
+        super().__init__()
+        self.dim = dim
+        self.mode = mode
+
+        # layers
+        if mode == 'upsample2d':
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2., 2.), mode='nearest-exact'),
+                nn.Conv2d(dim, dim // 2, 3, padding=1))
+        elif mode == 'upsample3d':
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2., 2.), mode='nearest-exact'),
+                nn.Conv2d(dim, dim // 2, 3, padding=1))
+            self.time_conv = CausalConv3d(
+                dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
+
+        elif mode == 'downsample2d':
+            self.resample = nn.Sequential(
+                nn.ZeroPad2d((0, 1, 0, 1)),
+                nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+        elif mode == 'downsample3d':
+            self.resample = nn.Sequential(
+                nn.ZeroPad2d((0, 1, 0, 1)),
+                nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+            self.time_conv = CausalConv3d(
+                dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
+
+        else:
+            self.resample = nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        b, c, t, h, w = x.size()
+        if self.mode == 'upsample3d':
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = 'Rep'
+                    feat_idx[0] += 1
+                else:
+
+                    cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                    if cache_x.shape[2] < 2 and feat_cache[
+                            idx] is not None and feat_cache[idx] != 'Rep':
+                        # cache last frame of last two chunk
+                        cache_x = torch.cat([
+                            feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                                cache_x.device), cache_x
+                        ],
+                                            dim=2)
+                    if cache_x.shape[2] < 2 and feat_cache[
+                            idx] is not None and feat_cache[idx] == 'Rep':
+                        cache_x = torch.cat([
+                            torch.zeros_like(cache_x).to(cache_x.device),
+                            cache_x
+                        ],
+                                            dim=2)
+                    if feat_cache[idx] == 'Rep':
+                        x = self.time_conv(x)
+                    else:
+                        x = self.time_conv(x, feat_cache[idx])
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+
+                    x = x.reshape(b, 2, c, t, h, w)
+                    x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]),
+                                    3)
+                    x = x.reshape(b, c, t * 2, h, w)
+        t = x.shape[2]
+        x = rearrange(x, 'b c t h w -> (b t) c h w')
+        x = self.resample(x)
+        x = rearrange(x, '(b t) c h w -> b c t h w', t=t)
+
+        if self.mode == 'downsample3d':
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = x.clone()
+                    feat_idx[0] += 1
+                else:
+
+                    cache_x = x[:, :, -1:, :, :].clone()
+                    # if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx]!='Rep':
+                    #     # cache last frame of last two chunk
+                    #     cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+
+                    x = self.time_conv(
+                        torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+        return x
+
+    def init_weight(self, conv):
+        conv_weight = conv.weight
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        one_matrix = torch.eye(c1, c2)
+        init_matrix = one_matrix
+        nn.init.zeros_(conv_weight)
+        #conv_weight.data[:,:,-1,1,1] = init_matrix * 0.5
+        conv_weight.data[:, :, 1, 0, 0] = init_matrix  #* 0.5
+        conv.weight.data.copy_(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+    def init_weight2(self, conv):
+        conv_weight = conv.weight.data
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        init_matrix = torch.eye(c1 // 2, c2)
+        #init_matrix = repeat(init_matrix, 'o ... -> (o 2) ...').permute(1,0,2).contiguous().reshape(c1,c2)
+        conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix
+        conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix
+        conv.weight.data.copy_(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+
+class ResidualBlock(nn.Module):
+
+    def __init__(self, in_dim, out_dim, dropout=0.0):
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+
+        # layers
+        self.residual = nn.Sequential(
+            RMS_norm(in_dim, images=False), nn.SiLU(),
+            CausalConv3d(in_dim, out_dim, 3, padding=1),
+            RMS_norm(out_dim, images=False), nn.SiLU(), nn.Dropout(dropout),
+            CausalConv3d(out_dim, out_dim, 3, padding=1))
+        self.shortcut = CausalConv3d(in_dim, out_dim, 1) \
+            if in_dim != out_dim else nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        h = self.shortcut(x)
+        for layer in self.residual:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    # cache last frame of last two chunk
+                    cache_x = torch.cat([
+                        feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                            cache_x.device), cache_x
+                    ],
+                                        dim=2)
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x + h
+
+
+class AttentionBlock(nn.Module):
+    """
+    Causal self-attention with a single head.
+    """
+
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+        # layers
+        self.norm = RMS_norm(dim)
+        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
+        self.proj = nn.Conv2d(dim, dim, 1)
+
+        # zero out the last layer params
+        nn.init.zeros_(self.proj.weight)
+
+    def forward(self, x):
+        identity = x
+        b, c, t, h, w = x.size()
+        x = rearrange(x, 'b c t h w -> (b t) c h w')
+        x = self.norm(x)
+        # compute query, key, value
+        q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3,
+                                         -1).permute(0, 1, 3,
+                                                     2).contiguous().chunk(
+                                                         3, dim=-1)
+
+        # apply attention
+        x = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+        )
+        x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
+
+        # output
+        x = self.proj(x)
+        x = rearrange(x, '(b t) c h w-> b c t h w', t=t)
+        return x + identity
+
+
+class Encoder3d(nn.Module):
+
+    def __init__(self,
+                 dim=128,
+                 z_dim=4,
+                 dim_mult=[1, 2, 4, 4],
+                 num_res_blocks=2,
+                 attn_scales=[],
+                 temperal_downsample=[True, True, False],
+                 dropout=0.0):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+
+        # dimensions
+        dims = [dim * u for u in [1] + dim_mult]
+        scale = 1.0
+
+        # init block
+        self.conv1 = CausalConv3d(3, dims[0], 3, padding=1)
+
+        # downsample blocks
+        downsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            # residual (+attention) blocks
+            for _ in range(num_res_blocks):
+                downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+                if scale in attn_scales:
+                    downsamples.append(AttentionBlock(out_dim))
+                in_dim = out_dim
+
+            # downsample block
+            if i != len(dim_mult) - 1:
+                mode = 'downsample3d' if temperal_downsample[
+                    i] else 'downsample2d'
+                downsamples.append(Resample(out_dim, mode=mode))
+                scale /= 2.0
+        self.downsamples = nn.Sequential(*downsamples)
+
+        # middle blocks
+        self.middle = nn.Sequential(
+            ResidualBlock(out_dim, out_dim, dropout), AttentionBlock(out_dim),
+            ResidualBlock(out_dim, out_dim, dropout))
+
+        # output blocks
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False), nn.SiLU(),
+            CausalConv3d(out_dim, z_dim, 3, padding=1))
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                # cache last frame of last two chunk
+                cache_x = torch.cat([
+                    feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                        cache_x.device), cache_x
+                ],
+                                    dim=2)
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        ## downsamples
+        for layer in self.downsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## middle
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## head
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    # cache last frame of last two chunk
+                    cache_x = torch.cat([
+                        feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                            cache_x.device), cache_x
+                    ],
+                                        dim=2)
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x
+
+
+class Decoder3d(nn.Module):
+
+    def __init__(self,
+                 dim=128,
+                 z_dim=4,
+                 dim_mult=[1, 2, 4, 4],
+                 num_res_blocks=2,
+                 attn_scales=[],
+                 temperal_upsample=[False, True, True],
+                 dropout=0.0):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_upsample = temperal_upsample
+
+        # dimensions
+        dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
+        scale = 1.0 / 2**(len(dim_mult) - 2)
+
+        # init block
+        self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)
+
+        # middle blocks
+        self.middle = nn.Sequential(
+            ResidualBlock(dims[0], dims[0], dropout), AttentionBlock(dims[0]),
+            ResidualBlock(dims[0], dims[0], dropout))
+
+        # upsample blocks
+        upsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            # residual (+attention) blocks
+            if i == 1 or i == 2 or i == 3:
+                in_dim = in_dim // 2
+            for _ in range(num_res_blocks + 1):
+                upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+                if scale in attn_scales:
+                    upsamples.append(AttentionBlock(out_dim))
+                in_dim = out_dim
+
+            # upsample block
+            if i != len(dim_mult) - 1:
+                mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d'
+                upsamples.append(Resample(out_dim, mode=mode))
+                scale *= 2.0
+        self.upsamples = nn.Sequential(*upsamples)
+
+        # output blocks
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False), nn.SiLU(),
+            CausalConv3d(out_dim, 3, 3, padding=1))
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        ## conv1
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                # cache last frame of last two chunk
+                cache_x = torch.cat([
+                    feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                        cache_x.device), cache_x
+                ],
+                                    dim=2)
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        ## middle
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## upsamples
+        for layer in self.upsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## head
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    # cache last frame of last two chunk
+                    cache_x = torch.cat([
+                        feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                            cache_x.device), cache_x
+                    ],
+                                        dim=2)
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x
+
+
+def count_conv3d(model):
+    count = 0
+    for m in model.modules():
+        if isinstance(m, CausalConv3d):
+            count += 1
+    return count
+
+
+class WanVAE_(nn.Module):
+
+    def __init__(self,
+                 dim=128,
+                 z_dim=4,
+                 dim_mult=[1, 2, 4, 4],
+                 num_res_blocks=2,
+                 attn_scales=[],
+                 temperal_downsample=[True, True, False],
+                 dropout=0.0):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+        self.temperal_upsample = temperal_downsample[::-1]
+
+        # modules
+        self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks,
+                                 attn_scales, self.temperal_downsample, dropout)
+        self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
+        self.conv2 = CausalConv3d(z_dim, z_dim, 1)
+        self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks,
+                                 attn_scales, self.temperal_upsample, dropout)
+
+    def forward(self, x):
+        mu, log_var = self.encode(x)
+        z = self.reparameterize(mu, log_var)
+        x_recon = self.decode(z)
+        return x_recon, mu, log_var
+
+    def encode(self, x, scale):
+        self.clear_cache()
+        ## cache
+        t = x.shape[2]
+        iter_ = 1 + (t - 1) // 4
+        for i in range(iter_):
+            self._enc_conv_idx = [0]
+            if i == 0:
+                out = self.encoder(
+                    x[:, :, :1, :, :],
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx)
+            else:
+                out_ = self.encoder(
+                    x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :],
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx)
+                out = torch.cat([out, out_], 2)
+        mu, log_var = self.conv1(out).chunk(2, dim=1)
+        if isinstance(scale[0], torch.Tensor):
+            mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(
+                1, self.z_dim, 1, 1, 1)
+        else:
+            mu = (mu - scale[0]) * scale[1]
+        self.clear_cache()
+        return mu
+
+    def decode(self, z, scale):
+        self.clear_cache()
+        # z: [b,c,t,h,w]
+        if isinstance(scale[0], torch.Tensor):
+            z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(
+                1, self.z_dim, 1, 1, 1)
+        else:
+            z = z / scale[1] + scale[0]
+        iter_ = z.shape[2]
+        x = self.conv2(z)
+        for i in range(iter_):
+            self._conv_idx = [0]
+            if i == 0:
+                out = self.decoder(
+                    x[:, :, i:i + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx)
+            else:
+                out_ = self.decoder(
+                    x[:, :, i:i + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx)
+                out = torch.cat([out, out_], 2)
+        self.clear_cache()
+        return out
+
+    def reparameterize(self, mu, log_var):
+        std = torch.exp(0.5 * log_var)
+        eps = torch.randn_like(std)
+        return eps * std + mu
+
+    def sample(self, imgs, deterministic=False):
+        mu, log_var = self.encode(imgs)
+        if deterministic:
+            return mu
+        std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0))
+        return mu + std * torch.randn_like(std)
+
+    def clear_cache(self):
+        self._conv_num = count_conv3d(self.decoder)
+        self._conv_idx = [0]
+        self._feat_map = [None] * self._conv_num
+        #cache encode
+        self._enc_conv_num = count_conv3d(self.encoder)
+        self._enc_conv_idx = [0]
+        self._enc_feat_map = [None] * self._enc_conv_num
+
+
+def _video_vae(pretrained_path=None, z_dim=None, device='cpu', **kwargs):
+    """
+    Autoencoder3d adapted from Stable Diffusion 1.x, 2.x and XL.
+    """
+    # params
+    cfg = dict(
+        dim=96,
+        z_dim=z_dim,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[False, True, True],
+        dropout=0.0)
+    cfg.update(**kwargs)
+
+    # init model
+    with torch.device('meta'):
+        model = WanVAE_(**cfg)
+
+    # load checkpoint
+    logging.info(f'loading {pretrained_path}')
+    model.load_state_dict(
+        torch.load(pretrained_path, map_location=device), assign=True)
+
+    return model
+
+
+class Wan2_1_VAE:
+
+    def __init__(self,
+                 z_dim=16,
+                 vae_pth='cache/vae_step_411000.pth',
+                 dtype=torch.float,
+                 device="cuda"):
+        self.dtype = dtype
+        self.device = device
+
+        mean = [
+            -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508,
+            0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921
+        ]
+        std = [
+            2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743,
+            3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160
+        ]
+        self.mean = torch.tensor(mean, dtype=dtype, device=device)
+        self.std = torch.tensor(std, dtype=dtype, device=device)
+        self.scale = [self.mean, 1.0 / self.std]
+
+        # init model
+        self.model = _video_vae(
+            pretrained_path=vae_pth,
+            z_dim=z_dim,
+        ).eval().requires_grad_(False).to(device)
+
+    def encode(self, videos):
+        """
+        videos: A list of videos each with shape [C, T, H, W].
+        """
+        with amp.autocast(dtype=self.dtype):
+            return [
+                self.model.encode(u.unsqueeze(0), self.scale).float().squeeze(0)
+                for u in videos
+            ]
+
+    def decode(self, zs):
+        with amp.autocast(dtype=self.dtype):
+            return [
+                self.model.decode(u.unsqueeze(0),
+                                  self.scale).float().clamp_(-1, 1).squeeze(0)
+                for u in zs
+            ]
+
+
+class LingBotWorld2WanVAE(nn.Module):
+    """FastVideo-facing wrapper around the exact LingBot World 2 Wan2.1 VAE computation."""
+
+    handles_latent_denorm = True
+
+    def __init__(self, config, checkpoint_path=None, dtype=torch.float):
+        """Load the official LingBot World 2 VAE weights and expose FastVideo VAE APIs."""
+        super().__init__()
+        self.config = config
+        self.dtype = dtype
+        z_dim = int(getattr(config, "z_dim", 16))
+        mean = torch.tensor(getattr(config, "latents_mean"), dtype=dtype)
+        std = torch.tensor(getattr(config, "latents_std"), dtype=dtype)
+        self.register_buffer("shift_factor", mean.view(1, z_dim, 1, 1, 1), persistent=False)
+        self.register_buffer("scaling_factor", (1.0 / std).view(1, z_dim, 1, 1, 1), persistent=False)
+        self.scale = [mean, 1.0 / std]
+
+        if checkpoint_path is None:
+            self.model = WanVAE_(dim=96, z_dim=z_dim, dim_mult=[1, 2, 4, 4],
+                                 num_res_blocks=2, attn_scales=[],
+                                 temperal_downsample=[False, True, True],
+                                 dropout=0.0)
+        else:
+            self.model = _video_vae(pretrained_path=checkpoint_path, z_dim=z_dim)
+        self.model.eval().requires_grad_(False)
+
+    def _scale_for(self, device: torch.device) -> list[torch.Tensor]:
+        """Return source VAE scale tensors on the active device."""
+        return [u.to(device) for u in self.scale]
+
+    def encode(self, videos: torch.Tensor):
+        """Encode `[B,C,T,H,W]` videos and return a FastVideo-style mean tensor."""
+        if videos.ndim != 5:
+            raise ValueError(f"LingBotWorld2WanVAE.encode expects 5D input, got {videos.shape}")
+        scale = self._scale_for(videos.device)
+        latents = []
+        with amp.autocast(dtype=self.dtype):
+            for video in videos:
+                latent = self.model.encode(video.unsqueeze(0), scale).float().squeeze(0)
+                latents.append(latent)
+        normalized = torch.stack(latents, dim=0)
+        return SimpleNamespace(mean=normalized)
+
+    def decode(self, latents: torch.Tensor) -> torch.Tensor:
+        """Decode normalized LingBot World 2 latents to clamped `[-1,1]` video tensors."""
+        if latents.ndim != 5:
+            raise ValueError(f"LingBotWorld2WanVAE.decode expects 5D input, got {latents.shape}")
+        scale = self._scale_for(latents.device)
+        videos = []
+        with amp.autocast(dtype=self.dtype):
+            for latent in latents:
+                video = self.model.decode(latent.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0)
+                videos.append(video)
+        return torch.stack(videos, dim=0)
+
+
+EntryClass = LingBotWorld2WanVAE

fastvideo/pipelines/__init__.py

@@ -35,7 +35,7 @@ def build_pipeline(fastvideo_args: FastVideoArgs,
     """
     # Get pipeline type
     model_path = fastvideo_args.model_path
-    model_path = maybe_download_model(model_path)
+    model_path = maybe_download_model(model_path, revision=fastvideo_args.revision)
     # fastvideo_args.downloaded_model_path = model_path
     logger.info("Model path: %s", model_path)
 

fastvideo/pipelines/basic/lingbotworld2/__init__.py

@@ -0,0 +1,5 @@
+from .causal_fast_pipeline import LingBotWorld2CausalFastPipeline
+
+__all__ = ["LingBotWorld2CausalFastPipeline"]
+
+EntryClass = LingBotWorld2CausalFastPipeline

fastvideo/pipelines/basic/lingbotworld2/causal_fast_pipeline.py

@@ -0,0 +1,365 @@
+# SPDX-License-Identifier: Apache-2.0
+"""LingBot World 2 causal-fast image-to-video pipeline."""
+
+import math
+import os
+
+from einops import rearrange
+import numpy as np
+import torch
+import torchvision.transforms.functional as TF
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.distributed.parallel_state import get_sp_world_size
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.dits.lingbotworld2.cam_utils import (
+    compute_relative_poses,
+    get_Ks_transformed,
+    get_plucker_embeddings,
+    interpolate_camera_poses,
+)
+from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import (
+    FlowUniPCMultistepScheduler, )
+from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages import (
+    ConditioningStage,
+    DecodingStage,
+    InputValidationStage,
+    TextEncodingStage,
+)
+from fastvideo.pipelines.stages.base import PipelineStage
+
+logger = init_logger(__name__)
+
+
+class LingBotWorld2TextEncodingStage(TextEncodingStage):
+    """Keep LingBot World 2 T5 attention masks so DiT context matches the source pipeline."""
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        """Initialize mask storage before running the shared text-encoding stage."""
+        if batch.prompt_attention_mask is None:
+            batch.prompt_attention_mask = []
+        return super().forward(batch, fastvideo_args)
+
+
+class LingBotWorld2CausalFastGenerationStage(PipelineStage):
+    """Prepare LingBot World 2 conditions and run the released causal-fast sampling loop."""
+
+    def __init__(self, transformer, scheduler, vae) -> None:
+        super().__init__()
+        self.transformer = transformer
+        self.scheduler = scheduler
+        self.vae = vae
+        self._cross_attn_initialized = False
+
+    def _convert_flow_pred_to_x0(
+        self,
+        flow_pred: torch.Tensor,
+        xt: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        """Convert LingBot World 2 flow prediction to x0 using the scheduler sigma."""
+        original_dtype = flow_pred.dtype
+        flow_pred, xt, sigmas, timesteps = map(
+            lambda x: x.double().to(flow_pred.device),
+            [flow_pred, xt, self.scheduler.sigmas, self.scheduler.timesteps],
+        )
+        timestep_id = torch.argmin((timesteps - timestep).abs())
+        sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        return (xt - sigma_t * flow_pred).to(original_dtype)
+
+    def _initialize_self_kv_cache(
+        self,
+        batch_size: int,
+        kv_size: int,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> list[dict]:
+        """Allocate per-block self-attention KV cache tensors."""
+        head_dim = self.transformer.dim // self.transformer.num_heads
+        num_heads = self.transformer.num_heads // get_sp_world_size()
+        shape = [batch_size, kv_size, num_heads, head_dim]
+        return [{
+            "k": torch.zeros(shape, dtype=dtype, device=device),
+            "v": torch.zeros(shape, dtype=dtype, device=device),
+            "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+            "local_end_index": torch.tensor([0], dtype=torch.long, device=device),
+        } for _ in range(self.transformer.num_layers)]
+
+    def _initialize_crossattn_cache(
+        self,
+        batch_size: int,
+        max_sequence_length: int,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> list[dict]:
+        """Allocate per-block text cross-attention KV cache tensors."""
+        head_dim = self.transformer.dim // self.transformer.num_heads
+        shape = [batch_size, max_sequence_length, self.transformer.num_heads, head_dim]
+        return [{
+            "k": torch.zeros(shape, dtype=dtype, device=device),
+            "v": torch.zeros(shape, dtype=dtype, device=device),
+            "is_init": torch.tensor([0], dtype=torch.bool, device=device),
+        } for _ in range(self.transformer.num_layers)]
+
+    @staticmethod
+    def _prompt_context(batch: ForwardBatch, device: torch.device) -> list[torch.Tensor]:
+        """Slice padded text encoder states back to LingBot World 2's unpadded context list."""
+        assert batch.prompt_embeds
+        context_tensor = batch.prompt_embeds[0].to(device)
+        if batch.prompt_attention_mask:
+            mask = batch.prompt_attention_mask[0].to(device)
+            seq_lens = mask.gt(0).sum(dim=1).long()
+            return [u[:v] for u, v in zip(context_tensor, seq_lens, strict=True)]
+        return [u for u in context_tensor]
+
+    def _prepare_image_tensor(self, batch: ForwardBatch, device: torch.device) -> torch.Tensor:
+        """Return the source-style normalized image tensor `[C,H,W]`."""
+        image = batch.pil_image
+        if image is None:
+            raise ValueError("LingBot World 2 causal-fast requires `image_path` or `pil_image`.")
+        if isinstance(image, torch.Tensor):
+            if image.ndim == 5:
+                return image[0, :, 0].to(device)
+            if image.ndim == 4:
+                return image[0].to(device)
+            return image.to(device)
+        return TF.to_tensor(image).sub_(0.5).div_(0.5).to(device)
+
+    def _prepare_camera(
+        self,
+        action_path: str,
+        c2ws: np.ndarray,
+        h: int,
+        w: int,
+        lat_f: int,
+        lat_h: int,
+        lat_w: int,
+        chunk_size: int,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> torch.Tensor:
+        """Build the LingBot World 2 camera Plucker tensor for latent chunks."""
+        Ks = torch.from_numpy(np.load(os.path.join(action_path, "intrinsics.npy"))).float()
+        Ks = get_Ks_transformed(
+            Ks,
+            height_org=480,
+            width_org=832,
+            height_resize=h,
+            width_resize=w,
+            height_final=h,
+            width_final=w,
+        )
+        Ks = Ks[0]
+        len_c2ws = len(c2ws)
+        len_c2ws_ = int((len_c2ws - 1) // 4) + 1
+        len_c2ws_ = int(len_c2ws_ - (len_c2ws_ % chunk_size))
+        c2ws_infer = interpolate_camera_poses(
+            src_indices=np.linspace(0, len_c2ws - 1, len_c2ws),
+            src_rot_mat=c2ws[:, :3, :3],
+            src_trans_vec=c2ws[:, :3, 3],
+            tgt_indices=np.linspace(0, len_c2ws - 1, len_c2ws_),
+        )
+        c2ws_infer = compute_relative_poses(c2ws_infer, framewise=True)
+        Ks = Ks.repeat(len(c2ws_infer), 1)
+        c2ws_plucker_emb = get_plucker_embeddings(c2ws_infer.to(device), Ks.to(device), h, w)
+        c2ws_plucker_emb = rearrange(
+            c2ws_plucker_emb,
+            "f (h c1) (w c2) c -> (f h w) (c c1 c2)",
+            c1=int(h // lat_h),
+            c2=int(w // lat_w),
+        )
+        c2ws_plucker_emb = c2ws_plucker_emb[None, ...]
+        return rearrange(
+            c2ws_plucker_emb,
+            "b (f h w) c -> b c f h w",
+            f=lat_f,
+            h=lat_h,
+            w=lat_w,
+        ).to(device=device, dtype=dtype)
+
+    def _encode_condition_video(
+        self,
+        img: torch.Tensor,
+        h: int,
+        w: int,
+        frames: int,
+        mask: torch.Tensor,
+        fastvideo_args: FastVideoArgs,
+    ) -> torch.Tensor:
+        """Encode the first-frame conditioning video and prepend mask channels."""
+        device = get_local_torch_device()
+        self.vae = self.vae.to(device)
+        video_condition = torch.concat(
+            [
+                torch.nn.functional.interpolate(
+                    img[None].cpu(),
+                    size=(h, w),
+                    mode="bicubic",
+                ).transpose(0, 1),
+                torch.zeros(3, frames - 1, h, w),
+            ],
+            dim=1,
+        ).to(device)
+        vae_dtype = torch.float32
+        with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=False):
+            encoder_output = self.vae.encode(video_condition.unsqueeze(0).to(torch.float32))
+        latent_condition = encoder_output.mean
+        if not bool(getattr(self.vae, "handles_latent_denorm", False)):
+            if hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None:
+                latent_condition -= self.vae.shift_factor.to(latent_condition.device, latent_condition.dtype)
+            latent_condition = latent_condition * self.vae.scaling_factor.to(latent_condition.device,
+                                                                             latent_condition.dtype)
+        if fastvideo_args.vae_cpu_offload:
+            self.vae.to("cpu")
+        return torch.concat([mask, latent_condition[0]], dim=0)
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        """Execute LingBot World 2 causal-fast generation and store final latents on the batch."""
+        device = get_local_torch_device()
+        cfg = fastvideo_args.pipeline_config.dit_config.arch_config
+        chunk_size = int(cfg.chunk_size)
+        max_sequence_length = int(batch.max_sequence_length or cfg.text_len)
+        action_path = batch.action_path
+        if action_path is None:
+            raise ValueError("LingBot World 2 causal-fast requires `action_path`.")
+
+        c2ws = np.load(os.path.join(action_path, "poses.npy"))
+        len_c2ws = ((len(c2ws) - 1) // 4) * 4 + 1
+        frame_num = ((int(batch.num_frames) - 1) // 4) * 4 + 1
+        frame_num = min(frame_num, len_c2ws)
+        c2ws = c2ws[:frame_num]
+
+        img = self._prepare_image_tensor(batch, device)
+        h0, w0 = img.shape[1:]
+        aspect_ratio = h0 / w0
+        lat_h = round(np.sqrt(cfg.max_area * aspect_ratio) // 8 // cfg.patch_size[1] * cfg.patch_size[1])
+        lat_w = round(np.sqrt(cfg.max_area / aspect_ratio) // 8 // cfg.patch_size[2] * cfg.patch_size[2])
+        h = lat_h * 8
+        w = lat_w * 8
+        lat_f = (frame_num - 1) // 4 + 1
+        lat_f = int(lat_f - (lat_f % chunk_size))
+        frames = (lat_f - 1) * 4 + 1
+        batch.height = h
+        batch.width = w
+        batch.num_frames = frames
+
+        seed = int(batch.seed if batch.seed is not None else 42)
+        seed_g = torch.Generator(device=device)
+        seed_g.manual_seed(seed)
+        noise = torch.randn(16, lat_f, lat_h, lat_w, dtype=torch.float32, generator=seed_g, device=device)
+
+        mask = torch.ones(1, frames, lat_h, lat_w, device=device)
+        mask[:, 1:] = 0
+        mask = torch.concat([torch.repeat_interleave(mask[:, 0:1], repeats=4, dim=1), mask[:, 1:]], dim=1)
+        mask = mask.view(1, mask.shape[1] // 4, 4, lat_h, lat_w).transpose(1, 2)[0]
+
+        self.scheduler.set_timesteps(cfg.num_train_timesteps, shift=cfg.sample_shift)
+        timesteps = self.scheduler.timesteps[list(cfg.timesteps_index)].to(device)
+        context = self._prompt_context(batch, device)
+        c2ws_plucker_emb = self._prepare_camera(
+            action_path,
+            c2ws,
+            h,
+            w,
+            lat_f,
+            lat_h,
+            lat_w,
+            chunk_size,
+            torch.bfloat16,
+            device,
+        )
+        y = self._encode_condition_video(img, h, w, frames, mask, fastvideo_args).to(device=device,
+                                                                                     dtype=torch.bfloat16)
+
+        transformer_dtype = torch.bfloat16
+        frame_seqlen = int(noise.shape[-2] * noise.shape[-1] // 4)
+        kv_size = frame_seqlen * cfg.local_attn_size if cfg.local_attn_size > -1 else frame_seqlen * lat_f
+        self_kv_cache = self._initialize_self_kv_cache(1, kv_size, transformer_dtype, device)
+        cross_kv_cache = self._initialize_crossattn_cache(1, max_sequence_length, transformer_dtype, device)
+
+        self.transformer = self.transformer.to(device)
+        self._cross_attn_initialized = False
+        pred_latent_chunks = []
+        latents_chunk = noise.split(chunk_size, dim=1)
+        condition_chunk = y.split(chunk_size, dim=1)
+        c2ws_plucker_emb_chunk = c2ws_plucker_emb.split(chunk_size, dim=2)
+        max_seq_len = int(math.ceil(chunk_size * lat_h * lat_w // 4))
+
+        with torch.amp.autocast("cuda", dtype=transformer_dtype):
+            for chunk_id, current_latent in enumerate(latents_chunk):
+                current_condition = condition_chunk[chunk_id]
+                current_c2ws_plucker_emb = c2ws_plucker_emb_chunk[chunk_id]
+                dit_cond_dict = {"c2ws_plucker_emb": current_c2ws_plucker_emb.chunk(1, dim=0)}
+                kwargs = {
+                    "context": [context[0]],
+                    "seq_len": max_seq_len,
+                    "y": [current_condition],
+                    "dit_cond_dict": dit_cond_dict,
+                    "kv_cache": self_kv_cache,
+                    "crossattn_cache": cross_kv_cache,
+                    "current_start": chunk_id * chunk_size * frame_seqlen,
+                    "max_attention_size": kv_size,
+                    "frame_seqlen": frame_seqlen,
+                }
+                x0 = current_latent
+                for timestep_idx, timestep_value in enumerate(timesteps):
+                    timestep = torch.stack([timestep_value]).to(device)
+                    noise_pred = self.transformer(
+                        x=[current_latent.to(device)],
+                        t=timestep,
+                        cross_attn_first_call=not self._cross_attn_initialized,
+                        **kwargs,
+                    )[0]
+                    self._cross_attn_initialized = True
+                    x0 = self._convert_flow_pred_to_x0(noise_pred, current_latent, timestep_value)
+                    if timestep_idx < len(timesteps) - 1:
+                        next_timestep = timesteps[timestep_idx + 1].reshape(1)
+                        current_latent = self.scheduler.add_noise(
+                            x0,
+                            torch.randn(x0.shape, generator=seed_g, device=x0.device, dtype=x0.dtype),
+                            next_timestep,
+                        )
+                pred_latent_chunks.append(x0)
+                context_timestep = torch.stack([timesteps[-1] * 0.0]).to(device)
+                self.transformer(x=[x0], t=context_timestep, cross_attn_first_call=False, **kwargs)
+
+        batch.latents = torch.cat(pred_latent_chunks, dim=1).unsqueeze(0)
+        return batch
+
+
+class LingBotWorld2CausalFastPipeline(LoRAPipeline, ComposedPipelineBase):
+    """FastVideo pipeline for LingBot World 2 14B causal-fast I2V generation."""
+
+    _required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "scheduler"]
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        if "scheduler" not in self.modules:
+            self.modules["scheduler"] = FlowUniPCMultistepScheduler(shift=1.0, use_dynamic_shifting=False)
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        """Set up the LingBot World 2 causal-fast pipeline stages."""
+        self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
+        self.add_stage(
+            stage_name="prompt_encoding_stage",
+            stage=LingBotWorld2TextEncodingStage(
+                text_encoders=[self.get_module("text_encoder")],
+                tokenizers=[self.get_module("tokenizer")],
+            ),
+        )
+        self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
+        self.add_stage(
+            stage_name="lingbotworld2_causal_fast_generation_stage",
+            stage=LingBotWorld2CausalFastGenerationStage(
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler"),
+                vae=self.get_module("vae"),
+            ),
+        )
+        self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")))
+
+
+EntryClass = LingBotWorld2CausalFastPipeline

fastvideo/pipelines/basic/lingbotworld2/presets.py

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+"""LingBotWorld2 causal-fast pipeline preset."""
+
+from fastvideo.api.presets import InferencePreset, PresetStageSpec
+
+_DENOISE_STAGE = PresetStageSpec(
+    name="denoise",
+    kind="denoising",
+    description="Causal-fast denoising pass",
+    allowed_overrides=frozenset({
+        "num_inference_steps",
+        "guidance_scale",
+    }),
+)
+
+LINGBOTWORLD2_CAUSAL_FAST_I2V = InferencePreset(
+    name="lingbotworld2_causal_fast_i2v",
+    version=1,
+    model_family="lingbotworld2",
+    description="LingBot World 2 14B causal-fast I2V",
+    workload_type="i2v",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults={
+        "guidance_scale": 1.0,
+        "num_inference_steps": 4,
+        "fps": 16,
+        "seed": 42,
+        "num_frames": 65,
+        "height": 480,
+        "width": 832,
+        "negative_prompt": "",
+    },
+)
+
+ALL_PRESETS = (LINGBOTWORLD2_CAUSAL_FAST_I2V, )

fastvideo/pipelines/basic/zimage/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Z-Image pipeline."""
+
+from fastvideo.pipelines.basic.zimage.zimage_pipeline import ZImagePipeline
+
+__all__ = ["ZImagePipeline"]

fastvideo/pipelines/basic/zimage/presets.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Z-Image inference presets."""
+
+from fastvideo.api.presets import InferencePreset, PresetStageSpec
+
+_DENOISE_STAGE = PresetStageSpec(
+    name="denoise",
+    kind="denoising",
+    description="Z-Image denoising pass",
+    allowed_overrides=frozenset({
+        "num_inference_steps",
+        "guidance_scale",
+        "cfg_normalization",
+        "cfg_truncation",
+    }),
+)
+
+ZIMAGE_TURBO = InferencePreset(
+    name="zimage_turbo",
+    version=1,
+    model_family="zimage",
+    description="Z-Image-Turbo text-to-image generation",
+    workload_type="t2i",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults={
+        "height": 1024,
+        "width": 1024,
+        "num_frames": 1,
+        "fps": 1,
+        "seed": 42,
+        "guidance_scale": 0.0,
+        "num_inference_steps": 8,
+        "negative_prompt": "",
+        "max_sequence_length": 512,
+        "cfg_normalization": False,
+        "cfg_truncation": 1.0,
+    },
+)
+
+ALL_PRESETS = (ZIMAGE_TURBO, )

fastvideo/pipelines/basic/zimage/stages.py

@@ -0,0 +1,337 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pipeline stages for the native Z-Image text-to-image path."""
+
+from __future__ import annotations
+
+import inspect
+
+import torch
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.hooks.activation_trace import trace_step
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+from fastvideo.pipelines.stages.input_validation import InputValidationStage
+from fastvideo.utils import PRECISION_TO_TYPE
+
+
+class ZImageInputValidationStage(InputValidationStage):
+    """Validate the image geometry and reproduce the official device RNG."""
+
+    def _generate_seeds(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> None:
+        del fastvideo_args
+        assert batch.seed is not None
+        batch.seeds = [batch.seed]
+        device = get_local_torch_device()
+        batch.generator = torch.Generator(device=device).manual_seed(batch.seed)
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if batch.do_classifier_free_guidance and batch.negative_prompt is None and not batch.negative_prompt_embeds:
+            batch.negative_prompt = ""
+        batch = super().forward(batch, fastvideo_args)
+        if batch.num_frames != 1:
+            raise ValueError(f"Z-Image is text-to-image and requires num_frames=1, got {batch.num_frames}")
+        if batch.height is None or batch.width is None:
+            raise ValueError("Z-Image requires height and width")
+        if batch.height % 16 or batch.width % 16:
+            raise ValueError("Z-Image height and width must be divisible by 16; "
+                             f"got {batch.height}x{batch.width}")
+        return batch
+
+
+class ZImageConditioningStage(PipelineStage):
+    """Trim padded Qwen states and materialize variable-length CFG streams."""
+
+    @staticmethod
+    def _trim_embeddings(
+        embeds: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+    ) -> list[torch.Tensor]:
+        if attention_mask is None:
+            return list(embeds.unbind(0))
+        return [
+            sample[mask.to(device=sample.device, dtype=torch.bool)]
+            for sample, mask in zip(embeds, attention_mask, strict=True)
+        ]
+
+    @staticmethod
+    def _repeat(items: list[torch.Tensor], count: int) -> list[torch.Tensor]:
+        return [item for item in items for _ in range(count)]
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        del fastvideo_args
+        if len(batch.prompt_embeds) != 1:
+            raise ValueError(f"Z-Image expects one text encoder, got {len(batch.prompt_embeds)}")
+
+        prompt_mask = batch.prompt_attention_mask[0] if batch.prompt_attention_mask else None
+        prompt_embeds = self._trim_embeddings(batch.prompt_embeds[0], prompt_mask)
+        batch.extra["zimage_prompt_embeds"] = self._repeat(prompt_embeds, batch.num_videos_per_prompt)
+
+        if batch.do_classifier_free_guidance:
+            if not batch.negative_prompt_embeds:
+                raise ValueError("Z-Image CFG requires negative prompt embeddings")
+            negative_mask = batch.negative_attention_mask[0] if batch.negative_attention_mask else None
+            negative_embeds = self._trim_embeddings(batch.negative_prompt_embeds[0], negative_mask)
+            batch.extra["zimage_negative_prompt_embeds"] = self._repeat(
+                negative_embeds,
+                batch.num_videos_per_prompt,
+            )
+        else:
+            batch.extra["zimage_negative_prompt_embeds"] = []
+        return batch
+
+
+class ZImageLatentPreparationStage(PipelineStage):
+    """Create the official fp32 image latents on the transformer device."""
+
+    def __init__(self, transformer) -> None:
+        self.transformer = transformer
+
+    @staticmethod
+    def _randn(
+        shape: tuple[int, ...],
+        generators: torch.Generator | list[torch.Generator] | None,
+        device: torch.device,
+    ) -> torch.Tensor:
+        if isinstance(generators, list):
+            if len(generators) != shape[0]:
+                raise ValueError(f"generator list length {len(generators)} does not match batch size {shape[0]}")
+            sample_shape = (1, *shape[1:])
+            return torch.cat([
+                torch.randn(sample_shape, generator=generator, device=device, dtype=torch.float32)
+                for generator in generators
+            ])
+        return torch.randn(shape, generator=generators, device=device, dtype=torch.float32)
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if batch.height is None or batch.width is None:
+            raise ValueError("Z-Image requires height and width before latent preparation")
+
+        prompt_embeds = batch.extra.get("zimage_prompt_embeds")
+        if not isinstance(prompt_embeds, list) or not prompt_embeds:
+            raise ValueError("Z-Image conditioning must run before latent preparation")
+
+        channels = int(getattr(self.transformer, "in_channels", 16))
+        spatial_ratio = int(fastvideo_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio)
+        shape = (
+            len(prompt_embeds),
+            channels,
+            1,
+            batch.height // spatial_ratio,
+            batch.width // spatial_ratio,
+        )
+        device = get_local_torch_device()
+        if batch.latents is None:
+            latents = self._randn(shape, batch.generator, device)
+        else:
+            latents = batch.latents
+            if latents.ndim == 4:
+                latents = latents.unsqueeze(2)
+            if tuple(latents.shape) != shape:
+                raise ValueError(f"Expected Z-Image latents with shape {shape}, got {tuple(latents.shape)}")
+            latents = latents.to(device=device, dtype=torch.float32)
+
+        batch.latents = latents
+        batch.raw_latent_shape = shape
+        return batch
+
+
+class ZImageTimestepPreparationStage(PipelineStage):
+    """Apply the native scheduler's zero endpoint and discrete schedule."""
+
+    def __init__(self, scheduler) -> None:
+        self.scheduler = scheduler
+
+    @staticmethod
+    def _calculate_shift(
+        image_seq_len: int,
+        base_seq_len: int,
+        max_seq_len: int,
+        base_shift: float,
+        max_shift: float,
+    ) -> float:
+        slope = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+        return image_seq_len * slope + base_shift - slope * base_seq_len
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if batch.latents is None:
+            raise ValueError("Z-Image latents must be prepared before timesteps")
+        if batch.timesteps is not None and batch.sigmas is not None:
+            raise ValueError("Only one of timesteps or sigmas may be supplied")
+
+        scheduler = self.scheduler
+        sigma_min = float(fastvideo_args.pipeline_config.scheduler_sigma_min)
+        use_reference_timesteps = bool(fastvideo_args.pipeline_config.scheduler_use_reference_discrete_timesteps)
+        scheduler.sigma_min = sigma_min
+        scheduler.register_to_config(
+            sigma_min=sigma_min,
+            use_reference_discrete_timesteps=use_reference_timesteps,
+        )
+        config = scheduler.config
+        image_seq_len = (batch.latents.shape[-2] // 2) * (batch.latents.shape[-1] // 2)
+        mu = self._calculate_shift(
+            image_seq_len,
+            int(config.get("base_image_seq_len", 256)),
+            int(config.get("max_image_seq_len", 4096)),
+            float(config.get("base_shift", 0.5)),
+            float(config.get("max_shift", 1.15)),
+        )
+        device = get_local_torch_device()
+
+        if batch.timesteps is not None:
+            if "timesteps" not in inspect.signature(scheduler.set_timesteps).parameters:
+                raise ValueError(f"{type(scheduler).__name__} does not accept custom timesteps")
+            scheduler.set_timesteps(timesteps=batch.timesteps, device=device, mu=mu)
+        elif batch.sigmas is not None:
+            if "sigmas" not in inspect.signature(scheduler.set_timesteps).parameters:
+                raise ValueError(f"{type(scheduler).__name__} does not accept custom sigmas")
+            scheduler.set_timesteps(sigmas=batch.sigmas, device=device, mu=mu)
+        else:
+            scheduler.set_timesteps(batch.num_inference_steps, device=device, mu=mu)
+
+        batch.timesteps = scheduler.timesteps
+        batch.num_inference_steps = len(batch.timesteps)
+        return batch
+
+
+class ZImageDenoisingStage(PipelineStage):
+    """Run the native Z-Image flow-matching loop."""
+
+    performance_component_metric = "transformer_time_s"
+
+    def __init__(self, transformer, scheduler) -> None:
+        self.transformer = transformer
+        self.scheduler = scheduler
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if batch.latents is None or batch.timesteps is None:
+            raise ValueError("Z-Image denoising requires latents and timesteps")
+
+        latents = batch.latents.float()
+        positive = batch.extra.get("zimage_prompt_embeds")
+        negative = batch.extra.get("zimage_negative_prompt_embeds", [])
+        if not isinstance(positive, list) or not positive:
+            raise ValueError("Z-Image denoising requires prompt embeddings")
+
+        device = get_local_torch_device()
+        target_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
+        batch_size = latents.shape[0]
+
+        for index, timestep_value in enumerate(batch.timesteps):
+            if timestep_value.item() == 0 and index == len(batch.timesteps) - 1:
+                continue
+
+            timestep = timestep_value.expand(batch_size).to(device=device, dtype=torch.float32)
+            timestep = (1000.0 - timestep) / 1000.0
+            current_guidance_scale = float(batch.guidance_scale)
+            if (batch.do_classifier_free_guidance and batch.cfg_truncation is not None and batch.cfg_truncation <= 1.0
+                    and timestep[0].item() > batch.cfg_truncation):
+                current_guidance_scale = 0.0
+            apply_cfg = batch.do_classifier_free_guidance and current_guidance_scale > 0.0
+
+            if apply_cfg:
+                if not isinstance(negative, list) or len(negative) != batch_size:
+                    raise ValueError("Z-Image CFG requires one negative embedding per image")
+                model_latents = latents.to(target_dtype).repeat(2, 1, 1, 1, 1)
+                model_embeddings = positive + negative
+                model_timestep = timestep.repeat(2)
+            else:
+                model_latents = latents.to(target_dtype)
+                model_embeddings = positive
+                model_timestep = timestep
+
+            with (
+                    torch.autocast(
+                        device_type=device.type,
+                        enabled=False,
+                    ),
+                    trace_step(index),
+                    set_forward_context(
+                        current_timestep=index,
+                        attn_metadata=None,
+                        forward_batch=batch,
+                    ),
+            ):
+                model_outputs = self.transformer(
+                    hidden_states=model_latents,
+                    encoder_hidden_states=model_embeddings,
+                    timestep=model_timestep,
+                )[0]
+
+            if apply_cfg:
+                positive_outputs = model_outputs[:batch_size]
+                negative_outputs = model_outputs[batch_size:]
+                guided_outputs = []
+                for positive_output, negative_output in zip(
+                        positive_outputs,
+                        negative_outputs,
+                        strict=True,
+                ):
+                    positive_fp32 = positive_output.float()
+                    prediction = positive_fp32 + current_guidance_scale * (positive_fp32 - negative_output.float())
+                    if batch.cfg_normalization:
+                        positive_norm = torch.linalg.vector_norm(positive_fp32)
+                        prediction_norm = torch.linalg.vector_norm(prediction)
+                        if prediction_norm > positive_norm:
+                            prediction = prediction * (positive_norm / prediction_norm)
+                    guided_outputs.append(prediction)
+                noise_pred = torch.stack(guided_outputs)
+            else:
+                noise_pred = torch.stack([output.float() for output in model_outputs])
+
+            noise_pred = -noise_pred
+            latents = self.scheduler.step(
+                noise_pred,
+                timestep_value,
+                latents,
+                return_dict=False,
+            )[0].float()
+
+        batch.latents = latents
+        return batch
+
+
+class ZImageDecodingStage(PipelineStage):
+    """Apply the official latent transform and decode one image frame."""
+
+    performance_component_metric = "vae_decode_time_s"
+
+    def __init__(self, vae) -> None:
+        self.vae = vae
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if batch.latents is None:
+            raise ValueError("Z-Image decoding requires latents")
+        if fastvideo_args.output_type == "latent":
+            # FastVideo standardizes image and video latents as [B,C,T,H,W].
+            # Tongyi's image-only API returns the equivalent tensor with T squeezed.
+            batch.output = batch.latents
+            return batch
+
+        latents = batch.latents
+        if latents.ndim != 5 or latents.shape[2] != 1:
+            raise ValueError(f"Expected Z-Image latents [B,C,1,H,W], got {tuple(latents.shape)}")
+
+        device = get_local_torch_device()
+        self.vae = self.vae.to(device)
+        vae_dtype = getattr(self.vae, "dtype", None)
+        if vae_dtype is None:
+            vae_dtype = next(self.vae.parameters()).dtype
+        config = self.vae.config
+        scaling_factor = float(config.scaling_factor)
+        shift_factor = float(config.shift_factor or 0.0)
+        latents_2d = latents.squeeze(2).to(device=device, dtype=vae_dtype)
+        latents_2d = latents_2d / scaling_factor + shift_factor
+        decoded = self.vae.decode(latents_2d, return_dict=False)[0]
+        decoded = (decoded / 2 + 0.5).clamp(0, 1)
+        batch.output = decoded.unsqueeze(2).float()
+
+        if fastvideo_args.vae_cpu_offload:
+            self.vae.to("cpu")
+        return batch

fastvideo/pipelines/basic/zimage/zimage_pipeline.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from fastvideo.configs.pipelines.zimage import ZImagePipelineConfig
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.stages.text_encoding import TextEncodingStage
+
+from .stages import (
+    ZImageConditioningStage,
+    ZImageDecodingStage,
+    ZImageDenoisingStage,
+    ZImageInputValidationStage,
+    ZImageLatentPreparationStage,
+    ZImageTimestepPreparationStage,
+)
+
+
+class ZImagePipeline(ComposedPipelineBase):
+    """Native Z-Image text-to-image pipeline."""
+
+    pipeline_config_cls: type[ZImagePipelineConfig] = ZImagePipelineConfig
+    _required_config_modules = [
+        "scheduler",
+        "text_encoder",
+        "tokenizer",
+        "transformer",
+        "vae",
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        scheduler = self.get_module("scheduler")
+        transformer = self.get_module("transformer")
+
+        self.add_stage("input_validation_stage", ZImageInputValidationStage())
+        self.add_stage(
+            "text_encoding_stage",
+            TextEncodingStage(
+                text_encoders=[self.get_module("text_encoder")],
+                tokenizers=[self.get_module("tokenizer")],
+            ),
+        )
+        self.add_stage("zimage_conditioning_stage", ZImageConditioningStage())
+        self.add_stage(
+            "latent_preparation_stage",
+            ZImageLatentPreparationStage(transformer=transformer),
+        )
+        self.add_stage(
+            "timestep_preparation_stage",
+            ZImageTimestepPreparationStage(scheduler=scheduler),
+        )
+        self.add_stage(
+            "denoising_stage",
+            ZImageDenoisingStage(transformer=transformer, scheduler=scheduler),
+        )
+        self.add_stage(
+            "decoding_stage",
+            ZImageDecodingStage(vae=self.get_module("vae")),
+        )
+
+
+EntryClass = ZImagePipeline

fastvideo/pipelines/composed_pipeline_base.py

@@ -303,7 +303,8 @@ class ComposedPipelineBase(ABC):
         self.modules[module_name] = module
 
     def _load_config(self, model_path: str) -> dict[str, Any]:
-        model_path = maybe_download_model(self.model_path)
+        revision = getattr(self.fastvideo_args, "revision", None)
+        model_path = maybe_download_model(self.model_path, revision=revision)
         self.model_path = model_path
         # fastvideo_args.downloaded_model_path = model_path
         logger.info("Model path: %s", model_path)

fastvideo/pipelines/pipeline_batch_info.py

@@ -147,8 +147,9 @@ class ForwardBatch:
     camera_trajectory: str | None = None  # Camera trajectory file/identifier
     action_list: list[str] | None = None  # List of actions (e.g., ['forward', 'left'])
     action_speed_list: list[float] | None = None  # Speed for each action
-    # Camera control inputs (LingBotWorld)
+    # Camera control inputs (LingBotWorld and LingBotWorld2)
     c2ws_plucker_emb: torch.Tensor | None = None  # Plucker embedding: [B, C, F_lat, H_lat, W_lat]
+    action_path: str | None = None  # Directory containing poses.npy and intrinsics.npy
 
     # Camera control inputs (GEN3C)
     trajectory_type: str | None = None
@@ -178,6 +179,8 @@ class ForwardBatch:
     num_inference_steps_sr: int = 50
     guidance_scale: float = 1.0
     guidance_scale_2: float | None = None
+    cfg_normalization: bool = False
+    cfg_truncation: float | None = 1.0
     guidance_rescale: float = 0.0
     eta: float = 0.0
     sigmas: list[float] | None = None

fastvideo/pipelines/stages/decoding.py

@@ -73,25 +73,36 @@ class DecodingStage(PipelineStage):
 
         cfg = getattr(self.vae, "config", None)
 
-        # Matrix-Game 2.0-style: z = z * std + mean
-        if (cfg is not None and hasattr(cfg, "latents_mean") and hasattr(cfg, "latents_std")):
-            latents_mean = torch.tensor(cfg.latents_mean, device=latents.device,
+        # Matrix-Game 2.0-style: z = z * std + mean. Some configs declare these
+        # fields as None, so test the values rather than mere attribute presence
+        # -- otherwise torch.tensor(None) raises instead of falling through.
+        latents_mean_value = getattr(cfg, "latents_mean", None) if cfg is not None else None
+        latents_std_value = getattr(cfg, "latents_std", None) if cfg is not None else None
+        if latents_mean_value is not None and latents_std_value is not None:
+            latents_mean = torch.tensor(latents_mean_value, device=latents.device,
                                         dtype=latents.dtype).view(1, -1, 1, 1, 1)
-            latents_std = torch.tensor(cfg.latents_std, device=latents.device, dtype=latents.dtype).view(1, -1, 1, 1, 1)
+            latents_std = torch.tensor(latents_std_value, device=latents.device,
+                                       dtype=latents.dtype).view(1, -1, 1, 1, 1)
             return latents * latents_std + latents_mean
 
-        # Diffusers-style: scaling_factor (+ optional shift_factor)
-        if hasattr(self.vae, "scaling_factor"):
-            if isinstance(self.vae.scaling_factor, torch.Tensor):
-                latents = latents / self.vae.scaling_factor.to(latents.device, latents.dtype)
+        # Diffusers-style: scaling_factor (+ optional shift_factor), read from
+        # the module only. Nearly every VAE *config* also carries a
+        # scaling_factor whose decode() already accounts for it, so sourcing it
+        # from cfg here would newly divide latents for VAEs that have always
+        # passed through untouched (flux2, hunyuanvideo, hunyuanvideo15, ltx2).
+        scaling_factor = getattr(self.vae, "scaling_factor", None)
+        if scaling_factor is not None:
+            if isinstance(scaling_factor, torch.Tensor):
+                latents = latents / scaling_factor.to(latents.device, latents.dtype)
             else:
-                latents = latents / self.vae.scaling_factor
+                latents = latents / scaling_factor
 
-            if hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None:
-                if isinstance(self.vae.shift_factor, torch.Tensor):
-                    latents = latents + self.vae.shift_factor.to(latents.device, latents.dtype)
+            shift_factor = getattr(self.vae, "shift_factor", None)
+            if shift_factor is not None:
+                if isinstance(shift_factor, torch.Tensor):
+                    latents = latents + shift_factor.to(latents.device, latents.dtype)
                 else:
-                    latents = latents + self.vae.shift_factor
+                    latents = latents + shift_factor
 
         return latents
 

fastvideo/pipelines/stages/text_encoding.py

@@ -73,6 +73,7 @@ class TextEncodingStage(PipelineStage):
             fastvideo_args,
             encoder_index=all_indices,
             return_attention_mask=True,
+            max_length=batch.max_sequence_length,
         )
         if self._last_audio_embeds is not None:
             batch.extra["ltx2_audio_prompt_embeds"] = self._last_audio_embeds
@@ -91,6 +92,7 @@ class TextEncodingStage(PipelineStage):
                 fastvideo_args,
                 encoder_index=all_indices,
                 return_attention_mask=True,
+                max_length=batch.max_sequence_length,
             )
             if self._last_audio_embeds is not None:
                 batch.extra["ltx2_audio_negative_embeds"] = self._last_audio_embeds
@@ -221,6 +223,20 @@ class TextEncodingStage(PipelineStage):
                 encoder_device = torch.device(target_device)
                 moved_for_forward = True
 
+            # An explicit `device=` wins. Otherwise follow the encoder's real
+            # param device. Once it has been moved for the forward that is the
+            # target device, and an HF-passthrough encoder's
+            # _fastvideo_input_device (stamped at load, e.g. "cpu" under
+            # text_encoder_cpu_offload) is stale -- honouring it would feed cpu
+            # token ids to cuda weights. The marker only speaks when nothing
+            # moved and the module has no parameters to speak for it.
+            if device is not None:
+                input_device = torch.device(target_device)
+            elif moved_for_forward:
+                input_device = encoder_device
+            else:
+                input_device = getattr(text_encoder, "_fastvideo_input_device", encoder_device)
+
             tok_kwargs = dict(encoder_config.tokenizer_kwargs)
             if max_length is not None:
                 tok_kwargs["max_length"] = max_length
@@ -264,7 +280,7 @@ class TextEncodingStage(PipelineStage):
                     # pre-format prompts into message lists upstream and rely on
                     # the inner tokenizer + full tokenizer_kwargs (which include
                     # add_generation_prompt). Preserve that original path exactly.
-                    text_inputs = tok.apply_chat_template(processed_texts, **tok_kwargs).to(encoder_device)
+                    text_inputs = tok.apply_chat_template(processed_texts, **tok_kwargs).to(input_device)
                 else:
                     # Two-step approach matching Diffusers: format with chat
                     # template first, then tokenize the resulting strings.
@@ -275,12 +291,12 @@ class TextEncodingStage(PipelineStage):
                             messages,
                             tokenize=False,
                             add_generation_prompt=True,
-                            enable_thinking=False,
+                            enable_thinking=encoder_config.chat_template_enable_thinking,
                         )
                         formatted_texts.append(formatted)
-                    text_inputs = tokenizer(formatted_texts, **tok_kwargs).to(encoder_device)
+                    text_inputs = tokenizer(formatted_texts, **tok_kwargs).to(input_device)
             else:
-                text_inputs = tok(processed_texts, **tok_kwargs).to(encoder_device)
+                text_inputs = tok(processed_texts, **tok_kwargs).to(input_device)
 
             input_ids = text_inputs["input_ids"]
             attention_mask = text_inputs["attention_mask"]

fastvideo/registry.py

@@ -30,6 +30,7 @@ from fastvideo.configs.pipelines.hunyuan15 import (Hunyuan15T2V480PConfig, Hunyu
 from fastvideo.configs.pipelines.hyworld import HYWorldConfig
 from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5I2VConfig, Kandinsky5T2VConfig
 from fastvideo.configs.pipelines.lingbotworld import LingBotWorldI2V480PConfig
+from fastvideo.configs.pipelines.lingbotworld2 import LingBotWorld2CausalFastI2V480PConfig
 from fastvideo.configs.pipelines.longcat import LongCatT2V480PConfig
 from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
 from fastvideo.configs.pipelines.flux_2 import (
@@ -62,6 +63,7 @@ from fastvideo.configs.pipelines.glm_image import GlmImageConfig
 from fastvideo.configs.pipelines.flux import FluxPipelineConfig
 from fastvideo.configs.pipelines.sd35 import SD35Config
 from fastvideo.configs.pipelines.stable_audio import (StableAudioOpenSmallConfig, StableAudioT2AConfig)
+from fastvideo.configs.pipelines.zimage import ZImagePipelineConfig
 from fastvideo.api.sampling_param import SamplingParam
 from fastvideo.api.matrixgame2 import MatrixGame2SamplingParam
 from fastvideo.api.matrixgame3 import MatrixGame3SamplingParam
@@ -488,6 +490,22 @@ def _register_configs() -> None:
         model_family="gamecraft",
         default_preset="gamecraft_i2v",
     )
+    # LingBotWorld2 causal-fast
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=LingBotWorld2CausalFastI2V480PConfig,
+        workload_types=(WorkloadType.I2V, ),
+        hf_model_paths=[
+            "robbyant/lingbot-world-v2-14b-causal-fast",
+        ],
+        model_detectors=[
+            lambda path:
+            ("lingbot-world-v2-14b-causal-fast" in path.lower() or "lingbotworld2causalfastpipeline" in path.lower())
+        ],
+        model_family="lingbotworld2",
+        default_preset="lingbotworld2_causal_fast_i2v",
+    )
+
     # LingBotWorld
     register_configs(
         sampling_param_cls=None,
@@ -496,7 +514,10 @@ def _register_configs() -> None:
         hf_model_paths=[
             "FastVideo/LingBot-World-Base-Cam-Diffusers",
         ],
-        model_detectors=[lambda path: ("lingbotworld" in path.lower() or "lingbot-world" in path.lower())],
+        model_detectors=[
+            lambda path: (("lingbotworld" in path.lower() or "lingbot-world" in path.lower()) and "causal-fast" not in
+                          path.lower() and "causalfast" not in path.lower())
+        ],
         model_family="lingbotworld",
         default_preset="lingbotworld_i2v",
     )
@@ -1104,6 +1125,19 @@ def _register_configs() -> None:
         ],
     )
 
+    # Z-Image-Turbo
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=ZImagePipelineConfig,
+        workload_types=(WorkloadType.T2I, ),
+        hf_model_paths=["Tongyi-MAI/Z-Image-Turbo"],
+        model_detectors=[
+            lambda path: "zimagepipeline" in path or "z-image" in path or "z_image" in path,
+        ],
+        model_family="zimage",
+        default_preset="zimage_turbo",
+    )
+
 
 # --- Part 3: Main Resolver ---
 
@@ -1210,6 +1244,8 @@ def _register_presets() -> None:
         ALL_PRESETS as KANDINSKY5_PRESETS, )
     from fastvideo.pipelines.basic.lingbotworld.presets import (
         ALL_PRESETS as LINGBOTWORLD_PRESETS, )
+    from fastvideo.pipelines.basic.lingbotworld2.presets import (
+        ALL_PRESETS as LINGBOTWORLD2_PRESETS, )
     from fastvideo.pipelines.basic.longcat.presets import (
         ALL_PRESETS as LONGCAT_PRESETS, )
     from fastvideo.pipelines.basic.ltx2.presets import (
@@ -1226,6 +1262,8 @@ def _register_presets() -> None:
         ALL_PRESETS as TURBODIFFUSION_PRESETS, )
     from fastvideo.pipelines.basic.wan.presets import (
         ALL_PRESETS as WAN_PRESETS, )
+    from fastvideo.pipelines.basic.zimage.presets import (
+        ALL_PRESETS as ZIMAGE_PRESETS, )
     from fastvideo.pipelines.basic.flux_2.presets import (
         ALL_PRESETS as FLUX2_PRESETS, )
 
@@ -1240,6 +1278,7 @@ def _register_presets() -> None:
         HYWORLD_PRESETS,
         KANDINSKY5_PRESETS,
         LINGBOTWORLD_PRESETS,
+        LINGBOTWORLD2_PRESETS,
         LONGCAT_PRESETS,
         LTX2_PRESETS,
         MATRIXGAME2_PRESETS,
@@ -1248,6 +1287,7 @@ def _register_presets() -> None:
         STABLE_AUDIO_PRESETS,
         TURBODIFFUSION_PRESETS,
         WAN_PRESETS,
+        ZIMAGE_PRESETS,
     )
     for group in all_preset_groups:
         for preset in group:

fastvideo/tests/api/test_parser.py

@@ -165,12 +165,14 @@ def test_load_run_config_supports_yaml_roundtrip(tmp_path) -> None:
                 "keyboard_cond": None,
                 "grid_sizes": None,
                 "c2ws_plucker_emb": None,
+                "action_path": None,
                 "refine_from": None,
                 "stage1_video": None,
             },
             "sampling": {
                 "num_videos_per_prompt": 1,
                 "seed": 1024,
+                "max_sequence_length": None,
                 "num_frames": 16,
                 "height": 720,
                 "width": 1280,
@@ -181,6 +183,8 @@ def test_load_run_config_supports_yaml_roundtrip(tmp_path) -> None:
                 "num_inference_steps_sr": 50,
                 "guidance_scale": 1.0,
                 "guidance_scale_2": None,
+                "cfg_normalization": False,
+                "cfg_truncation": 1.0,
                 "guidance_rescale": 0.0,
                 "true_cfg_scale": None,
                 "use_embedded_guidance": None,

fastvideo/tests/modal/pr_test.py

@@ -518,7 +518,7 @@ def run_performance_tests():
         "fi; "
         "(nvidia-smi --query-gpu=index,timestamp,clocks.sm,clocks.max.sm,power.draw,power.limit,temperature.gpu "
         "--format=csv -l 10 > /tmp/gpu_telemetry.csv 2>/dev/null &); "
-        "pytest ./fastvideo/tests/performance -vs; "
+        "pytest ./fastvideo/tests/performance -v; "
         "PYTEST_RC=$?; "
         "PERF_RC=0; "
         "if [ $PYTEST_RC -eq 0 ] || [ \"$PERF_UPLOAD_POLICY\" = 'always' ]; then "

fastvideo/tests/ssim/reference_videos_cli.py

@@ -18,7 +18,7 @@ VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv", ".webm", ".flv")
 # pixel-space models (e.g. LTX-2 distilled). Keeping them in the same upload
 # flow means seeding a new test only requires one HF round-trip.
 LATENT_EXTENSIONS = (".pt",)
-REFERENCE_EXTENSIONS = VIDEO_EXTENSIONS + LATENT_EXTENSIONS
+REFERENCE_EXTENSIONS = VIDEO_EXTENSIONS + LATENT_EXTENSIONS + (".png",)
 HF_TOKEN_ENV_KEYS = ("HF_API_KEY", "HUGGINGFACE_HUB_TOKEN", "HF_TOKEN")
 HF_REPO_ENV_KEY = "FASTVIDEO_SSIM_REFERENCE_HF_REPO"
 HF_REPO_TYPE_ENV_KEY = "FASTVIDEO_SSIM_REFERENCE_HF_REPO_TYPE"
@@ -51,10 +51,10 @@ def _default_repo_type() -> str:
 
 
 def _iter_reference_files(root: Path) -> Iterable[Path]:
-    """Yield video and latent (.pt) references under `root`.
+    """Yield supported references under `root`.
 
     Used by copy-local and the "has local references" marker probe so that
-    latent-only tests (no mp4) still satisfy readiness checks.
+    image- and latent-only tests (no mp4) still satisfy readiness checks.
     """
     for path in root.rglob("*"):
         if path.is_file() and path.suffix.lower() in REFERENCE_EXTENSIONS:

fastvideo/tests/ssim/test_reference_videos_cli.py

@@ -44,6 +44,22 @@ def _fake_snapshot_download(*, allow_patterns, cache_dir, **_kwargs):
     return str(snapshot_root)
 
 
+def test_png_references_are_copied_and_mark_local_cache_ready(tmp_path):
+    generated_dir = tmp_path / "generated"
+    generated = generated_dir / "model-a" / "flash" / "sample.png"
+    generated.parent.mkdir(parents=True)
+    generated.write_bytes(b"image")
+    reference_dir = tmp_path / "reference_videos" / "default" / "L40S_reference_videos"
+
+    assert reference_videos_cli.copy_generated_to_reference(
+        generated_dir,
+        reference_dir,
+    ) == 1
+    marker = reference_dir.parent / ".download_complete_default"
+    marker.touch()
+    assert reference_videos_cli._has_local_reference_videos(tmp_path, "default")
+
+
 @pytest.fixture(autouse=True)
 def _fake_hf(monkeypatch):
     _FakeHfApi.instances = []

fastvideo/tests/ssim/test_zimage_similarity.py

@@ -0,0 +1,124 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PNG regression for Tongyi-MAI/Z-Image@26f23eda's native Turbo recipe."""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+import torch
+
+from fastvideo.logger import init_logger
+from fastvideo.tests.ssim.inference_similarity_utils import (
+    run_text_to_video_similarity_test,
+)
+from fastvideo.tests.ssim.reference_utils import (
+    get_cuda_device_name,
+    resolve_device_reference_folder,
+)
+
+logger = init_logger(__name__)
+
+REQUIRED_GPUS = 1
+ZIMAGE_MIN_SSIM = 0.98
+
+ZIMAGE_MODEL_PATH = os.getenv(
+    "ZIMAGE_MODEL_DIR",
+    "Tongyi-MAI/Z-Image-Turbo",
+)
+ZIMAGE_MODEL_REVISION = os.getenv(
+    "ZIMAGE_MODEL_REVISION",
+    "f332072aa78be7aecdf3ee76d5c247082da564a6",
+)
+
+TEST_PROMPTS = [
+    "Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. "
+    "Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. "
+    "Neon lightning-bolt lamp (⚡️), bright yellow glow, above extended left palm. Soft-lit outdoor night background, "
+    "silhouetted tiered pagoda (西安大雁塔), blurred colorful distant lights.",
+]
+
+device_reference_folder = resolve_device_reference_folder(
+    (
+        ("A40", "A40"),
+        ("L40S", "L40S"),
+        ("H100", "H100"),
+        ("H200", "H200"),
+        ("B200", "B200"),
+    ),
+    device_name=get_cuda_device_name(),
+    fallback_device_prefix="L40S",
+    logger=logger,
+)
+
+ZIMAGE_MODEL_TO_PARAMS: dict[str, dict[str, object]] = {
+    "Tongyi-MAI__Z-Image-Turbo": {
+        "num_gpus": 1,
+        "model_path": ZIMAGE_MODEL_PATH,
+        "sp_size": 1,
+        "tp_size": 1,
+        "height": 1024,
+        "width": 1024,
+        "num_frames": 1,
+        "fps": 1,
+        "num_inference_steps": 8,
+        "guidance_scale": 0.0,
+        "seed": 42,
+    },
+}
+
+ZIMAGE_FULL_QUALITY_MODEL_TO_PARAMS = {
+    model_id: dict(params)
+    for model_id, params in ZIMAGE_MODEL_TO_PARAMS.items()
+}
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="Z-Image SSIM test requires CUDA",
+)
+@pytest.mark.parametrize("prompt", TEST_PROMPTS)
+@pytest.mark.parametrize("attention_backend_name", ["TORCH_SDPA"])
+@pytest.mark.parametrize("model_id", list(ZIMAGE_MODEL_TO_PARAMS))
+def test_zimage_similarity(
+    prompt: str,
+    attention_backend_name: str,
+    model_id: str,
+) -> None:
+    is_hf_repo = "/" in ZIMAGE_MODEL_PATH and not ZIMAGE_MODEL_PATH.startswith("/")
+    if not is_hf_repo and not os.path.isdir(ZIMAGE_MODEL_PATH):
+        pytest.skip(
+            f"Z-Image weights not found at {ZIMAGE_MODEL_PATH} "
+            "(set ZIMAGE_MODEL_DIR to override)"
+        )
+
+    run_text_to_video_similarity_test(
+        logger=logger,
+        script_dir=os.path.dirname(os.path.abspath(__file__)),
+        device_reference_folder=device_reference_folder,
+        prompt=prompt,
+        attention_backend_name=attention_backend_name,
+        model_id=model_id,
+        default_params_map=ZIMAGE_MODEL_TO_PARAMS,
+        full_quality_params_map=ZIMAGE_FULL_QUALITY_MODEL_TO_PARAMS,
+        min_acceptable_ssim=ZIMAGE_MIN_SSIM,
+        media_extension=".png",
+        init_kwargs_override={
+            "revision": ZIMAGE_MODEL_REVISION,
+            "workload_type": "t2i",
+            "trust_remote_code": True,
+            "use_fsdp_inference": False,
+            "dit_cpu_offload": False,
+            "text_encoder_cpu_offload": False,
+            "vae_cpu_offload": False,
+            "image_encoder_cpu_offload": False,
+            "pin_cpu_memory": False,
+        },
+        generation_kwargs_override={
+            "save_video": True,
+            "negative_prompt": "",
+            "max_sequence_length": 512,
+            "cfg_normalization": False,
+            "cfg_truncation": 1.0,
+        },
+    )

fastvideo/tests/stages/test_text_encoding.py

@@ -2,6 +2,7 @@ import torch
 import types
 import pytest
 
+from fastvideo.configs.models.encoders.clip import CLIPTextArchConfig, CLIPTextConfig
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.configs.pipelines.base import PipelineConfig
@@ -13,6 +14,9 @@ class TensorDict(dict):
         return TensorDict({k: v.to(device) for k, v in self.items()})
 
 class FakeTokenizer:
+    def __init__(self):
+        self.last_chat_template_kwargs = None
+
     def __call__(self, texts, **kwargs):
         B = len(texts)
         seq_len = int(kwargs.get("max_length", 4))
@@ -21,6 +25,10 @@ class FakeTokenizer:
             "attention_mask": torch.ones(B, seq_len, dtype=torch.long),
         })
 
+    def apply_chat_template(self, messages, **kwargs):
+        self.last_chat_template_kwargs = kwargs
+        return "formatted prompt"
+
 
 class FakeChatTokenizer:
     def __init__(self):
@@ -40,17 +48,22 @@ class FakeChatTokenizer:
             "attention_mask": torch.ones(B, seq_len, dtype=torch.long),
         })
 
-
 class FakeTextEncoder(torch.nn.Module):
     def __init__(self, hidden_size=8):
         super().__init__()
         self.hidden_size = hidden_size
+        self.last_input_device = None
         self.last_output_hidden_states = None
 
     def forward(self, input_ids, attention_mask, output_hidden_states=False):
+        self.last_input_device = input_ids.device
         self.last_output_hidden_states = bool(output_hidden_states)
         B, T = input_ids.shape
-        last_hidden_state = torch.arange(B * T * self.hidden_size, dtype=torch.float32).view(B, T, self.hidden_size)
+        last_hidden_state = torch.arange(
+            B * T * self.hidden_size,
+            dtype=torch.float32,
+            device=input_ids.device,
+        ).view(B, T, self.hidden_size)
         hidden_states = (last_hidden_state, ) if output_hidden_states else None
         return types.SimpleNamespace(last_hidden_state=last_hidden_state,
                                      hidden_states=hidden_states)
@@ -213,3 +226,104 @@ def test_chat_list_preprocess_output_is_not_stripped():
         {"role": "user", "content": "a robotic arm welding a metal structure"},
     ]]
     assert tokenizer.last_kwargs["return_tensors"] == "pt"
+
+
+def test_chat_template_thinking_is_keyword_only_and_preserves_subclass_positions():
+    arch = CLIPTextArchConfig()
+    config = CLIPTextConfig(
+        arch,
+        "legacy-prefix",
+        None,
+        None,
+        False,
+        False,
+        7,
+        False,
+        True,
+        False,
+    )
+
+    assert config.num_hidden_layers_override == 7
+    assert config.require_post_norm is False
+    assert config.enable_scale is True
+    assert config.is_causal is False
+    assert config.chat_template_enable_thinking is False
+
+    configured = CLIPTextConfig(chat_template_enable_thinking=True)
+    assert configured.chat_template_enable_thinking is True
+
+    with pytest.raises(TypeError):
+        CLIPTextConfig(
+            arch,
+            "legacy-prefix",
+            None,
+            None,
+            False,
+            False,
+            7,
+            False,
+            True,
+            False,
+            True,
+        )
+
+
+@pytest.mark.parametrize("enable_thinking", [False, True])
+def test_encode_text_forwards_chat_template_thinking_config(enable_thinking):
+    fastvideo_args, hidden = make_args(num_encoders=1, text_len=4, hidden_size=8)
+    encoder_config = fastvideo_args.pipeline_config.text_encoder_configs[0]
+    encoder_config.is_chat_model = True
+    encoder_config.chat_template_enable_thinking = enable_thinking
+
+    stage = make_stage(num_encoders=1, hidden_size=hidden)
+    stage.encode_text("a", fastvideo_args, encoder_index=[0])
+
+    assert stage.tokenizers[0].last_chat_template_kwargs == {
+        "tokenize": False,
+        "add_generation_prompt": True,
+        "enable_thinking": enable_thinking,
+    }
+
+
+def test_encode_text_uses_hf_passthrough_input_device(monkeypatch):
+    fastvideo_args, hidden = make_args(num_encoders=1, text_len=4, hidden_size=8)
+    stage = make_stage(num_encoders=1, hidden_size=hidden)
+    stage.text_encoders[0]._fastvideo_input_device = torch.device("cpu")
+
+    monkeypatch.setattr(
+        "fastvideo.pipelines.stages.text_encoding.get_local_torch_device",
+        lambda: torch.device("meta"),
+    )
+
+    batch = ForwardBatch(
+        data_type="video",
+        prompt="a",
+        do_classifier_free_guidance=False,
+        prompt_embeds=[],
+        negative_prompt_embeds=None,
+        prompt_attention_mask=[],
+        negative_attention_mask=None,
+    )
+    output = stage.forward(batch, fastvideo_args)
+
+    # The marker governs where the encoder *receives* its tokens...
+    assert stage.text_encoders[0].last_input_device == torch.device("cpu")
+    # ...while the stage still normalizes the embeds it returns onto the
+    # caller's target device.
+    assert output.prompt_embeds[0].device.type == "meta"
+
+
+def test_encode_text_explicit_device_overrides_hf_passthrough_marker():
+    fastvideo_args, hidden = make_args(num_encoders=1, text_len=4, hidden_size=8)
+    stage = make_stage(num_encoders=1, hidden_size=hidden)
+    stage.text_encoders[0]._fastvideo_input_device = torch.device("cpu")
+
+    output = stage.encode_text(
+        "a",
+        fastvideo_args,
+        encoder_index=[0],
+        device=torch.device("meta"),
+    )
+
+    assert stage.text_encoders[0].last_input_device == torch.device("meta")
+    assert output[0].device.type == "meta"

fastvideo/utils.py

@@ -493,7 +493,12 @@ def import_pynvml():
     return pynvml
 
 
-def maybe_download_model(model_name_or_path: str, local_dir: str | None = None, download: bool = True) -> str:
+def maybe_download_model(
+    model_name_or_path: str,
+    local_dir: str | None = None,
+    download: bool = True,
+    revision: str | None = None,
+) -> str:
     """
     Check if the model path is a Hugging Face Hub model ID and download it if needed.
 
@@ -510,6 +515,7 @@ def maybe_download_model(model_name_or_path: str, local_dir: str | None = None,
             ``org/repo/subfolder`` umbrella-repo reference.
         local_dir: Local directory to save the model
         download: Whether to download the model from Hugging Face Hub
+        revision: Optional immutable Hub revision.
 
     Returns:
         Local path to the model (or to the subfolder inside the snapshot).
@@ -550,6 +556,7 @@ def maybe_download_model(model_name_or_path: str, local_dir: str | None = None,
                     repo_id=repo_id,
                     allow_patterns=[f"{subfolder}/**"],
                     local_dir=local_dir,
+                    revision=revision,
                 )
             # Defense-in-depth: ensure the resolved subfolder path stays
             # inside the snapshot root and that snapshot_download actually
@@ -569,7 +576,8 @@ def maybe_download_model(model_name_or_path: str, local_dir: str | None = None,
         with get_lock(model_name_or_path):
             local_path = snapshot_download(repo_id=model_name_or_path,
                                            ignore_patterns=["*.onnx", "*.msgpack"],
-                                           local_dir=local_dir)
+                                           local_dir=local_dir,
+                                           revision=revision)
         logger.info("Downloaded model to %s", local_path)
         return str(local_path)
     except Exception as e:

scripts/checkpoint_conversion/convert_lingbotworld2_causal_fast.py

@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Create a FastVideo-native LingBot World 2 causal-fast bundle."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from pathlib import Path
+
+
+LATENTS_MEAN = [
+    -0.7571,
+    -0.7089,
+    -0.9113,
+    0.1075,
+    -0.1745,
+    0.9653,
+    -0.1517,
+    1.5508,
+    0.4134,
+    -0.0715,
+    0.5517,
+    -0.3632,
+    -0.1922,
+    -0.9497,
+    0.2503,
+    -0.2921,
+]
+LATENTS_STD = [
+    2.8184,
+    1.4541,
+    2.3275,
+    2.6558,
+    1.2196,
+    1.7708,
+    2.6052,
+    2.0743,
+    3.2687,
+    2.1526,
+    2.8652,
+    1.5579,
+    1.6382,
+    1.1253,
+    2.8251,
+    1.9160,
+]
+
+
+def _write_json(path: Path, payload: dict) -> None:
+    """Write deterministic JSON config files for a local bundle."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
+
+
+def _symlink_or_replace(source: Path, target: Path) -> None:
+    """Create a relative symlink, replacing stale symlinks only."""
+    target.parent.mkdir(parents=True, exist_ok=True)
+    if target.is_symlink() or target.exists():
+        if target.is_symlink() and target.resolve() == source.resolve():
+            return
+        raise FileExistsError(f"Refusing to replace existing path: {target}")
+    target.symlink_to(os.path.relpath(source, target.parent))
+
+
+def convert_vae(source_dir: Path, output_dir: Path) -> None:
+    """Link the exact LingBot World 2 Wan VAE checkpoint consumed by LingBotWorld2WanVAE."""
+    _symlink_or_replace(source_dir / "Wan2.1_VAE.pth", output_dir / "vae" / "Wan2.1_VAE.pth")
+
+
+def build_bundle(source_dir: Path, output_dir: Path) -> None:
+    """Create configs and symlinks for FastVideo-native LingBot World 2 loading."""
+    output_dir.mkdir(parents=True, exist_ok=True)
+    _write_json(
+        output_dir / "model_index.json",
+        {
+            "_class_name": "LingBotWorld2CausalFastPipeline",
+            "_diffusers_version": "0.30.0",
+            "scheduler": ["diffusers", "FlowUniPCMultistepScheduler"],
+            "text_encoder": ["transformers", "LingBotWorld2T5EncoderModel"],
+            "tokenizer": ["transformers", "AutoTokenizer"],
+            "transformer": ["diffusers", "LingBotWorld2CausalFastTransformer3DModel"],
+            "vae": ["diffusers", "LingBotWorld2WanVAE"],
+        },
+    )
+    _write_json(
+        output_dir / "transformer" / "config.json",
+        {
+            "_class_name": "LingBotWorld2CausalFastTransformer3DModel",
+            "_diffusers_version": "0.30.0",
+            "model_type": "i2v",
+            "patch_size": [1, 2, 2],
+            "text_len": 512,
+            "in_dim": 36,
+            "dim": 5120,
+            "ffn_dim": 13824,
+            "freq_dim": 256,
+            "text_dim": 4096,
+            "out_dim": 16,
+            "num_heads": 40,
+            "num_layers": 40,
+            "local_attn_size": 18,
+            "sink_size": 6,
+            "qk_norm": True,
+            "cross_attn_norm": True,
+            "eps": 1e-6,
+        },
+    )
+    _write_json(
+        output_dir / "text_encoder" / "config.json",
+        {
+            "architectures": ["LingBotWorld2T5EncoderModel"],
+            "vocab_size": 256384,
+            "dim": 4096,
+            "dim_attn": 4096,
+            "dim_ffn": 10240,
+            "num_heads": 64,
+            "num_layers": 24,
+            "num_buckets": 32,
+            "text_len": 512,
+            "hidden_size": 4096,
+            "dropout": 0.1,
+        },
+    )
+    _write_json(
+        output_dir / "vae" / "config.json",
+        {
+            "_class_name": "LingBotWorld2WanVAE",
+            "base_dim": 96,
+            "z_dim": 16,
+            "dim_mult": [1, 2, 4, 4],
+            "num_res_blocks": 2,
+            "attn_scales": [],
+            "temperal_downsample": [False, True, True],
+            "dropout": 0.0,
+            "latents_mean": LATENTS_MEAN,
+            "latents_std": LATENTS_STD,
+        },
+    )
+    _write_json(
+        output_dir / "scheduler" / "scheduler_config.json",
+        {
+            "_class_name": "FlowUniPCMultistepScheduler",
+            "num_train_timesteps": 1000,
+            "shift": 1.0,
+            "use_dynamic_shifting": False,
+            "prediction_type": "flow_prediction",
+            "solver_order": 2,
+            "solver_type": "bh2",
+            "final_sigmas_type": "zero",
+        },
+    )
+    _write_json(output_dir / "scheduler" / "config.json", json.loads((output_dir / "scheduler" / "scheduler_config.json").read_text()))
+
+    _symlink_or_replace(source_dir / "models_t5_umt5-xxl-enc-bf16.pth", output_dir / "text_encoder" / "pytorch_model.pt")
+    _symlink_or_replace(source_dir / "google" / "umt5-xxl", output_dir / "tokenizer")
+    for shard in sorted((source_dir / "transformers").glob("*.safetensors")):
+        _symlink_or_replace(shard, output_dir / "transformer" / shard.name)
+    index_path = source_dir / "transformers" / "diffusion_pytorch_model.safetensors.index.json"
+    if index_path.exists():
+        _symlink_or_replace(index_path, output_dir / "transformer" / "diffusion_pytorch_model.safetensors.index.json")
+    convert_vae(source_dir, output_dir)
+
+
+def main() -> None:
+    """Parse CLI arguments and build the local converted bundle."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--source-dir", required=True, type=Path)
+    parser.add_argument("--output-dir", required=True, type=Path)
+    args = parser.parse_args()
+    build_bundle(args.source_dir, args.output_dir)
+
+
+if __name__ == "__main__":
+    main()

tests/local_tests/lingbotworld2/PORT_STATUS.md

@@ -0,0 +1,82 @@
+# LingBot World 2 Port Status
+
+## Summary
+
+- model_family: `lingbotworld2`
+- workload_types: I2V causal-fast
+- official_ref: `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2`
+- official_ref_dir: `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2`
+- hf_weights_path: `robbyant/lingbot-world-v2-14b-causal-fast`
+- local_weights_dir: `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/ckpts/lingbot-world-v2-14b-causal-fast`
+- source_layout: `raw_official`
+- local_tests_readme: `/mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot/tests/local_tests/lingbotworld2/README.md`
+
+## Current Phase
+
+- phase: port verified
+- status: complete
+- owner: orchestrator
+- last_updated: 2026-07-10
+
+## Component Matrix
+
+| Component      | Type      | Reuse/Port | Official Definition                                                                                                                 | Official Instantiation                                                                               | FastVideo Target                                                                                                             | Prototype        | Conversion                       | Parity                    | Open Issues |
+| -------------- | --------- | ---------- | ---------------------------------------------------------------------------------------------------                                 | ---------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------- | ---------------- | -------------------------------  | ------------------------- | ----------- |
+| Transformer    | DiT       | ported     | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/modules/model_fast.py`                                                         | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/image2video.py`                                 | `/mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot/fastvideo/models/dits/lingbotworld2/causal_fast.py`                   | import/meta pass | symlink raw shards               | exact generated parity    | none        |
+| VAE            | VAE       | reused     | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/modules/vae2_1.py`                                                             | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/image2video.py`                                 | `/mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot/fastvideo/models/vaes/lingbotworld2_wanvae.py`                        | reuse proven     | symlink raw `.pth`               | exact encode/decode pass  | none        |
+| Text encoder   | encoder   | ported     | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/modules/t5.py`                                                                 | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/image2video.py`                                 | `/mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot/fastvideo/models/encoders/lingbotworld2_t5.py`                        | import pass      | symlink official `.pth` as `.pt` | strict-load pass          | none        |
+| Camera utility | generic   | ported     | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/utils/cam_utils.py`                                                            | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/image2video.py`                                 | `/mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot/fastvideo/models/dits/lingbotworld2/cam_utils.py`                     | import pass      | stateless                        | covered by pipeline smoke | none        |
+| Scheduler      | scheduler | reused     | `FlowUniPCMultistepScheduler` in official sampling loop                                                                             | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/image2video.py`                                 | `/mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py`       | reuse proven     | config emitted                   | exact generated parity    | none        |
+| Pipeline       | pipeline  | ported     | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/generate.py`; `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/image2video.py` | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/run_fast.sh`                                        | `/mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot/fastvideo/pipelines/basic/lingbotworld2/causal_fast_pipeline.py`      | import pass      | model_index emitted              | exact generated parity    | none        |
+
+## Conversion State
+
+- conversion_script: `/mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot/scripts/checkpoint_conversion/convert_lingbotworld2_causal_fast.py`
+- converted_weights_dir: `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/ckpts/lingbot-world-v2-14b-causal-fast-fastvideo`
+- source_layout: `raw_official`
+- strict_load_status: VAE pass; text encoder pass; transformer full 14B load pass
+- passthrough_components: transformer shards, text encoder checkpoint, tokenizer files, raw VAE checkpoint
+- retry_history: generic `AutoencoderKLWan` VAE conversion strict-loaded but was not numerically exact; final bundle uses `LingBotWorld2WanVAE`.
+
+## Parity Commands
+
+| Scope                    | Command                                                                                                                                                                                                                                                                                                   | Last Result | Notes                                                                         |
+| ------------------------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------                                                                              | ----------- | -----------------------------------------------                               |
+| Preflight and parity     | `/mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/python -m pytest tests/local_tests/lingbotworld2/test_lingbotworld2_causal_fast_preflight.py tests/local_tests/pipelines/test_lingbotworld2_causal_fast_pipeline_smoke.py tests/local_tests/pipelines/test_lingbotworld2_causal_fast_pipeline_parity.py -q -rs` | pass        | 2026-07-10: 5 passed, 1 skipped; skip is opt-in 14B heavy smoke.              |
+| Reference matrix         | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/scripts/run_reference_matrix.sh`                                                                                                                                                                                                                         | pass        | `5` and `9` frame cases recorded as official_unsupported.                     |
+| FastVideo matrix         | `/mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/python /mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/scripts/run_fastvideo_matrix.py`                                                                                                                                                                         | pass        | Matches official-supported cases and records `5`/`9` as official_unsupported. |
+| Generated video parity   | `/mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/python /mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/scripts/compare_matrix_outputs.py`                                                                                                                                                                       | pass        | 2026-07-09: all 9 supported rows exact; max abs 0, mean abs 0.0.              |
+| Heavy FastVideo smoke    | `LINGBOTWORLD2_RUN_HEAVY_SMOKE=1 /mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/python -m pytest tests/local_tests/pipelines/test_lingbotworld2_causal_fast_pipeline_smoke.py -q -rs`                                                                                                                           | pass        | 2026-07-10: 2 passed; 8-GPU 14B latent generation passed in 3:15.             |
+
+## Open Questions
+
+| ID   | Question                                                                                   | Owner        | Needed By Phase | Status   | Resolution                                                             |
+| ---- | ------------------------------------------------------------------------------------------ | ------------ | --------------- | -------- | ----------                                                             |
+| Q001 | Should final reported matrix include source-impossible requested `5` and `9` frame cases?  | user/author  | final parity    | resolved | Record them as official-source failures; do not change `chunk_size=4`. |
+
+## Issues And Blockers
+
+| ID   | Phase              | Component | Severity | Issue                                                                  | Evidence                                                                                                           | Owner        | Status | Resolution                                                                  |
+| ---- | ------------------ | --------- | -------- | ---------------------------------------------------------------------  | ------------------------------------------------------------------------------------------------------------------ | ------------ | ------ | ----------                                                                  |
+| I001 | pipeline parity    | pipeline  | high     | Requested `5` and `9` frames cannot produce official reference videos. | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/image2video.py` sets `lat_f=0`, then `F=-3`.                  | orchestrator | closed | Matrix scripts record unsupported cases; supported cases remain `17/33/65`. |
+
+## Escape Hatches
+
+| ID | Phase | Decision Type | Question | Recommended Option | Status | Resolution |
+| -- | ----- | ------------- | -------- | ------------------ | ------ | ---------- |
+
+## Decisions
+
+| Date       | Decision                                                                                  | Rationale                                                                                  | Impact                                                                                                                   |
+| ---------- | ------------------------------------------------------------------------                  | ------------------------------------------------------------------------------------------ | ------------------------------------------------------------                                                             |
+| 2026-07-09 | Keep the FastVideo causal-fast port native and do not import source repo.                 | User required a FastVideo-native model/pipeline implementation.                            | Production code owns DiT, T5 wrapper, camera utilities, and pipeline loop.                                               |
+| 2026-07-09 | Preserve official `chunk_size=4` behavior.                                                | `run_fast.sh` does not override chunk size and `generate.py` defaults to `4`.              | Requested `5` and `9` frame cases are recorded as official failures.                                                     |
+| 2026-07-09 | Store generated outputs under the LingBot World 2 source tree.                            | User required outputs to stay within the named directories.                                | Matrix outputs live under `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/outputs`.                                     |
+| 2026-07-09 | Use the exact LingBot World 2 Wan VAE implementation in FastVideo.                        | Generic `AutoencoderKLWan` strict-loaded but decoded output was not numerically exact.     | `LingBotWorld2WanVAE` loads the original `Wan2.1_VAE.pth` and restores exact video parity.                               |
+| 2026-07-10 | Use `lingbotworld2` for the causal-fast model family and `LingBotWorld2` for its classes. | FastVideo already has a separate LingBot World v1 family under `lingbotworld`.             | V2 modules, configs, presets, registry entries, tests, converter, dataset, and example no longer share the v1 namespace. |
+
+## Handoff Notes
+
+- The final comparison report is `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/outputs/verification/matrix_compare.md`.
+- Exact generated-output parity is only meaningful for cases where the official reference writes a video.
+- Do not loosen parity by changing `chunk_size` unless the user explicitly changes the reference config.

tests/local_tests/lingbotworld2/README.md

@@ -0,0 +1,163 @@
+# LingBot World 2 Local Tests
+
+Local-only smoke, conversion, and generated-output parity notes for the
+`lingbotworld2` causal-fast FastVideo port.
+
+## Reference Assets
+
+| Field               | Value                                                                                                             |
+| ------------------- | ----------------------------------------------------------------------------------------------------------------- |
+| Model family        | `lingbotworld2`                                                                                                   |
+| First scope         | `robbyant/lingbot-world-v2-14b-causal-fast`, I2V causal-fast inference                                            |
+| Workload types      | I2V: text + first image + camera/action path to video                                                             |
+| Official reference  | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2`                                                                 |
+| HF weights          | `robbyant/lingbot-world-v2-14b-causal-fast`                                                                       |
+| Raw weights         | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/ckpts/lingbot-world-v2-14b-causal-fast`                          |
+| FastVideo bundle    | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/ckpts/lingbot-world-v2-14b-causal-fast-fastvideo`                |
+| Source layout       | `raw_official`                                                                                                    |
+| Needs conversion    | `yes`                                                                                                             |
+| Token env var       | none recorded                                                                                                     |
+
+Do not write token values in this file.
+
+## Shared Environment Setup
+
+Use the existing venv and do not rebuild FlashAttention:
+
+```bash
+/mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/pip install -r /mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/requirements.txt
+/mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/pip install --no-deps -e /mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot
+```
+
+The current environment already has the FastVideo editable install pointing at:
+
+```text
+/mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot
+```
+
+Runtime cache and temporary paths used by the matrix scripts:
+
+```text
+/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/.cache
+```
+
+## Conversion
+
+The converter writes a FastVideo-loadable bundle with native component configs,
+the exact LingBot World 2 Wan VAE checkpoint, and symlinked raw LingBot World 2 transformer/T5
+assets:
+
+```bash
+cd /mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot
+/mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/python scripts/checkpoint_conversion/convert_lingbotworld2_causal_fast.py \
+  --source-dir /mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/ckpts/lingbot-world-v2-14b-causal-fast \
+  --output-dir /mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/ckpts/lingbot-world-v2-14b-causal-fast-fastvideo
+```
+
+Strict-load checks already run locally:
+
+| Component    | Status | Notes                                                                  |
+| ------------ | ------ | ---------------------------------------------------------------------  |
+| VAE          | pass   | `LingBotWorld2WanVAE` loads the original `Wan2.1_VAE.pth` checkpoint.  |
+| Text encoder | pass   | Native `LingBotWorld2T5EncoderModel` strict-load from symlinked `.pt`. |
+| Transformer  | pass   | Full 14B load and 8-GPU causal-fast generation pass.                   |
+
+## Official Frame-Count Constraint
+
+The requested verification frame counts are `5, 9, 17, 33, 65`. The official
+causal-fast source uses default `chunk_size=4`, then truncates latent frames to a
+multiple of the chunk size. This makes `5` and `9` source-impossible and changes
+the generated frame counts for the remaining values.
+
+| Requested frames | Official effective frames | Status                         |
+| ---------------- | ------------------------- | ------------------------------ |
+| `5`              | `-3`                      | official source fails          |
+| `9`              | `-3`                      | official source fails          |
+| `17`             | `13`                      | runs with truncated output     |
+| `33`             | `29`                      | runs with truncated output     |
+| `65`             | `61`                      | runs with truncated output     |
+
+Evidence: `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/image2video.py`
+rounds `frame_num` to `4n+1`, computes latent frames, truncates `lat_f` by
+`chunk_size`, and resets `F = (lat_f - 1) * 4 + 1`.
+
+## Local Tests
+
+```bash
+cd /mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot
+/mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/python -m pytest \
+  tests/local_tests/lingbotworld2/test_lingbotworld2_causal_fast_preflight.py \
+  tests/local_tests/pipelines/test_lingbotworld2_causal_fast_pipeline_smoke.py \
+  tests/local_tests/pipelines/test_lingbotworld2_causal_fast_pipeline_parity.py \
+  -q -rs
+```
+
+Current result:
+
+```text
+2026-07-10: 5 passed, 1 skipped.
+Skipped test: opt-in 14B heavy smoke gated by LINGBOTWORLD2_RUN_HEAVY_SMOKE=1.
+```
+
+Heavy smoke is opt-in because it loads the 14B model:
+
+```bash
+cd /mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot
+LINGBOTWORLD2_RUN_HEAVY_SMOKE=1 /mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/python -m pytest \
+  tests/local_tests/pipelines/test_lingbotworld2_causal_fast_pipeline_smoke.py \
+  -q -rs
+```
+
+Current heavy result:
+
+```text
+2026-07-10: 2 passed in 3:15 on 8 H200 GPUs.
+```
+
+## Matrix Verification
+
+Reference generation:
+
+```bash
+/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/scripts/run_reference_matrix.sh
+```
+
+FastVideo generation:
+
+```bash
+/mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/python \
+  /mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/scripts/run_fastvideo_matrix.py
+```
+
+Decoded-frame comparison:
+
+```bash
+/mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/python \
+  /mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/scripts/compare_matrix_outputs.py
+```
+
+Current comparison result:
+
+```text
+2026-07-09: all 9 official-supported FastVideo videos are exactly equal to the
+reference videos after decoded-frame comparison: max abs 0, mean abs 0.0.
+The requested 5-frame and 9-frame cases are recorded as official_unsupported for
+both reference and FastVideo.
+```
+
+Expected output roots:
+
+| Output type  | Path                                                                       |
+| ------------ | -------------------------------------------------------------------------- |
+| Reference    | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/outputs/reference`        |
+| FastVideo    | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/outputs/fastvideo`        |
+| Verification | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/outputs/verification`     |
+
+Generated-output parity test:
+
+```bash
+cd /mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot
+/mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/python -m pytest \
+  tests/local_tests/pipelines/test_lingbotworld2_causal_fast_pipeline_parity.py \
+  -q -rs
+```

tests/local_tests/lingbotworld2/__init__.py

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

tests/local_tests/lingbotworld2/test_lingbotworld2_causal_fast_preflight.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Preflight checks for the LingBot World 2 causal-fast FastVideo port."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+
+
+ROOT = Path("/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2")
+FASTVIDEO_ROOT = Path(__file__).resolve().parents[3]
+RAW_MODEL_DIR = ROOT / "ckpts" / "lingbot-world-v2-14b-causal-fast"
+FASTVIDEO_MODEL_DIR = ROOT / "ckpts" / "lingbot-world-v2-14b-causal-fast-fastvideo"
+ACTION_PATH = FASTVIDEO_ROOT / "examples" / "dataset" / "lingbotworld2"
+
+
+def _source_effective_frames(frame_num: int, pose_count: int, chunk_size: int = 4) -> int:
+    """Mirror the released `wan/image2video.py` frame-count truncation."""
+    rounded = ((frame_num - 1) // 4) * 4 + 1
+    rounded = min(rounded, ((pose_count - 1) // 4) * 4 + 1)
+    latent_frames = (rounded - 1) // 4 + 1
+    latent_frames = int(latent_frames - (latent_frames % chunk_size))
+    return (latent_frames - 1) * 4 + 1
+
+
+def test_lingbotworld2_bundle_declares_native_components() -> None:
+    assert RAW_MODEL_DIR.exists()
+    assert FASTVIDEO_MODEL_DIR.exists()
+    assert (FASTVIDEO_MODEL_DIR / "model_index.json").exists()
+    assert (FASTVIDEO_MODEL_DIR / "transformer" / "config.json").exists()
+    assert (FASTVIDEO_MODEL_DIR / "text_encoder" / "config.json").exists()
+    assert (FASTVIDEO_MODEL_DIR / "text_encoder" / "pytorch_model.pt").exists()
+    assert (FASTVIDEO_MODEL_DIR / "vae" / "config.json").exists()
+    assert (FASTVIDEO_MODEL_DIR / "vae" / "Wan2.1_VAE.pth").exists()
+    assert (FASTVIDEO_MODEL_DIR / "scheduler" / "scheduler_config.json").exists()
+
+    model_index = json.loads((FASTVIDEO_MODEL_DIR / "model_index.json").read_text(encoding="utf-8"))
+    assert model_index["_class_name"] == "LingBotWorld2CausalFastPipeline"
+    assert model_index["transformer"][1] == "LingBotWorld2CausalFastTransformer3DModel"
+    assert model_index["text_encoder"][1] == "LingBotWorld2T5EncoderModel"
+    assert model_index["vae"][1] == "LingBotWorld2WanVAE"
+
+
+def test_lingbotworld2_source_frame_requests_truncate_as_expected() -> None:
+    pose_count = int(np.load(ACTION_PATH / "poses.npy").shape[0])
+    effective = {
+        frame_num: _source_effective_frames(frame_num, pose_count)
+        for frame_num in (5, 9, 17, 33, 65)
+    }
+    assert effective == {
+        5: -3,
+        9: -3,
+        17: 13,
+        33: 29,
+        65: 61,
+    }
+
+
+def test_lingbotworld2_request_preserves_action_path() -> None:
+    from fastvideo.api.compat import normalize_generation_request, request_to_sampling_param
+
+    request = normalize_generation_request({
+        "prompt": "camera path schema check",
+        "inputs": {
+            "image_path": str(ACTION_PATH / "image.jpg"),
+            "action_path": str(ACTION_PATH),
+        },
+        "sampling": {
+            "height": 480,
+            "width": 832,
+            "num_frames": 17,
+            "num_inference_steps": 4,
+            "guidance_scale": 1.0,
+        },
+        "output": {
+            "save_video": False,
+            "return_frames": False,
+        },
+    })
+    sampling_param = request_to_sampling_param(request, model_path=str(FASTVIDEO_MODEL_DIR))
+    assert sampling_param.action_path == str(ACTION_PATH)
+    assert sampling_param.image_path == str(ACTION_PATH / "image.jpg")
+    assert sampling_param.num_frames == 17

tests/local_tests/pipelines/test_lingbotworld2_causal_fast_pipeline_parity.py

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Generated-output parity for LingBot World 2 causal-fast matrix runs."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+import subprocess
+import sys
+
+import pytest
+
+
+ROOT = Path("/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2")
+REFERENCE_STATUS = ROOT / "outputs" / "reference" / "status.tsv"
+FASTVIDEO_STATUS = ROOT / "outputs" / "fastvideo" / "status.tsv"
+COMPARE_SCRIPT = ROOT / "scripts" / "compare_matrix_outputs.py"
+COMPARE_JSON = ROOT / "outputs" / "verification" / "matrix_compare.json"
+
+
+def test_lingbotworld2_generated_matrix_matches_reference_exactly() -> None:
+    """Compare generated matrix videos after the source-side launch scripts run."""
+    if not REFERENCE_STATUS.exists() or not FASTVIDEO_STATUS.exists():
+        pytest.skip("Run lingbot-world-v2/scripts/run_reference_matrix.sh and run_fastvideo_matrix.py first.")
+
+    subprocess.run([sys.executable, str(COMPARE_SCRIPT)], check=True)
+    results = json.loads(COMPARE_JSON.read_text(encoding="utf-8"))
+    compared = [item for item in results if item.get("comparison") and "exact" in item["comparison"]]
+    if not compared:
+        pytest.fail("No LingBot World 2 matrix cases had both reference and FastVideo videos to compare.")
+
+    failures = [
+        item
+        for item in compared
+        if not item["comparison"].get("exact")
+    ]
+    assert failures == []

tests/local_tests/pipelines/test_lingbotworld2_causal_fast_pipeline_smoke.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Smoke tests for the LingBot World 2 causal-fast pipeline."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any, cast
+
+import pytest
+import torch
+
+
+MODEL_DIR = Path("/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/ckpts/lingbot-world-v2-14b-causal-fast-fastvideo")
+FASTVIDEO_ROOT = Path(__file__).resolve().parents[3]
+ACTION_PATH = FASTVIDEO_ROOT / "examples" / "dataset" / "lingbotworld2"
+
+
+def test_lingbotworld2_checkpoint_selects_expected_pipeline_and_defaults() -> None:
+    import fastvideo.registry as registry
+    from fastvideo.api.presets import get_preset, get_presets_for_family
+    from fastvideo.configs.pipelines.lingbotworld2 import LingBotWorld2CausalFastI2V480PConfig
+    from fastvideo.fastvideo_args import WorkloadType
+    from fastvideo.pipelines.basic.lingbotworld2.causal_fast_pipeline import (
+        EntryClass,
+        LingBotWorld2CausalFastPipeline,
+    )
+
+    assert EntryClass is LingBotWorld2CausalFastPipeline
+    assert LingBotWorld2CausalFastPipeline._required_config_modules == [
+        "text_encoder",
+        "tokenizer",
+        "vae",
+        "transformer",
+        "scheduler",
+    ]
+
+    default_preset, model_family = registry.get_preset_selection(str(MODEL_DIR))
+    assert model_family == "lingbotworld2"
+    assert default_preset == "lingbotworld2_causal_fast_i2v"
+
+    info = registry.get_model_info(
+        str(MODEL_DIR),
+        workload_type=WorkloadType.I2V,
+        override_pipeline_cls_name="LingBotWorld2CausalFastPipeline",
+    )
+    assert info.pipeline_cls is LingBotWorld2CausalFastPipeline
+    assert info.pipeline_config_cls is LingBotWorld2CausalFastI2V480PConfig
+
+    names = {preset.name for preset in get_presets_for_family("lingbotworld2")}
+    assert "lingbotworld2_causal_fast_i2v" in names
+    preset = get_preset("lingbotworld2_causal_fast_i2v", "lingbotworld2")
+    assert preset.defaults["num_inference_steps"] == 4
+    assert preset.defaults["height"] == 480
+    assert preset.defaults["width"] == 832
+    assert preset.defaults["num_frames"] == 65
+    assert preset.defaults["guidance_scale"] == 1.0
+
+    cfg = LingBotWorld2CausalFastI2V480PConfig()
+    arch = cfg.dit_config.arch_config
+    assert arch.model_type == "i2v"
+    assert arch.in_dim == 36
+    assert arch.local_attn_size == 18
+    assert arch.sink_size == 6
+    assert arch.chunk_size == 4
+    assert tuple(arch.timesteps_index) == (0, 250, 500, 750)
+
+
+@pytest.mark.skipif(
+    os.getenv("LINGBOTWORLD2_RUN_HEAVY_SMOKE") != "1",
+    reason="Set LINGBOTWORLD2_RUN_HEAVY_SMOKE=1 to load the 14B causal-fast model.",
+)
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="LingBot World 2 heavy smoke requires CUDA")
+def test_lingbotworld2_14b_generates_finite_latents_on_8_gpus() -> None:
+    from fastvideo import VideoGenerator
+
+    generator = VideoGenerator.from_pretrained(
+        str(MODEL_DIR),
+        num_gpus=8,
+        sp_size=8,
+        hsdp_shard_dim=8,
+        use_fsdp_inference=True,
+        dit_layerwise_offload=False,
+        dit_cpu_offload=False,
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=False,
+        pin_cpu_memory=True,
+        output_type="latent",
+        override_pipeline_cls_name="LingBotWorld2CausalFastPipeline",
+    )
+    try:
+        result = generator.generate_video(
+            "A serene lakeside scene with a lone tree standing in calm water.",
+            image_path=str(ACTION_PATH / "image.jpg"),
+            action_path=str(ACTION_PATH),
+            output_path="/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/outputs/fastvideo/heavy_smoke",
+            save_video=False,
+            return_frames=True,
+            height=480,
+            width=832,
+            num_frames=17,
+            num_inference_steps=4,
+            guidance_scale=1.0,
+            negative_prompt="",
+            fps=16,
+            seed=42,
+        )
+    finally:
+        generator.shutdown()
+
+    samples = cast(dict[str, Any], result)["samples"]
+    assert torch.is_tensor(samples)
+    assert samples.ndim == 5
+    assert samples.shape[1] == 16
+    assert torch.isfinite(samples).all()

tests/local_tests/pipelines/test_zimage_pipeline_parity.py

@@ -0,0 +1,328 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Native Z-Image stage contract and pinned-repository pipeline parity."""
+
+from __future__ import annotations
+
+import gc
+import importlib
+import os
+import subprocess
+import sys
+from contextlib import contextmanager
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch.testing import assert_close
+
+from fastvideo.configs.pipelines.zimage import ZImagePipelineConfig
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler,
+)
+from fastvideo.pipelines.basic.zimage.stages import (
+    ZImageConditioningStage,
+    ZImageDecodingStage,
+    ZImageDenoisingStage,
+    ZImageInputValidationStage,
+    ZImageLatentPreparationStage,
+    ZImageTimestepPreparationStage,
+)
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+REFERENCE_REVISION = "26f23eda626ffadda020b04ff79488e1d72004cd"
+REFERENCE_REPO = Path(os.getenv("ZIMAGE_REFERENCE_REPO", REPO_ROOT / "Z-Image"))
+MODEL_DIR = Path(os.getenv("ZIMAGE_MODEL_DIR", REPO_ROOT / "official_weights" / "Z-Image"))
+
+
+class _ConstantTransformer:
+    in_channels = 16
+
+    def __init__(self) -> None:
+        self.calls: list[tuple[torch.Tensor, list[torch.Tensor], torch.Tensor]] = []
+
+    def __call__(
+        self,
+        *,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: list[torch.Tensor],
+        timestep: torch.Tensor,
+    ) -> tuple[list[torch.Tensor], dict]:
+        self.calls.append((hidden_states, encoder_hidden_states, timestep))
+        return [torch.ones_like(sample) for sample in hidden_states.unbind(0)], {}
+
+
+class _RecordingVAE(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.anchor = torch.nn.Parameter(torch.zeros(()))
+        self.config = SimpleNamespace(scaling_factor=0.3611, shift_factor=0.1159)
+        self.decode_input: torch.Tensor | None = None
+
+    def decode(self, latents: torch.Tensor, return_dict: bool = False) -> tuple[torch.Tensor]:
+        assert return_dict is False
+        self.decode_input = latents.detach().clone()
+        batch, _, height, width = latents.shape
+        return (torch.zeros(batch, 3, height * 8, width * 8, device=latents.device),)
+
+
+def _stage_args(*, output_type: str = "pil") -> SimpleNamespace:
+    return SimpleNamespace(
+        pipeline_config=ZImagePipelineConfig(),
+        output_type=output_type,
+        vae_cpu_offload=False,
+        disable_autocast=False,
+    )
+
+
+def test_zimage_native_default_stage_math(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Exercise the pinned native 1-step formulas without model assets."""
+    autocast_calls: list[dict[str, object]] = []
+
+    @contextmanager
+    def record_autocast(**kwargs):
+        autocast_calls.append(kwargs)
+        yield
+
+    monkeypatch.setattr(
+        "fastvideo.pipelines.basic.zimage.stages.get_local_torch_device",
+        lambda: torch.device("cpu"),
+    )
+    monkeypatch.setattr("fastvideo.pipelines.basic.zimage.stages.torch.autocast", record_autocast)
+    args = _stage_args()
+    transformer = _ConstantTransformer()
+    scheduler = FlowMatchEulerDiscreteScheduler(shift=3.0)
+    vae = _RecordingVAE()
+    batch = ForwardBatch(
+        data_type="image",
+        prompt="test prompt",
+        negative_prompt="",
+        prompt_embeds=[torch.arange(32, dtype=torch.float32).reshape(1, 4, 8)],
+        prompt_attention_mask=[torch.tensor([[True, True, False, False]])],
+        height=64,
+        width=64,
+        num_frames=1,
+        num_inference_steps=1,
+        guidance_scale=0.0,
+        seed=42,
+        latents=torch.zeros(1, 16, 1, 8, 8),
+    )
+
+    ZImageInputValidationStage().forward(batch, args)
+    ZImageConditioningStage().forward(batch, args)
+    ZImageLatentPreparationStage(transformer).forward(batch, args)
+    ZImageTimestepPreparationStage(scheduler).forward(batch, args)
+
+    assert len(batch.extra["zimage_prompt_embeds"]) == 1
+    assert batch.extra["zimage_prompt_embeds"][0].shape == (2, 8)
+    assert_close(batch.timesteps, torch.tensor([1000.0]))
+    assert scheduler.sigma_min == 0.0
+    assert scheduler.config.use_reference_discrete_timesteps is True
+
+    ZImageDenoisingStage(transformer, scheduler).forward(batch, args)
+    assert autocast_calls == [{"device_type": "cpu", "enabled": False}]
+    assert_close(batch.latents, torch.ones_like(batch.latents))
+    _, encoded, normalized_timestep = transformer.calls[0]
+    assert encoded[0].shape == (2, 8)
+    assert_close(normalized_timestep, torch.tensor([0.0]))
+
+    ZImageDecodingStage(vae).forward(batch, args)
+    assert vae.decode_input is not None
+    expected_vae_input = torch.full((1, 16, 8, 8), 1 / 0.3611 + 0.1159)
+    assert_close(vae.decode_input, expected_vae_input)
+    assert batch.output is not None
+    assert batch.output.shape == (1, 3, 1, 64, 64)
+    assert_close(batch.output, torch.full_like(batch.output, 0.5))
+
+
+def test_zimage_batch_rng_matches_one_native_generator(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(
+        "fastvideo.pipelines.basic.zimage.stages.get_local_torch_device",
+        lambda: torch.device("cpu"),
+    )
+    args = _stage_args(output_type="latent")
+    transformer = _ConstantTransformer()
+    batch = ForwardBatch(
+        data_type="image",
+        prompt=["one", "two"],
+        negative_prompt="",
+        height=64,
+        width=64,
+        num_frames=1,
+        num_inference_steps=1,
+        guidance_scale=0.0,
+        seed=42,
+    )
+    ZImageInputValidationStage().forward(batch, args)
+    batch.extra["zimage_prompt_embeds"] = [torch.zeros(1, 8), torch.zeros(1, 8)]
+    ZImageLatentPreparationStage(transformer).forward(batch, args)
+
+    expected = torch.randn(
+        (2, 16, 1, 8, 8),
+        generator=torch.Generator("cpu").manual_seed(42),
+    )
+    assert_close(batch.latents, expected)
+
+
+def test_zimage_text_stage_honors_request_max_length(monkeypatch: pytest.MonkeyPatch) -> None:
+    from fastvideo.pipelines.stages.text_encoding import TextEncodingStage
+
+    args = _stage_args()
+    stage = TextEncodingStage(text_encoders=[object()], tokenizers=[object()])
+    observed: list[int | None] = []
+
+    def fake_encode_text(
+        text,
+        fastvideo_args,
+        encoder_index,
+        return_attention_mask,
+        max_length=None,
+    ):
+        del text, fastvideo_args, encoder_index, return_attention_mask
+        observed.append(max_length)
+        return [torch.zeros(1, 4, 8)], [torch.ones(1, 4, dtype=torch.long)]
+
+    monkeypatch.setattr(stage, "encode_text", fake_encode_text)
+    batch = ForwardBatch(
+        data_type="image",
+        prompt="test",
+        negative_prompt="",
+        max_sequence_length=64,
+        guidance_scale=0.0,
+    )
+    stage.forward(batch, args)
+    assert observed == [64]
+
+
+def test_zimage_cfg_none_uses_native_empty_negative_prompt(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(
+        "fastvideo.pipelines.basic.zimage.stages.get_local_torch_device",
+        lambda: torch.device("cpu"),
+    )
+    batch = ForwardBatch(
+        data_type="image",
+        prompt="test",
+        negative_prompt=None,
+        height=64,
+        width=64,
+        num_frames=1,
+        num_inference_steps=1,
+        guidance_scale=2.0,
+        seed=42,
+    )
+    ZImageInputValidationStage().forward(batch, _stage_args())
+    assert batch.negative_prompt == ""
+
+
+def _require_pinned_reference() -> Path:
+    if not REFERENCE_REPO.is_dir():
+        pytest.skip(f"Pinned Tongyi Z-Image clone not found: {REFERENCE_REPO}")
+    result = subprocess.run(
+        ["git", "-C", str(REFERENCE_REPO), "rev-parse", "HEAD"],
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+    actual = result.stdout.strip()
+    assert actual == REFERENCE_REVISION, (
+        f"Tongyi Z-Image oracle must be pinned at {REFERENCE_REVISION}, got {actual}"
+    )
+    return REFERENCE_REPO / "src"
+
+
+def _import_reference_modules() -> tuple[object, object]:
+    source_root = _require_pinned_reference().resolve()
+    if str(source_root) not in sys.path:
+        sys.path.insert(0, str(source_root))
+    pipeline_module = importlib.import_module("zimage.pipeline")
+    utils_module = importlib.import_module("utils")
+    for module in (pipeline_module, utils_module):
+        module_path = Path(module.__file__ or "").resolve()
+        assert module_path.is_relative_to(source_root), (
+            f"Z-Image oracle import escaped the pinned repository: {module_path}"
+        )
+    return pipeline_module, utils_module
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Full Z-Image pipeline parity requires CUDA")
+def test_zimage_pipeline_latents_match_pinned_native_repo(tmp_path: Path) -> None:
+    """Compare one complete native step; the oracle never imports Diffusers' pipeline."""
+    if not (MODEL_DIR / "model_index.json").is_file():
+        pytest.skip(f"Pinned Z-Image weights not found: {MODEL_DIR}")
+    pipeline_module, utils_module = _import_reference_modules()
+    prompt = "Young Chinese woman in red Hanfu, intricate embroidery."
+
+    previous_fastvideo_backend = os.environ.get("FASTVIDEO_ATTENTION_BACKEND")
+    try:
+        # Exercise the shared PyTorch SDPA path with each process's native
+        # kernel selection; the official and FastVideo implementations both
+        # pass the same padding mask into scaled_dot_product_attention.
+        utils_module.set_attention_backend(None)
+        components = utils_module.load_from_local_dir(
+            MODEL_DIR,
+            device="cuda",
+            dtype=torch.bfloat16,
+            compile=False,
+        )
+        reference = pipeline_module.generate(
+            prompt=prompt,
+            **components,
+            height=64,
+            width=64,
+            num_inference_steps=1,
+            guidance_scale=0.0,
+            negative_prompt="",
+            max_sequence_length=512,
+            generator=torch.Generator("cuda").manual_seed(42),
+            output_type="latent",
+        ).detach().float().cpu()
+        del components
+        gc.collect()
+        torch.cuda.empty_cache()
+
+        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "TORCH_SDPA"
+        from fastvideo import VideoGenerator
+
+        generator = VideoGenerator.from_pretrained(
+            str(MODEL_DIR),
+            num_gpus=1,
+            tp_size=1,
+            sp_size=1,
+            use_fsdp_inference=False,
+            dit_cpu_offload=False,
+            dit_layerwise_offload=False,
+            text_encoder_cpu_offload=False,
+            vae_cpu_offload=False,
+            pin_cpu_memory=False,
+            output_type="latent",
+        )
+        try:
+            result = generator.generate_video(
+                prompt,
+                negative_prompt="",
+                output_path=str(tmp_path),
+                save_video=False,
+                return_frames=True,
+                height=64,
+                width=64,
+                num_frames=1,
+                fps=1,
+                num_inference_steps=1,
+                guidance_scale=0.0,
+                max_sequence_length=512,
+                seed=42,
+            )
+        finally:
+            generator.shutdown()
+        actual = result["samples"].squeeze(2).float()
+        assert actual.shape == reference.shape
+        diff = (actual - reference).abs()
+        print(f"Z-Image pipeline latent parity: max={diff.max():.6f}, mean={diff.mean():.6f}")
+        assert_close(actual, reference, atol=1e-2, rtol=1e-2)
+    finally:
+        if previous_fastvideo_backend is None:
+            os.environ.pop("FASTVIDEO_ATTENTION_BACKEND", None)
+        else:
+            os.environ["FASTVIDEO_ATTENTION_BACKEND"] = previous_fastvideo_backend

tests/local_tests/pipelines/test_zimage_pipeline_smoke.py

@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Import/registry preflight and optional real Z-Image-Turbo smoke."""
+
+from __future__ import annotations
+
+import os
+from contextlib import nullcontext
+from pathlib import Path
+from typing import Any, cast
+
+import pytest
+import torch
+
+
+MODEL_DIR_ENV = os.getenv("ZIMAGE_MODEL_DIR")
+
+
+def test_zimage_typed_surface_preflight() -> None:
+    from fastvideo.api import GeneratorConfig, PipelineSelection
+    from fastvideo.api.compat import generator_config_to_fastvideo_args
+    import fastvideo.registry as registry
+    from fastvideo.api.presets import get_preset, get_presets_for_family
+    from fastvideo.configs.models.dits.zimage import ZImageDiTConfig
+    from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig
+    from fastvideo.configs.models.vaes.autoencoder_kl import AutoencoderKLVAEConfig
+    from fastvideo.configs.pipelines.zimage import ZImagePipelineConfig
+    from fastvideo.fastvideo_args import WorkloadType
+    from fastvideo.pipelines.basic.zimage.zimage_pipeline import EntryClass, ZImagePipeline
+
+    assert ZImagePipeline.__name__ == "ZImagePipeline"
+    assert EntryClass is ZImagePipeline
+    assert ZImagePipeline.pipeline_config_cls is ZImagePipelineConfig
+    assert ZImagePipeline._required_config_modules == [
+        "scheduler",
+        "text_encoder",
+        "tokenizer",
+        "transformer",
+        "vae",
+    ]
+
+    default_preset, model_family = registry.get_preset_selection("Tongyi-MAI/Z-Image-Turbo")
+    assert (default_preset, model_family) == ("zimage_turbo", "zimage")
+    info = registry.get_model_info(
+        "Tongyi-MAI/Z-Image-Turbo",
+        workload_type=WorkloadType.T2I,
+        override_pipeline_cls_name="ZImagePipeline",
+    )
+    assert info.pipeline_cls is ZImagePipeline
+    assert info.pipeline_config_cls is ZImagePipelineConfig
+
+    assert {preset.name for preset in get_presets_for_family("zimage")} == {"zimage_turbo"}
+    preset = get_preset("zimage_turbo", "zimage")
+    assert preset.defaults == {
+        "height": 1024,
+        "width": 1024,
+        "num_frames": 1,
+        "fps": 1,
+        "seed": 42,
+        "guidance_scale": 0.0,
+        "num_inference_steps": 8,
+        "negative_prompt": "",
+        "max_sequence_length": 512,
+        "cfg_normalization": False,
+        "cfg_truncation": 1.0,
+    }
+
+    config = ZImagePipelineConfig()
+    assert isinstance(config.dit_config, ZImageDiTConfig)
+    assert isinstance(config.vae_config, AutoencoderKLVAEConfig)
+    assert isinstance(config.text_encoder_configs[0], Qwen3TextConfig)
+    assert config.text_encoder_configs[0].chat_template_enable_thinking is True
+    assert config.text_encoder_configs[0].output_hidden_states is True
+    assert config.text_encoder_configs[0].tokenizer_kwargs["max_length"] == 512
+    assert config.scheduler_sigma_min == 0.0
+    assert config.scheduler_use_reference_discrete_timesteps is True
+
+    args = generator_config_to_fastvideo_args(
+        GeneratorConfig(
+            model_path="Tongyi-MAI/Z-Image-Turbo",
+            pipeline=PipelineSelection(workload_type="t2i"),
+        ))
+    assert isinstance(args.pipeline_config, ZImagePipelineConfig)
+    assert args.workload_type is WorkloadType.T2I
+
+
+def test_model_download_honors_pinned_revision(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+    import fastvideo.utils as utils
+
+    observed: dict[str, object] = {}
+
+    def fake_snapshot_download(**kwargs):
+        observed.update(kwargs)
+        return str(tmp_path)
+
+    monkeypatch.setattr(utils, "snapshot_download", fake_snapshot_download)
+    monkeypatch.setattr(utils, "get_lock", lambda _key: nullcontext())
+    resolved = utils.maybe_download_model(
+        "Tongyi-MAI/Z-Image-Turbo",
+        revision="f332072aa78be7aecdf3ee76d5c247082da564a6",
+    )
+
+    assert resolved == str(tmp_path)
+    assert observed["repo_id"] == "Tongyi-MAI/Z-Image-Turbo"
+    assert observed["revision"] == "f332072aa78be7aecdf3ee76d5c247082da564a6"
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Z-Image load/generate smoke requires CUDA")
+def test_zimage_pipeline_load_generate_smoke() -> None:
+    if MODEL_DIR_ENV is None:
+        pytest.skip("Set ZIMAGE_MODEL_DIR to activate the real load/generate smoke")
+    model_dir = Path(MODEL_DIR_ENV)
+    if not (model_dir / "model_index.json").is_file():
+        pytest.skip(f"Z-Image model_index.json not found under {model_dir}")
+
+    from fastvideo import VideoGenerator
+
+    generator = VideoGenerator.from_pretrained(
+        str(model_dir),
+        num_gpus=1,
+        tp_size=1,
+        sp_size=1,
+        use_fsdp_inference=False,
+        dit_cpu_offload=False,
+        dit_layerwise_offload=False,
+        text_encoder_cpu_offload=False,
+        vae_cpu_offload=False,
+        pin_cpu_memory=False,
+        output_type="latent",
+    )
+    try:
+        result = generator.generate_video(
+            prompt="a red panda reading a book",
+            negative_prompt="",
+            output_path="outputs/zimage/smoke",
+            save_video=False,
+            return_frames=True,
+            height=64,
+            width=64,
+            num_frames=1,
+            fps=1,
+            num_inference_steps=1,
+            guidance_scale=0.0,
+            max_sequence_length=64,
+            cfg_normalization=False,
+            cfg_truncation=1.0,
+            seed=42,
+        )
+    finally:
+        generator.shutdown()
+
+    result_dict = cast(dict[str, Any], result)
+    samples = result_dict["samples"]
+    assert torch.is_tensor(samples)
+    assert samples.ndim in (4, 5)
+    assert torch.isfinite(samples).all()

tests/local_tests/zimage/PORT_STATUS.md

@@ -0,0 +1,162 @@
+# Z-Image Port Status
+
+## Summary
+
+- model_family: `zimage`
+- variant: `Z-Image-Turbo`
+- workload_types: `T2I`
+- PR scope: required components plus same-PR FastVideo-native pipeline, example/parity, SSIM test, and L40S seed workflow
+- official_ref: `https://github.com/Tongyi-MAI/Z-Image`
+- official_ref_dir: `<repo_root>/Z-Image/src`
+- official_ref_commit: `26f23eda626ffadda020b04ff79488e1d72004cd`
+- hf_weights_path: `Tongyi-MAI/Z-Image-Turbo@f332072aa78be7aecdf3ee76d5c247082da564a6`
+- local_weights_dir: `<repo_root>/official_weights/Z-Image/`
+- source_layout: official repo-native `src/zimage/`; HF component subfolders provide assets only
+- weight_layout: HF per-component subfolders; these are assets, not the reference implementation
+- local_tests_readme: `tests/local_tests/zimage/README.md`
+
+The pinned text encoder has 36 decoder blocks, hidden size 2560,
+intermediate size 9728, 32 attention heads, 8 KV heads, and head dimension
+128. FastVideo obtains those values from the pinned `config.json`.
+
+## Current Phase
+
+- phase: Phase 11 (post-parity review and handoff); quality reference published
+- status: external_ci_gate
+- owner: infra/reviewer
+- last_updated: 2026-07-15
+
+## Pinned Official Pipeline Context
+
+- official_pipeline: `Tongyi-MAI/Z-Image@26f23eda626ffadda020b04ff79488e1d72004cd`,
+  `src/zimage/pipeline.py::generate` and `src/config/inference.py`
+- oracle_policy: official Z-Image source only; Diffusers is not the implementation oracle
+- workload_and_output: `T2I`; prompt plus optional negative prompt to PIL images, or latents when requested
+- model_index_contract: `_class_name=ZImagePipeline`; required modules are `transformer`, `vae`,
+  `text_encoder`, `tokenizer`, and `scheduler`
+- official_defaults: `height=1024`, `width=1024`, `num_inference_steps=8`,
+  `guidance_scale=0.0`, `cfg_truncation=1.0`, `max_sequence_length=512`
+- reproducible_rng: `torch.Generator("cuda").manual_seed(42)`; the API accepts a generator and has no fixed seed default
+- scheduler_contract: `use_reference_discrete_timesteps=True`, runtime `sigma_min=0.0`, endpoint-preserving
+  `num_steps + 1` schedule, and no DiT call for the final zero timestep
+- decode_contract: apply `(latents / 0.3611) + 0.1159` before official VAE decode
+- implementation_status: native pipeline/config/preset/registry/example and the 8-step 1024x1024 PNG SSIM test are
+  implemented; one-step native GPU parity is bit-exact; the L40S candidate generated at inference head `9a6dbcadc`
+  was author-approved and published at dataset revision `7b27d00004b723ef0cf4923e77ae57128573f01c`
+
+## Component Matrix
+
+| Component | Type | Reuse/Port | Official Definition / Instantiation | FastVideo Target | Prototype | Conversion | Parity | Open Issues |
+|---|---|---|---|---|---|---|---|---|
+| Production text encoder (Qwen3 body) | text_encoder | reused production passthrough | Official `src/utils/loader.py` uses Transformers with typed load then device-only move; official pipeline consumes `hidden_states[-2]` | `TextEncoderLoader` -> `Qwen3ForCausalLM.from_pretrained_local` -> body-only Transformers `AutoModel` | DONE | not_needed for current HF subfolder | PINNED GB200 PASS; reference/production outputs exact | none |
+| Native Qwen quality | implementation_subcomponent | shared native implementation | Same pinned text-encoder weights and official `hidden_states[-2]` boundary; not loaded by the pipeline | direct `fastvideo.models.encoders.qwen3.Qwen3ForCausalLM` construction | DONE | not_needed | PINNED GB200 PASS; fp32 `atol=2e-4`, `rtol=1e-4`; bf16 thresholds unchanged | none |
+| Tokenizer | tokenizer | reused | Official loader's `AutoTokenizer` from `<weights>/tokenizer/`; official thinking template and length 512 | `TokenizerLoader` | DONE | not_needed | PINNED GB200 PASS (2 passed) | none |
+| VAE | vae | reused with accepted wrapper exception | Pinned official `src/zimage/autoencoder.py`; raw decode is the component boundary | Existing Diffusers-backed `fastvideo.models.vaes.autoencoder_kl.AutoencoderKL` through production `VAELoader`, plus direct raw-decode coverage | DONE | not_needed | PINNED GB200 PASS (2 passed) | none |
+| Scheduler | scheduler | reused with Z-Image extension | Pinned official `src/zimage/scheduler.py`; official pipeline sets `sigma_min=0.0` | `FlowMatchEulerDiscreteScheduler(use_reference_discrete_timesteps=True, sigma_min=0.0)` | DONE | not_needed | PINNED GB200 PASS (5 passed) | none |
+| Transformer (`ZImageTransformer2DModel`) | dit | native port | Pinned official `src/zimage/transformer.py`; HF `transformer/` supplies weights only | `fastvideo.models.dits.zimage.ZImageTransformer2DModel` through `TransformerLoader` | DONE | not_needed; exact 521-key surface | PINNED GB200 PASS (5 passed; real-weight full forward exact) | none |
+| Pipeline | pipeline | native same-PR implementation | Pinned official `src/zimage/pipeline.py::generate` | FastVideo `ZImagePipeline` plus config/preset/registry/example/parity | DONE | not_needed | RTX 5090 ONE-STEP LATENTS BIT-EXACT; L40S REFERENCE PUBLISHED | I005 (CI verification only) |
+
+## Conversion State
+
+- conversion_script: none for the reused components or native transformer
+- converted_weights_dir: none
+- source_layout: official repo-native Python (`Tongyi-MAI/Z-Image`)
+- hf_weight_layout: per-component subfolders; not a Diffusers implementation oracle
+- transformer_conversion: not_needed; official, HF index, and native names align exactly
+- transformer_key_surface: 521 tensors; sorted-key SHA256 `3a9216f208c1873b2cf06394411a53e1e95e10fae3b01dca0f7223556e47c354`
+- strict_load_status: transformer real-weight strict load passed; native Qwen strictness regressions passed
+- passthrough_components: body-only Qwen3 `AutoModel`, tokenizer assets
+- direct_load_components: native Qwen3 parity target, VAE, scheduler
+- native_qwen_exclusion: `lm_head.weight` only; the native encoder owns the body, not an LM head
+- native_qwen_fused_contract: every Q/K/V and gate/up shard is required unless the exact fused destination is present
+- native_qwen_auxiliary_contract: mapped quantization auxiliaries, including scale parameters, must load and may not be silently discarded
+- retry_history: none
+
+## Parity Commands
+
+| Scope | Command | Last Result | Notes |
+|---|---|---|---|
+| Scheduler | `pytest tests/local_tests/zimage/test_zimage_scheduler_parity.py -v -s` | 5 PASSED, no skips, GB200, 2026-07-15 | Enforces official-source SHA/import origin, fp32 scalar contract, exact zero endpoint, and positional/default regressions. |
+| Tokenizer | `pytest tests/local_tests/zimage/test_zimage_tokenizer_parity.py -v -s` | 2 PASSED as part of the pinned full GB200 run, 2026-07-15 | Uses `enable_thinking=True`, `max_length=512`, and fails if chat templating is unavailable. |
+| VAE | `pytest tests/local_tests/zimage/test_zimage_vae_parity.py -v -s` | 2 PASSED, no skips, GB200, 2026-07-15 | Direct raw decode plus production `VAELoader`; config values are checked, while latent scale/shift application is pipeline behavior. |
+| Production text encoder | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_parity_forward[fp32] -v -s` | REFERENCE/PRODUCTION CHECKS EXACT, GB200, 2026-07-15 | The required pipeline path matches an independent official-loader materialization exactly. |
+| Native Qwen fp32 quality | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_parity_forward[fp32] -v -s` | PASS, GB200, 2026-07-15 | Z-Image-only `atol=2e-4`, `rtol=1e-4`; production equality and shared Qwen math are unchanged. |
+| Native Qwen bf16 quality | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_parity_forward[bf16] -v -s` | PASS, GB200, 2026-07-15 | Existing mean/median distribution thresholds pass unchanged against an independent materialization of the official loader contract. |
+| Native Qwen per-layer bf16 diagnostic | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_per_layer_bf16_diagnostic -v -s` | PASS, GB200, 2026-07-15 | Pinned architecture has 36 blocks and 37 hidden-state entries under the tested contract. |
+| Transformer parity | `ZIMAGE_OFFICIAL_REF_DIR=<pinned-clone> pytest tests/local_tests/zimage/test_zimage_transformer_parity.py -v -s` | 5 PASSED, no skips, GB200, 2026-07-15 | Production 521-key/shape surface, tiny CPU parity, mask behavior, SP guard, real-weight strict load, and 6.15B full forward pass; real-weight output was exact (`max=0`, `mean=0`). |
+| Complete component closure | `ZIMAGE_OFFICIAL_REF_DIR=<pinned-clone> ZIMAGE_TRANSFORMER_DIR=<pinned-transformer> pytest tests/local_tests/zimage/ -v -s` | 34 PASSED, no skips, GB200, 2026-07-15, `f310aacb1` | Official source `26f23eda...`, HF snapshot `f332072a...`; production/native encoder, scheduler, tokenizer, transformer, and VAE all passed. |
+| Pipeline typed surface | `pytest tests/local_tests/pipelines/test_zimage_pipeline_smoke.py::test_zimage_typed_surface_preflight -v` | PASS, hardware-free, 2026-07-14 | Covers exact `EntryClass`, config, preset, registry, model detection, and official defaults. |
+| Native stage contract | `pytest tests/local_tests/pipelines/test_zimage_pipeline_parity.py::test_zimage_native_default_stage_math -v` | PASS, hardware-free, 2026-07-14 | Covers conditioning trim, endpoint-preserving timestep setup, denoising sign/timestep math, and VAE scale/shift. |
+| Full pipeline parity | `ZIMAGE_REFERENCE_REPO=<pinned-clone> ZIMAGE_MODEL_DIR=<pinned-weights> DISABLE_SP=1 pytest tests/local_tests/pipelines/test_zimage_pipeline_parity.py::test_zimage_pipeline_latents_match_pinned_native_repo -v -s` | PASS, RTX 5090, `fec2baa6c` | Real FastVideo one-step latents match pinned native Tongyi exactly (`max=0`, `mean=0`). Later commits through `f310aacb1` change only API inventory/tests, not inference math. |
+| PNG SSIM quality | `pytest fastvideo/tests/ssim/test_zimage_similarity.py -v -s` | L40S REFERENCE UPLOADED; END-TO-END SSIM CI PENDING | Dataset revision `7b27d00004b723ef0cf4923e77ae57128573f01c`; remote SHA-256 `6c2be648...f7a6`; upload scoped to `Tongyi-MAI__Z-Image-Turbo/TORCH_SDPA`. |
+
+## Open Questions
+
+| ID | Question | Owner | Needed By Phase | Status | Resolution |
+|---|---|---|---|---|---|
+| Q001 | Which exact Z-Image source revision is the parity oracle? | prep | Phase 1 | resolved | Pinned `Tongyi-MAI/Z-Image@26f23eda626ffadda020b04ff79488e1d72004cd`; scheduler/VAE tests fail on a wrong HEAD or import origin. |
+| Q002 | Which immutable HF snapshot supplies component weights? | prep | Phase 1 | resolved | `Tongyi-MAI/Z-Image-Turbo@f332072aa78be7aecdf3ee76d5c247082da564a6`. |
+| Q003 | What prompt/tokenization/hidden-state contract does the official pipeline use? | pipeline | Phase 6 | resolved | `apply_chat_template(tokenize=False, add_generation_prompt=True, enable_thinking=True)`, tokenization length 512, and valid-token `hidden_states[-2]`. |
+| Q004 | Does the native transformer require conversion? | conversion | Phase 5 | resolved | No. Pinned HF index, official production meta model, and FastVideo production meta model expose the same 521 keys; official/native key shapes match exactly. |
+
+## Issues And Blockers
+
+| ID | Phase | Component | Severity | Issue | Evidence | Owner | Status | Resolution |
+|---|---|---|---|---|---|---|---|---|
+| I001 | Phase 6 | native_qwen_quality | medium | The optional native fp32 `hidden_states[-2]` tail narrowly exceeded `atol=1e-4` after smooth cross-kernel accumulation. | Exact layer trace: embeddings exact, layer-0 max `3.815e-6`, first misses at layer 33, largest failing boundary difference `1.686811e-4`; final RMSNorm passes. The required production `AutoModel` path is exact. | parity | resolved | Author approved Z-Image-only `_FP32_ATOL=2e-4` with `_FP32_RTOL=1e-4`; complete pinned GB200 suite passed 34/34 with bf16 and all other gates unchanged. |
+| I002 | Phase 7 | scheduler/pipeline | high | Published `scheduler_config.json` omits the reference-schedule flag and official zero endpoint. | The native stage applies both values; hardware-free stage coverage and pinned native GPU parity pass. | pipeline | resolved | The Z-Image pipeline applies `use_reference_discrete_timesteps=True` and `sigma_min=0.0`; one-step latents match official Tongyi exactly. |
+| I003 | Phase 6 | tokenizer/text_encoder | medium | Prompt formatting and target length previously differed from the official call. | Pinned source uses thinking chat template, 512 tokens, and `hidden_states[-2]`. | parity | resolved | Corrected behavior passed the pinned GB200 run. |
+| I004 | Phase 4 | transformer | high | `ZImageTransformer2DModel` was not ported. | Native config/model/registry and deterministic parity test now exist. | port | resolved | Production meta surface matches all 521 pinned HF keys; tiny full-forward parity passes exactly. |
+| I005 | Phase 10 | pipeline/quality | high | Verify the published L40S PNG through end-to-end SSIM CI. | Author approved the reviewed candidate; the guarded upload published exactly one file at dataset revision `7b27d000...f01c`, and remote bytes match SHA-256 `6c2be648...f7a6`. | infra | external_ci_gate | Reference publication is resolved. The full-suite launcher remains blocked by the external-fork documentation workflow's `action_required` state; merge is separate. |
+| I006 | Phase 6 | text_encoder | high | HF passthrough loading inherited an ambient target-device context, and the earlier test harness incorrectly cast float32 RoPE buffers during bf16 placement. | HF loading now occurs outside the ambient context; typed load is followed by device-only movement, matching official `src/utils/loader.py`; production/independent outputs match exactly. | encoder | resolved | Hardware-free placement/default-dtype tests and the pinned GB200 production path pass. |
+| I007 | Phase 6 | text_encoder | high | Lenient native loading could miss one fused source shard or an auxiliary quantization scale. | Loader completeness covers every destination, required split shard, exact fused destination, and mapped scale parameter. | encoder | resolved | Strictness regressions and pinned checkpoint loading pass. |
+| I008 | Phase 6 | vae | medium | Direct decode alone did not cover production registry/config/device/strict-load behavior; the shared target is a runtime Diffusers wrapper. | VAE parity now covers direct raw decode and production `VAELoader` against official `src/zimage/autoencoder.py`. | parity | resolved | Pinned GB200 subset passed 2/2. The wrapper is an implementation exception, not the oracle; official latent normalization remains Phase 7 pipeline scope. |
+| I009 | Phase 2 | reference assets | high | Adding a nonexistent clone path to `sys.path` could import another installed `zimage` package and yield false parity. | Tests now verify exact clone HEAD and module paths under `Z-Image/src`; `/Z-Image/` is ignored. | prep | resolved | Wrong SHA/path/import is a failure; absent local assets may skip. |
+| I010 | Phase 6 | tokenizer | medium | Missing `apply_chat_template` previously skipped and could hide an incompatible tokenizer. | Tokenizer parity now asserts the API and exact official kwargs. | parity | resolved | Missing chat-template behavior fails; pinned asset-backed parity passed. |
+| I011 | Phase 6 | transformer | high | Production real-weight strict load and full-forward parity were not verified. | Pinned 24.6 GB weights and official source were exercised on GB200. | parity | resolved | Strict load and full forward passed exactly (`max=0`, `mean=0`). |
+| I012 | Phase 6 | required_components | high | Record a complete non-skip run after the loader/oracle corrections and approved native-Qwen tolerance. | Exact PR head `f310aacb1`, pinned official source `26f23eda...`, immutable HF snapshot `f332072a...`, GB200. | parity | resolved | Complete suite passed 34/34 in 94.21s with no skips; real-weight transformer output was exact (`max=0`, `mean=0`). |
+
+## Escape Hatches
+
+| ID | Phase | Decision Type | Question | Recommended Option | Status | Resolution |
+|---|---|---|---|---|---|---|
+| E001 | Phase 6 | cost | Approve one standard 8-hour 4×GB200 hold to run the complete pinned component and real-weight transformer parity suite at PR #1339? | Approve `model/r9`; reuse a suitable existing allocation if available. | resolved | Author approved. Existing job `1545932` was reused, so no new hold was created. Excludes Modal, SSIM seeding/upload, and merge. |
+| E002 | Phase 6 | numerical | Accept the traced Z-Image-only native encoder fp32 tail without weakening shared Qwen or bf16 gates? | Approve `model/r10`: change only `_FP32_ATOL` from `1e-4` to `2e-4`, retain `_FP32_RTOL=1e-4`, then rerun native parity at the exact resulting head. | resolved | Author approved. The one-line test threshold changed at `f310aacb1`; complete pinned GB200 suite passed 34/34 with production equality, bf16, strict-load, and official-oracle checks unchanged. |
+
+## Decisions
+
+| Date | Decision | Rationale | Impact |
+|---|---|---|---|
+| 2026-07-15 | Publish the visually approved L40S PNG reference to `FastVideo/ssim-reference-videos`. | The author confirmed the staged image looks good and explicitly requested the HF upload. | Guarded no-force upload added one Z-Image PNG at dataset revision `7b27d000...f01c`; remote bytes match the reviewed artifact. |
+| 2026-07-15 | Accept the traced Z-Image-only native encoder fp32 absolute tolerance of `2e-4` while retaining `rtol=1e-4`. | Exact tracing showed gradual fused/split-kernel accumulation rather than a semantic break; the largest failing absolute tail was `1.686811e-4`. | One test constant changed. The full pinned suite passed 34/34; production equality, bf16 thresholds, strict loading, and official Tongyi oracle are unchanged. |
+| 2026-07-14 | Carry the native FastVideo pipeline, example/parity, SSIM test, and L40S seed workflow on PR #1339. | The author explicitly requested the pipeline and quality regression on the same PR rather than a follow-up. | Implementation, component closure, native GPU parity, and reference publication are complete. Later commits do not change inference math. |
+| 2026-07-14 | Treat the exact production `AutoModel` loader as the required pipeline text encoder and direct native-Qwen parity as a separate quality row. | `TextEncoderLoader` calls `Qwen3ForCausalLM.from_pretrained_local`, which returns body-only Transformers `AutoModel`; the pipeline does not direct-construct FastVideo-native Qwen. | The production component is exact. The separately approved native-Qwen tolerance changed no production math. |
+| 2026-07-13 | Use pinned `Tongyi-MAI/Z-Image@26f23eda...` as the implementation oracle; treat HF subfolders only as weight/config assets. | The official repo defines transformer, VAE, scheduler, pipeline, and loader semantics. A runtime component may inherit Diffusers without making Diffusers the reference. | Tests verify official clone SHA/import origin; docs distinguish the oracle from storage/runtime implementation details. |
+| 2026-07-13 | Keep official latent scale/shift application out of the VAE component assertion. | `(latents / 0.3611) + 0.1159` is pipeline-stage behavior, while VAE component parity compares raw decode and production loading. | The Z-Image stage implements the exact formula; pinned native GPU parity passed. |
+| 2026-07-13 | Preserve the published transformer key surface and use raw masked SDPA at SP=1. | The checkpoint already matches the official class, while padded variable-length attention needs a key mask unavailable in current distributed wrappers. | No converter is needed; runtime rejects SP>1 until a mask-aware distributed path has parity evidence. |
+| 2026-07-04 | Pin both external sources and treat all earlier GPU results as historical until rerun. | Earlier evidence did not record an immutable HF revision, and source imports were not origin-checked. | README has executable clone/download commands; tests enforce the source SHA; all asset-backed parity rows were revalidated. |
+| 2026-07-04 | Use a body-only Transformers `AutoModel` for production and a separate `AutoModel` object to materialize the official loader contract independently. | The pinned official Z-Image source is the oracle. Every FastVideo consumer needs hidden states, not full-vocabulary logits or an LM head, and comparing an object with itself is not independent evidence. | Production avoids the LM head while parity covers production loading and native implementation separately. |
+| 2026-07-04 | Make fused-source and quantization-scale completeness part of native Qwen strict loading. | A loaded destination name alone did not prove that every required Q/K/V or gate/up shard and auxiliary scale arrived. | Missing split shards or mapped scale parameters fail; an exact fused destination remains valid. |
+| 2026-07-04 | Reuse the existing Diffusers-backed AutoencoderKL wrapper for the Z-Image VAE component. | The target is established shared code with production-loader parity against pinned official `src/zimage/autoencoder.py`. | The wrapper remains an implementation detail for this VAE only; official Z-Image source defines pipeline behavior and numerical parity. |
+| 2026-07-04 | Treat absent chat-template behavior as incompatibility, not an optional test environment. | Thinking-format prompt construction is part of the official numerical contract. | Tokenizer parity fails when `apply_chat_template` is unavailable. |
+| 2026-07-04 | Record the pinned Qwen architecture as 36 blocks, hidden size 2560, intermediate size 9728, 32 attention heads, and 8 KV heads. | Earlier notes described a different architecture and could mis-size the native target. | Config/review evidence now matches `f332072aa78be7aecdf3ee76d5c247082da564a6`; expected hidden-state length is 37. |
+| 2026-06-21 | Expand generated Qwen `position_ids` to `[batch_size, seq_len]`. | The rotary layer flattens positions; `[1, seq_len]` misaligned RoPE for batch sizes greater than one. | Batch-one behavior is unchanged; pinned batch-two parity passed. |
+| 2026-06-21 | Reuse the shared native `Qwen3ForCausalLM` implementation rather than maintain a second Z-Image encoder. | The shared implementation supports the same config-driven GQA, fused QKV/gate-up, and hidden-state output contract. | Registry/native parity uses the shared class; architecture values come from the pinned config. |
+| 2026-05-12 | Use mean and median drift checks for deep bf16 encoder parity. | Cross-kernel fused/unfused bf16 arithmetic produces a long max-error tail; fp32 and per-layer diagnostics distinguish this from structural divergence. | Existing thresholds passed unchanged on the pinned 36-block snapshot. |
+| 2026-05-12 | Forward the complete scheduler config except loader metadata. | Hand-picking three fields could silently omit future scheduler behavior. | Scheduler parity consumes all on-disk constructor fields and explicitly supplies Z-Image runtime overrides. |
+
+## Handoff Notes
+
+- The implementation oracle is the pinned official `Tongyi-MAI/Z-Image` clone,
+  never Diffusers. Transformers is used for text components only because the
+  official loader does so; HF component subfolders provide immutable assets.
+- The required production text encoder is the exact body-only `AutoModel` path.
+  `model/r10` is resolved: only the Z-Image test's fp32 absolute tolerance is
+  `2e-4`; relative tolerance, shared Qwen math, bf16, and production loading are unchanged.
+- I012 is closed by the exact-head 34/34 non-skip GB200 run at `f310aacb1`.
+- Native pipeline/config/preset/registry/example, parity tests, and the 8-step
+  1024x1024 PNG SSIM test are implemented. Native RTX 5090 one-step parity is
+  bit-exact. The reviewed L40S PNG is live at dataset revision
+  `7b27d00004b723ef0cf4923e77ae57128573f01c` with matching remote bytes.
+- End-to-end SSIM/full-suite execution remains an infra/reviewer gate; reference
+  publication does not authorize PR approval or merge.

tests/local_tests/zimage/README.md

@@ -0,0 +1,200 @@
+# Z-Image Local Tests
+
+Local component coverage and quality handoff for the `zimage` FastVideo port
+(`T2I`, Z-Image-Turbo only). The same PR is authorized to carry the components,
+FastVideo-native `ZImagePipeline`, example, pipeline parity, SSIM test, and L40S
+reference seed. The native pipeline surface and 8-step 1024x1024 PNG SSIM test
+are implemented. Native one-step GPU parity is bit-exact, and the L40S candidate
+has been published to `FastVideo/ssim-reference-videos`.
+
+> **Status:** Implementation and quality-reference seeding complete; Phase 11
+> final handoff. The
+> Phase 6 exact-head closure is complete: the full pinned component suite passed
+> 34/34 with no skips on GB200 at `f310aacb1`. The required production text
+> encoder is the official-loader-compatible Transformers `AutoModel` path, which
+> matches an independently materialized reference exactly. Scheduler, tokenizer,
+> VAE, native Qwen fp32/bf16, and the 6.15B real-weight transformer all pass.
+> Native one-step pipeline parity is bit-exact on RTX 5090. The reviewed
+> 1024x1024 L40S PNG generated at inference head `9a6dbcadc` is published at
+> dataset revision `7b27d00004b723ef0cf4923e77ae57128573f01c`. See
+> [`PORT_STATUS.md`](./PORT_STATUS.md) for the live state.
+
+## Reference assets and scope
+
+| Field | Value |
+|---|---|
+| Model family / variant | `zimage` / `Z-Image-Turbo` |
+| Workload | `T2I` |
+| PR scope | required components, native pipeline, example/parity, SSIM test, and L40S seed |
+| Required component scope | production `AutoModel` text encoder, tokenizer, VAE, scheduler, native transformer |
+| Additional quality scope | direct FastVideo-native Qwen parity (`model/r10`; not used by the pipeline loader) |
+| Official reference | `https://github.com/Tongyi-MAI/Z-Image` |
+| Local reference dir | `<repo_root>/Z-Image/src` |
+| Official commit | `26f23eda626ffadda020b04ff79488e1d72004cd` |
+| HF weights | `Tongyi-MAI/Z-Image-Turbo@f332072aa78be7aecdf3ee76d5c247082da564a6` |
+| Local weights dir | `<repo_root>/official_weights/Z-Image/` (component subfolders; transformer is 24.6 GB) |
+| Reference implementation layout | Official repo-native `src/zimage/` Python; Diffusers is not the parity oracle |
+| Weight layout | HF per-component subfolders (`text_encoder/`, `vae/`, `scheduler/`, `transformer/`) |
+| Conversion | not needed; the native and official production transformer expose the same 521 checkpoint keys/shapes |
+| HF token env | `HF_TOKEN` (the pinned repository is public; never record a token value) |
+
+The pinned text-encoder config declares `Qwen3ForCausalLM` with 36 decoder
+blocks, hidden size 2560, intermediate size 9728, 32 attention heads, 8 KV
+heads, and head dimension 128. FastVideo loads those values from
+`text_encoder/config.json`; they are not hard-coded in the Z-Image integration.
+
+## Shared environment setup
+
+Run from the FastVideo repository root in the same environment used for
+FastVideo. Keep the reference clone and weights inside the ignored paths shown
+below.
+
+Clone and pin the official reference:
+
+```bash
+python .agents/skills/add-model-01-prep/scripts/clone_reference_repo.py \
+  https://github.com/Tongyi-MAI/Z-Image.git \
+  Z-Image \
+  --commit 26f23eda626ffadda020b04ff79488e1d72004cd \
+  --update-gitignore
+git -C Z-Image rev-parse HEAD
+```
+
+The final command must print
+`26f23eda626ffadda020b04ff79488e1d72004cd`. The scheduler and VAE tests also
+enforce that exact HEAD and verify that imported `zimage` modules resolve under
+`<repo_root>/Z-Image/src`. A missing clone may skip local parity; a wrong SHA or
+wrong import origin fails rather than silently testing another installation.
+
+Download the immutable HF snapshot:
+
+```bash
+python .agents/skills/add-model-01-prep/scripts/download_hf_weights.py \
+  Tongyi-MAI/Z-Image-Turbo \
+  official_weights/Z-Image \
+  --revision f332072aa78be7aecdf3ee76d5c247082da564a6 \
+  --allow-pattern 'text_encoder/*' \
+  --allow-pattern 'tokenizer/*' \
+  --allow-pattern 'vae/*' \
+  --allow-pattern 'scheduler/*'
+```
+
+The real-weight transformer test additionally needs `transformer/*` from the
+same pinned revision (24.6 GB). Do not download it without approving that cost;
+point `ZIMAGE_TRANSFORMER_DIR` at an existing pinned transformer subfolder.
+
+Do not change core dependency versions (`torch`, `diffusers`, `transformers`,
+`flash-attn`, `triton`, or CUDA packages) without explicit approval.
+
+```text
+dependency_changes: none
+official_env_status: official source pinned and imported; GB200 component validation run completed
+private_dep_stubs: none
+blocked_on: repo-wide full-suite/SSIM launch gate and PR review; model implementation work is complete
+```
+
+## Run the tests
+
+```bash
+pytest tests/local_tests/zimage/ -v -s
+```
+
+| Component | Coverage scope and contract | Status |
+|---|---|---|
+| Scheduler ([`test_zimage_scheduler_parity.py`](./test_zimage_scheduler_parity.py)) | `implementation_subcomponent`; exact official `sigma_min=0.0` plus `use_reference_discrete_timesteps=True`; positional/default-path regressions remain asset-free | PINNED GB200 PASS (5 passed) |
+| Tokenizer ([`test_zimage_tokenizer_parity.py`](./test_zimage_tokenizer_parity.py)) | `production_loader`; exact `apply_chat_template(tokenize=False, add_generation_prompt=True, enable_thinking=True)` and `max_length=512`; absence of `apply_chat_template` is a failure, not a skip | PINNED GB200 PASS (2 passed) |
+| VAE ([`test_zimage_vae_parity.py`](./test_zimage_vae_parity.py)) | `both`; direct raw decode plus production `VAELoader`, config values, and raw decode | PINNED GB200 PASS (2 passed) |
+| Production text encoder ([`test_zimage_encoder_parity.py`](./test_zimage_encoder_parity.py)) | `production_loader`; `TextEncoderLoader` materializes a body-only Transformers `AutoModel` matching official `src/utils/loader.py`; required by the pipeline | PINNED GB200 PASS (reference/production outputs exact) |
+| Native Qwen quality ([`test_zimage_encoder_parity.py`](./test_zimage_encoder_parity.py)) | `implementation_subcomponent`; direct FastVideo-native construction, strict loading, and fp32/bf16 diagnostics; not loaded by the pipeline | PINNED GB200 PASS; fp32 `atol=2e-4`, `rtol=1e-4`; bf16 thresholds unchanged |
+| Transformer ([`test_zimage_transformer_parity.py`](./test_zimage_transformer_parity.py)) | `both`; pinned production meta key/shape surface, additive-mask attention parity, deterministic tiny CPU full-forward parity, explicit `sp_size=1` guard, and real-weight production-loader/full-forward parity | PINNED GB200 PASS (5 passed; real-weight full forward exact) |
+
+## Pinned official pipeline context
+
+| Field | Contract |
+|---|---|
+| Oracle | `Tongyi-MAI/Z-Image@26f23eda626ffadda020b04ff79488e1d72004cd`, `src/zimage/pipeline.py::generate` and `src/config/inference.py`; never Diffusers |
+| Pipeline/modules | `_class_name=ZImagePipeline`; `transformer`, `vae`, `text_encoder`, `tokenizer`, `scheduler` |
+| Inputs/outputs | text-to-image prompt with optional negative prompt; PIL images by default or latents when requested |
+| Official defaults | `1024x1024`, `num_inference_steps=8`, `guidance_scale=0.0`, `cfg_truncation=1.0`, `max_sequence_length=512` |
+| Reproducible RNG | `torch.Generator("cuda").manual_seed(42)`; `generate()` accepts the generator rather than fixing a seed internally |
+| Scheduler | `use_reference_discrete_timesteps=True`, runtime `sigma_min=0.0`, and the endpoint-preserving `num_steps + 1` schedule with the final zero-timestep DiT call skipped |
+| Decode | apply `(latents / 0.3611) + 0.1159` before official VAE decode |
+| FastVideo status | Native pipeline/config/preset/registry/example and PNG SSIM test are implemented. One-step RTX 5090 parity is bit-exact; the reviewed L40S PNG generated at `9a6dbcadc` is published at dataset revision `7b27d000...f01c`. Later PR changes affect only tests/docs. |
+
+## Pipeline and quality verification
+
+| Scope | Command | Status |
+|---|---|---|
+| Hardware-free pipeline surface | `pytest tests/local_tests/pipelines/test_zimage_pipeline_smoke.py::test_zimage_typed_surface_preflight -v` | PASS |
+| Hardware-free native stage contract | `pytest tests/local_tests/pipelines/test_zimage_pipeline_parity.py::test_zimage_native_default_stage_math -v` | PASS |
+| Pinned native full GPU pipeline parity | `ZIMAGE_REFERENCE_REPO=<pinned-clone> ZIMAGE_MODEL_DIR=<pinned-weights> DISABLE_SP=1 pytest tests/local_tests/pipelines/test_zimage_pipeline_parity.py::test_zimage_pipeline_latents_match_pinned_native_repo -v -s` | PASS on RTX 5090; one-step latents bit-exact (`max=0`, `mean=0`) |
+| 8-step 1024x1024 PNG SSIM | `pytest fastvideo/tests/ssim/test_zimage_similarity.py -v -s` | Reviewed L40S PNG uploaded; remote SHA-256 `6c2be648...f7a6`; end-to-end SSIM CI pending |
+
+## Component contracts
+
+### Text encoder and tokenizer
+
+- The production `TextEncoderLoader` path deliberately loads Transformers
+  `AutoModel`, even when the checkpoint advertises `Qwen3ForCausalLM`. FastVideo
+  consumes hidden states only, so the production path is body-only and does not
+  materialize or execute an LM head.
+- The parity oracle is the pinned official Z-Image source and pipeline contract.
+  A separate Transformers `AutoModel` instance materializes the official
+  loader behavior independently; it is not the object returned by FastVideo's
+  production loader. The native `Qwen3ForCausalLM` implementation is compared
+  separately.
+- CPU/MPS text-encoder offload remains on the requested device. The loader records
+  `_fastvideo_input_device`, and `TextEncodingStage` sends token tensors there.
+- Native loading must account for every destination parameter and every required
+  fused source shard. Q/K/V and gate/up shards are all required unless the exact
+  destination is supplied already fused; quantized auxiliary parameters such as
+  `scale_weight` are not silently ignored.
+- The pinned official prompt path enables thinking, tokenizes to 512, and consumes
+  `hidden_states[-2]` at valid-token positions. A 36-block Qwen model exposes 37
+  hidden-state entries in this contract (embedding/intermediate entries plus the
+  final normalized state).
+- `TextEncoderConfig.chat_template_enable_thinking` is keyword-only and defaults
+  to `False` for existing model families. The Z-Image pipeline config opts in
+  with `True`.
+
+### VAE
+
+The test covers both direct implementation parity and production `VAELoader`
+resolution/strict loading, confirms the official config's scaling/shift values,
+and compares raw decode outputs. The official pipeline's
+`(latents / 0.3611) + 0.1159` transformation is pipeline behavior, so it remains
+a Phase 7 Z-Image pipeline-stage and pipeline-parity requirement. This port
+reuses the
+existing shared `fastvideo.models.vaes.autoencoder_kl.AutoencoderKL` wrapper,
+which subclasses Diffusers `AutoencoderKL`, as a narrowly scoped exception to
+the native-component boundary. That runtime inheritance does not make Diffusers
+the reference implementation: parity is against pinned
+`Tongyi-MAI/Z-Image/src/zimage/autoencoder.py`. This is not precedent for the
+transformer or pipeline implementation.
+
+### Scheduler
+
+The pinned official pipeline mutates `scheduler.sigma_min = 0.0` before building
+its `num_steps + 1` schedule. The Z-Image pipeline applies both
+`use_reference_discrete_timesteps=True` and `sigma_min=0.0`; the published
+`scheduler_config.json` alone does not encode the complete runtime contract.
+
+### Transformer
+
+- The FastVideo-native transformer preserves all 521 published checkpoint keys.
+  The sorted key digest from the pinned HF safetensors index is
+  `3a9216f208c1873b2cf06394411a53e1e95e10fae3b01dca0f7223556e47c354`.
+- Tiny CPU parity uses the pinned official class, identical deterministic state,
+  variable-length image/text batches, and concrete full-forward output checks.
+- Attention uses raw torch SDPA because the padded variable-length stream needs
+  a key mask not exposed by FastVideo's distributed wrappers. Runtime rejects
+  initialized sequence-parallel world sizes above one.
+- The production-loader/full-forward test runs when the pinned transformer
+  subfolder and CUDA are available. The pinned GB200 run strict-loaded the
+  6.15B model and matched the official full-forward output exactly.
+
+## Remaining work
+
+- Run the end-to-end SSIM/full suite after the external-fork documentation
+  approval gate permits the repo-wide launcher.
+- Keep PR approval and merge separate from the completed reference upload.

tests/local_tests/zimage/__init__.py

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

tests/local_tests/zimage/test_zimage_encoder_parity.py

@@ -0,0 +1,769 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Parity test for Z-Image Qwen3 text encoder support in FastVideo.
+
+This compares:
+1) direct transformers AutoModel output, and
+2) FastVideo's production ``TextEncoderLoader`` passthrough, and
+3) FastVideo's shared native Qwen3 encoder (``Qwen3ForCausalLM``, added for Flux2
+   Klein and reused here — Z-Image-Turbo's ``Qwen3Model`` checkpoint routes
+   to it via the model registry),
+
+using identical local checkpoint, tokenization, and inputs.
+
+Usage:
+    pytest tests/local_tests/zimage/test_zimage_encoder_parity.py -v -s
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from types import SimpleNamespace
+import gc
+import json
+import os
+
+import pytest
+import torch
+
+
+def _reclaim_vram() -> None:
+    """Actually reclaim VRAM after the caller has ``del``'d its model refs.
+
+    HF transformer models hold reference cycles, so ``del`` alone does not free
+    them — a ``gc.collect()`` is required before the CUDA caching allocator will
+    release the blocks. Without this the parity test keeps two full encoder
+    copies resident and OOMs on a 44 GB GPU.
+    """
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
+def _gpu_mem_mb() -> float:
+    if torch.cuda.is_available():
+        return torch.cuda.memory_allocated() / (1024 ** 2)
+    return 0.0
+from torch.testing import assert_close
+from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer
+from safetensors.torch import safe_open
+
+from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig
+from fastvideo.distributed.parallel_state import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
+from fastvideo.layers.quantization.absmax_fp8 import AbsMaxFP8MergedParameter
+from fastvideo.models.encoders.qwen3 import Qwen3ForCausalLM
+from fastvideo.models.loader.component_loader import TextEncoderLoader
+from fastvideo.models.loader.utils import set_default_torch_dtype
+from fastvideo.models.registry import ModelRegistry
+
+PARITY_SCOPE = "both"
+
+# Strict-load contract: the shared Qwen3 encoder is body-only (embed_tokens +
+# layers + norm, no lm_head). Z-Image-Turbo ships a full Qwen3 checkpoint, so
+# `lm_head.weight` is the only key that may go unmatched; anything else means a
+# real silent drop. Enforced here in the test since the shared encoder's loader
+# is intentionally lenient (it's used by multiple models).
+_ALLOWED_UNEXPECTED_KEYS = {"lm_head.weight"}
+
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+ZIMAGE_TEXT_ENCODER_DIR = REPO_ROOT / "official_weights" / "Z-Image" / "text_encoder"
+ZIMAGE_TOKENIZER_DIR = REPO_ROOT / "official_weights" / "Z-Image" / "tokenizer"
+
+
+@pytest.fixture(scope="module", autouse=True)
+def _init_dist_and_tp_groups():
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "29531")
+    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 _load_json(path: Path) -> dict:
+    with path.open("r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+def _iter_safetensors(path: Path):
+    with safe_open(str(path), framework="pt", device="cpu") as f:
+        for k in f.keys():
+            yield k, f.get_tensor(k)
+
+
+def _iter_pretrained_safetensors(model_dir: Path):
+    single = model_dir / "model.safetensors"
+    if single.exists():
+        yield from _iter_safetensors(single)
+        return
+
+    index = model_dir / "model.safetensors.index.json"
+    if index.exists():
+        idx = _load_json(index)
+        shard_names = sorted(set(idx["weight_map"].values()))
+        for shard in shard_names:
+            yield from _iter_safetensors(model_dir / shard)
+        return
+
+    raise FileNotFoundError(
+        f"Missing safetensors checkpoint in {model_dir} (expected model.safetensors or model.safetensors.index.json)"
+    )
+
+
+def _pretrained_safetensor_keys(model_dir: Path) -> set[str]:
+    return {name for name, _ in _iter_pretrained_safetensors(model_dir)}
+
+
+def _load_qwen3_config() -> Qwen3TextConfig:
+    cfg_raw = _load_json(ZIMAGE_TEXT_ENCODER_DIR / "config.json")
+    for key in ("_name_or_path", "transformers_version", "model_type", "torch_dtype"):
+        cfg_raw.pop(key, None)
+    config = Qwen3TextConfig()
+    config.update_model_arch(cfg_raw)
+    assert config.num_hidden_layers == 36
+    assert config.hidden_size == 2560
+    assert config.intermediate_size == 9728
+    assert config.num_attention_heads == 32
+    assert config.num_key_value_heads == 8
+    return config
+
+
+def _loader_args(cpu_offload: bool) -> SimpleNamespace:
+    return SimpleNamespace(
+        text_encoder_cpu_offload=cpu_offload,
+        override_text_encoder_quant=None,
+        override_text_encoder_safetensors=None,
+        pin_cpu_memory=False,
+    )
+
+
+def _precision_name(dtype: torch.dtype) -> str:
+    if dtype == torch.float32:
+        return "fp32"
+    if dtype == torch.bfloat16:
+        return "bf16"
+    raise ValueError(f"Unsupported parity dtype: {dtype}")
+
+
+def _tokenize_official_prompts(tokenizer, prompts: list[str]):
+    formatted = [
+        tokenizer.apply_chat_template(
+            [{"role": "user", "content": prompt}],
+            tokenize=False,
+            add_generation_prompt=True,
+            enable_thinking=True,
+        )
+        for prompt in prompts
+    ]
+    return tokenizer(
+        formatted,
+        padding="max_length",
+        truncation=True,
+        max_length=512,
+        return_tensors="pt",
+    )
+
+
+def _assert_native_load_surface(
+    model,
+    checkpoint_keys: set[str],
+    loaded_params: set[str],
+) -> None:
+    param_names = {name for name, _ in model.named_parameters()}
+    missing = param_names - loaded_params
+    assert not missing, f"FastVideo parameters not loaded: {sorted(missing)}"
+
+    normalized_keys = {
+        name[len("model."):] if name.startswith("model.") else name
+        for name in checkpoint_keys
+    }
+    mapped_keys: set[str] = set()
+    for checkpoint_name in normalized_keys:
+        for param_name, weight_name, _ in model.config.arch_config.stacked_params_mapping:
+            if weight_name in checkpoint_name:
+                mapped_keys.add(checkpoint_name.replace(weight_name, param_name))
+                break
+        else:
+            mapped_keys.add(checkpoint_name)
+
+    unexpected = mapped_keys - param_names - {
+        "rotary_emb.inv_freq",
+        "rotary_emb.cos_cached",
+        "rotary_emb.sin_cached",
+    }
+    assert unexpected <= _ALLOWED_UNEXPECTED_KEYS, (
+        "Unexpected checkpoint keys not in allowlist: "
+        f"{sorted(unexpected - _ALLOWED_UNEXPECTED_KEYS)}"
+    )
+
+
+def _fake_qwen3_weight_target(include_quant_scales: bool = False):
+    stacked = Qwen3TextConfig().arch_config.stacked_params_mapping
+    params = {
+        "layers.0.self_attn.qkv_proj.weight": torch.nn.Parameter(torch.empty(1)),
+        "layers.0.mlp.gate_up_proj.weight": torch.nn.Parameter(torch.empty(1)),
+    }
+    if include_quant_scales:
+        qkv_scale = AbsMaxFP8MergedParameter(torch.zeros(3), requires_grad=False)
+        qkv_scale.output_partition_sizes = [1, 1, 1]
+        gate_up_scale = AbsMaxFP8MergedParameter(torch.zeros(2), requires_grad=False)
+        gate_up_scale.output_partition_sizes = [1, 1]
+        params.update({
+            "layers.0.self_attn.qkv_proj.scale_weight": qkv_scale,
+            "layers.0.mlp.gate_up_proj.scale_weight": gate_up_scale,
+        })
+    for param in params.values():
+        if not hasattr(param, "weight_loader"):
+            param.weight_loader = (
+                lambda target, loaded, *args: target.data.copy_(loaded)
+            )
+
+    model = SimpleNamespace(
+        config=SimpleNamespace(
+            arch_config=SimpleNamespace(stacked_params_mapping=stacked)
+        ),
+        named_parameters=lambda: params.items(),
+    )
+    return model, params, stacked
+
+
+# The provisional bf16 thresholds below come from the historical unpinned
+# 35-layer run. Per add-model-02-parity's calibration block, bf16 max-based
+# asserts are meaningless for deep encoders — fused-vs-unfused linear order
+# alone produces a long tail, so the test uses mean + median instead. The
+# pinned snapshot has 36 blocks and must be rerun before these thresholds are
+# treated as current evidence.
+#
+# Historical empirical numbers (unrecorded snapshot, 35 layers, CUDA):
+#   last_hidden_state (post-norm):  mean=0.0152, median=0.0117
+#   hidden_states[-2] (pre-norm):   mean=0.0754, median=0.0625
+# Thresholds below add ~1.6x headroom over observed.
+_BF16_MEAN_DRIFT_POST_NORM = 0.025
+_BF16_MEAN_DRIFT_PRE_NORM = 0.120
+_BF16_MEDIAN_DRIFT_POST_NORM = 0.020
+_BF16_MEDIAN_DRIFT_PRE_NORM = 0.100
+_FP32_ATOL = 2e-4
+_FP32_RTOL = 1e-4
+
+
+def _print_diag(label: str, ref: torch.Tensor, fv: torch.Tensor) -> torch.Tensor:
+    diff = (ref.float() - fv.float()).abs()
+    flat = diff.flatten()
+    p99 = flat.kthvalue(max(1, int(0.99 * flat.numel()))).values
+    print(
+        f"[{label}] max_diff={diff.max():.4f}  mean_diff={diff.mean():.4f}  "
+        f"median_diff={diff.median():.4f}  p99_diff={p99:.4f}",
+        flush=True,
+    )
+    return diff
+
+
+@pytest.mark.parametrize(
+    ("cpu_offload", "target_device", "expected_device"),
+    [
+        (True, torch.device("cuda"), torch.device("cpu")),
+        (False, torch.device("meta"), torch.device("meta")),
+    ],
+)
+def test_qwen3_production_loader_avoids_device_context_and_honors_placement(
+    monkeypatch: pytest.MonkeyPatch,
+    cpu_offload: bool,
+    target_device: torch.device,
+    expected_device: torch.device,
+):
+    """HF loading stays ambient while final placement honors target/offload."""
+    import fastvideo.platforms as platforms
+
+    placements: list[torch.device] = []
+    auto_model_calls: list[dict] = []
+
+    class FakeHFModel:
+        def eval(self):
+            return self
+
+        def to(self, device: torch.device):
+            placements.append(torch.device(device))
+            return self
+
+    fake_model = FakeHFModel()
+
+    ambient_device = torch.empty(0).device
+    default_devices: list[torch.device] = []
+
+    def fake_auto_model_from_pretrained(*args, **kwargs):
+        auto_model_calls.append(kwargs)
+        default_devices.append(torch.empty(0).device)
+        return fake_model
+
+    def fail_causal_lm_from_pretrained(*args, **kwargs):
+        raise AssertionError("Qwen text encoding must not load AutoModelForCausalLM")
+
+    monkeypatch.setattr(
+        AutoModel,
+        "from_pretrained",
+        staticmethod(fake_auto_model_from_pretrained),
+    )
+    monkeypatch.setattr(
+        AutoModelForCausalLM,
+        "from_pretrained",
+        staticmethod(fail_causal_lm_from_pretrained),
+    )
+    monkeypatch.setattr(
+        platforms,
+        "_current_platform",
+        SimpleNamespace(
+            is_mps=lambda: False,
+            verify_model_arch=lambda _arch: None,
+        ),
+    )
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(
+        "fastvideo.distributed.get_local_torch_device",
+        lambda: torch.device("cuda"),
+    )
+
+    config = Qwen3TextConfig()
+    # The pinned checkpoint names the causal wrapper even though the official
+    # Z-Image loader intentionally asks Transformers for the body-only model.
+    config.arch_config.architectures = ["Qwen3ForCausalLM"]
+    loaded = TextEncoderLoader().load_model(
+        "unused-by-mock",
+        config,
+        target_device,
+        _loader_args(cpu_offload=cpu_offload),
+        dtype="fp32",
+        use_text_encoder_override=True,
+    )
+
+    assert loaded is fake_model
+    assert auto_model_calls == [{
+        "local_files_only": True,
+        "dtype": torch.float32,
+        "low_cpu_mem_usage": True,
+    }]
+    assert default_devices == [ambient_device]
+    assert placements == [expected_device]
+    assert loaded._fastvideo_input_device == expected_device
+
+
+def test_set_default_torch_dtype_restores_after_exception():
+    original = torch.get_default_dtype()
+
+    with pytest.raises(RuntimeError, match="expected failure"):
+        with set_default_torch_dtype(torch.bfloat16):
+            raise RuntimeError("expected failure")
+
+    assert torch.get_default_dtype() == original
+
+
+@pytest.mark.parametrize("checkpoint_prefix", ["", "model."])
+def test_qwen3_native_load_accepts_already_fused_weights_and_scales(
+    checkpoint_prefix: str,
+):
+    """Direct FastVideo/state-dict keys satisfy each fused target atomically."""
+    fake_model, params, _ = _fake_qwen3_weight_target(include_quant_scales=True)
+    expected = {
+        name: torch.arange(1, param.numel() + 1, dtype=param.dtype)
+        for name, param in params.items()
+    }
+    weights = [
+        (f"{checkpoint_prefix}{name}", expected[name])
+        for name in params
+    ]
+
+    loaded = Qwen3ForCausalLM.load_weights(fake_model, weights)
+
+    assert loaded == set(params)
+    for name, param in params.items():
+        assert torch.equal(param.data, expected[name])
+    _assert_native_load_surface(
+        fake_model,
+        {name for name, _ in weights},
+        loaded,
+    )
+
+
+def test_qwen3_native_load_rejects_incompatible_fused_quant_scale_shape():
+    fake_model, params, _ = _fake_qwen3_weight_target(include_quant_scales=True)
+    weights = [
+        (name, torch.ones_like(param))
+        for name, param in params.items()
+        if name != "layers.0.self_attn.qkv_proj.scale_weight"
+    ]
+    weights.append(("layers.0.self_attn.qkv_proj.scale_weight", torch.ones(1)))
+
+    with pytest.raises(AssertionError, match="Attempted to load weight"):
+        Qwen3ForCausalLM.load_weights(fake_model, weights)
+
+
+def test_qwen3_native_load_uses_real_shard_loader_for_split_quant_scales():
+    fake_model, params, stacked = _fake_qwen3_weight_target(include_quant_scales=True)
+    weights = [
+        (name, torch.ones_like(param))
+        for name, param in params.items()
+        if name.endswith(".weight")
+    ]
+    weights.extend(
+        (f"model.layers.0.self_attn{weight_name}.scale_weight", torch.tensor(float(index + 1)))
+        for index, (_, weight_name, _) in enumerate(stacked[:3])
+    )
+    weights.extend(
+        (f"model.layers.0.mlp{weight_name}.scale_weight", torch.tensor(float(index + 4)))
+        for index, (_, weight_name, _) in enumerate(stacked[3:])
+    )
+
+    loaded = Qwen3ForCausalLM.load_weights(fake_model, weights)
+
+    assert loaded == set(params)
+    assert torch.equal(
+        params["layers.0.self_attn.qkv_proj.scale_weight"].data,
+        torch.tensor([1.0, 2.0, 3.0]),
+    )
+    assert torch.equal(
+        params["layers.0.mlp.gate_up_proj.scale_weight"].data,
+        torch.tensor([4.0, 5.0]),
+    )
+
+
+@pytest.mark.parametrize(
+    ("missing_weight_name", "missing_shard_id"),
+    [
+        (".q_proj", "q"),
+        (".k_proj", "k"),
+        (".v_proj", "v"),
+        (".gate_proj", 0),
+        (".up_proj", 1),
+    ],
+)
+def test_qwen3_native_load_rejects_missing_fused_weight_shard(
+    missing_weight_name: str,
+    missing_shard_id: str | int,
+):
+    fake_model, _, stacked = _fake_qwen3_weight_target()
+    weights = [
+        (f"model.layers.0.self_attn{weight_name}.weight", torch.empty(1))
+        for _, weight_name, _ in stacked[:3]
+        if weight_name != missing_weight_name
+    ]
+    weights.extend(
+        (f"model.layers.0.mlp{weight_name}.weight", torch.empty(1))
+        for _, weight_name, _ in stacked[3:]
+        if weight_name != missing_weight_name
+    )
+
+    with pytest.raises(ValueError, match="Missing required stacked checkpoint shards") as exc_info:
+        Qwen3ForCausalLM.load_weights(fake_model, weights)
+
+    destination = (
+        "layers.0.self_attn.qkv_proj.weight"
+        if isinstance(missing_shard_id, str)
+        else "layers.0.mlp.gate_up_proj.weight"
+    )
+    assert f"{destination}[{missing_shard_id}]" in str(exc_info.value)
+
+
+@pytest.mark.parametrize(
+    ("missing_weight_name", "missing_shard_id"),
+    [
+        (".q_proj", "q"),
+        (".k_proj", "k"),
+        (".v_proj", "v"),
+        (".gate_proj", 0),
+        (".up_proj", 1),
+    ],
+)
+def test_qwen3_native_load_rejects_missing_fused_quant_scale_shard(
+    missing_weight_name: str,
+    missing_shard_id: str | int,
+):
+    fake_model, params, stacked = _fake_qwen3_weight_target(include_quant_scales=True)
+    weights = [
+        (name, torch.empty(1))
+        for name in params
+        if name.endswith(".weight")
+    ]
+    weights.extend(
+        (f"model.layers.0.self_attn{weight_name}.scale_weight", torch.empty(1))
+        for _, weight_name, _ in stacked[:3]
+        if weight_name != missing_weight_name
+    )
+    weights.extend(
+        (f"model.layers.0.mlp{weight_name}.scale_weight", torch.empty(1))
+        for _, weight_name, _ in stacked[3:]
+        if weight_name != missing_weight_name
+    )
+
+    with pytest.raises(ValueError, match="Missing required stacked checkpoint shards") as exc_info:
+        Qwen3ForCausalLM.load_weights(fake_model, weights)
+
+    destination = (
+        "layers.0.self_attn.qkv_proj.scale_weight"
+        if isinstance(missing_shard_id, str)
+        else "layers.0.mlp.gate_up_proj.scale_weight"
+    )
+    assert f"{destination}[{missing_shard_id}]" in str(exc_info.value)
+
+
+@pytest.mark.skipif(not ZIMAGE_TEXT_ENCODER_DIR.exists(), reason="Z-Image text encoder checkpoint required")
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        pytest.param(torch.float32, id="fp32"),
+        pytest.param(
+            torch.bfloat16,
+            id="bf16",
+            marks=pytest.mark.skipif(
+                not (torch.cuda.is_available() and torch.cuda.is_bf16_supported()),
+                reason="bf16 parity requires a bf16-capable CUDA device",
+            ),
+        ),
+    ],
+)
+def test_zimage_qwen3_encoder_parity_forward(dtype: torch.dtype):
+    if not ZIMAGE_TOKENIZER_DIR.exists():
+        pytest.skip(f"Z-Image tokenizer dir not found: {ZIMAGE_TOKENIZER_DIR}")
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    torch.manual_seed(11)
+
+    tokenizer = AutoTokenizer.from_pretrained(str(ZIMAGE_TOKENIZER_DIR), local_files_only=True)
+
+    prompts = [
+        "A cinematic shot of a rainy neon street at night.",
+        "A watercolor illustration of a fox in autumn leaves.",
+    ]
+    toks = _tokenize_official_prompts(tokenizer, prompts)
+    input_ids = toks["input_ids"].to(device)
+    attention_mask = toks["attention_mask"].to(device)
+    config = _load_qwen3_config()
+
+    ref = AutoModel.from_pretrained(
+        str(ZIMAGE_TEXT_ENCODER_DIR),
+        local_files_only=True,
+        trust_remote_code=True,
+        dtype=dtype,
+        low_cpu_mem_usage=True,
+    ).eval().to(device=device)
+
+    with torch.no_grad():
+        ref_out = ref(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True,
+            return_dict=True,
+        )
+        ref_last = ref_out.last_hidden_state.detach().float().cpu()
+        assert ref_out.hidden_states is not None
+        ref_hs_m2 = ref_out.hidden_states[-2].detach().float().cpu()
+
+    print(f"[mem] after HF ref forward: {_gpu_mem_mb():.0f} MiB", flush=True)
+    # Keep only CPU outputs so each full encoder is resident one at a time.
+    del ref, ref_out
+    _reclaim_vram()
+    print(f"[mem] after freeing HF ref:  {_gpu_mem_mb():.0f} MiB", flush=True)
+
+    production = TextEncoderLoader().load_model(
+        str(ZIMAGE_TEXT_ENCODER_DIR),
+        config,
+        device,
+        _loader_args(cpu_offload=False),
+        dtype=_precision_name(dtype),
+        use_text_encoder_override=True,
+    )
+    assert production.__class__.__name__ == "Qwen3Model"
+    assert not hasattr(production, "lm_head")
+    assert production._fastvideo_input_device == device
+
+    with torch.no_grad():
+        production_out = production(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True,
+            return_dict=True,
+        )
+        production_last = production_out.last_hidden_state.detach().float().cpu()
+        assert production_out.hidden_states is not None
+        production_hs_m2 = production_out.hidden_states[-2].detach().float().cpu()
+
+    # Production loading must preserve the independent Transformers oracle
+    # exactly; native-kernel drift is assessed separately below.
+    assert_close(ref_last, production_last, atol=0.0, rtol=0.0)
+    assert_close(ref_hs_m2, production_hs_m2, atol=0.0, rtol=0.0)
+    del production, production_out, production_last, production_hs_m2
+    _reclaim_vram()
+
+    fv_cls, _ = ModelRegistry.resolve_model_cls("Qwen3Model")
+    fv = fv_cls(config).eval()
+    loaded = fv.load_weights(_iter_pretrained_safetensors(ZIMAGE_TEXT_ENCODER_DIR))
+    assert loaded, "No Qwen3 weights were loaded into FastVideo model"
+    fv = fv.to(device=device, dtype=dtype)
+    print(f"[mem] after FastVideo load:  {_gpu_mem_mb():.0f} MiB", flush=True)
+
+    _assert_native_load_surface(
+        fv,
+        _pretrained_safetensor_keys(ZIMAGE_TEXT_ENCODER_DIR),
+        loaded,
+    )
+
+    with torch.no_grad():
+        fv_out = fv(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True,
+        )
+        fv_last = fv_out.last_hidden_state.detach().float().cpu()
+        assert fv_out.hidden_states is not None
+        fv_hs_m2 = fv_out.hidden_states[-2].detach().float().cpu()
+
+    # Z-Image uses only valid token embeddings selected by attention mask.
+    # Compare parity on that exact path to avoid padded-token artifacts.
+    mask_cpu = attention_mask.detach().bool().cpu()
+    is_bf16 = dtype == torch.bfloat16
+
+    for i in range(mask_cpu.shape[0]):
+        valid = mask_cpu[i]
+        ref_last_v = ref_last[i][valid]
+        fv_last_v = fv_last[i][valid]
+        ref_hs_m2_v = ref_hs_m2[i][valid]
+        fv_hs_m2_v = fv_hs_m2[i][valid]
+
+        last_diff = _print_diag(f"batch{i} last_hidden_state {dtype}", ref_last_v, fv_last_v)
+        hs_m2_diff = _print_diag(f"batch{i} hidden_states[-2] {dtype}", ref_hs_m2_v, fv_hs_m2_v)
+
+        if is_bf16:
+            # Distribution checks instead of element-wise assert_close: cross-
+            # kernel bf16 produces a long max tail (fused vs unfused linears),
+            # but a real op bug pushes mean *and* median high. Mean catches
+            # systematic drift; median catches "wrong weights / swapped
+            # layers" where most elements diverge. Pre-norm tensors get a
+            # looser bound because the final RMSNorm compresses drift ~5x.
+            assert last_diff.mean().item() < _BF16_MEAN_DRIFT_POST_NORM, (
+                f"batch{i} last_hidden_state bf16 mean drift "
+                f"{last_diff.mean():.4f} >= {_BF16_MEAN_DRIFT_POST_NORM}"
+            )
+            assert last_diff.median().item() < _BF16_MEDIAN_DRIFT_POST_NORM, (
+                f"batch{i} last_hidden_state bf16 median drift "
+                f"{last_diff.median():.4f} >= {_BF16_MEDIAN_DRIFT_POST_NORM}"
+            )
+            assert hs_m2_diff.mean().item() < _BF16_MEAN_DRIFT_PRE_NORM, (
+                f"batch{i} hidden_states[-2] bf16 mean drift "
+                f"{hs_m2_diff.mean():.4f} >= {_BF16_MEAN_DRIFT_PRE_NORM}"
+            )
+            assert hs_m2_diff.median().item() < _BF16_MEDIAN_DRIFT_PRE_NORM, (
+                f"batch{i} hidden_states[-2] bf16 median drift "
+                f"{hs_m2_diff.median():.4f} >= {_BF16_MEDIAN_DRIFT_PRE_NORM}"
+            )
+        else:
+            assert_close(ref_last_v, fv_last_v, atol=_FP32_ATOL, rtol=_FP32_RTOL)
+            assert_close(ref_hs_m2_v, fv_hs_m2_v, atol=_FP32_ATOL, rtol=_FP32_RTOL)
+
+
+@pytest.mark.skipif(not ZIMAGE_TEXT_ENCODER_DIR.exists(), reason="Z-Image text encoder checkpoint required")
+@pytest.mark.skipif(
+    not (torch.cuda.is_available() and torch.cuda.is_bf16_supported()),
+    reason="bf16 per-layer diagnostic requires a bf16-capable CUDA device",
+)
+def test_zimage_qwen3_encoder_per_layer_bf16_diagnostic():
+    """Diagnostic-only: prints per-layer FastVideo-vs-HF drift for the bf16
+    forward to distinguish monotonic accumulation (bf16 tail) from a
+    single-layer spike (real op mismatch).
+
+    This test always passes; review the captured stdout to interpret.
+    Expected pattern for healthy bf16 accumulation: max/mean/median grow
+    smoothly with layer depth. Suspect pattern: one layer's diff is
+    >>2x the previous, suggesting a real bug at that block.
+    """
+    if not ZIMAGE_TOKENIZER_DIR.exists():
+        pytest.skip(f"Z-Image tokenizer dir not found: {ZIMAGE_TOKENIZER_DIR}")
+
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+    torch.manual_seed(11)
+
+    tokenizer = AutoTokenizer.from_pretrained(str(ZIMAGE_TOKENIZER_DIR), local_files_only=True)
+    toks = _tokenize_official_prompts(
+        tokenizer,
+        ["A cinematic shot of a rainy neon street at night."],
+    )
+    input_ids = toks["input_ids"].to(device)
+    attention_mask = toks["attention_mask"].to(device)
+
+    ref = AutoModel.from_pretrained(
+        str(ZIMAGE_TEXT_ENCODER_DIR),
+        local_files_only=True,
+        trust_remote_code=True,
+        dtype=dtype,
+        low_cpu_mem_usage=True,
+    ).eval().to(device=device)
+
+    with torch.no_grad():
+        ref_out = ref(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True,
+            return_dict=True,
+        )
+        # Cache the reference hidden states on CPU, then free the HF model so
+        # only one encoder is resident when FastVideo runs (avoids OOM).
+        ref_hidden_cpu = [h.detach().float().cpu() for h in ref_out.hidden_states]
+    del ref, ref_out
+    _reclaim_vram()
+
+    fv_cls, _ = ModelRegistry.resolve_model_cls("Qwen3Model")
+    cfg_raw = _load_json(ZIMAGE_TEXT_ENCODER_DIR / "config.json")
+    for k in ("_name_or_path", "transformers_version", "model_type", "torch_dtype"):
+        cfg_raw.pop(k, None)
+    cfg = Qwen3TextConfig()
+    cfg.update_model_arch(cfg_raw)
+    fv = fv_cls(cfg).eval()
+    fv.load_weights(_iter_pretrained_safetensors(ZIMAGE_TEXT_ENCODER_DIR))
+    fv = fv.to(device=device, dtype=dtype)
+
+    with torch.no_grad():
+        fv_out = fv(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True,
+        )
+        fv_hidden_cpu = [h.detach().float().cpu() for h in fv_out.hidden_states]
+
+    assert fv_out.hidden_states is not None
+    assert len(ref_hidden_cpu) == 37, (
+        f"expected embeddings plus 36 block outputs, got {len(ref_hidden_cpu)} entries"
+    )
+    assert len(ref_hidden_cpu) == len(fv_hidden_cpu), (
+        f"len mismatch: HF={len(ref_hidden_cpu)} FV={len(fv_hidden_cpu)}"
+    )
+
+    mask_cpu = attention_mask.detach().bool().cpu()[0]
+    print(
+        f"\n[per-layer bf16 diagnostic] num_hidden_states={len(ref_hidden_cpu)}  "
+        f"valid_tokens={mask_cpu.sum().item()}/{mask_cpu.numel()}",
+        flush=True,
+    )
+    print(
+        f"{'idx':>3}  {'kind':<10}  {'max':>8}  {'mean':>8}  {'median':>8}  {'p99':>8}",
+        flush=True,
+    )
+    for idx, (ref_h, fv_h) in enumerate(zip(ref_hidden_cpu, fv_hidden_cpu)):
+        ref_v = ref_h[0][mask_cpu]
+        fv_v = fv_h[0][mask_cpu]
+        diff = (ref_v - fv_v).abs()
+        flat = diff.flatten()
+        p99 = torch.quantile(flat, 0.99).item()
+        if idx == 0:
+            kind = "embedding"
+        elif idx == len(ref_hidden_cpu) - 1:
+            kind = "post-norm"
+        else:
+            kind = f"layer{idx - 1}-out"
+        print(
+            f"{idx:>3}  {kind:<10}  "
+            f"{diff.max().item():>8.4f}  {diff.mean().item():>8.4f}  "
+            f"{diff.median().item():>8.4f}  {p99:>8.4f}",
+            flush=True,
+        )

tests/local_tests/zimage/test_zimage_scheduler_parity.py

@@ -0,0 +1,233 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Parity tests for Z-Image FlowMatchEulerDiscreteScheduler.
+
+These tests compare FastVideo's scheduler implementation against the
+reference scheduler shipped in the local Z-Image repository.
+
+Usage:
+    pytest tests/local_tests/zimage/test_zimage_scheduler_parity.py -v
+"""
+
+from __future__ import annotations
+
+import importlib
+import json
+import subprocess
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from torch.testing import assert_close
+
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler as FastVideoScheduler,
+)
+
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+ZIMAGE_REPO = REPO_ROOT / "Z-Image"
+ZIMAGE_SRC = REPO_ROOT / "Z-Image" / "src"
+ZIMAGE_SCHEDULER_CFG = (
+    REPO_ROOT / "official_weights" / "Z-Image" / "scheduler" / "scheduler_config.json"
+)
+ZIMAGE_REFERENCE_REVISION = "26f23eda626ffadda020b04ff79488e1d72004cd"
+PARITY_SCOPE = "implementation_subcomponent"
+
+
+def _require_pinned_reference_module(module_name: str, source_file: Path):
+    if not ZIMAGE_REPO.exists():
+        pytest.skip(f"Pinned Z-Image reference clone not found: {ZIMAGE_REPO}")
+    if not source_file.is_file():
+        pytest.fail(f"Z-Image reference clone is incomplete; missing {source_file}")
+
+    try:
+        result = subprocess.run(
+            ["git", "-C", str(ZIMAGE_REPO), "rev-parse", "HEAD"],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+    except (FileNotFoundError, subprocess.CalledProcessError) as exc:
+        pytest.fail(f"Cannot verify Z-Image reference revision: {exc}")
+
+    actual_revision = result.stdout.strip()
+    assert actual_revision == ZIMAGE_REFERENCE_REVISION, (
+        "Z-Image reference clone is not at the pinned revision: "
+        f"expected {ZIMAGE_REFERENCE_REVISION}, got {actual_revision}"
+    )
+
+    if str(ZIMAGE_SRC) not in sys.path:
+        sys.path.insert(0, str(ZIMAGE_SRC))
+    try:
+        module = importlib.import_module(module_name)
+    except Exception as exc:
+        pytest.fail(f"Cannot import pinned Z-Image module {module_name}: {exc}")
+
+    module_file = Path(module.__file__ or "").resolve()
+    assert module_file.is_relative_to(ZIMAGE_SRC.resolve()), (
+        f"{module_name} resolved outside the pinned clone: {module_file}"
+    )
+    return module
+
+
+@pytest.fixture(scope="module")
+def reference_scheduler_cls():
+    module = _require_pinned_reference_module(
+        "zimage.scheduler",
+        ZIMAGE_SRC / "zimage" / "scheduler.py",
+    )
+    return module.FlowMatchEulerDiscreteScheduler
+
+
+@pytest.fixture(scope="module")
+def scheduler_config() -> dict:
+    if not ZIMAGE_SCHEDULER_CFG.exists():
+        pytest.skip(f"Z-Image scheduler config not found: {ZIMAGE_SCHEDULER_CFG}")
+    with ZIMAGE_SCHEDULER_CFG.open("r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+# Diffusers-style scheduler-config keys we do NOT want to forward — they are
+# loader/state metadata, not constructor kwargs. Anything else in
+# scheduler_config.json (including future keys like `time_shift_type`,
+# `invert_sigmas`, etc.) must be forwarded so parity actually exercises the
+# config-on-disk and not a hand-curated subset.
+_SCHEDULER_CONFIG_LOADER_KEYS = frozenset({"_class_name", "_diffusers_version", "_name_or_path"})
+
+
+def _scheduler_kwargs_from_config(scheduler_config: dict) -> dict:
+    return {k: v for k, v in scheduler_config.items() if k not in _SCHEDULER_CONFIG_LOADER_KEYS}
+
+
+def _make_scheduler_pair(reference_scheduler_cls, scheduler_config: dict, **overrides):
+    kwargs = _scheduler_kwargs_from_config(scheduler_config)
+    kwargs.update(overrides)
+
+    ref = reference_scheduler_cls(**kwargs)
+    # The pinned Z-Image pipeline mutates this immediately before scheduling.
+    ref.sigma_min = 0.0
+
+    # These values must be serialized by the future production pipeline config.
+    # setdefault also allows this test to consume that config once it lands.
+    fv_kwargs = dict(kwargs)
+    fv_kwargs.setdefault("use_reference_discrete_timesteps", True)
+    fv_kwargs.setdefault("sigma_min", 0.0)
+    fv = FastVideoScheduler(**fv_kwargs)
+    return ref, fv
+
+
+@pytest.fixture(scope="module")
+def scheduler_pair(reference_scheduler_cls, scheduler_config: dict):
+    return _make_scheduler_pair(reference_scheduler_cls, scheduler_config)
+
+
+def _run_step_loop(ref_scheduler, fv_scheduler, sample_shape=(2, 4, 32, 32)):
+    torch.manual_seed(123)
+
+    ref_sample = torch.randn(sample_shape, dtype=torch.float32)
+    fv_sample = ref_sample.clone()
+
+    for t in ref_scheduler.timesteps:
+        model_output = torch.randn_like(ref_sample)
+
+        ref_next = ref_scheduler.step(
+            model_output,
+            t,
+            ref_sample,
+            return_dict=False,
+        )[0]
+        fv_next = fv_scheduler.step(
+            model_output,
+            t,
+            fv_sample,
+            return_dict=False,
+        )[0]
+
+        assert_close(ref_next, fv_next, atol=1e-6, rtol=1e-6)
+
+        ref_sample = ref_next
+        fv_sample = fv_next
+
+
+def test_zimage_scheduler_parity_default_schedule_and_step(scheduler_pair):
+    ref, fv = scheduler_pair
+
+    num_inference_steps = 8
+    ref.set_timesteps(num_inference_steps=num_inference_steps, device="cpu")
+    fv.set_timesteps(num_inference_steps=num_inference_steps, device="cpu")
+
+    assert_close(ref.timesteps, fv.timesteps, atol=1e-4, rtol=1e-6)
+    assert_close(ref.sigmas, fv.sigmas, atol=1e-7, rtol=1e-6)
+
+    _run_step_loop(ref, fv)
+
+
+def test_scheduler_positional_args_keep_existing_bindings_and_default_schedule():
+    scheduler = FastVideoScheduler(1000, 1.0, False, 0.7)
+
+    assert scheduler.config.base_shift == 0.7
+    assert scheduler.config.use_reference_discrete_timesteps is False
+
+    scheduler.set_timesteps(num_inference_steps=8, device="cpu")
+    assert_close(
+        scheduler.sigmas[-2],
+        torch.tensor(scheduler.sigma_min, dtype=torch.float32),
+    )
+
+
+def test_scheduler_honors_explicit_zero_sigma_min():
+    scheduler = FastVideoScheduler(
+        sigma_min=0.0,
+        use_reference_discrete_timesteps=True,
+    )
+
+    assert scheduler.sigma_min == 0.0
+    scheduler.set_timesteps(num_inference_steps=8, device="cpu")
+    assert_close(
+        scheduler.timesteps[-1],
+        torch.tensor(125.0, dtype=torch.float32),
+    )
+
+
+def test_reference_schedule_preserves_float64_linspace_rounding():
+    scheduler = FastVideoScheduler(
+        sigma_min=0.0,
+        use_reference_discrete_timesteps=True,
+    )
+    scheduler.set_timesteps(num_inference_steps=9, device="cpu")
+
+    expected = torch.from_numpy(
+        np.linspace(1000.0, 0.0, 10)[:-1] / 1000.0,
+    ).to(torch.float32) * 1000.0
+    float32_regression = torch.from_numpy(
+        np.linspace(1000.0, 0.0, 10, dtype=np.float32)[:-1] / np.float32(1000.0),
+    ) * 1000.0
+
+    # The fifth timestep differs by one float32 ULP depending on where the
+    # linspace is rounded, so this assertion fails if dtype=np.float32 returns.
+    assert expected[4].item() != float32_regression[4].item()
+    assert scheduler.timesteps[4].item() == expected[4].item()
+
+
+def test_zimage_scheduler_parity_dynamic_shifting_with_mu(
+    reference_scheduler_cls,
+    scheduler_config: dict,
+):
+    ref, fv = _make_scheduler_pair(
+        reference_scheduler_cls,
+        scheduler_config,
+        use_dynamic_shifting=True,
+    )
+
+    mu = 0.75
+    num_inference_steps = 8
+    ref.set_timesteps(num_inference_steps=num_inference_steps, device="cpu", mu=mu)
+    fv.set_timesteps(num_inference_steps=num_inference_steps, device="cpu", mu=mu)
+
+    assert_close(ref.timesteps, fv.timesteps, atol=1e-4, rtol=1e-6)
+    assert_close(ref.sigmas, fv.sigmas, atol=1e-7, rtol=1e-6)
+
+    _run_step_loop(ref, fv)

tests/local_tests/zimage/test_zimage_tokenizer_parity.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Parity test for Z-Image tokenizer loading and tokenization behavior.
+
+This compares:
+1) direct transformers AutoTokenizer loading from local tokenizer dir, and
+2) FastVideo TokenizerLoader loading path,
+
+using the exact chat-template and tokenization settings from the pinned
+Z-Image pipeline.
+
+Usage:
+    pytest tests/local_tests/zimage/test_zimage_tokenizer_parity.py -v
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import pytest
+import torch
+from torch.testing import assert_close
+from transformers import AutoTokenizer
+
+from fastvideo.models.loader.component_loader import TokenizerLoader
+
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+ZIMAGE_TOKENIZER_DIR = REPO_ROOT / "official_weights" / "Z-Image" / "tokenizer"
+PARITY_SCOPE = "production_loader"
+OFFICIAL_MAX_SEQUENCE_LENGTH = 512
+OFFICIAL_CHAT_TEMPLATE_KWARGS = {
+    "tokenize": False,
+    "add_generation_prompt": True,
+    "enable_thinking": True,
+}
+
+
+@dataclass
+class _DummyPipelineConfig:
+    text_encoder_configs: tuple = field(default_factory=tuple)
+
+
+@dataclass
+class _DummyFastVideoArgs:
+    pipeline_config: _DummyPipelineConfig = field(default_factory=_DummyPipelineConfig)
+
+
+def _load_reference_tokenizer():
+    if not ZIMAGE_TOKENIZER_DIR.exists():
+        pytest.skip(f"Z-Image tokenizer dir not found: {ZIMAGE_TOKENIZER_DIR}")
+    return AutoTokenizer.from_pretrained(str(ZIMAGE_TOKENIZER_DIR), local_files_only=True)
+
+
+def _load_fastvideo_tokenizer():
+    if not ZIMAGE_TOKENIZER_DIR.exists():
+        pytest.skip(f"Z-Image tokenizer dir not found: {ZIMAGE_TOKENIZER_DIR}")
+    loader = TokenizerLoader()
+    return loader.load(str(ZIMAGE_TOKENIZER_DIR), _DummyFastVideoArgs())
+
+
+def test_zimage_tokenizer_loader_and_tokenization_parity():
+    ref_tok = _load_reference_tokenizer()
+    fv_tok = _load_fastvideo_tokenizer()
+
+    prompts = [
+        "A close-up portrait of a cat in warm studio lighting.",
+        "An oil painting of a lighthouse at sunset.",
+    ]
+
+    tok_kwargs = {
+        "padding": "max_length",
+        "max_length": OFFICIAL_MAX_SEQUENCE_LENGTH,
+        "truncation": True,
+        "return_tensors": "pt",
+    }
+
+    ref_out = ref_tok(prompts, **tok_kwargs)
+    fv_out = fv_tok(prompts, **tok_kwargs)
+
+    assert_close(ref_out["input_ids"], fv_out["input_ids"], atol=0, rtol=0)
+    assert_close(ref_out["attention_mask"], fv_out["attention_mask"], atol=0, rtol=0)
+
+    # Compare key tokenizer attributes that affect generation-time behavior.
+    assert ref_tok.padding_side == fv_tok.padding_side
+    assert ref_tok.pad_token_id == fv_tok.pad_token_id
+    assert ref_tok.eos_token_id == fv_tok.eos_token_id
+
+
+@pytest.mark.skipif(not ZIMAGE_TOKENIZER_DIR.exists(), reason="Z-Image tokenizer assets are required")
+def test_zimage_tokenizer_chat_template_parity():
+    ref_tok = _load_reference_tokenizer()
+    fv_tok = _load_fastvideo_tokenizer()
+
+    assert callable(getattr(ref_tok, "apply_chat_template", None)), (
+        "Pinned Z-Image tokenizer must expose apply_chat_template"
+    )
+    assert callable(getattr(fv_tok, "apply_chat_template", None)), (
+        "FastVideo TokenizerLoader dropped the required apply_chat_template API"
+    )
+
+    messages = [{"role": "user", "content": "Describe a futuristic city skyline."}]
+
+    ref_text = ref_tok.apply_chat_template(messages, **OFFICIAL_CHAT_TEMPLATE_KWARGS)
+    fv_text = fv_tok.apply_chat_template(messages, **OFFICIAL_CHAT_TEMPLATE_KWARGS)
+
+    assert ref_text == fv_text
+
+    # Z-Image deliberately enables Qwen's thinking template. Guard against a
+    # pipeline silently using the generic non-thinking prompt path.
+    non_thinking_text = fv_tok.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        enable_thinking=False,
+    )
+    assert fv_text != non_thinking_text
+
+    tokenization_kwargs = {
+        "padding": "max_length",
+        "max_length": OFFICIAL_MAX_SEQUENCE_LENGTH,
+        "truncation": True,
+        "return_tensors": "pt",
+    }
+    ref_tokens = ref_tok(ref_text, **tokenization_kwargs)
+    fv_tokens = fv_tok(fv_text, **tokenization_kwargs)
+
+    assert_close(ref_tokens["input_ids"], fv_tokens["input_ids"], atol=0, rtol=0)
+    assert_close(ref_tokens["attention_mask"], fv_tokens["attention_mask"], atol=0, rtol=0)

tests/local_tests/zimage/test_zimage_transformer_parity.py

@@ -0,0 +1,344 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Z-Image transformer parity.
+
+Coverage scope: both. The tiny CPU tests compare the pinned implementation and
+FastVideo model directly. The production-loader test activates when the real
+transformer subfolder is available.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import importlib.util
+import json
+import os
+from pathlib import Path
+import subprocess
+import sys
+from types import ModuleType
+
+import pytest
+import torch
+from torch.testing import assert_close
+
+from fastvideo.configs.models.dits import ZImageDiTConfig
+from fastvideo.models.dits.zimage import ZImageAttention, ZImageTransformer2DModel, _prepare_attention_mask
+
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+OFFICIAL_REF_DIR = Path(os.getenv("ZIMAGE_OFFICIAL_REF_DIR", REPO_ROOT / "Z-Image"))
+OFFICIAL_SRC = OFFICIAL_REF_DIR / "src"
+TRANSFORMER_DIR = Path(
+    os.getenv("ZIMAGE_TRANSFORMER_DIR", REPO_ROOT / "official_weights" / "Z-Image" / "transformer"))
+REFERENCE_REVISION = "26f23eda626ffadda020b04ff79488e1d72004cd"
+HF_REVISION = "f332072aa78be7aecdf3ee76d5c247082da564a6"
+HF_KEY_COUNT = 521
+HF_KEY_SHA256 = "3a9216f208c1873b2cf06394411a53e1e95e10fae3b01dca0f7223556e47c354"
+PARITY_SCOPE = "both"
+
+PRODUCTION_CONFIG = {
+    "all_patch_size": (2, ),
+    "all_f_patch_size": (1, ),
+    "in_channels": 16,
+    "dim": 3840,
+    "n_layers": 30,
+    "n_refiner_layers": 2,
+    "n_heads": 30,
+    "n_kv_heads": 30,
+    "norm_eps": 1e-5,
+    "qk_norm": True,
+    "cap_feat_dim": 2560,
+    "rope_theta": 256.0,
+    "t_scale": 1000.0,
+    "axes_dims": (32, 48, 48),
+    "axes_lens": (1536, 512, 512),
+}
+
+TINY_CONFIG = {
+    "all_patch_size": (2, ),
+    "all_f_patch_size": (1, ),
+    "in_channels": 2,
+    "dim": 12,
+    "n_layers": 1,
+    "n_refiner_layers": 1,
+    "n_heads": 2,
+    "n_kv_heads": 2,
+    "norm_eps": 1e-5,
+    "qk_norm": True,
+    "cap_feat_dim": 8,
+    "rope_theta": 256.0,
+    "t_scale": 1000.0,
+    "axes_dims": (2, 2, 2),
+    "axes_lens": (192, 16, 16),
+}
+
+
+def _load_source_module(name: str, path: Path):
+    spec = importlib.util.spec_from_file_location(name, path)
+    if spec is None or spec.loader is None:
+        pytest.fail(f"Cannot build import spec for {path}")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+@pytest.fixture(scope="module")
+def reference_transformer_cls():
+    source_file = OFFICIAL_SRC / "zimage" / "transformer.py"
+    if not OFFICIAL_REF_DIR.exists():
+        pytest.skip(f"Pinned Z-Image reference clone not found: {OFFICIAL_REF_DIR}")
+    if not source_file.is_file():
+        pytest.fail(f"Z-Image reference clone is incomplete; missing {source_file}")
+
+    try:
+        revision = subprocess.run(
+            ["git", "-C", str(OFFICIAL_REF_DIR), "rev-parse", "HEAD"],
+            check=True,
+            capture_output=True,
+            text=True,
+        ).stdout.strip()
+    except (FileNotFoundError, subprocess.CalledProcessError) as exc:
+        pytest.fail(f"Cannot verify Z-Image reference revision: {exc}")
+    assert revision == REFERENCE_REVISION, (
+        f"expected Z-Image reference {REFERENCE_REVISION}, got {revision}")
+
+    module_names = (
+        "config",
+        "config.inference",
+        "config.model",
+        "utils",
+        "utils.import_utils",
+        "utils.attention",
+        "_zimage_reference_transformer",
+    )
+    saved_modules = {name: sys.modules.get(name) for name in module_names}
+    added_path = str(OFFICIAL_SRC) not in sys.path
+    if added_path:
+        sys.path.insert(0, str(OFFICIAL_SRC))
+
+    try:
+        for name in module_names:
+            sys.modules.pop(name, None)
+        utils_package = ModuleType("utils")
+        utils_package.__path__ = [str(OFFICIAL_SRC / "utils")]
+        sys.modules["utils"] = utils_package
+        _load_source_module("utils.import_utils", OFFICIAL_SRC / "utils" / "import_utils.py")
+        _load_source_module("utils.attention", OFFICIAL_SRC / "utils" / "attention.py")
+        module = _load_source_module("_zimage_reference_transformer", source_file)
+        assert Path(module.__file__).resolve() == source_file.resolve()
+        yield module.ZImageTransformer2DModel
+    finally:
+        for name in module_names:
+            sys.modules.pop(name, None)
+            if saved_modules[name] is not None:
+                sys.modules[name] = saved_modules[name]
+        if added_path:
+            sys.path.remove(str(OFFICIAL_SRC))
+
+
+def _build_fastvideo(config_values: dict) -> ZImageTransformer2DModel:
+    config = ZImageDiTConfig()
+    config.update_model_arch(config_values)
+    return ZImageTransformer2DModel(
+        config=config,
+        hf_config={"_class_name": "ZImageTransformer2DModel", **config_values},
+    )
+
+
+def _state_shapes(model: torch.nn.Module) -> dict[str, tuple[int, ...]]:
+    return {name: tuple(tensor.shape) for name, tensor in model.state_dict().items()}
+
+
+def test_zimage_transformer_production_meta_key_surface(reference_transformer_cls):
+    from fastvideo.models.registry import ModelRegistry
+
+    model_cls, _ = ModelRegistry.resolve_model_cls("ZImageTransformer2DModel")
+    assert model_cls is ZImageTransformer2DModel
+
+    with torch.device("meta"):
+        reference = reference_transformer_cls(**PRODUCTION_CONFIG)
+        fastvideo = _build_fastvideo(PRODUCTION_CONFIG)
+
+    reference_shapes = _state_shapes(reference)
+    fastvideo_shapes = _state_shapes(fastvideo)
+    assert fastvideo_shapes == reference_shapes
+    keys = sorted(reference_shapes)
+    key_digest = hashlib.sha256(("\n".join(keys) + "\n").encode()).hexdigest()
+    assert len(keys) == HF_KEY_COUNT
+    assert key_digest == HF_KEY_SHA256
+    print(f"Z-Image production transformer key/shape surface: {len(reference_shapes)} tensors")
+
+
+def test_zimage_transformer_tiny_cpu_parity(reference_transformer_cls):
+    reference = reference_transformer_cls(**TINY_CONFIG).eval()
+    generator = torch.Generator(device="cpu").manual_seed(123)
+    with torch.no_grad():
+        for parameter in reference.parameters():
+            parameter.copy_(torch.randn(parameter.shape, generator=generator, dtype=parameter.dtype) * 0.02)
+
+    fastvideo = _build_fastvideo(TINY_CONFIG).eval()
+    assert _state_shapes(fastvideo) == _state_shapes(reference)
+    fastvideo.load_state_dict(reference.state_dict(), strict=True)
+
+    generator.manual_seed(456)
+    images = [
+        torch.randn((2, 1, 4, 4), generator=generator),
+        torch.randn((2, 1, 8, 10), generator=generator),
+    ]
+    caption_features = [
+        torch.randn((3, 8), generator=generator),
+        torch.randn((35, 8), generator=generator),
+    ]
+    timestep = torch.tensor([0.25, 0.75], dtype=torch.float32)
+
+    with torch.inference_mode(), torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH):
+        reference_outputs = reference(images, timestep, caption_features)[0]
+        positional_outputs = fastvideo(images, caption_features, timestep)[0]
+        keyword_outputs = fastvideo(
+            hidden_states=images,
+            encoder_hidden_states=caption_features,
+            timestep=timestep,
+        )[0]
+
+    assert len(positional_outputs) == len(keyword_outputs) == len(reference_outputs) == 2
+    for positional, actual, expected in zip(positional_outputs, keyword_outputs, reference_outputs):
+        assert actual.shape == expected.shape
+        assert expected.abs().max() > 0
+        assert_close(positional, expected, atol=1e-6, rtol=1e-6)
+        assert_close(actual, expected, atol=1e-6, rtol=1e-6)
+
+
+def test_zimage_attention_additive_mask_parity(reference_transformer_cls):
+    reference_module = sys.modules[reference_transformer_cls.__module__]
+    reference = reference_module.ZImageAttention(12, 2, 2).eval()
+    generator = torch.Generator(device="cpu").manual_seed(321)
+    with torch.no_grad():
+        for parameter in reference.parameters():
+            parameter.copy_(torch.randn(parameter.shape, generator=generator, dtype=parameter.dtype) * 0.02)
+
+    fastvideo = ZImageAttention(12, 2, 2).eval()
+    fastvideo.load_state_dict(reference.state_dict(), strict=True)
+    hidden_states = torch.randn((2, 5, 12), generator=generator)
+    attention_mask = torch.tensor([[True, True, True, False, False], [True, True, True, True, True]])
+    additive_mask = _prepare_attention_mask(attention_mask, hidden_states.dtype)
+    assert additive_mask is not None
+    assert additive_mask.dtype == hidden_states.dtype
+    assert additive_mask.shape == (2, 1, 1, 5)
+    assert additive_mask[0, 0, 0, 0] == 0
+    assert torch.isneginf(additive_mask[0, 0, 0, -1])
+
+    with torch.inference_mode(), torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH):
+        expected = reference(hidden_states, attention_mask=attention_mask)
+        actual = fastvideo(hidden_states, attention_mask=attention_mask)
+    assert_close(actual, expected, atol=1e-6, rtol=1e-6)
+
+
+def test_zimage_transformer_rejects_sequence_parallelism(monkeypatch: pytest.MonkeyPatch):
+    model = _build_fastvideo(TINY_CONFIG)
+    monkeypatch.setattr("fastvideo.models.dits.zimage.model_parallel_is_initialized", lambda: True)
+    monkeypatch.setattr("fastvideo.models.dits.zimage.get_sp_world_size", lambda: 2)
+
+    with pytest.raises(NotImplementedError, match="sp_size=1"):
+        model(
+            hidden_states=[torch.zeros(2, 1, 4, 4)],
+            encoder_hidden_states=[torch.zeros(3, 8)],
+            timestep=torch.tensor([0.5]),
+        )
+
+
+def _has_real_weights() -> bool:
+    return any(TRANSFORMER_DIR.glob("*.safetensors"))
+
+
+def _load_real_config() -> dict:
+    with (TRANSFORMER_DIR / "config.json").open(encoding="utf-8") as file:
+        config = json.load(file)
+    assert config.pop("_class_name") == "ZImageTransformer2DModel"
+    config.pop("_diffusers_version", None)
+    config.pop("_name_or_path", None)
+    for key, expected in PRODUCTION_CONFIG.items():
+        actual = config[key]
+        if isinstance(expected, tuple):
+            actual = tuple(actual)
+        assert actual == expected, f"unexpected production config {key}={actual!r}"
+    assert set(config) == set(PRODUCTION_CONFIG)
+    return config
+
+
+def _run_fastvideo_production(inputs: tuple, dtype: torch.dtype) -> list[torch.Tensor]:
+    from fastvideo.configs.pipelines.base import PipelineConfig
+    from fastvideo.fastvideo_args import FastVideoArgs
+    from fastvideo.models.loader.component_loader import TransformerLoader
+
+    precision = "bf16" if dtype == torch.bfloat16 else "fp16"
+    args = FastVideoArgs(
+        model_path=str(TRANSFORMER_DIR),
+        dit_cpu_offload=False,
+        dit_layerwise_offload=False,
+        use_fsdp_inference=False,
+        pipeline_config=PipelineConfig(dit_config=ZImageDiTConfig(), dit_precision=precision),
+    )
+    model = TransformerLoader().load(str(TRANSFORMER_DIR), args).eval()
+    assert isinstance(model, ZImageTransformer2DModel)
+    assert args.model_paths["transformer"] == str(TRANSFORMER_DIR)
+    assert next(model.parameters()).device.type == "cuda"
+    with torch.inference_mode(), torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH):
+        outputs = model(
+            hidden_states=inputs[0],
+            encoder_hidden_states=inputs[1],
+            timestep=inputs[2],
+        )[0]
+    result = [output.detach().float().cpu() for output in outputs]
+    del model
+    torch.cuda.empty_cache()
+    return result
+
+
+def _run_reference_production(reference_transformer_cls, config: dict, inputs: tuple,
+                              dtype: torch.dtype) -> list[torch.Tensor]:
+    from safetensors.torch import load_file
+
+    state_dict = {}
+    for shard in sorted(TRANSFORMER_DIR.glob("*.safetensors")):
+        state_dict.update(load_file(str(shard), device="cpu"))
+    with torch.device("meta"):
+        model = reference_transformer_cls(**config)
+    model.load_state_dict(state_dict, strict=True, assign=True)
+    del state_dict
+    model = model.to(device="cuda", dtype=dtype).eval()
+    with torch.inference_mode(), torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH):
+        outputs = model(inputs[0], inputs[2], inputs[1])[0]
+    result = [output.detach().float().cpu() for output in outputs]
+    del model
+    torch.cuda.empty_cache()
+    return result
+
+
+@pytest.mark.skipif(
+    not _has_real_weights(),
+    reason=(
+        "Pinned Z-Image transformer weights are absent; set ZIMAGE_TRANSFORMER_DIR "
+        f"to Tongyi-MAI/Z-Image-Turbo@{HF_REVISION}/transformer"),
+)
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Real-weight Z-Image transformer parity requires CUDA")
+def test_zimage_transformer_production_loader_forward_parity(reference_transformer_cls):
+    config = _load_real_config()
+    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+    generator = torch.Generator(device="cuda").manual_seed(789)
+    inputs = (
+        [torch.randn((16, 1, 4, 4), device="cuda", dtype=dtype, generator=generator)],
+        [torch.randn((3, 2560), device="cuda", dtype=dtype, generator=generator)],
+        torch.tensor([0.25], device="cuda", dtype=dtype),
+    )
+
+    fastvideo_outputs = _run_fastvideo_production(inputs, dtype)
+    reference_outputs = _run_reference_production(reference_transformer_cls, config, inputs, dtype)
+    assert len(fastvideo_outputs) == len(reference_outputs) == 1
+
+    actual, expected = fastvideo_outputs[0], reference_outputs[0]
+    diff = (actual - expected).abs()
+    mean_scale = expected.abs().mean().clamp_min(1e-6)
+    print(f"Z-Image real-weight parity: max={diff.max():.6f}, mean={diff.mean():.6f}")
+    assert diff.mean() / mean_scale < 0.01
+    assert_close(actual, expected, atol=1e-2, rtol=1e-2)

tests/local_tests/zimage/test_zimage_vae_parity.py

@@ -0,0 +1,198 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Parity test for Z-Image VAE (AutoencoderKL decode path).
+
+This compares FastVideo's AutoencoderKL wrapper against the Z-Image
+reference AutoencoderKL implementation using the same local weights.
+
+Usage:
+    pytest tests/local_tests/zimage/test_zimage_vae_parity.py -v
+"""
+
+from __future__ import annotations
+
+import importlib
+import json
+import subprocess
+import sys
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import load_file as safetensors_load_file
+from torch.testing import assert_close
+
+from fastvideo.configs.models.vaes.autoencoder_kl import AutoencoderKLVAEConfig
+from fastvideo.models.loader.component_loader import VAELoader
+from fastvideo.models.vaes.autoencoder_kl import AutoencoderKL as FastVideoAutoencoderKL
+
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+ZIMAGE_REPO = REPO_ROOT / "Z-Image"
+ZIMAGE_SRC = REPO_ROOT / "Z-Image" / "src"
+ZIMAGE_VAE_DIR = REPO_ROOT / "official_weights" / "Z-Image" / "vae"
+ZIMAGE_VAE_CFG = ZIMAGE_VAE_DIR / "config.json"
+ZIMAGE_VAE_WEIGHTS = ZIMAGE_VAE_DIR / "diffusion_pytorch_model.safetensors"
+ZIMAGE_REFERENCE_REVISION = "26f23eda626ffadda020b04ff79488e1d72004cd"
+ZIMAGE_VAE_SCALING_FACTOR = 0.3611
+ZIMAGE_VAE_SHIFT_FACTOR = 0.1159
+PARITY_SCOPE = "both"
+
+
+def _require_pinned_reference_module(module_name: str, source_file: Path):
+    if not ZIMAGE_REPO.exists():
+        pytest.skip(f"Pinned Z-Image reference clone not found: {ZIMAGE_REPO}")
+    if not source_file.is_file():
+        pytest.fail(f"Z-Image reference clone is incomplete; missing {source_file}")
+
+    try:
+        result = subprocess.run(
+            ["git", "-C", str(ZIMAGE_REPO), "rev-parse", "HEAD"],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+    except (FileNotFoundError, subprocess.CalledProcessError) as exc:
+        pytest.fail(f"Cannot verify Z-Image reference revision: {exc}")
+
+    actual_revision = result.stdout.strip()
+    assert actual_revision == ZIMAGE_REFERENCE_REVISION, (
+        "Z-Image reference clone is not at the pinned revision: "
+        f"expected {ZIMAGE_REFERENCE_REVISION}, got {actual_revision}"
+    )
+
+    if str(ZIMAGE_SRC) not in sys.path:
+        sys.path.insert(0, str(ZIMAGE_SRC))
+    try:
+        module = importlib.import_module(module_name)
+    except Exception as exc:
+        pytest.fail(f"Cannot import pinned Z-Image module {module_name}: {exc}")
+
+    module_file = Path(module.__file__ or "").resolve()
+    assert module_file.is_relative_to(ZIMAGE_SRC.resolve()), (
+        f"{module_name} resolved outside the pinned clone: {module_file}"
+    )
+    return module
+
+
+@pytest.fixture(scope="module")
+def reference_autoencoder_cls():
+    module = _require_pinned_reference_module(
+        "zimage.autoencoder",
+        ZIMAGE_SRC / "zimage" / "autoencoder.py",
+    )
+    return module.AutoencoderKL
+
+
+def _load_cfg() -> dict:
+    if not ZIMAGE_VAE_CFG.exists():
+        pytest.skip(f"Z-Image VAE config not found: {ZIMAGE_VAE_CFG}")
+    with ZIMAGE_VAE_CFG.open("r", encoding="utf-8") as f:
+        cfg = json.load(f)
+    cfg.pop("_class_name", None)
+    cfg.pop("_diffusers_version", None)
+    cfg.pop("_name_or_path", None)
+    return cfg
+
+
+def _load_weights() -> dict[str, torch.Tensor]:
+    if not ZIMAGE_VAE_WEIGHTS.exists():
+        pytest.skip(f"Z-Image VAE weights not found: {ZIMAGE_VAE_WEIGHTS}")
+    return safetensors_load_file(str(ZIMAGE_VAE_WEIGHTS), device="cpu")
+
+
+def _build_reference(reference_autoencoder_cls, cfg: dict, weights: dict[str, torch.Tensor]) -> torch.nn.Module:
+    ref = reference_autoencoder_cls(
+        in_channels=cfg["in_channels"],
+        out_channels=cfg["out_channels"],
+        down_block_types=tuple(cfg["down_block_types"]),
+        up_block_types=tuple(cfg["up_block_types"]),
+        block_out_channels=tuple(cfg["block_out_channels"]),
+        layers_per_block=cfg["layers_per_block"],
+        latent_channels=cfg["latent_channels"],
+        norm_num_groups=cfg["norm_num_groups"],
+        scaling_factor=cfg["scaling_factor"],
+        shift_factor=cfg["shift_factor"],
+        use_quant_conv=cfg["use_quant_conv"],
+        use_post_quant_conv=cfg["use_post_quant_conv"],
+        mid_block_add_attention=cfg["mid_block_add_attention"],
+    ).eval()
+    ref.load_state_dict(weights, strict=True)
+    return ref
+
+
+def _build_fastvideo(cfg: dict, weights: dict[str, torch.Tensor]) -> torch.nn.Module:
+    fv_cfg = AutoencoderKLVAEConfig()
+    fv_cfg.update_model_arch(cfg)
+    fv = FastVideoAutoencoderKL(fv_cfg).eval()
+    fv.load_state_dict(weights, strict=True)
+    return fv
+
+
+def _load_fastvideo_production_vae(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(
+        "fastvideo.models.loader.component_loader.get_local_torch_device",
+        lambda: torch.device("cpu"),
+    )
+    pipeline_config = SimpleNamespace(
+        vae_config=AutoencoderKLVAEConfig(),
+        vae_precision="fp32",
+    )
+    fastvideo_args = SimpleNamespace(
+        model_paths={},
+        pipeline_config=pipeline_config,
+        vae_cpu_offload=False,
+    )
+    vae = VAELoader().load(str(ZIMAGE_VAE_DIR), fastvideo_args)
+    assert isinstance(vae, FastVideoAutoencoderKL)
+    assert fastvideo_args.model_paths["vae"] == str(ZIMAGE_VAE_DIR)
+    return vae
+
+
+def test_zimage_vae_decode_parity(reference_autoencoder_cls):
+    torch.manual_seed(7)
+
+    cfg = _load_cfg()
+    weights = _load_weights()
+
+    ref = _build_reference(reference_autoencoder_cls, cfg, weights)
+    fv = _build_fastvideo(cfg, weights)
+
+    # Decode-only parity is the critical path for Z-Image pipeline integration.
+    latents = torch.randn(1, cfg["latent_channels"], 8, 8, dtype=torch.float32)
+
+    with torch.no_grad():
+        ref_out = ref.decode(latents, return_dict=False)[0].detach().float()
+        fv_out = fv.decode(latents, return_dict=False)[0].detach().float()
+
+    assert ref_out.shape == fv_out.shape
+    assert_close(ref_out, fv_out, atol=1e-4, rtol=1e-4)
+
+
+def test_zimage_vae_production_loader_and_raw_decode_parity(
+    reference_autoencoder_cls,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    torch.manual_seed(17)
+
+    cfg = _load_cfg()
+    weights = _load_weights()
+    assert cfg["scaling_factor"] == ZIMAGE_VAE_SCALING_FACTOR
+    assert cfg["shift_factor"] == ZIMAGE_VAE_SHIFT_FACTOR
+    ref = _build_reference(reference_autoencoder_cls, cfg, weights)
+    fv = _load_fastvideo_production_vae(monkeypatch)
+
+    assert ref.config.scaling_factor == cfg["scaling_factor"]
+    assert fv.config.scaling_factor == cfg["scaling_factor"]
+    assert ref.config.shift_factor == cfg["shift_factor"]
+    assert fv.config.shift_factor == cfg["shift_factor"]
+
+    decode_latents = torch.randn(1, cfg["latent_channels"], 8, 8, dtype=torch.float32)
+
+    with torch.no_grad():
+        ref_out = ref.decode(decode_latents, return_dict=False)[0].detach().float()
+        fv_out = fv.decode(decode_latents, return_dict=False)[0].detach().float()
+
+    assert ref_out.shape == fv_out.shape
+    assert_close(ref_out, fv_out, atol=1e-4, rtol=1e-4)