[test]: cover Z-Image registry resolution

[feat]: add native Z-Image transformer prototype
[docs]: support image SSIM reference seeding
2026-07-13 17:47:18 -07:00 · 2026-07-13 17:46:16 -07:00 · 2026-07-13 17:44:50 -07:00 · 2026-07-13 17:35:21 -07:00 · 2026-07-13 17:26:46 -07:00 · 2026-07-13 17:26:46 -07:00
27 changed files with 3098 additions and 97 deletions
@@ -1,20 +1,24 @@
 ---
 name: seed-ssim-references
-description: Seed HF reference artefacts for a single newly-added SSIM test (pixel `.mp4` for `run_text_to_video_similarity_test`-style tests, or latent `.pt` for `run_text_to_latent_similarity_test`-style tests). Runs the test on Modal L40S, downloads the generated artefacts via `modal volume get`, pauses for the user to verify (visual eyeball for mp4, numerics dump for pt), then uploads only that test's files to `FastVideo/ssim-reference-videos`. Use when a new `fastvideo/tests/ssim/test_*_similarity.py` has just been added and has no references on HF yet.
+description: Seed HF reference artefacts for a single newly-added SSIM test (pixel video/image or latent `.pt`). Runs the test on Modal L40S, downloads the generated artefacts via `modal volume get`, pauses for the user to verify (visual review for videos/images, numerics dump for pt), then uploads only that test's files to `FastVideo/ssim-reference-videos`. Use when a new `fastvideo/tests/ssim/test_*_similarity.py` has just been added and has no references on HF yet.
 ---

-# Seed SSIM Reference Artefacts (mp4 or pt)
+# Seed SSIM Reference Artefacts (video, image, or pt)

 ## Purpose

 A brand-new SSIM test in `fastvideo/tests/ssim/` fails forever until its
 reference artefacts exist on the HF dataset
-(`FastVideo/ssim-reference-videos`). The dataset hosts two kinds of artefacts
+(`FastVideo/ssim-reference-videos`). The dataset hosts three kinds of artefacts
 side-by-side per `(model_id, backend, prompt)`:

 - **`.mp4`** — pixel ground-truth for tests that call
  `run_text_to_video_similarity_test` / `run_image_to_video_similarity_test`
-  in `inference_similarity_utils.py`. Compared via SSIM.
+  in `inference_similarity_utils.py` with the default media extension.
+  Compared via SSIM.
+- **`.png` / `.jpg` / `.jpeg`** — pixel ground-truth for image-generation
+  tests that call `run_text_to_video_similarity_test` with the matching
+  `media_extension`. Compared via SSIM.
 - **`.pt`** — pre-VAE latent bundle (fp16 full latent + fp32 slice +
  metadata + `slice_spec` + `format_version`) for tests that call
  `run_text_to_latent_similarity_test` in `latent_similarity_utils.py`.
@@ -22,11 +26,12 @@ side-by-side per `(model_id, backend, prompt)`:

 This skill:

-1. Detects which artefact type the test produces (pixel vs latent).
+1. Detects which artefact type the test produces (video, image, or latent).
 2. Runs the test on Modal's L40S pool to generate the artefacts.
 3. Downloads them to the local repo via `modal volume get`.
 4. Pauses so the user can verify quality:
-   - **mp4**: visual eyeball in a video player.
+   - **video**: visual eyeball in a video player.
+   - **image**: visual eyeball in an image viewer.
   - **pt**: numerics dump (shape, slice stats, NaN/Inf check, metadata).
 5. Uploads only the new test's files to HF, with a guard that refuses to
   overwrite anything already present.
@@ -67,8 +72,8 @@ Everything else is fixed:
  seeded by this skill.
 - HF repo: `FastVideo/ssim-reference-videos` (dataset).
 - Multi-model test files: all model ids in `*_MODEL_TO_PARAMS` are seeded
-  together; the Modal run produces one mp4 per (model, prompt, backend) and
-  the upload scopes by `--model-id`, looping if there is more than one.
+  together; the Modal run produces one artefact per (model, prompt, backend)
+  and the upload scopes by `--model-id`, looping if there is more than one.

 ## Prerequisites

@@ -77,7 +82,7 @@ The user has confirmed:
 - `modal` CLI authenticated.
 - `HF_API_KEY` (or `HUGGINGFACE_HUB_TOKEN` / `HF_TOKEN`) exported with write
  access to `FastVideo/ssim-reference-videos`.
- The test file runs locally end-to-end (generates an mp4; SSIM assertion
+- The test file runs locally end-to-end (generates its declared artefact; SSIM assertion
  failure due to missing reference is expected and fine).

 Fail fast if the token env var is missing.
@@ -100,16 +105,22 @@ Detect artefact type by inspecting the file's imports / helper call:
 - **latent** (`.pt`) — file imports `run_text_to_latent_similarity_test`
  from `fastvideo.tests.ssim.latent_similarity_utils` (or any other helper
  that ends with `_latent_similarity_test`).
- **pixel** (`.mp4`) — file imports
+- **image** (`.png`, `.jpg`, or `.jpeg`) — a pixel helper call resolves
+  `media_extension` to one of those image suffixes. For example, Z-Image calls
+  `run_text_to_video_similarity_test(..., media_extension=".png")`.
+- **video** (`.mp4` by default) — file imports
  `run_text_to_video_similarity_test` / `run_image_to_video_similarity_test`
-  from `fastvideo.tests.ssim.inference_similarity_utils`, OR uses the
-  legacy custom-inline helper pattern (see `test_gamecraft`,
-  `test_longcat`, etc.). Default to pixel when both heuristics fail.
+  from `fastvideo.tests.ssim.inference_similarity_utils` without an image
+  `media_extension`, OR uses the legacy custom-inline helper pattern (see
+  `test_gamecraft`, `test_longcat`, etc.). Default to video when all
+  heuristics fail.

-Record `ARTEFACT_TYPE ∈ {pixel, latent}` for use in step 4. Steps 2, 3, 5,
+Record `ARTEFACT_TYPE ∈ {video, image, latent}` for use in step 4. Steps 2, 3, 5,
 and 6 are artefact-type-agnostic — `_iter_reference_files`,
 `copy_generated_to_reference`, and `upload_reference_videos` already walk
-both `.mp4` and `.pt` (see `reference_videos_cli.py`).
+videos, images, and `.pt` bundles (see `REFERENCE_EXTENSIONS` in
+`reference_videos_cli.py`). The Modal exporter likewise counts all three via
+`GENERATED_ARTIFACT_EXTENSIONS` before copying the whole generated tree.

 If either check fails, stop and tell the user what's wrong.

@@ -166,10 +177,10 @@ Flag rationale:
 - `--skip-reference-download`: no refs exist yet, so conftest must not try to
  pull them.
 - `--no-fail-fast`: lets the test finish generation before `_assert_similarity`
-  raises `FileNotFoundError: Reference video folder does not exist`. The
-  expected failure is what we want — the mp4 has already been written.
+  raises for the missing reference. The expected failure is what we want —
+  the generated artefact has already been written.
 - `--sync-generated-to-volume` + `--generated-volume-subdir`: copies the
-  generated mp4s to the `hf-model-weights` Modal volume under
+  generated artefacts to the `hf-model-weights` Modal volume under
  `ssim_generated_videos/default/<SUBDIR>/generated_videos/` so we can pull
  them locally.

@@ -177,7 +188,7 @@ The Modal run will end with a nonzero exit (expected) and print a
 `modal volume get hf-model-weights ssim_generated_videos/default/<SUBDIR>/generated_videos ./generated_videos_modal/default`
 command. Capture that `<SUBDIR>` — you need it for step 3.

-### 3. Download generated videos locally
+### 3. Download generated artefacts locally

 ```bash
 modal volume get --force hf-model-weights \
@@ -189,8 +200,8 @@ modal volume get --force hf-model-weights \
 already exists; without it, `modal volume get` errors with `[Errno 21] Is a
 directory`. Safe to pass on the first run too.

-After this, the mp4s live at
-`./generated_videos_modal/default/generated_videos/L40S_reference_videos/<model_id>/<backend>/<prompt>.mp4`.
+After this, the artefacts live at
+`./generated_videos_modal/default/generated_videos/L40S_reference_videos/<model_id>/<backend>/<prompt>.<ext>`.
 The extra `generated_videos/` level comes from the volume layout in
 `_sync_generated_videos_to_volume` (`ssim_test.py`) — the command copies
 `<repo>/fastvideo/tests/ssim/generated_videos/<tier>` to
@@ -201,11 +212,16 @@ get` preserves that trailing `generated_videos/` segment.

 Type-aware verification.

-**For `ARTEFACT_TYPE = pixel`** — list the downloaded mp4s and ask the user to
+**For `ARTEFACT_TYPE = video`** — list the downloaded videos and ask the user to
 open them in a video player:

 > "Generated videos downloaded to `./generated_videos_modal/default/generated_videos/L40S_reference_videos/`. Please open them and confirm the quality looks correct. Reply **`upload`** to continue, or anything else to abort."

+**For `ARTEFACT_TYPE = image`** — list the downloaded `.png` / `.jpg` /
+`.jpeg` files and ask the user to open them in an image viewer:
+
+> "Generated images downloaded to `./generated_videos_modal/default/generated_videos/L40S_reference_videos/`. Please open them and confirm the quality looks correct. Reply **`upload`** to continue, or anything else to abort."
+
 **For `ARTEFACT_TYPE = latent`** — `.pt` files are not human-watchable. Print
 a numerics dump for each `.pt` so the user can sanity-check shape, distribution,
 and metadata:
@@ -250,8 +266,9 @@ everything on disk so they can inspect further — no cleanup.

 ### 5. Copy into the local reference layout

-Scoped copy — only the new test's artefacts. Single command works for both
-artefact types because `_iter_reference_files` walks `.mp4` and `.pt`:
+Scoped copy — only the new test's artefacts. One command works for all three
+artefact types because `_iter_reference_files` walks video, image, and latent
+extensions:

 ```bash
 python fastvideo/tests/ssim/reference_videos_cli.py copy-local \
@@ -261,12 +278,13 @@ python fastvideo/tests/ssim/reference_videos_cli.py copy-local \
 ```

 (The `--generated-dir` points at the device-folder root inside the
-downloaded tree; `copy-local` walks all `<model>/<backend>/*.{mp4,pt}`
+downloaded tree; `copy-local` walks all supported reference extensions
 underneath it. Since the Modal run was scoped to a single test file via
 `--test-files`, only that test's model(s) are present — so the copy is
 implicitly per-test.)

-Result for pixel: `fastvideo/tests/ssim/reference_videos/default/L40S_reference_videos/<model_id>/<backend>/<prompt>.mp4`.
+Result for video: `fastvideo/tests/ssim/reference_videos/default/L40S_reference_videos/<model_id>/<backend>/<prompt>.mp4`.
+Result for image: same path with `.png`, `.jpg`, or `.jpeg` extension.
 Result for latent: same path with `.pt` extension.

 ### 6. Upload to HF — scoped per model_id, with overwrite guard
@@ -307,7 +325,7 @@ it will auto-download the refs they just uploaded.
 - **Modal run fails before generation.** No artefacts on the volume — nothing
  to download. Fix the test locally (`pytest fastvideo/tests/ssim/<test_file>`)
  and retry from step 2.
- **`./generated_videos_modal/default/L40S_reference_videos/` missing after
+- **`./generated_videos_modal/default/generated_videos/L40S_reference_videos/` missing after
  `modal volume get`.** The run didn't produce artefacts (most likely the
  test crashed before writing, or `REQUIRED_GPUS` exceeded the partition
  capacity — see Modal logs).
@@ -340,14 +358,14 @@ it will auto-download the refs they just uploaded.
 - The overwrite guard in `reference_videos_cli.py upload` is default-on
  specifically because this skill exists. Re-seeding is a distinct operation
  that requires explicit `--force`.
- Both artefact types share the same Modal flow: the orchestrator sets
+- All three artefact types share the same Modal flow: the orchestrator sets
  `--skip-reference-download` + `--no-fail-fast`, runs pytest, the test's
-  helper writes the artefact (`.mp4` via `imageio` for pixel,
+  helper writes the artefact (`.mp4` / `.png` / `.jpg` / `.jpeg` for pixel,
  `save_latent_reference` → `torch.save` for latent) BEFORE the
  missing-reference assertion raises. `_sync_generated_videos_to_volume` in
  `ssim_test.py` does a `shutil.copytree` of the whole `generated_videos/`
-  tree, picking up `.mp4`, `.pt`, and the `*_ssim.json` / `*_latent.json`
-  metric files alongside.
+  tree. Its artifact count recognizes videos, images, and `.pt` bundles; JSON
+  metric files are copied alongside but are not prepared as references.

 ## References

@@ -357,7 +375,7 @@ it will auto-download the refs they just uploaded.
 - `fastvideo/tests/ssim/reference_videos_cli.py` — `copy-local`, `upload`
  (with `--model-id`, `--force`), `download`, `ensure` subcommands.
  Extension allowlist is `REFERENCE_EXTENSIONS = VIDEO_EXTENSIONS +
-  LATENT_EXTENSIONS` (`.pt`).
+  IMAGE_EXTENSIONS + LATENT_EXTENSIONS`.
 - `fastvideo/tests/ssim/README.md` — reference layout, HF repo conventions.
 - `fastvideo/tests/ssim/inference_similarity_utils.py` — pixel helpers
  (`run_text_to_video_similarity_test`,
@@ -376,3 +394,4 @@ it will auto-download the refs they just uploaded.
 | 2026-04-21 | Rewrite: single-test scope, explicit user-review pause, per-`model_id` upload, HF overwrite guard. Dropped `scripts/seed_ssim.sh`. |
 | 2026-04-21 | Post-first-run fixes: `modal volume get` needs `--force` when parent exists; download tree has an extra `generated_videos/` level so `--generated-dir` must reflect it. |
 | 2026-05-01 | Latent (`*.pt`) artefact support: artefact-type detection in step 1, type-aware verification (visual eyeball for mp4, numerics dump for pt) in step 4, FSDP+inference_mode failure-mode added, design notes for the unified Modal flow. Triggered by PR #1253 (LTX-2 latent migration + Stable Audio latent test). |
+| 2026-07-13 | Image (`.png`, `.jpg`, `.jpeg`) artefact support: image-aware detection and preview guidance, plus current Modal export and reference CLI extension coverage. |
@@ -34,6 +34,7 @@ env
 *.log
 weights/
 logs/
+/Z-Image/
 official_weights/
 converted_weights/

@@ -12,6 +12,7 @@ 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

@@ -19,5 +20,5 @@ __all__ = [
    "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "DreamXWorldConfig",
    "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "FluxDiTConfig", "Flux2Config",
    "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig",
-    "StableAudioConfig", "GlmImageDiTConfig"
+    "StableAudioConfig", "GlmImageDiTConfig", "ZImageDiTConfig"
 ]
@@ -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"
@@ -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
@@ -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
@@ -327,14 +327,12 @@ 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,
@@ -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

@@ -401,6 +401,10 @@ class TextEncoderLoader(ComponentLoader):
                        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()

                model = model_cls(model_config)  # type: ignore
@@ -64,6 +64,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 = {
@@ -78,6 +79,9 @@ _TEXT_ENCODER_MODELS = {
    "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":
@@ -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)
@@ -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

@@ -221,6 +221,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 +278,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 +289,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"]
@@ -79,6 +79,18 @@ RAW_GENERATED_VOLUME_ROOT = "ssim_generated_videos"
 DEFAULT_OUTPUT_QUALITY_TIER = "default"
 FULL_OUTPUT_QUALITY_TIER = "full_quality"
 MODAL_DEVICE_REFERENCE_FOLDER = "L40S_reference_videos"
+GENERATED_ARTIFACT_EXTENSIONS = (
+    ".mp4",
+    ".avi",
+    ".mov",
+    ".mkv",
+    ".webm",
+    ".flv",
+    ".png",
+    ".jpg",
+    ".jpeg",
+    ".pt",
+)
 SSIM_COMMON_KWARGS = dict(
    image=image,
    timeout=5400,
@@ -277,11 +289,11 @@ def _print_local_reference_copy_command(quality_tier: str) -> None:
    )


-def _count_video_files(root: str) -> int:
+def _count_generated_artifacts(root: str) -> int:
    count = 0
    for current_root, _, filenames in os.walk(root):
        for filename in filenames:
-            if filename.lower().endswith((".mp4", ".avi", ".mov", ".mkv", ".webm", ".flv")):
+            if filename.lower().endswith(GENERATED_ARTIFACT_EXTENSIONS):
                count += 1
    return count

@@ -316,8 +328,11 @@ def _sync_generated_videos_to_volume(
    shutil.copytree(generated_root, absolute_dst)
    model_vol.commit()

-    num_videos = _count_video_files(absolute_dst)
-    print(f"Raw generated videos exported to Modal volume path: {relative_dst} ({num_videos} video files).")
+    num_artifacts = _count_generated_artifacts(absolute_dst)
+    print(
+        f"Raw generated artifacts exported to Modal volume path: "
+        f"{relative_dst} ({num_artifacts} artifact files)."
+    )
    print(
        "Download command:\n"
        f"  modal volume get hf-model-weights {relative_dst} "
@@ -1002,7 +1017,7 @@ def run_ssim_tests(
            quality_tier=quality_tier,
        )
        print(
-            "Raw generated videos will be saved to Modal volume path: "
+            "Raw generated artifacts will be saved to Modal volume path: "
            f"{generated_volume_path}"
        )
    else:
@@ -1042,7 +1057,7 @@ def run_ssim_tests(
        )
        local_download_dir = _build_local_generated_download_dir(quality_tier)
        print(
-            "To download raw generated videos locally, run:\n"
+            "To download raw generated artifacts locally, run:\n"
            f"  modal volume get hf-model-weights {download_src} "
            f"{local_download_dir}"
        )
@@ -99,3 +99,11 @@ def test_spawn_ssim_task_passes_hf_api_key_to_pytest_env(monkeypatch, tmp_path):
    assert env["FASTVIDEO_SSIM_MODEL_ID"] == "model-a"
    assert running_task.process.__class__ is FakeProcess
    running_task.log_handle.close()
+
+
+def test_count_generated_artifacts_includes_images(monkeypatch, tmp_path):
+    module = _load_ssim_test_module(monkeypatch)
+    for filename in ("sample.mp4", "sample.png", "sample.jpg", "sample.jpeg", "sample.pt", "metrics.json"):
+        (tmp_path / filename).write_text(filename, encoding="utf-8")
+
+    assert module._count_generated_artifacts(str(tmp_path)) == 5
@@ -1,17 +1,20 @@
-The reference videos are used as part of e2e SSIM regression tests.
+Reference media are used as part of e2e SSIM regression tests.
 Wan inference coverage now lives in `test_wan_t2v_similarity.py` and
 `test_wan_i2v_similarity.py` alongside the other model-specific SSIM files.
-These tests compare newly generated videos against references to detect quality
+These tests compare newly generated videos or images against references to detect quality
 regressions.

 reference layout:
 - `reference_videos/default/<GPU>_reference_videos/<model>/<backend>/...`
 - `reference_videos/full_quality/<GPU>_reference_videos/<model>/<backend>/...`

+Reference artifacts may be videos, images (`.png`, `.jpg`, `.jpeg`), or latent
+bundles (`.pt`).
+
 `A40_reference_videos` are generated on A40s and so on. (legacy default layout
 `<ssim_dir>/<GPU>_reference_videos/...` is still read as fallback.)

-Before SSIM tests run, missing reference videos are auto-downloaded from a
+Before SSIM tests run, missing reference artifacts are auto-downloaded from a
 public HF repo (configured by `FASTVIDEO_SSIM_REFERENCE_HF_REPO`, default:
 `FastVideo/ssim-reference-videos`).

@@ -43,7 +46,7 @@ to skip auto-download of missing refs:
 to bootstrap draft refs for a new model:
 `pytest fastvideo/tests/ssim/ -vs --ssim-bootstrap-mode`

-generated videos are written under:
+generated artifacts are written under:
 - default params: `generated_videos/default/<GPU>_reference_videos/<model_id>/<ATTENTION_BACKEND>/`
 - full-quality params: `generated_videos/full_quality/<GPU>_reference_videos/<model_id>/<ATTENTION_BACKEND>/`

@@ -86,7 +89,7 @@ Normal SSIM runs are strict: missing references fail the test. For a new model
 PR, add `[new-model]` to the PR title or set `FASTVIDEO_SSIM_BOOTSTRAP_MODE=1`
 for the Buildkite job. CI then passes `--ssim-bootstrap-mode` to pytest.

-In bootstrap mode, a missing pixel `.mp4` or latent `.pt` reference is treated
+In bootstrap mode, a missing pixel `.mp4`/`.png`/`.jpg`/`.jpeg` or latent `.pt` reference is treated
 as an expected draft-reference case after the generated artifact has been
 written. The test uploads that generated artifact to the HF reference repo under
 `drafts/...` and marks the case `xfail`. Bootstrap mode does not weaken normal
@@ -13,12 +13,13 @@ from pathlib import Path


 VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv", ".webm", ".flv")
+IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg")
 # Additional artefact types stored under the same reference folders. Latent
 # tensors (`.pt`) back the latent-slice regression tests used by flaky
 # 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 + IMAGE_EXTENSIONS + LATENT_EXTENSIONS
 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 +52,10 @@ def _default_repo_type() -> str:


 def _iter_reference_files(root: Path) -> Iterable[Path]:
-    """Yield video and latent (.pt) references under `root`.
+    """Yield video, image, and latent 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 still satisfy readiness checks.
    """
    for path in root.rglob("*"):
        if path.is_file() and path.suffix.lower() in REFERENCE_EXTENSIONS:
@@ -824,7 +825,7 @@ def main(argv: Sequence[str] | None = None) -> int:
            reference_dir=reference_dir,
            dry_run=args.dry_run,
        )
-        print(f"Done. {'Would copy' if args.dry_run else 'Copied'} {copied} video files.")
+        print(f"Done. {'Would copy' if args.dry_run else 'Copied'} {copied} reference files.")
        if not args.dry_run and copied > 0:
            marker = (
                _ssim_dir() / REFERENCE_VIDEOS_DIRNAME / args.quality_tier / f".download_complete_{args.quality_tier}"
@@ -55,6 +55,46 @@ def _fake_hf(monkeypatch):
    )


+def test_image_references_are_copied_ready_and_uploaded(tmp_path):
+    generated_dir = tmp_path / "generated"
+    generated_backend_dir = generated_dir / "model-a" / "flash"
+    generated_backend_dir.mkdir(parents=True)
+    for filename in ("sample.png", "sample.jpg", "sample.jpeg", "ignored.json"):
+        (generated_backend_dir / filename).write_text(filename, encoding="utf-8")
+
+    ssim_dir = tmp_path / "ssim"
+    reference_dir = (
+        ssim_dir
+        / "reference_videos"
+        / "default"
+        / "L40S_reference_videos"
+    )
+
+    copied = reference_videos_cli.copy_generated_to_reference(
+        generated_dir=generated_dir,
+        reference_dir=reference_dir,
+    )
+
+    assert copied == 3
+    assert not (reference_dir / "model-a" / "flash" / "ignored.json").exists()
+    (reference_dir.parent / ".download_complete_default").touch()
+    assert reference_videos_cli._has_local_reference_videos(ssim_dir, "default")
+
+    reference_videos_cli.upload_reference_videos(
+        repo_id="FastVideo/ssim-reference-videos",
+        repo_type="dataset",
+        reference_dirs_by_tier=[("default", reference_dir)],
+        token="hf_token",
+        private=False,
+        model_id="model-a",
+    )
+    uploaded_folder = Path(_FakeHfApi.instances[-1].uploaded_folders[0]["folder_path"])
+    assert all(
+        (uploaded_folder / "flash" / filename).exists()
+        for filename in ("sample.png", "sample.jpg", "sample.jpeg")
+    )
+
+
 def test_upload_draft_reference_artifact_uses_drafts_prefix(tmp_path):
    ssim_dir = tmp_path / "ssim"
    generated = (
@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+
+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
+# Calibrate after the first repeated L40S reference runs.
+ZIMAGE_MIN_SSIM = 0.98
+
+ZIMAGE_MODEL_PATH = os.getenv(
+    "ZIMAGE_MODEL_DIR",
+    "Tongyi-MAI/Z-Image-Turbo",
+)
+
+TEST_PROMPTS = [
+    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
+]
+
+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": 9,
+        "guidance_scale": 0.0,
+        "seed": 42,
+    },
+}
+
+ZIMAGE_FULL_QUALITY_MODEL_TO_PARAMS: dict[str, dict[str, object]] = {
+    "Tongyi-MAI__Z-Image-Turbo": dict(
+        ZIMAGE_MODEL_TO_PARAMS["Tongyi-MAI__Z-Image-Turbo"]
+    ),
+}
+
+
+@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={
+            "workload_type": "t2i",
+            "use_fsdp_inference": False,
+        },
+        generation_kwargs_override={
+            "save_video": True,
+        },
+    )
@@ -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"
@@ -0,0 +1,125 @@
+# Z-Image Port Status
+
+## Summary
+
+- model_family: `zimage`
+- variant: `Z-Image-Turbo`
+- workload_types: `T2I`
+- PR scope: components plus native transformer prototype/parity
+- 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: `diffusers`
+- 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 6 (transformer parity activation and component revalidation)
+- status: in_progress
+- owner: dit/parity
+- last_updated: 2026-07-13
+
+## Component Matrix
+
+| Component | Type | Reuse/Port | Official Definition / Instantiation | FastVideo Target | Prototype | Conversion | Parity | Open Issues |
+|---|---|---|---|---|---|---|---|---|
+| Text encoder (Qwen3 body) | text_encoder | reused production body + native implementation | Pinned `Qwen3ForCausalLM` checkpoint; Transformers `AutoModel` body; official pipeline consumes `hidden_states[-2]` | `TextEncoderLoader` -> `Qwen3ForCausalLM.from_pretrained_local` -> independent body-only `AutoModel`, plus native `fastvideo.models.encoders.qwen3.Qwen3ForCausalLM` | DONE | not_needed for current HF subfolder | REVALIDATION REQUIRED (historical unpinned native fp32/bf16 PASS only) | I001, I006, I007 |
+| Tokenizer | tokenizer | reused | `AutoTokenizer` from `<weights>/tokenizer/`; official thinking template and length 512 | `TokenizerLoader` | DONE | not_needed | REVALIDATION REQUIRED (historical unpinned PASS) | I003, I010 |
+| VAE | vae | reused with accepted wrapper exception | Pinned `zimage.AutoencoderKL`; official decode uses `(latents / scaling_factor) + shift_factor` | Existing Diffusers-backed `fastvideo.models.vaes.autoencoder_kl.AutoencoderKL` through production `VAELoader`, plus direct implementation coverage | DONE | not_needed | REVALIDATION REQUIRED (historical unpinned direct-decode PASS only) | I008 |
+| Scheduler | scheduler | reused with Z-Image extension | Pinned `zimage.FlowMatchEulerDiscreteScheduler`; official pipeline sets `sigma_min=0.0` | `FlowMatchEulerDiscreteScheduler(use_reference_discrete_timesteps=True, sigma_min=0.0)` | DONE | not_needed | REVALIDATION REQUIRED; asset-free regressions remain separately runnable | I002, I009 |
+| Transformer (`ZImageTransformer2DModel`) | dit | native port | Pinned `src/zimage/transformer.py`; published Diffusers transformer subfolder | `fastvideo.models.dits.zimage.ZImageTransformer2DModel` through `TransformerLoader` | DONE (production meta + tiny CPU forward) | not_needed; exact 521-key surface | TINY NON-SKIP PASS; production real-weight forward PENDING | I011 |
+| Pipeline | pipeline | planned | Pinned `ZImagePipeline` | future PR | NOT_STARTED | depends_on_transformer | NOT_STARTED | I005 |
+
+## Conversion State
+
+- conversion_script: none for the reused components in this component-only PR
+- converted_weights_dir: none
+- source_layout: diffusers
+- 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 pending; other components revalidation_required
+- production_passthrough: 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` | RERUN REQUIRED; historical unpinned 2/2 PASS on A40, 2026-05-12 | Now enforces reference SHA/import origin and exact zero endpoint; includes asset-free positional/default regressions. |
+| Tokenizer | `pytest tests/local_tests/zimage/test_zimage_tokenizer_parity.py -v -s` | RERUN REQUIRED; historical unpinned 2/2 PASS on A40, 2026-05-12 | Now 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` | RERUN REQUIRED; historical unpinned 1/1 direct-decode PASS on A40, 2026-05-12 | Scope is now `both`: direct decode plus `VAELoader` and official latent normalization. |
+| Text encoder fp32 | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_parity_forward[fp32] -v -s` | RERUN REQUIRED; historical unpinned native PASS on L40S, 2026-06-21 | New run must cover independent reference, body-only production path, and native implementation. |
+| Text encoder bf16 | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_parity_forward[bf16] -v -s` | RERUN REQUIRED; historical unpinned native PASS on L40S, 2026-06-21 | Historical thresholds retained pending pinned-snapshot validation. |
+| Per-layer bf16 diagnostic | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_per_layer_bf16_diagnostic -v -s` | RERUN REQUIRED; historical unpinned informational run on L40S, 2026-06-21 | Pinned architecture has 36 blocks and 37 hidden-state entries under the tested contract. |
+| Transformer meta/tiny parity | `ZIMAGE_OFFICIAL_REF_DIR=<pinned-clone> pytest tests/local_tests/zimage/test_zimage_transformer_parity.py -v -s` | 4 PASSED, 1 SKIPPED on CPU, 2026-07-13 | Production 521-key/shape surface and pinned HF key digest match; additive-mask attention and tiny variable-length full-forward parity are exact; SP>1 guard passes. Real-weight test skipped because the 24.6 GB transformer subfolder/CUDA were absent. |
+
+## 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 | text_encoder | medium | The full checkpoint includes `lm_head.weight`, while the native FastVideo encoder owns only the body. | Native parity allowlists exactly `lm_head.weight`; production uses body-only Transformers `AutoModel`. | parity | in_progress | Implementation is complete; resolve only after the pinned fp32/bf16 rerun. |
+| I002 | Phase 7 | scheduler/pipeline | high | Published `scheduler_config.json` omits the reference-schedule flag and official zero endpoint. | Pinned pipeline mutates `sigma_min=0.0`; parity requires both scheduler values. | pipeline | open | Future pipeline config must serialize `use_reference_discrete_timesteps=True` and `sigma_min=0.0`. |
+| 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 | in_progress | Test/config behavior corrected; pinned asset rerun required before resolving. |
+| 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 7 | pipeline | high | No FastVideo pipeline/config/preset/registry/example or quality regression exists. | Component-only PR scope. | pipeline | open | Separate future PR after component gates pass. |
+| I006 | Phase 6 | text_encoder | high | Production loading previously used a causal-LM object and could override CPU placement, while parity exercised only direct native construction. | Production path is now body-only `AutoModel`; loader records `_fastvideo_input_device`; stage moves tokens to that device; parity uses a distinct reference instance. | encoder | in_progress | Implementation complete; pinned production-loader rerun required. |
+| I007 | Phase 6 | text_encoder | high | Lenient native loading could miss one fused source shard or an auxiliary quantization scale. | Loader completeness now covers every destination, required split shard, exact fused destination, and mapped scale parameter. | encoder | in_progress | Asset-free strictness regressions and pinned checkpoint load must pass before resolving. |
+| 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 scope is now `both` and includes `VAELoader` plus official normalization. | parity | in_progress | Existing Diffusers `AutoencoderKL` wrapper exception is accepted only for this component-only PR; pinned rerun required. |
+| 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. Asset-backed numerical rerun remains I003. |
+| I011 | Phase 6 | transformer | high | Production real-weight strict load and full-forward parity are not verified. | The pinned transformer is 24.6 GB and is not present locally; the CUDA test scaffold therefore skipped. | parity | open | Run `test_zimage_transformer_production_loader_forward_parity` non-skip with the pinned HF transformer subfolder on CUDA and record drift. |
+
+## Escape Hatches
+
+| ID | Phase | Decision Type | Question | Recommended Option | Status | Resolution |
+|---|---|---|---|---|---|---|
+
+## Decisions
+
+| Date | Decision | Rationale | Impact |
+|---|---|---|---|
+| 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 are revalidation-required. |
+| 2026-07-04 | Use a body-only Transformers `AutoModel` for production and a separate `AutoModel` object as the parity oracle. | Every FastVideo consumer needs hidden states, not full-vocabulary logits or an LM head; 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 | Accept the existing Diffusers-backed AutoencoderKL wrapper only for this component-only PR. | The target is established shared code and this PR adds production-loader evidence rather than a new VAE implementation. | Exception is limited to the VAE here and is not precedent for the transformer or future full pipeline. |
+| 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 still requires rerun. |
+| 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 remain provisional until the pinned 36-block snapshot is rerun. |
+| 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
+
+- No asset-backed component currently has pinned-snapshot `non_skip_pass`
+  evidence. Do not promote a historical row to PASS without recording the exact
+  command and non-skip result against both immutable pins.
+- Text-encoder revalidation must cover the independent `AutoModel` oracle,
+  body-only production loader, native implementation, requested input device,
+  fused/split source completeness, and auxiliary quantization scales.
+- VAE revalidation must cover both direct implementation decode and production
+  `VAELoader` behavior with official scaling/shift normalization.
+- Transformer production key/shape and tiny output parity pass. The real-weight
+  production-loader/full-forward CUDA test remains I011 and must pass non-skip.
+- Next work after I011: pipeline/config/preset/registry/example, pipeline parity,
+  and image-quality regression.
@@ -0,0 +1,161 @@
+# Z-Image Local Tests
+
+Local-only component coverage for the `zimage` FastVideo port (`T2I`,
+Z-Image-Turbo only). Coverage includes the Qwen3 text encoder, tokenizer,
+AutoencoderKL VAE, FlowMatchEulerDiscreteScheduler, and a native
+`ZImageTransformer2DModel`. Pipeline wiring, pipeline parity, examples, and
+image-quality regression remain out of scope.
+
+> **Status:** DRAFT port, in progress. Every test that consumes the reference
+> clone or HF assets requires a fresh non-skip rerun against the immutable pins
+> below. Earlier A40/L40S results used an unrecorded HF snapshot and are retained
+> only as historical diagnostics, not current verification evidence. 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` |
+| Component scope | text encoder, tokenizer, VAE, scheduler, native transformer |
+| 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) |
+| Source layout | Diffusers-style per-component subfolders |
+| 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: transformer source import and tiny CPU parity verified; asset-backed components pending
+private_dep_stubs: none
+blocked_on: pinned-snapshot component reruns and transformer real-weight CUDA parity
+```
+
+## 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 | REVALIDATION REQUIRED |
+| 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 | REVALIDATION REQUIRED (historical unpinned PASS) |
+| VAE ([`test_zimage_vae_parity.py`](./test_zimage_vae_parity.py)) | `both`; direct implementation decode plus production `VAELoader`; production check applies `(latents / scaling_factor) + shift_factor` before decode | REVALIDATION REQUIRED (historical unpinned direct-decode PASS only) |
+| Text encoder ([`test_zimage_encoder_parity.py`](./test_zimage_encoder_parity.py)) | `both`; independent Transformers `AutoModel` reference, body-only production loader, and FastVideo-native implementation; fp32/bf16 output checks plus fused/split/quant-scale strictness | REVALIDATION REQUIRED (historical unpinned native PASS only) |
+| 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 scaffold | TINY PARITY PASS (4 passed); REAL-WEIGHT PARITY PENDING (24.6 GB weights absent) |
+
+## 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 a separate `AutoModel` instance; it is not the object
+  returned by the 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. A future Z-Image pipeline config must
+  opt in with `True`.
+
+### VAE
+
+The test covers both direct implementation parity and production `VAELoader`
+resolution/strict loading. This component-only PR explicitly accepts the
+existing shared `fastvideo.models.vaes.autoencoder_kl.AutoencoderKL` wrapper,
+which subclasses Diffusers `AutoencoderKL`, as a narrowly scoped exception to
+the native-component boundary. This is not precedent for the future transformer
+or full Z-Image pipeline, and the exception must be revisited if that scope grows.
+
+### Scheduler
+
+The pinned official pipeline mutates `scheduler.sigma_min = 0.0` before building
+its `num_steps + 1` schedule. A future Z-Image pipeline config must set 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; its non-skip result is still required.
+
+## Remaining work
+
+- Run every asset-backed component test non-skip against the pinned clone and HF
+  snapshot before claiming component parity.
+- Run the transformer production-loader/full-forward parity test non-skip with
+  the pinned 24.6 GB transformer weights on CUDA.
+- Add the Z-Image pipeline/config/preset/registry/example, pipeline parity, and
+  image-quality regression after all component gates pass.
@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0
@@ -0,0 +1,743 @@
+# 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.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 = 1e-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
+
+
+def test_qwen3_production_loader_uses_body_only_model_and_honors_cpu_offload(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """The production passthrough must avoid the LM head and preserve CPU placement."""
+    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()
+
+    def fake_auto_model_from_pretrained(*args, **kwargs):
+        auto_model_calls.append(kwargs)
+        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,
+        torch.device("cuda"),
+        _loader_args(cpu_offload=True),
+        dtype="fp32",
+        use_text_encoder_override=True,
+    )
+
+    assert loaded is fake_model
+    assert auto_model_calls == [{
+        "local_files_only": True,
+        "torch_dtype": torch.float32,
+        "low_cpu_mem_usage": True,
+    }]
+    assert placements == [torch.device("cpu")]
+    assert loaded._fastvideo_input_device == torch.device("cpu")
+
+
+@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, dtype=dtype)
+
+    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, dtype=dtype)
+
+    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,
+        )
@@ -0,0 +1,227 @@
+# 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))
+
+
+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))
+
+
+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)
@@ -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)
@@ -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)
@@ -0,0 +1,226 @@
+# 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
+from fastvideo.pipelines.stages.decoding import DecodingStage
+
+
+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 _official_decode_latents(latents: torch.Tensor, vae: torch.nn.Module) -> torch.Tensor:
+    scaling_factor = vae.config.scaling_factor
+    shift_factor = vae.config.shift_factor or 0.0
+    dtype = next(vae.parameters()).dtype
+    return (latents.to(dtype=dtype) / scaling_factor) + shift_factor
+
+
+def test_zimage_decode_normalization_uses_autoencoder_config_values():
+    vae = SimpleNamespace(
+        config=SimpleNamespace(
+            latents_mean=None,
+            latents_std=None,
+            scaling_factor=ZIMAGE_VAE_SCALING_FACTOR,
+            shift_factor=ZIMAGE_VAE_SHIFT_FACTOR,
+        )
+    )
+    latents = torch.tensor([[[[0.0, 0.5], [-0.5, 1.0]]]], dtype=torch.float32)
+
+    actual = DecodingStage(vae)._denormalize_latents(latents)
+    expected = latents / ZIMAGE_VAE_SCALING_FACTOR + ZIMAGE_VAE_SHIFT_FACTOR
+
+    assert_close(actual, expected, atol=0, rtol=0)
+
+
+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_decode_normalization(
+    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"]
+
+    pipeline_latents = torch.randn(1, cfg["latent_channels"], 8, 8, dtype=torch.float32)
+    ref_decode_latents = _official_decode_latents(pipeline_latents, ref)
+    fv_decode_latents = DecodingStage(fv)._denormalize_latents(pipeline_latents)
+    assert_close(ref_decode_latents, fv_decode_latents, atol=0, rtol=0)
+
+    with torch.no_grad():
+        ref_out = ref.decode(ref_decode_latents, return_dict=False)[0].detach().float()
+        fv_out = fv.decode(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)
Author	SHA1	Message	Date
SolitaryThinker	f27573c59e	[test]: cover Z-Image registry resolution	2026-07-13 17:47:18 -07:00
SolitaryThinker	fac19aa6fa	[feat]: add native Z-Image transformer prototype	2026-07-13 17:46:16 -07:00
SolitaryThinker	5421a77e71	[docs]: support image SSIM reference seeding	2026-07-13 17:44:50 -07:00
SolitaryThinker	91389697bf	[test]: add Z-Image PNG SSIM coverage	2026-07-13 17:35:21 -07:00
Raghav	70bd4bb756	[bugfix]: don't let a stale HF-passthrough marker override a moved encoder (#1339 ) _fastvideo_input_device is stamped at load time with the loader's target device, which is "cpu" when text_encoder_cpu_offload is on. encode_text then moves such an encoder to the compute device for the forward, after which the marker is stale. Honouring it fed cpu token ids to cuda weights: RuntimeError: Expected all tensors to be on the same device, but got index is on cpu, different from other tensors on cuda:0 Seen on FLUX.2 Klein, whose Qwen3 text encoder takes the HF-passthrough path. Once the encoder has been moved, its param device is authoritative; the marker only speaks when nothing moved and the module has no parameters of its own.	2026-07-13 17:26:46 -07:00
RaghavandClaude Opus 4.8	35e55aa7f6	[bugfix]: keep DecodingStage scaling_factor module-only (#1339 ) _denormalize_latents previously ran its scaling branch only when the VAE module exposed scaling_factor -- in practice only gamecraftvae does. Falling back to vae.config.scaling_factor turned that branch on for every VAE whose config declares one (flux2, hunyuanvideo, hunyuanvideo15, ltx2), newly dividing latents that had always passed through untouched. Surfaced as FLUX.2 4B/9B SSIM failures. Keep the latents_mean/latents_std guard, which only tests the values instead of attribute presence so a config declaring them as None falls through rather than raising in torch.tensor(None). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>	2026-07-13 17:26:46 -07:00
SolitaryThinker	67c023fca9	[fix]: address Z-Image review findings (#1339 )	2026-07-13 17:26:46 -07:00
Raghav	bf839dde18	[fix]: Z-Image reuses the shared Qwen3 encoder; fix its batched RoPE positions main gained a config-driven Qwen3 text encoder (Qwen3ForCausalLM + Qwen3TextConfig) via the Flux2 Klein port (#1349). Reuse it for Z-Image instead of the parallel ~540-line bespoke Qwen3Model this PR originally added: drop the duplicate encoder + config, and map Z-Image-Turbo's "Qwen3Model" architecture string to the shared encoder in the registry. update_model_arch populates Z-Image's dims from config.json. Validating the reuse on L40S surfaced a real, batch-only divergence in the shared encoder: it built position_ids as [1, seq_len], but the rotary layer flattens positions to num_tokens and reshapes q/k to (num_tokens, -1, head_dim). For batch>1 that folded the batch dim into the head dim and misaligned RoPE (fp32 mean_diff 0.22 at batch=2; batch=1 was fine). Fix: expand position_ids to [batch_size, seq_len]. batch=1 is byte-identical, so Flux2 Klein is unaffected; this also fixes a latent batch bug in the shared encoder. Encoder parity now PASSES on L40S (Z-Image-Turbo): fp32 bit-exact against the shared encoder (both batch elements, max=0.0000); bf16 within the existing thresholds (last_hidden mean ~0.016, pre-norm mean ~0.07-0.08). Also in this PR (unchanged): the gated scheduler option use_reference_discrete_timesteps (default False) for Z-Image timestep parity, the Z-Image component parity tests (encoder/scheduler/tokenizer/VAE), and the encoder parity test's strict-load allowlist + OOM-safe two-model handling (free HF ref + gc.collect before loading FastVideo).	2026-07-13 17:26:46 -07:00
Raghav	b1d954c894	[docs]: pin GPU SKU + final empirical numbers in PORT_STATUS Component parity validated on Tongyi-MAI/Z-Image-Turbo on NVIDIA A40 (46068 MiB, driver 565.57.01) as of 2026-05-12. All thresholds pass with 1.5-1.6x headroom over observed: last_hidden_state bf16 mean (worst): 0.0168 < 0.025 last_hidden_state bf16 median (worst): 0.0127 < 0.020 hidden_states[-2] bf16 mean (worst): 0.0739 < 0.120 hidden_states[-2] bf16 median (worst): 0.0625 < 0.100 Handoff notes updated to reflect that component parity is done; next port-stack steps (DiT, pipeline, conversion, SSIM) are scoped to a future PR, not #1339.	2026-07-13 17:26:46 -07:00
Raghav	e4eef6ed20	[test]: calibrate bf16 encoder parity thresholds against observed Z-Image-Turbo Per-layer bf16 diagnostic test on Z-Image-Turbo (Qwen3 35 layers, not the 24 originally assumed) confirmed the divergence between FastVideo's Qwen3 port and HF Qwen3 is a textbook bf16-tail signature, not a code bug: - fp32 forward is bit-exact (last_hidden_state max=0.0000) - bf16 per-layer mean drift grows smoothly from 5e-4 to 7.5e-2 over 35 layers; no single-layer spike (consecutive layer ratio ≤ 1.5x) - embedding (idx 0) is bit-exact in bf16 too — weight load is right Root cause is well-understood: FastVideo's TP-aware fused linears (QKVParallel, MergedColumnParallel, SiluAndMul) execute matmul ops in a different order than HF's unfused equivalents. In fp32 the per-element differences cancel statistically (~zero mean drift); in bf16 the rounding-order accumulates into a long max tail. Per add-model-02-parity calibration block: replace element-wise assert_close (meaningless when the max tail is 4.0) with distribution checks (mean + median). Thresholds calibrated to empirical observation: last_hidden_state (post-norm): mean < 0.025, median < 0.020 hidden_states[-2] (pre-norm): mean < 0.120, median < 0.100 Headroom is ~1.6x over observed (0.0152, 0.0117, 0.0754, 0.0625). fp32 path keeps strict assert_close at atol=1e-4. Also revert `torch_dtype=dtype` -> `dtype=dtype`: transformers 4.57.3 emits a deprecation warning recommending `dtype=`. The earlier switch (Copilot review) was wrong-direction for the pinned transformers version. The two config-key strips that mention `torch_dtype` stay — they're stripping the legacy serialized config key, not the constructor kwarg. PORT_STATUS.md updated: pinned Z-Image SHA, pinned HF id (Z-Image-Turbo), empirical bf16 signature recorded in the decisions table, 35-layer correction reflected in the parity table.	2026-07-13 17:26:46 -07:00
Raghav	7390f528c6	[test]: add per-layer bf16 diagnostic for Qwen3 encoder Diagnostic-only test that prints abs-diff (max/mean/median/p99) for each of the 25 hidden states (embedding + 24 layer outputs + post-norm) in bf16. Lets a reviewer distinguish: - Healthy bf16 tail: monotonic growth per layer, median ≪ atol - Real op mismatch: single-layer spike where one layer's diff is >>2x the previous Always passes; output is captured via `-s` and inspected manually. Run after the parametrized parity test (which currently fails on the strict mean-drift threshold) to decide whether to relax thresholds or investigate a specific layer.	2026-07-13 17:26:46 -07:00
Raghav	92168e5020	[fix]: address Gemini + Copilot review on PR #1339 scheduler: - Restore numpy default (float64) on the non-reference timestep linspace path. The earlier `dtype=np.float32` silently changed rounded timestep values for every model already using this scheduler. Float32 is now only applied on the new `use_reference_discrete_timesteps=True` branch. - Document `use_reference_discrete_timesteps` in the class docstring with default + when-to-enable + backward-compat notes. qwen3 encoder: - `torch.arange(...)` for position_ids now passes `dtype=torch.long` explicitly so embedding/RoPE indices are unambiguous. - Drop the redundant `isinstance(x.device.type, str)` check; `device.type` is always a str. Keep the mps -> cpu fallback for torch.autocast. encoder parity test: - `AutoModel.from_pretrained` uses `torch_dtype=` to match repo convention (the rest of the repo and existing transformers callers all use this keyword). transformers 4.57.3 accepts both, but `torch_dtype` keeps the test green across the older transformers versions that some downstream forks pin.	2026-07-13 17:26:46 -07:00
Raghav	d4b8f08674	[docs]: add tests/local_tests/zimage/README.md + PORT_STATUS.md Required by `.agents/skills/add-model/contracts/port_state.md` and `add-model-10-pr-review` "Quality and evidence lane" before any handoff. README captures: reference assets, weight-dir layout, run commands, per-component status, and live blockers (lm_head allowlist, scheduler config-flag pin requirement, tokenizer `text_len` reconciliation). PORT_STATUS captures: component matrix, parity commands + last result, open questions (pin Z-Image clone SHA, pin HF id), issues I001-I005 (lm_head allowlist resolved, scheduler-config pin still open, transformer + pipeline NOT_STARTED), and the bf16 / scheduler-config-widening decisions made in this PR.	2026-07-13 17:26:46 -07:00
Raghav	13daf39e91	[test]: forward full scheduler_config.json into parity, not 3 keys The fixture only forwarded `num_train_timesteps`, `shift`, and `use_dynamic_shifting` from `scheduler_config.json`. Any future field in the on-disk config (e.g. `time_shift_type`, `invert_sigmas`, etc.) would be silently dropped on both sides — parity would pass while the produced schedulers diverged from the file Z-Image actually ships. Spread the full dict (minus Diffusers loader keys) into the constructor on both sides. The dynamic-shifting test follows the same shape.	2026-07-13 17:26:46 -07:00
Raghav	2d0e552a7e	[test]: add tests/local_tests/zimage/__init__.py All sibling family dirs (`encoders/`, `pipelines/`, `transformers/`, `vaes/`, `sd35/`) ship an `__init__.py`. Z-Image was the only port-in- progress dir missing one; conftest setups that walk packages will skip the dir without it.	2026-07-13 17:26:46 -07:00
Raghav	46d1f91c7e	[test]: Qwen3 encoder parity in bf16 with calibrated tolerance + diagnostics The original test forced fp32 because the strict `atol=1e-4` failed in bf16 across 24 transformer blocks. Per add-model-02-parity's calibration block, deep encoders in bf16 accumulate per-GEMM epsilon into a long tail — the right fix is to keep both dtypes covered, log per-batch max/mean/median/p99 diagnostics, and require median≈0 + abs-mean drift below threshold so a real bug (mean ≫ atol) cannot pass. Also: assert that the unexpected-key surface from the production loader equals `Qwen3Model.ALLOWED_UNEXPECTED_KEYS` so future `Qwen3ForCausalLM` checkpoint changes are caught. Fix stale docstring filename. Fix typo `torch.bloat16` (removed entirely with the dtype rewrite).	2026-07-13 17:26:46 -07:00
Mrinaal Dogra	9307aeb590	[feat]: Add Z-Image Qwen3 Encoder and parity test	2026-07-13 17:26:46 -07:00
Mrinaal Dogra	926f0c9ab4	[test]: add Z-Image VAE decode parity test	2026-07-13 17:26:46 -07:00
Mrinaal Dogra	b2ba79a4cd	[test]: add Z-Image scheduler parity and reference timestep option	2026-07-13 17:26:46 -07:00