[ci]: add exact-identity perf comparator statuses

Comparator verdicts per issue #1532: PASS / REGRESSION /
CALIBRATION_NEEDED / RECIPE_MISMATCH / INFRA_ERROR (QUALITY_BLOCKED
reserved for the promoted-baseline workflow). REGRESSION and
RECIPE_MISMATCH fail CI; CALIBRATION_NEEDED is loud but no longer
auto-seeds a passing baseline.

.buildkite/pipeline.yml

@@ -436,7 +436,7 @@ steps:
                 - "pyproject.toml"
                 - "docker/Dockerfile"
               config:
-                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                command: "timeout 25m .buildkite/scripts/pr_test.sh"
                 label: ":test_tube: LoRA Training Tests"
                 env:
                   - TEST_TYPE=training_lora

.gitignore

@@ -72,8 +72,7 @@ docs/distillation/examples/
 # Python pickle files
 *.pkl
 
-# Reference videos
-!fastvideo/tests/ssim/reference_videos/**/*.mp4
+# Reference videos (negations must come after the catch-all on line below)
 
 # Static images
 !docs/assets/images/**/*.png
@@ -127,6 +126,8 @@ apps/dreamverse/web/.env.production.local
 .sisyphus/
 openspec/
 fastvideo/tests/ssim/reference_videos/**
+!fastvideo/tests/ssim/reference_videos/**/*.mp4
+!fastvideo/tests/ssim/reference_videos/**/*.png
 
 # Editor logs and local Python version pins (accidentally committed)
 *.nvimlog

docs/contributing/performance_benchmarks.md

@@ -232,13 +232,25 @@ by the harness and is not config-declarable.)
 
 Recipe fingerprinting, hardware/software profile IDs, exact-identity
 comparison, and dashboard cohort grouping land with this change: v2 records
-compare only within their identity cohort, and a record that opens a NEW
-cohort is marked `baseline_status: "initialized_new_cohort"` (regression
-gating starts once that cohort accumulates history). Legacy v1 configs still
-run and are normalized for reporting, but their records skip rolling-baseline
-comparison entirely (`baseline_status: "skipped_missing_identity"`, never
-baseline eligible); only static thresholds gate them. Metric-specific
-threshold policies and promoted baselines remain separate follow-ups.
+compare only within their identity cohort and are never compared against v1
+records (or vice versa). Legacy records missing any comparison identity field
+still run and are normalized for reporting, but they skip rolling-baseline
+comparison entirely (`baseline_status: "skipped_missing_identity"`, comparator
+verdict `PASS`, never baseline eligible); only static thresholds gate them.
+Each comparison reports an explicit `comparator_status` on the normalized
+record and in the Markdown summary:
+
+| Status | Meaning | CI |
+|---|---|---|
+| `PASS` | Comparable baseline found with no gated regression, or comparison was skipped for a record without the v2 identity block. | passes |
+| `REGRESSION` | A gated metric regressed past both its percent and absolute floors. | fails |
+| `CALIBRATION_NEEDED` | No comparable baseline exists. Gating is inactive and the record does **not** seed a baseline — seed new cohorts explicitly via the reseed workflow. The record also carries `baseline_status: "initialized_new_cohort"`. | passes |
+| `RECIPE_MISMATCH` | Baseline history exists for the same variant/hardware/software cohort but under a different `recipe_fingerprint`. Represent recipe changes as a new `variant_id`, or reseed. | fails |
+| `INFRA_ERROR` | The comparison itself failed (for example, baseline history could not be loaded). | fails |
+| `QUALITY_BLOCKED` | Reserved for the promoted-baseline workflow; never emitted by the comparator. | n/a |
+
+Metric-specific threshold policies and promoted baselines remain separate
+follow-ups.
 
 ### Raw record (`results/perf_*.json`)
 
@@ -367,6 +379,8 @@ result, used as the rolling-baseline source of truth.
   "build_id": "<buildkite-build-id>",
   "job_id": "<buildkite-job-id>",
   "quality_metadata": { "quality_status": "canonical" },
+  "baseline_status": "compared",
+  "comparator_status": "PASS",
   "success": true
 }
 ```
@@ -487,10 +501,10 @@ When the rolling-baseline phase runs, it emits:
 2. The pytest test auto-discovers all configs — no test code needed. CI
    picks it up on the next `/test performance` run.
 
-3. The first persisted main-branch run with no HF history initializes the
-   baseline (passes automatically). Subsequent runs compare against it. Local
-   and pull-request runs with no HF history also pass, but they do not seed the
-   shared baseline.
+3. The first run with no HF history reports `CALIBRATION_NEEDED`: it passes,
+   but it does not seed the shared baseline. Seed the cohort explicitly with
+   the `reseed-performance-baseline` skill; subsequent scheduled-main runs
+   then compare against it and keep the rolling baseline advancing.
 
 4. If the benchmark targets a GPU not currently in `thresholds`, either add
    that GPU as a key or rely on the `default` block. Note that `default` is
@@ -510,8 +524,13 @@ When the rolling-baseline phase runs, it emits:
 
 ## Troubleshooting
 
-**"No baseline for ... Initializing"** — first run for this comparison cohort.
-Run will pass and (if persisting) seed the first record.
+**`CALIBRATION_NEEDED: no comparable baseline`** — first run for this cohort.
+The run passes but does not seed a baseline; seed it explicitly with the
+`reseed-performance-baseline` skill.
+
+**`RECIPE_MISMATCH`** — the benchmark recipe changed without a new
+`variant_id`. Either bump `variant_id` to open a new cohort, or reseed the
+baseline if the existing variant should adopt the new recipe.
 
 **Persistent failure right after a torch / kernel / image upgrade** —
 genuine regression *or* baseline drift. Compare the failing normalized record

docs/design/inference_schema_parity_inventory.yaml

@@ -458,6 +458,8 @@ surfaces:
       guidance_scale: request.sampling.guidance_scale
       guidance_scale_2: request.sampling.guidance_scale_2
       guidance_rescale: request.sampling.guidance_rescale
+      use_embedded_guidance: request.sampling.use_embedded_guidance
+      true_cfg_scale: request.sampling.true_cfg_scale
       boundary_ratio: request.sampling.boundary_ratio
       sigmas: request.sampling.sigmas
       enable_teacache: request.runtime.enable_teacache

examples/inference/basic/basic_flux_dev.py

@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import argparse
+import contextlib
+import os
+import re
+
+DEFAULT_PROMPTS = [
+    "a photo of a cat",
+    (
+        "a cinematic photo of a red panda wearing a tiny backpack, standing on a "
+        "rainy neon-lit street at night, shallow depth of field, sharp focus, "
+        "35mm, bokeh"
+    ),
+]
+
+
+def _safe_filename(text: str, max_len: int = 100) -> str:
+    """Make a stable, filesystem-friendly filename base."""
+    s = text[:max_len].strip()
+    s = s.replace(os.sep, "_")
+    if os.altsep:
+        s = s.replace(os.altsep, "_")
+    s = re.sub(r"\s+", " ", s)
+    s = re.sub(r"[^A-Za-z0-9 .,_-]", "_", s)
+    s = s.strip(" .")
+    return s or "prompt"
+
+
+def _remove_existing_outputs(out_dir: str, filename_base: str) -> None:
+    """Delete prior outputs so reruns do not get _1, _2 suffixes."""
+    if not os.path.isdir(out_dir):
+        return
+
+    pattern = re.compile(rf"^{re.escape(filename_base)}(_\d+)?\.(mp4|png)$")
+    for fn in os.listdir(out_dir):
+        if pattern.match(fn):
+            with contextlib.suppress(FileNotFoundError):
+                os.remove(os.path.join(out_dir, fn))
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(
+        description="Run FLUX.1-dev text-to-image with FastVideo VideoGenerator.",
+    )
+    p.add_argument(
+        "--model-path",
+        default="official_weights/FLUX.1-dev",
+        help="Local Diffusers checkpoint dir or HF repo id.",
+    )
+    p.add_argument(
+        "--out-dir",
+        "--outdir",
+        default="outputs/flux_dev/samples",
+        help="Directory for saved PNG outputs.",
+    )
+    p.add_argument(
+        "--prompt",
+        action="append",
+        default=None,
+        help="Prompt. Repeat for multiple images.",
+    )
+    p.add_argument(
+        "--backend",
+        default=None,
+        help="Set FASTVIDEO_ATTENTION_BACKEND (e.g. TORCH_SDPA).",
+    )
+    p.add_argument("--seed", type=int, default=42, help="Base seed; each prompt uses seed + index.")
+    p.add_argument("--height", type=int, default=1024, help="Output height.")
+    p.add_argument("--width", type=int, default=1024, help="Output width.")
+    p.add_argument("--steps", type=int, default=28, help="Number of inference steps.")
+    p.add_argument("--guidance", type=float, default=3.5, help="Guidance scale.")
+    p.add_argument("--num-gpus", type=int, default=1, help="GPU count.")
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    prompts: list[str] = args.prompt if args.prompt else DEFAULT_PROMPTS
+
+    if args.backend:
+        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = args.backend
+
+    from fastvideo import VideoGenerator
+
+    os.makedirs(args.out_dir, exist_ok=True)
+
+    init_kwargs = {
+        "num_gpus": args.num_gpus,
+        "workload_type": "t2i",
+        "sp_size": 1,
+        "tp_size": 1,
+        "dit_cpu_offload": False,
+        "dit_layerwise_offload": False,
+        "text_encoder_cpu_offload": False,
+        "vae_cpu_offload": False,
+        "image_encoder_cpu_offload": False,
+        "pin_cpu_memory": False,
+        "use_fsdp_inference": False,
+    }
+
+    generator = VideoGenerator.from_pretrained(
+        model_path=args.model_path,
+        **init_kwargs,
+    )
+    try:
+        for i, prompt in enumerate(prompts):
+            seed = args.seed + i
+            filename_base = (
+                f"flux_dev_{i:02d}_seed{seed}_{_safe_filename(prompt, max_len=80)}"
+            )
+            _remove_existing_outputs(args.out_dir, filename_base)
+            output_path = os.path.join(args.out_dir, f"{filename_base}.png")
+            print(f"[flux] prompt_idx={i} seed={seed} output_path={output_path}")
+
+            generation_kwargs = {
+                "output_path": output_path,
+                "height": args.height,
+                "width": args.width,
+                "num_frames": 1,
+                "fps": 1,
+                "num_inference_steps": args.steps,
+                "guidance_scale": args.guidance,
+                "use_embedded_guidance": True,
+                "true_cfg_scale": 1.0,
+                "seed": seed,
+                "save_video": True,
+            }
+
+            generator.generate_video(prompt, **generation_kwargs)
+
+        print(f"[flux] done. outputs written to: {args.out_dir}")
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_glm_image.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Run GLM-Image text-to-image generation through FastVideo.
+
+User story:
+    "I have the HF `zai-org/GLM-Image` checkpoint and want a minimal
+    text-to-image generation command, saved as a PNG."
+"""
+import argparse
+from pathlib import Path
+
+from PIL import Image
+
+from fastvideo import VideoGenerator
+from fastvideo.api import (
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    OutputConfig,
+    ParallelismConfig,
+    PipelineSelection,
+    SamplingConfig,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run GLM-Image text-to-image generation.")
+    parser.add_argument(
+        "--model-path",
+        default="zai-org/GLM-Image",
+        help="HF id or local diffusers-format GLM-Image weights directory.",
+    )
+    parser.add_argument(
+        "--output",
+        default="image_output/landscape.png",
+        help="Output PNG path.",
+    )
+    parser.add_argument(
+        "--prompt",
+        default=("A beautiful landscape photography with rolling hills, "
+                 "a winding river, and a vibrant sunset in the background. "
+                 "Warm golden light, photorealistic style."),
+        help="Text prompt.",
+    )
+    parser.add_argument("--height", type=int, default=1024)
+    parser.add_argument("--width", type=int, default=1024)
+    parser.add_argument("--steps", type=int, default=50)
+    parser.add_argument("--guidance-scale", type=float, default=1.5)
+    parser.add_argument("--seed", type=int, default=1024)
+    parser.add_argument("--num-gpus", type=int, default=1)
+    parser.add_argument("--tp-size", type=int, default=None)
+    parser.add_argument("--sp-size", type=int, default=None)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    tp_size = args.tp_size if args.tp_size is not None else (args.num_gpus if args.num_gpus > 1 else 1)
+    sp_size = args.sp_size if args.sp_size is not None else (1 if args.num_gpus > 1 else args.num_gpus)
+
+    # GLM-Image needs trust_remote_code for its AR encoder; offload and the
+    # pipeline class come from the model's registered defaults — don't override.
+    generator_config = GeneratorConfig(
+        model_path=args.model_path,
+        trust_remote_code=True,
+        engine=EngineConfig(
+            num_gpus=args.num_gpus,
+            parallelism=ParallelismConfig(tp_size=tp_size, sp_size=sp_size),
+        ),
+        pipeline=PipelineSelection(workload_type="t2i"),
+    )
+
+    generator = VideoGenerator.from_config(generator_config)
+    try:
+        request = GenerationRequest(
+            prompt=args.prompt,
+            sampling=SamplingConfig(
+                height=args.height,
+                width=args.width,
+                num_frames=1,
+                fps=1,
+                num_inference_steps=args.steps,
+                guidance_scale=args.guidance_scale,
+                seed=args.seed,
+            ),
+            output=OutputConfig(
+                output_path=str(output.parent),
+                save_video=False,
+                return_frames=True,
+            ),
+        )
+        result = generator.generate(request)
+        if isinstance(result, list):
+            result = result[0]
+
+        frames = result.frames
+        if frames is not None and len(frames):
+            Image.fromarray(frames[0]).save(output)
+            print(f"Saved image to {output}")
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/edit_glm_image.py

@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Run GLM-Image image-to-image (edit) generation through FastVideo.
+
+User story:
+    "I have the HF `zai-org/GLM-Image` checkpoint and a condition image, and
+    want a minimal edit command (text + image -> edited image), saved as a PNG."
+
+GLM-Image is a single unified pipeline: passing a condition image switches it
+from text-to-image to the edit path (the condition enters the DiT via a KV-cache
+write pass), so the generator config is identical to `basic_glm_image.py` — the
+`inputs.pil_image` on the request is what selects the edit mode.
+"""
+import argparse
+from pathlib import Path
+
+from PIL import Image
+
+from fastvideo import VideoGenerator
+from fastvideo.api import (
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    InputConfig,
+    OutputConfig,
+    ParallelismConfig,
+    PipelineSelection,
+    SamplingConfig,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run GLM-Image image-to-image (edit) generation.")
+    parser.add_argument(
+        "--model-path",
+        default="zai-org/GLM-Image",
+        help="HF id or local diffusers-format GLM-Image weights directory.",
+    )
+    parser.add_argument(
+        "--image",
+        default="assets/images/couple.jpg",
+        help="Condition image to edit.",
+    )
+    parser.add_argument(
+        "--output",
+        default="image_output/edited.png",
+        help="Output PNG path.",
+    )
+    parser.add_argument(
+        "--prompt",
+        default="Change the background to a snowy mountain landscape at golden hour.",
+        help="Edit instruction.",
+    )
+    parser.add_argument("--height", type=int, default=1024)
+    parser.add_argument("--width", type=int, default=1024)
+    parser.add_argument("--steps", type=int, default=50)
+    parser.add_argument("--guidance-scale", type=float, default=1.5)
+    parser.add_argument("--seed", type=int, default=1024)
+    parser.add_argument("--num-gpus", type=int, default=1)
+    parser.add_argument("--tp-size", type=int, default=None)
+    parser.add_argument("--sp-size", type=int, default=None)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    condition = Image.open(args.image).convert("RGB")
+    tp_size = args.tp_size if args.tp_size is not None else (args.num_gpus if args.num_gpus > 1 else 1)
+    sp_size = args.sp_size if args.sp_size is not None else (1 if args.num_gpus > 1 else args.num_gpus)
+
+    # GLM-Image needs trust_remote_code for its AR encoder; offload and the
+    # pipeline class come from the model's registered defaults — don't override.
+    # The pipeline is registered as t2i; passing inputs.pil_image below switches
+    # it to the edit path.
+    generator_config = GeneratorConfig(
+        model_path=args.model_path,
+        trust_remote_code=True,
+        engine=EngineConfig(
+            num_gpus=args.num_gpus,
+            parallelism=ParallelismConfig(tp_size=tp_size, sp_size=sp_size),
+        ),
+        pipeline=PipelineSelection(workload_type="t2i"),
+    )
+
+    generator = VideoGenerator.from_config(generator_config)
+    try:
+        request = GenerationRequest(
+            prompt=args.prompt,
+            inputs=InputConfig(pil_image=condition),
+            sampling=SamplingConfig(
+                height=args.height,
+                width=args.width,
+                num_frames=1,
+                fps=1,
+                num_inference_steps=args.steps,
+                guidance_scale=args.guidance_scale,
+                seed=args.seed,
+            ),
+            output=OutputConfig(
+                output_path=str(output.parent),
+                save_video=False,
+                return_frames=True,
+            ),
+        )
+        result = generator.generate(request)
+        if isinstance(result, list):
+            result = result[0]
+
+        frames = result.frames
+        if frames is not None and len(frames):
+            Image.fromarray(frames[0]).save(output)
+            print(f"Saved image to {output}")
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

fastvideo/api/flux.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from fastvideo.api.sampling_param import SamplingParam
+
+
+@dataclass
+class FluxSamplingParam(SamplingParam):
+
+    prompt: str | None = "a photo of a cat"
+    negative_prompt: str = ""
+
+    num_videos_per_prompt: int = 1
+    seed: int = 0
+
+    num_frames: int = 1
+    height: int = 1024
+    width: int = 1024
+    fps: int = 1
+
+    num_inference_steps: int = 28
+    guidance_scale: float = 3.5
+    use_embedded_guidance: bool = True
+    true_cfg_scale: float = 1.0

fastvideo/api/sampling_param.py

@@ -90,6 +90,10 @@ class SamplingParam:
     num_inference_steps_sr: int = 50
     guidance_scale: float = 1.0
     guidance_scale_2: float | None = None
+    # Embedded guidance (FLUX): do not treat ``guidance_scale > 1`` as classic CFG.
+    use_embedded_guidance: bool = False
+    # Diffusers-style true CFG for FLUX when > 1 (requires negative prompt encoding).
+    true_cfg_scale: float = 1.0
     guidance_rescale: float = 0.0
     boundary_ratio: float | None = None
     sigmas: list[float] | None = None
@@ -325,6 +329,18 @@ class SamplingParam:
             default=SamplingParam.guidance_rescale,
             help="Guidance rescale factor",
         )
+        parser.add_argument(
+            "--use-embedded-guidance",
+            action="store_true",
+            default=SamplingParam.use_embedded_guidance,
+            help="Use embedded guidance scale (FLUX-style) instead of classic CFG",
+        )
+        parser.add_argument(
+            "--true-cfg-scale",
+            type=float,
+            default=SamplingParam.true_cfg_scale,
+            help="True CFG scale for FLUX when > 1 (requires negative prompt encoding)",
+        )
         parser.add_argument(
             "--boundary-ratio",
             type=float,

fastvideo/api/schema.py

@@ -150,6 +150,7 @@ class SamplingConfig:
     guidance_scale_2: float | None = None
     guidance_rescale: float = 0.0
     true_cfg_scale: float | None = None
+    use_embedded_guidance: bool | None = None
     boundary_ratio: float | None = None
     sigmas: list[float] | None = None
 

fastvideo/attention/backends/sdpa.py

@@ -15,8 +15,14 @@ class SDPABackend(AttentionBackend):
     accept_output_buffer: bool = True
 
     @staticmethod
-    def get_supported_head_sizes() -> list[int]:
-        return [32, 64, 96, 128, 160, 192, 224, 256]
+    def get_supported_head_sizes() -> list[int] | None:
+        # torch.nn.functional.scaled_dot_product_attention is head-size
+        # agnostic: the math backend handles any size and the fused kernels
+        # fall back internally when they cannot. The previous list
+        # ([32, 64, ..., 256]) was copied from FlashAttentionBackend in the
+        # v1 refactor (#270) and wrongly excluded sizes such as 80 (e.g.
+        # CLIP vision encoders). None means "no restriction".
+        return None
 
     @staticmethod
     def get_name() -> str:

fastvideo/configs/models/dits/__init__.py

@@ -1,7 +1,9 @@
 from fastvideo.configs.models.dits.cosmos import CosmosVideoConfig
 from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig
 from fastvideo.configs.models.dits.dreamx_world import DreamXWorldARConfig, DreamXWorldConfig
+from fastvideo.configs.models.dits.flux import FluxDiTConfig
 from fastvideo.configs.models.dits.flux_2 import Flux2Config
+from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig
 from fastvideo.configs.models.dits.hunyuangamecraft import HunyuanGameCraftConfig
 from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
 from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config
@@ -15,6 +17,7 @@ from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig
 
 __all__ = [
     "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "DreamXWorldConfig",
-    "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig",
-    "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig", "StableAudioConfig", "Flux2Config"
+    "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "FluxDiTConfig", "Flux2Config",
+    "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig",
+    "StableAudioConfig", "GlmImageDiTConfig"
 ]

fastvideo/configs/models/dits/flux.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+@dataclass
+class FluxTransformer2DArchConfig(DiTArchConfig):
+
+    patch_size: int = 1
+    in_channels: int = 64
+    out_channels: int | None = None
+    num_layers: int = 19
+    num_single_layers: int = 38
+    attention_head_dim: int = 128
+    num_attention_heads: int = 24
+    joint_attention_dim: int = 4096
+    pooled_projection_dim: int = 768
+    guidance_embeds: bool = True
+    axes_dims_rope: tuple[int, int, int] = (16, 56, 56)
+
+
+@dataclass
+class FluxDiTConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=FluxTransformer2DArchConfig)
+    prefix: str = "flux"

fastvideo/configs/models/dits/glm_image.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+def is_blocks(n: str, m) -> bool:
+    return "transformer_blocks" in n and str.isdigit(n.split(".")[-1])
+
+
+@dataclass
+class GlmImageDiTArchConfig(DiTArchConfig):
+
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
+
+    hidden_size: int = 4096
+    num_attention_heads: int = 32
+    attention_head_dim: int = 128
+    in_channels: int = 16
+    out_channels: int = 16
+    num_layers: int = 30
+
+    text_embed_dim: int = 1472
+    time_embed_dim: int = 512
+    condition_dim: int = 256
+
+    prior_vq_quantizer_codebook_size: int = 16384
+
+    patch_size: int = 2
+
+    max_height: int = 2048
+    max_width: int = 2048
+
+    qk_norm: str = "layer_norm"
+    eps: float = 1e-5
+
+    exclude_lora_layers: list[str] = field(
+        default_factory=lambda: ["image_projector", "glyph_projector", "prior_token_embedding"])
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^glyph_projector\.net\.0\.proj\.(.*)$": r"glyph_projector.fc_in.\1",
+            r"^glyph_projector\.net\.2\.(.*)$": r"glyph_projector.fc_out.\1",
+            r"^prior_projector\.net\.0\.proj\.(.*)$": r"prior_projector.fc_in.\1",
+            r"^prior_projector\.net\.2\.(.*)$": r"prior_projector.fc_out.\1",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.0\.proj\.(.*)$": r"transformer_blocks.\1.ff.fc_in.\2",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.2\.(.*)$": r"transformer_blocks.\1.ff.fc_out.\2",
+        })
+
+    reverse_param_names_mapping: dict = field(default_factory=dict)
+    lora_param_names_mapping: dict = field(default_factory=dict)
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.num_channels_latents = self.out_channels
+
+
+@dataclass
+class GlmImageDiTConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=GlmImageDiTArchConfig)
+    prefix: str = "GlmImage"

fastvideo/configs/models/vaes/__init__.py

@@ -2,6 +2,7 @@ from fastvideo.configs.models.vaes.cosmosvae import CosmosVAEConfig
 from fastvideo.configs.models.vaes.cosmos2_5vae import Cosmos25VAEConfig
 from fastvideo.configs.models.vaes.gamecraftvae import GameCraftVAEConfig
 from fastvideo.configs.models.vaes.gen3cvae import Gen3CVAEConfig
+from fastvideo.configs.models.vaes.glm_image import GlmImageVAEConfig
 from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
 from fastvideo.configs.models.vaes.hunyuan15vae import Hunyuan15VAEConfig
 from fastvideo.configs.models.vaes.ltx2vae import LTX2VAEConfig
@@ -21,4 +22,5 @@ __all__ = [
     "OobleckVAEArchConfig",
     "OobleckVAEConfig",
     "Flux2VAEConfig",
+    "GlmImageVAEConfig",
 ]

fastvideo/configs/models/vaes/glm_image.py

@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.vaes.autoencoder_kl import (AutoencoderKLArchConfig, AutoencoderKLVAEConfig)
+
+_GLM_IMAGE_LATENTS_MEAN: tuple[float, ...] = (
+    -0.2080078125,
+    1.875,
+    -0.470703125,
+    -1.265625,
+    -1.421875,
+    0.77734375,
+    -0.3671875,
+    -0.9453125,
+    0.318359375,
+    0.7734375,
+    -0.1884765625,
+    -0.022216796875,
+    -0.220703125,
+    -1.59375,
+    -0.81640625,
+    -0.255859375,
+)
+_GLM_IMAGE_LATENTS_STD: tuple[float, ...] = (
+    3.0625,
+    2.203125,
+    2.265625,
+    4.84375,
+    2.5,
+    3.9375,
+    2.203125,
+    3.03125,
+    2.1875,
+    2.046875,
+    2.71875,
+    2.390625,
+    2.390625,
+    2.453125,
+    2.25,
+    2.15625,
+)
+
+
+@dataclass
+class GlmImageVAEArchConfig(AutoencoderKLArchConfig):
+    act_fn: str = "silu"
+    block_out_channels: tuple[int, ...] = (128, 512, 1024, 1024)
+    down_block_types: tuple[str, ...] = (
+        "DownEncoderBlock2D",
+        "DownEncoderBlock2D",
+        "DownEncoderBlock2D",
+        "DownEncoderBlock2D",
+    )
+    up_block_types: tuple[str, ...] = (
+        "UpDecoderBlock2D",
+        "UpDecoderBlock2D",
+        "UpDecoderBlock2D",
+        "UpDecoderBlock2D",
+    )
+    force_upcast: bool = True
+    in_channels: int = 3
+    out_channels: int = 3
+    latent_channels: int = 16
+    latents_mean: tuple[float, ...] = _GLM_IMAGE_LATENTS_MEAN
+    latents_std: tuple[float, ...] = _GLM_IMAGE_LATENTS_STD
+    layers_per_block: int = 3
+    mid_block_add_attention: bool = False
+    norm_num_groups: int = 32
+    sample_size: int = 1024
+    scaling_factor: float = 0.18215
+    shift_factor: float | None = None
+    use_quant_conv: bool = False
+    use_post_quant_conv: bool = False
+
+    temporal_compression_ratio: int = 1
+    spatial_compression_ratio: int = 8
+
+
+@dataclass
+class GlmImageVAEConfig(AutoencoderKLVAEConfig):
+    arch_config: GlmImageVAEArchConfig = field(default_factory=GlmImageVAEArchConfig)
+
+    use_tiling: bool = True
+    use_temporal_tiling: bool = False
+    use_parallel_tiling: bool = False
+
+    tile_sample_min_height: int = 512
+    tile_sample_min_width: int = 512
+    tile_sample_stride_height: int = 384
+    tile_sample_stride_width: int = 384
+
+    load_encoder: bool = True
+    load_decoder: bool = True

fastvideo/configs/pipelines/flux.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import torch
+
+from fastvideo.configs.models import EncoderConfig
+from fastvideo.configs.models.dits.flux import FluxDiTConfig
+from fastvideo.configs.models.encoders import (
+    BaseEncoderOutput,
+    CLIPTextConfig,
+    T5LargeConfig,
+)
+from fastvideo.configs.models.vaes.autoencoder_kl import AutoencoderKLVAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig, preprocess_text
+
+
+def _flux_clip_pooled_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
+    """CLIP branch for FLUX: Diffusers uses pooled prompt embeddings only."""
+    if outputs.pooler_output is None:
+        raise RuntimeError(
+            "FLUX CLIP conditioning requires pooler_output. Ensure the CLIP text encoder returns pooled features.")
+    return outputs.pooler_output
+
+
+def _flux_t5_sequence_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
+    if outputs.last_hidden_state is None:
+        raise RuntimeError("FLUX T5 conditioning requires last_hidden_state.")
+    return outputs.last_hidden_state
+
+
+@dataclass
+class FluxPipelineConfig(PipelineConfig):
+    """Pipeline layout for Diffusers FLUX.1-dev (CLIP + T5 + packed DiT + FlowMatch)."""
+
+    scheduler_arch: str = "FlowMatchEulerDiscreteScheduler"
+    transformer_arch: str = "FluxTransformer2DModel"
+    vae_arch: str = "AutoencoderKL"
+    text_encoder_archs: tuple[str, ...] = ("CLIPTextModel", "T5EncoderModel")
+    tokenizer_archs: tuple[str, ...] = ("CLIPTokenizer", "T5TokenizerFast")
+
+    dit_config: FluxDiTConfig = field(default_factory=FluxDiTConfig)
+    vae_config: AutoencoderKLVAEConfig = field(default_factory=AutoencoderKLVAEConfig)
+
+    embedded_cfg_scale: float = 3.5
+    flow_shift: float | None = None
+
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (CLIPTextConfig(), T5LargeConfig()))
+    preprocess_text_funcs: tuple[Callable[[str], str],
+                                 ...] = field(default_factory=lambda: (preprocess_text, preprocess_text))
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], ...] = field(
+        default_factory=lambda: (_flux_clip_pooled_postprocess, _flux_t5_sequence_postprocess))
+
+    dit_precision: str = "bf16"
+    vae_precision: str = "fp32"
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32", "bf16"))
+
+    def __post_init__(self) -> None:
+        te_cfgs = list(self.text_encoder_configs)
+        if len(te_cfgs) >= 1:
+            te_cfgs[0].tokenizer_kwargs.setdefault("padding", "max_length")
+            te_cfgs[0].tokenizer_kwargs.setdefault("max_length", 77)
+            te_cfgs[0].tokenizer_kwargs.setdefault("truncation", True)
+            te_cfgs[0].tokenizer_kwargs.setdefault("return_tensors", "pt")
+        if len(te_cfgs) >= 2:
+            cap = 512
+            te_cfgs[1].tokenizer_kwargs["max_length"] = min(int(te_cfgs[1].tokenizer_kwargs.get("max_length", cap)),
+                                                            cap)
+            te_cfgs[1].tokenizer_kwargs.setdefault("padding", "max_length")
+            te_cfgs[1].tokenizer_kwargs.setdefault("truncation", True)
+            te_cfgs[1].tokenizer_kwargs.setdefault("return_tensors", "pt")

fastvideo/configs/pipelines/glm_image.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import torch
+
+from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
+from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig
+from fastvideo.configs.models.encoders import BaseEncoderOutput, T5Config
+from fastvideo.configs.models.vaes.glm_image import GlmImageVAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig
+
+
+def glm_image_t5_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
+    mask: torch.Tensor = outputs.attention_mask
+    hidden_state: torch.Tensor = outputs.last_hidden_state
+    seq_lens = mask.gt(0).sum(dim=1).long()
+
+    assert torch.isnan(hidden_state).sum() == 0, "T5 hidden states contain NaN"
+
+    max_len = 512
+    prompt_embeds = [u[:min(v, max_len)] for u, v in zip(hidden_state, seq_lens, strict=True)]
+    prompt_embeds_tensor: torch.Tensor = torch.stack(
+        [torch.cat([u, u.new_zeros(max_len - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0)
+
+    return prompt_embeds_tensor
+
+
+@dataclass
+class GlmImageConfig(PipelineConfig):
+
+    dit_config: DiTConfig = field(default_factory=GlmImageDiTConfig)
+    dit_precision: str = "bf16"
+
+    vae_config: VAEConfig = field(default_factory=GlmImageVAEConfig)
+    vae_precision: str = "fp32"
+    vae_tiling: bool = True
+    vae_sp: bool = False
+
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (T5Config(), ))
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32", ))
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
+                                  ...] = field(default_factory=lambda: (glm_image_t5_postprocess, ))
+
+    flow_shift: float | None = 1.0
+    embedded_cfg_scale: float = 7.5

fastvideo/layers/rotary_embedding.py

@@ -295,6 +295,7 @@ def get_1d_rotary_pos_embed(
     interpolation_factor: float = 1.0,
     dtype: torch.dtype = torch.float32,
     use_real: bool = True,
+    freqs_dtype: torch.dtype | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Precompute the frequency tensor for complex exponential (cis) with given dimensions.
@@ -319,12 +320,16 @@ def get_1d_rotary_pos_embed(
     if isinstance(pos, int):
         pos = torch.arange(pos).float()
 
+    # freqs_dtype is an alias for dtype (Diffusers-compatible calling convention).
+    if freqs_dtype is not None:
+        dtype = freqs_dtype
+
     # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
     # has some connection to NTK literature
     if theta_rescale_factor != 1.0:
         theta *= theta_rescale_factor**(dim / (dim - 2))
 
-    freqs = 1.0 / (theta**(torch.arange(0, dim, 2)[:(dim // 2)].to(dtype) / dim))  # [D/2]
+    freqs = 1.0 / (theta**(torch.arange(0, dim, 2, device=pos.device)[:(dim // 2)].to(dtype) / dim))  # [D/2]
     freqs = torch.outer(pos * interpolation_factor, freqs)  # [S, D/2]
     freqs_cos = freqs.cos()  # [S, D/2]
     freqs_sin = freqs.sin()  # [S, D/2]
@@ -445,6 +450,21 @@ def get_nd_rotary_pos_embed(
     return cos, sin
 
 
+_ROTARY_POS_EMBED_CACHE: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {}
+# Bound the table cache so long-running servers / causal models (which vary
+# start_frame per frame) cannot grow it without limit; entries are large float64
+# tensors. Least-recently-used eviction keeps the active resolution(s) hot while
+# capping memory.
+_ROTARY_POS_EMBED_CACHE_MAXSIZE = 16
+
+
+def _hashable(value: Any) -> Any:
+    """Return a hashable view of a scalar or sequence for use in a cache key."""
+    if isinstance(value, list | tuple):
+        return tuple(value)
+    return value
+
+
 def get_rotary_pos_embed(
     rope_sizes,
     hidden_size,
@@ -495,6 +515,31 @@ def get_rotary_pos_embed(
         sp_rank = 0
         sp_world_size = 1
 
+    # Memoize on every output-affecting argument; the table is constant across
+    # denoising steps, so this avoids recomputing the float64 cos/sin tables.
+    cache_key = (
+        _hashable(rope_sizes),
+        tuple(rope_dim_list),
+        rope_theta,
+        _hashable(theta_rescale_factor),
+        _hashable(interpolation_factor),
+        shard_dim,
+        sp_rank,
+        sp_world_size,
+        dtype,
+        start_frame,
+        use_real,
+    )
+    cached = _ROTARY_POS_EMBED_CACHE.get(cache_key)
+    if cached is not None:
+        # Move to most-recently-used position so the active table is not evicted
+        # when several resolutions / buckets share the process (LRU recency).
+        # Pop with a default: a concurrent eviction between the get() above and
+        # here would otherwise raise KeyError on the hit path.
+        if _ROTARY_POS_EMBED_CACHE.pop(cache_key, None) is not None:
+            _ROTARY_POS_EMBED_CACHE[cache_key] = cached
+        return cached
+
     freqs_cos, freqs_sin = get_nd_rotary_pos_embed(
         rope_dim_list,
         rope_sizes,
@@ -508,6 +553,14 @@ def get_rotary_pos_embed(
         start_frame=start_frame,
         use_real=use_real,
     )
+    # The returned tensors are shared cache entries: callers must never mutate
+    # them in place. Note .to(device) is an identity alias when the tensor is
+    # already on the target device (e.g. CPU runs), so it does NOT guarantee a
+    # copy — treat the tables as read-only and copy before any in-place op.
+    # Reached only on a miss, so evict the least-recently-used entry at capacity.
+    if len(_ROTARY_POS_EMBED_CACHE) >= _ROTARY_POS_EMBED_CACHE_MAXSIZE:
+        _ROTARY_POS_EMBED_CACHE.pop(next(iter(_ROTARY_POS_EMBED_CACHE)))
+    _ROTARY_POS_EMBED_CACHE[cache_key] = (freqs_cos, freqs_sin)
     return freqs_cos, freqs_sin
 
 

fastvideo/models/dits/flux.py

@@ -0,0 +1,578 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from contextlib import nullcontext
+from dataclasses import dataclass
+import math
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from fastvideo.layers.rotary_embedding import apply_rotary_emb, get_1d_rotary_pos_embed
+
+from fastvideo.attention import DistributedAttention
+from fastvideo.configs.models import DiTConfig
+from fastvideo.forward_context import get_forward_context, set_forward_context
+from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.layers.visual_embedding import Timesteps
+from fastvideo.models.dits.base import BaseDiT
+from fastvideo.models.dits.sd3 import (
+    CombinedTimestepTextProjEmbeddings,
+    SD3AdaLayerNormContinuous,
+    SD3AdaLayerNormZero,
+    SD3FeedForward,
+    SD3TextProjection,
+    SD3TimestepEmbedding,
+)
+from fastvideo.platforms import AttentionBackendEnum
+
+
+@dataclass
+class FluxTransformer2DModelOutput:
+    sample: torch.Tensor
+
+
+class FluxPosEmbed(nn.Module):
+    """1D RoPE axes concatenated per Diffusers `FluxPosEmbed`."""
+
+    def __init__(self, theta: int, axes_dim: list[int]) -> None:
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        n_axes = ids.shape[-1]
+        cos_out: list[torch.Tensor] = []
+        sin_out: list[torch.Tensor] = []
+        pos = ids.float()
+        is_mps = ids.device.type == "mps"
+        is_npu = ids.device.type == "npu"
+        freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64
+        for i in range(n_axes):
+            cos, sin = get_1d_rotary_pos_embed(
+                self.axes_dim[i],
+                pos[:, i],
+                theta=self.theta,
+                use_real=True,
+                freqs_dtype=freqs_dtype,
+            )
+            cos_out.append(cos)
+            sin_out.append(sin)
+        freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device)
+        freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device)
+        return freqs_cos, freqs_sin
+
+
+class FluxCombinedTimestepGuidanceTextProjEmbeddings(nn.Module):
+    def __init__(self, embedding_dim: int, pooled_projection_dim: int) -> None:
+        super().__init__()
+        self.time_proj = Timesteps(
+            num_channels=256,
+            flip_sin_to_cos=True,
+            downscale_freq_shift=0,
+        )
+        self.timestep_embedder = SD3TimestepEmbedding(
+            in_channels=256,
+            time_embed_dim=embedding_dim,
+            act_fn="silu",
+        )
+        self.guidance_embedder = SD3TimestepEmbedding(
+            in_channels=256,
+            time_embed_dim=embedding_dim,
+            act_fn="silu",
+        )
+        self.text_embedder = SD3TextProjection(
+            pooled_projection_dim,
+            embedding_dim,
+            act_fn="silu",
+        )
+
+    def forward(
+        self,
+        timestep: torch.Tensor,
+        guidance: torch.Tensor,
+        pooled_projection: torch.Tensor,
+    ) -> torch.Tensor:
+        timesteps_proj = self.time_proj(timestep)
+        timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype))
+        guidance_proj = self.time_proj(guidance)
+        guidance_emb = self.guidance_embedder(guidance_proj.to(dtype=pooled_projection.dtype))
+        time_guidance_emb = timesteps_emb + guidance_emb
+        pooled_projections = self.text_embedder(pooled_projection)
+        return time_guidance_emb + pooled_projections
+
+
+class FluxAdaLayerNormZeroSingle(nn.Module):
+    def __init__(self, embedding_dim: int, bias: bool = True) -> None:
+        super().__init__()
+        self.silu = nn.SiLU()
+        self.linear = ReplicatedLinear(embedding_dim, 3 * embedding_dim, bias=bias)
+        self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        emb: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        emb, _ = self.linear(self.silu(emb))
+        shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=1)
+        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        return x, gate_msa
+
+
+class FluxJointAttention(nn.Module):
+    """Joint attention: text tokens precede image tokens (Diffusers order)."""
+
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+    ) -> None:
+        super().__init__()
+        self.heads = num_attention_heads
+        self.head_dim = attention_head_dim
+        self.inner_dim = num_attention_heads * attention_head_dim
+
+        self.norm_q = nn.RMSNorm(attention_head_dim, eps=1e-6)
+        self.norm_k = nn.RMSNorm(attention_head_dim, eps=1e-6)
+        self.norm_added_q = nn.RMSNorm(attention_head_dim, eps=1e-6)
+        self.norm_added_k = nn.RMSNorm(attention_head_dim, eps=1e-6)
+
+        self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.to_k = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.to_v = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.add_q_proj = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.add_k_proj = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.add_v_proj = ReplicatedLinear(dim, self.inner_dim, bias=True)
+
+        self.to_out = nn.ModuleList(
+            [
+                ReplicatedLinear(self.inner_dim, dim, bias=True),
+                nn.Dropout(0.0),
+            ]
+        )
+        self.to_add_out = ReplicatedLinear(self.inner_dim, dim, bias=True)
+
+        self.attn = DistributedAttention(
+            num_heads=num_attention_heads,
+            head_size=attention_head_dim,
+            causal=False,
+            supported_attention_backends=supported_attention_backends,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        batch_size = hidden_states.shape[0]
+        text_seq_len = encoder_hidden_states.shape[1]
+        img_seq_len = hidden_states.shape[1]
+
+        q, _ = self.to_q(hidden_states)
+        k, _ = self.to_k(hidden_states)
+        v, _ = self.to_v(hidden_states)
+        q = q.view(batch_size, img_seq_len, self.heads, self.head_dim)
+        k = k.view(batch_size, img_seq_len, self.heads, self.head_dim)
+        v = v.view(batch_size, img_seq_len, self.heads, self.head_dim)
+        q = self.norm_q(q)
+        k = self.norm_k(k)
+
+        enc_q, _ = self.add_q_proj(encoder_hidden_states)
+        enc_k, _ = self.add_k_proj(encoder_hidden_states)
+        enc_v, _ = self.add_v_proj(encoder_hidden_states)
+        enc_q = enc_q.view(batch_size, text_seq_len, self.heads, self.head_dim)
+        enc_k = enc_k.view(batch_size, text_seq_len, self.heads, self.head_dim)
+        enc_v = enc_v.view(batch_size, text_seq_len, self.heads, self.head_dim)
+        enc_q = self.norm_added_q(enc_q)
+        enc_k = self.norm_added_k(enc_k)
+
+        q = torch.cat([enc_q, q], dim=1)
+        k = torch.cat([enc_k, k], dim=1)
+        v = torch.cat([enc_v, v], dim=1)
+
+        q = apply_rotary_emb(q, image_rotary_emb, sequence_dim=1)
+        k = apply_rotary_emb(k, image_rotary_emb, sequence_dim=1)
+
+        joint_out, _ = self.attn(q, k, v)
+        joint_out = joint_out.reshape(batch_size, text_seq_len + img_seq_len, self.inner_dim)
+
+        enc_out = joint_out[:, :text_seq_len]
+        img_out = joint_out[:, text_seq_len:]
+
+        img_out, _ = self.to_out[0](img_out)
+        img_out = self.to_out[1](img_out)
+        enc_out, _ = self.to_add_out(enc_out)
+        return img_out, enc_out
+
+
+class FluxSingleStreamAttention(nn.Module):
+    """Self-attention on concatenated text+image sequence (single blocks)."""
+
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+    ) -> None:
+        super().__init__()
+        self.heads = num_attention_heads
+        self.head_dim = attention_head_dim
+        self.inner_dim = num_attention_heads * attention_head_dim
+
+        self.norm_q = nn.RMSNorm(attention_head_dim, eps=1e-6)
+        self.norm_k = nn.RMSNorm(attention_head_dim, eps=1e-6)
+        self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.to_k = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.to_v = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.attn = DistributedAttention(
+            num_heads=num_attention_heads,
+            head_size=attention_head_dim,
+            causal=False,
+            supported_attention_backends=supported_attention_backends,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        batch_size, seq_len, _ = hidden_states.shape
+        q, _ = self.to_q(hidden_states)
+        k, _ = self.to_k(hidden_states)
+        v, _ = self.to_v(hidden_states)
+        q = q.view(batch_size, seq_len, self.heads, self.head_dim)
+        k = k.view(batch_size, seq_len, self.heads, self.head_dim)
+        v = v.view(batch_size, seq_len, self.heads, self.head_dim)
+        q = self.norm_q(q)
+        k = self.norm_k(k)
+        q = apply_rotary_emb(q, image_rotary_emb, sequence_dim=1)
+        k = apply_rotary_emb(k, image_rotary_emb, sequence_dim=1)
+        out, _ = self.attn(q, k, v)
+        return out.reshape(batch_size, seq_len, self.inner_dim)
+
+
+class FluxTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+    ) -> None:
+        super().__init__()
+        self.norm1 = SD3AdaLayerNormZero(dim)
+        self.norm1_context = SD3AdaLayerNormZero(dim)
+        self.attn = FluxJointAttention(
+            dim=dim,
+            num_attention_heads=num_attention_heads,
+            attention_head_dim=attention_head_dim,
+            supported_attention_backends=supported_attention_backends,
+        )
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff = SD3FeedForward(
+            dim=dim,
+            dim_out=dim,
+            activation_fn="gelu-approximate",
+        )
+        self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff_context = SD3FeedForward(
+            dim=dim,
+            dim_out=dim,
+            activation_fn="gelu-approximate",
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
+        joint_attention_kwargs: dict[str, Any] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        del joint_attention_kwargs
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
+        (norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp) = self.norm1_context(
+            encoder_hidden_states, emb=temb
+        )
+
+        attn_output, context_attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+        )
+
+        attn_output = gate_msa.unsqueeze(1) * attn_output
+        hidden_states = hidden_states + attn_output
+
+        norm_hidden_states = self.norm2(hidden_states)
+        norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        ff_output = self.ff(norm_hidden_states)
+        hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output
+
+        context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output
+        encoder_hidden_states = encoder_hidden_states + context_attn_output
+
+        norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
+        norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
+        context_ff_output = self.ff_context(norm_encoder_hidden_states)
+        encoder_hidden_states = encoder_hidden_states + (c_gate_mlp.unsqueeze(1) * context_ff_output)
+        if encoder_hidden_states.dtype == torch.float16:
+            encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+
+        return encoder_hidden_states, hidden_states
+
+
+class FluxSingleTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        mlp_ratio: float = 4.0,
+        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+    ) -> None:
+        super().__init__()
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.norm = FluxAdaLayerNormZeroSingle(dim)
+        self.proj_mlp = ReplicatedLinear(dim, mlp_hidden_dim, bias=True)
+        self.act_mlp = nn.GELU(approximate="tanh")
+        self.proj_out = ReplicatedLinear(dim + mlp_hidden_dim, dim, bias=True)
+        self.attn = FluxSingleStreamAttention(
+            dim=dim,
+            num_attention_heads=num_attention_heads,
+            attention_head_dim=attention_head_dim,
+            supported_attention_backends=supported_attention_backends,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
+        joint_attention_kwargs: dict[str, Any] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        del joint_attention_kwargs
+        text_seq_len = encoder_hidden_states.shape[1]
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+        residual = hidden_states
+        norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
+        mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states)[0])
+        attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+        )
+        hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
+        gate = gate.unsqueeze(1)
+        hidden_states = gate * self.proj_out(hidden_states)[0]
+        hidden_states = residual + hidden_states
+        if hidden_states.dtype == torch.float16:
+            hidden_states = hidden_states.clip(-65504, 65504)
+        encoder_hidden_states = hidden_states[:, :text_seq_len]
+        hidden_states = hidden_states[:, text_seq_len:]
+        return encoder_hidden_states, hidden_states
+
+
+class FluxTransformer2DModel(BaseDiT):
+    """FastVideo FLUX transformer; load Diffusers FLUX safetensors 1:1."""
+
+    _fsdp_shard_conditions = [
+        lambda n, m: (n.startswith("transformer_blocks.") or n.startswith("single_transformer_blocks."))
+        and n.split(".")[-1].isdigit(),
+    ]
+    _compile_conditions = _fsdp_shard_conditions
+    # HF weight names already match this module layout (cf. SGLang regex maps).
+    param_names_mapping: dict[str, Any] = {}
+    reverse_param_names_mapping: dict[str, Any] = {}
+    lora_param_names_mapping: dict[str, Any] = {}
+    _supported_attention_backends = (
+        AttentionBackendEnum.FLASH_ATTN,
+        AttentionBackendEnum.TORCH_SDPA,
+    )
+
+    def __init__(self, config: DiTConfig, hf_config: dict[str, Any], **kwargs) -> None:
+        del kwargs
+        super().__init__(config=config, hf_config=hf_config)
+        self.fastvideo_config = config
+        self.hf_config = hf_config
+        arch = config.arch_config
+
+        out_ch = arch.out_channels
+        self.out_channels = out_ch if out_ch is not None else arch.in_channels
+        self.inner_dim = arch.num_attention_heads * arch.attention_head_dim
+        self.hidden_size = self.inner_dim
+        self.num_attention_heads = arch.num_attention_heads
+        self.num_channels_latents = arch.in_channels
+
+        axes_list = list(arch.axes_dims_rope)
+        self.pos_embed = FluxPosEmbed(theta=10000, axes_dim=axes_list)
+        if arch.guidance_embeds:
+            self.time_text_embed = FluxCombinedTimestepGuidanceTextProjEmbeddings(
+                embedding_dim=self.inner_dim,
+                pooled_projection_dim=arch.pooled_projection_dim,
+            )
+        else:
+            self.time_text_embed = CombinedTimestepTextProjEmbeddings(
+                embedding_dim=self.inner_dim,
+                pooled_projection_dim=arch.pooled_projection_dim,
+            )
+        self.context_embedder = ReplicatedLinear(arch.joint_attention_dim, self.inner_dim)
+        self.x_embedder = ReplicatedLinear(arch.in_channels, self.inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                FluxTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=arch.num_attention_heads,
+                    attention_head_dim=arch.attention_head_dim,
+                    supported_attention_backends=self._supported_attention_backends,
+                )
+                for _ in range(arch.num_layers)
+            ]
+        )
+        self.single_transformer_blocks = nn.ModuleList(
+            [
+                FluxSingleTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=arch.num_attention_heads,
+                    attention_head_dim=arch.attention_head_dim,
+                    supported_attention_backends=self._supported_attention_backends,
+                )
+                for _ in range(arch.num_single_layers)
+            ]
+        )
+
+        self.norm_out = SD3AdaLayerNormContinuous(
+            self.inner_dim,
+            self.inner_dim,
+            elementwise_affine=False,
+            eps=1e-6,
+            bias=True,
+            norm_type="layer_norm",
+        )
+        self.proj_out = ReplicatedLinear(
+            self.inner_dim,
+            arch.patch_size * arch.patch_size * self.out_channels,
+            bias=True,
+        )
+        self.gradient_checkpointing = False
+        self.__post_init__()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | None = None,
+        pooled_projections: torch.Tensor | None = None,
+        timestep: torch.LongTensor | torch.Tensor | None = None,
+        img_ids: torch.Tensor | None = None,
+        txt_ids: torch.Tensor | None = None,
+        guidance: torch.Tensor | None = None,
+        joint_attention_kwargs: dict[str, Any] | None = None,
+        return_dict: bool = True,
+        controlnet_block_samples: Any | None = None,
+        controlnet_single_block_samples: Any | None = None,
+        controlnet_blocks_repeat: bool = False,
+        **kwargs: Any,
+    ) -> FluxTransformer2DModelOutput | tuple[torch.Tensor, ...]:
+        del kwargs
+        if encoder_hidden_states is None:
+            raise ValueError("encoder_hidden_states must be provided")
+        if pooled_projections is None:
+            raise ValueError("pooled_projections must be provided")
+        if timestep is None:
+            raise ValueError("timestep must be provided")
+        if img_ids is None or txt_ids is None:
+            raise ValueError("img_ids and txt_ids must be provided")
+
+        arch = self.fastvideo_config.arch_config
+        if arch.guidance_embeds and guidance is None:
+            raise ValueError("guidance must be provided when guidance_embeds=True")
+
+        if timestep.dim() == 0:
+            timestep = timestep[None]
+        if timestep.dim() > 1:
+            timestep = timestep.reshape(-1)
+        if timestep.shape[0] == 1 and hidden_states.shape[0] > 1:
+            timestep = timestep.expand(hidden_states.shape[0])
+
+        try:
+            get_forward_context()
+            forward_context = nullcontext()
+        except AssertionError:
+            if timestep.numel() == 0:
+                ts0 = 0
+            elif torch.is_floating_point(timestep):
+                ts0 = int(round(timestep[0].item() * 1000))
+            else:
+                ts0 = int(timestep[0].item())
+            forward_context = set_forward_context(current_timestep=ts0, attn_metadata=None)
+
+        with forward_context:
+            hidden_states, _ = self.x_embedder(hidden_states)
+
+            ts = timestep.to(hidden_states.dtype) * 1000
+            g = None if guidance is None else guidance.to(hidden_states.dtype) * 1000
+
+            if arch.guidance_embeds:
+                assert g is not None
+                temb = self.time_text_embed(ts, g, pooled_projections)
+            else:
+                temb = self.time_text_embed(timestep=ts, pooled_projection=pooled_projections)
+
+            encoder_hidden_states, _ = self.context_embedder(encoder_hidden_states)
+
+            if txt_ids.ndim == 3:
+                txt_ids = txt_ids[0]
+            if img_ids.ndim == 3:
+                img_ids = img_ids[0]
+
+            ids = torch.cat((txt_ids, img_ids), dim=0)
+            image_rotary_emb = self.pos_embed(ids)
+
+            jkwargs = joint_attention_kwargs or {}
+
+            for idx, block in enumerate(self.transformer_blocks):
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=jkwargs,
+                )
+                if controlnet_block_samples:
+                    interval = len(self.transformer_blocks) / len(controlnet_block_samples)
+                    interval = int(math.ceil(interval))
+                    if controlnet_blocks_repeat:
+                        hidden_states = hidden_states + controlnet_block_samples[idx % len(controlnet_block_samples)]
+                    else:
+                        hidden_states = hidden_states + controlnet_block_samples[idx // interval]
+
+            for idx, block in enumerate(self.single_transformer_blocks):
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=jkwargs,
+                )
+                if controlnet_single_block_samples:
+                    interval = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
+                    interval = int(math.ceil(interval))
+                    hidden_states = hidden_states + controlnet_single_block_samples[idx // interval]
+
+            hidden_states = self.norm_out(hidden_states, temb)
+            output, _ = self.proj_out(hidden_states)
+
+            if not return_dict:
+                return (output,)
+            return FluxTransformer2DModelOutput(sample=output)
+
+
+EntryClass = FluxTransformer2DModel

fastvideo/models/dits/glm_image.py

@@ -0,0 +1,776 @@
+# SPDX-License-Identifier: Apache-2.0
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from fastvideo.layers.mlp import MLP
+
+from fastvideo.attention import LocalAttention
+from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig
+from fastvideo.layers.layernorm import ScaleResidualLayerNormScaleShift
+from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.layers.rotary_embedding import _apply_rotary_emb
+from fastvideo.layers.visual_embedding import Timesteps
+from fastvideo.models.dits.base import BaseDiT
+from fastvideo.platforms import AttentionBackendEnum
+
+
+class GlmImageLayerKVCache:
+
+    def __init__(self):
+        self.k_cache = None
+        self.v_cache = None
+        self.mode: Optional[str] = None
+
+    def store(self, k: torch.Tensor, v: torch.Tensor):
+        # Append along seq (dim=1).
+        if self.k_cache is None:
+            self.k_cache = k
+            self.v_cache = v
+        else:
+            self.k_cache = torch.cat([self.k_cache, k], dim=1)
+            self.v_cache = torch.cat([self.v_cache, v], dim=1)
+
+    def get(self):
+        return self.k_cache, self.v_cache
+
+    def clear(self):
+        self.k_cache = None
+        self.v_cache = None
+        self.mode = None
+
+
+class GlmImageKVCache:
+
+    def __init__(self, num_layers: int):
+        self.num_layers = num_layers
+        self.caches = [GlmImageLayerKVCache() for _ in range(num_layers)]
+
+    def __getitem__(self, layer_idx: int) -> GlmImageLayerKVCache:
+        return self.caches[layer_idx]
+
+    def set_mode(self, mode: Optional[str]):
+        if mode is not None and mode not in ["write", "read", "skip"]:
+            raise ValueError(
+                f"Invalid mode: {mode}, must be one of 'write', 'read', 'skip'"
+            )
+        for cache in self.caches:
+            cache.mode = mode
+
+    def clear(self):
+        for cache in self.caches:
+            cache.clear()
+
+
+# =============================================================================
+# Timestep and Text Projection
+# =============================================================================
+class GlmImageTimestepEmbedding(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        time_embed_dim: int,
+        act_fn: str = "silu",
+        out_dim: int = None,
+    ):
+        super().__init__()
+        if out_dim is None:
+            out_dim = time_embed_dim
+        self.linear_1 = ReplicatedLinear(in_channels, time_embed_dim, bias=True)
+        if act_fn == "silu":
+            self.act = nn.SiLU()
+        elif act_fn == "gelu":
+            self.act = nn.GELU(approximate="tanh")
+        else:
+            self.act = nn.SiLU()
+        self.linear_2 = ReplicatedLinear(time_embed_dim, out_dim, bias=True)
+
+    def forward(self, sample: torch.Tensor) -> torch.Tensor:
+        sample, _ = self.linear_1(sample)
+        sample = self.act(sample)
+        sample, _ = self.linear_2(sample)
+        return sample
+
+
+class GlmImageTextProjection(nn.Module):
+
+    def __init__(
+        self,
+        in_features: int,
+        hidden_size: int,
+        out_features: int = None,
+        act_fn: str = "silu",
+    ):
+        super().__init__()
+        if out_features is None:
+            out_features = hidden_size
+        self.linear_1 = ReplicatedLinear(in_features, hidden_size, bias=True)
+        if act_fn == "silu":
+            self.act_1 = nn.SiLU()
+        elif act_fn == "gelu_tanh":
+            self.act_1 = nn.GELU(approximate="tanh")
+        else:
+            self.act_1 = nn.SiLU()
+        self.linear_2 = ReplicatedLinear(hidden_size, out_features, bias=True)
+
+    def forward(self, caption: torch.Tensor) -> torch.Tensor:
+        hidden_states, _ = self.linear_1(caption)
+        hidden_states = self.act_1(hidden_states)
+        hidden_states, _ = self.linear_2(hidden_states)
+        return hidden_states
+
+
+class GlmImageCombinedTimestepSizeEmbeddings(nn.Module):
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        condition_dim: int,
+        pooled_projection_dim: int,
+        timesteps_dim: int = 256,
+    ):
+        super().__init__()
+
+        self.time_proj = Timesteps(
+            num_channels=timesteps_dim, flip_sin_to_cos=True, downscale_freq_shift=0
+        )
+        self.condition_proj = Timesteps(
+            num_channels=condition_dim, flip_sin_to_cos=True, downscale_freq_shift=0
+        )
+        self.timestep_embedder = GlmImageTimestepEmbedding(
+            in_channels=timesteps_dim, time_embed_dim=embedding_dim
+        )
+        self.condition_embedder = GlmImageTextProjection(
+            pooled_projection_dim, embedding_dim, act_fn="silu"
+        )
+
+    def forward(
+        self,
+        timestep: torch.Tensor,
+        target_size: torch.Tensor,
+        crop_coords: torch.Tensor,
+        hidden_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        timesteps_proj = self.time_proj(timestep)
+
+        crop_coords_proj = self.condition_proj(crop_coords.flatten()).view(
+            crop_coords.size(0), -1
+        )
+        target_size_proj = self.condition_proj(target_size.flatten()).view(
+            target_size.size(0), -1
+        )
+
+        condition_proj = torch.cat([crop_coords_proj, target_size_proj], dim=1)
+
+        timesteps_emb = self.timestep_embedder(
+            timesteps_proj.to(dtype=hidden_dtype)
+        )
+        condition_emb = self.condition_embedder(
+            condition_proj.to(dtype=hidden_dtype)
+        )
+
+        conditioning = timesteps_emb + condition_emb
+        return conditioning
+
+
+# =============================================================================
+# Image Projector
+# =============================================================================
+class GlmImageImageProjector(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int = 16,
+        hidden_size: int = 2560,
+        patch_size: int = 2,
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+        self.proj = nn.Linear(in_channels * patch_size**2, hidden_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, channel, height, width = hidden_states.shape
+        post_patch_height = height // self.patch_size
+        post_patch_width = width // self.patch_size
+
+        hidden_states = hidden_states.reshape(
+            batch_size,
+            channel,
+            post_patch_height,
+            self.patch_size,
+            post_patch_width,
+            self.patch_size,
+        )
+        hidden_states = (
+            hidden_states.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2)
+        )
+        hidden_states = self.proj(hidden_states)
+
+        return hidden_states
+
+
+# =============================================================================
+# AdaLayerNorm
+# =============================================================================
+class GlmImageAdaLayerNormZero(nn.Module):
+
+    def __init__(self, embedding_dim: int, dim: int) -> None:
+        super().__init__()
+
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
+        self.norm_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
+        self.linear = ReplicatedLinear(embedding_dim, 12 * dim, bias=True)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+    ) -> Tuple[torch.Tensor, ...]:
+        dtype = hidden_states.dtype
+        norm_hidden_states = self.norm(hidden_states).to(dtype=dtype)
+        norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to(
+            dtype=dtype
+        )
+
+        emb, _ = self.linear(temb)
+        (
+            shift_msa,
+            c_shift_msa,
+            scale_msa,
+            c_scale_msa,
+            gate_msa,
+            c_gate_msa,
+            shift_mlp,
+            c_shift_mlp,
+            scale_mlp,
+            c_scale_mlp,
+            gate_mlp,
+            c_gate_mlp,
+        ) = emb.chunk(12, dim=1)
+
+        hidden_states = norm_hidden_states * (
+            1 + scale_msa.unsqueeze(1)
+        ) + shift_msa.unsqueeze(1)
+        encoder_hidden_states = norm_encoder_hidden_states * (
+            1 + c_scale_msa.unsqueeze(1)
+        ) + c_shift_msa.unsqueeze(1)
+
+        return (
+            hidden_states,
+            gate_msa,
+            shift_mlp,
+            scale_mlp,
+            gate_mlp,
+            encoder_hidden_states,
+            c_gate_msa,
+            c_shift_mlp,
+            c_scale_mlp,
+            c_gate_mlp,
+        )
+
+
+# =============================================================================
+# Attention
+# =============================================================================
+class GlmImageAttention(nn.Module):
+
+    def __init__(
+        self,
+        query_dim: int,
+        heads: int,
+        dim_head: int,
+        out_dim: int,
+        bias: bool = True,
+        qk_norm: str = "layer_norm",
+        elementwise_affine: bool = False,
+        eps: float = 1e-5,
+        supported_attention_backends: set[AttentionBackendEnum] | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+
+        self.heads = out_dim // dim_head if out_dim is not None else heads
+        self.num_kv_heads = self.heads
+        self.dim_head = dim_head
+        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+        self.inner_kv_dim = self.inner_dim
+        self.out_dim = out_dim if out_dim is not None else query_dim
+
+        self.to_q = ReplicatedLinear(query_dim, self.inner_dim, bias=bias)
+        self.to_k = ReplicatedLinear(query_dim, self.inner_kv_dim, bias=bias)
+        self.to_v = ReplicatedLinear(query_dim, self.inner_kv_dim, bias=bias)
+
+        self.to_out = nn.ModuleList(
+            [ReplicatedLinear(self.inner_dim, self.out_dim, bias=True)]
+        )
+
+        if qk_norm is None:
+            self.norm_q = None
+            self.norm_k = None
+        elif qk_norm == "layer_norm":
+            self.norm_q = nn.LayerNorm(
+                dim_head, eps=eps, elementwise_affine=elementwise_affine
+            )
+            self.norm_k = nn.LayerNorm(
+                dim_head, eps=eps, elementwise_affine=elementwise_affine
+            )
+        else:
+            raise ValueError(f"unknown qk_norm: {qk_norm}")
+
+        self.attn = LocalAttention(
+            num_heads=self.heads,
+            head_size=dim_head,
+            num_kv_heads=self.heads,
+            softmax_scale=None,
+            causal=False,
+            supported_attention_backends=supported_attention_backends,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        kv_cache: Optional[GlmImageLayerKVCache] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        dtype = encoder_hidden_states.dtype
+
+        batch_size, text_seq_length, embed_dim = encoder_hidden_states.shape
+        batch_size, image_seq_length, embed_dim = hidden_states.shape
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        # 1. QKV projections
+        query, _ = self.to_q(hidden_states)
+        key, _ = self.to_k(hidden_states)
+        value, _ = self.to_v(hidden_states)
+
+        query = query.unflatten(2, (self.heads, -1))
+        key = key.unflatten(2, (self.heads, -1))
+        value = value.unflatten(2, (self.heads, -1))
+
+        # 2. QK normalization
+        if self.norm_q is not None:
+            query = self.norm_q(query).to(dtype=dtype)
+        if self.norm_k is not None:
+            key = self.norm_k(key).to(dtype=dtype)
+
+        # 3. Rotational positional embeddings applied to latent stream
+        if image_rotary_emb is not None:
+            cos, sin = image_rotary_emb
+
+            query[:, text_seq_length:, :, :] = _apply_rotary_emb(
+                query[:, text_seq_length:, :, :], cos, sin, is_neox_style=True
+            )
+            key[:, text_seq_length:, :, :] = _apply_rotary_emb(
+                key[:, text_seq_length:, :, :], cos, sin, is_neox_style=True
+            )
+
+        # 4. KV Cache handling
+        if kv_cache is not None:
+            if kv_cache.mode == "write":
+                kv_cache.store(key, value)
+            elif kv_cache.mode == "read":
+                # Prepend cached condition k/v along seq (dim=1).
+                k_cache, v_cache = kv_cache.get()
+                key = torch.cat([k_cache, key], dim=1) if k_cache is not None else key
+                value = (
+                    torch.cat([v_cache, value], dim=1) if v_cache is not None else value
+                )
+            elif kv_cache.mode == "skip":
+                pass
+
+        hidden_states = self.attn(query, key, value)
+        hidden_states = hidden_states.flatten(2, 3)
+        hidden_states = hidden_states.to(query.dtype)
+
+        # 6. Output projection
+        hidden_states, _ = self.to_out[0](hidden_states)
+
+        encoder_hidden_states, hidden_states = hidden_states.split(
+            [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1
+        )
+        return hidden_states, encoder_hidden_states
+
+
+# =============================================================================
+# Transformer Block
+# =============================================================================
+class GlmImageTransformerBlock(nn.Module):
+
+    def __init__(
+        self,
+        dim: int = 2560,
+        num_attention_heads: int = 64,
+        attention_head_dim: int = 40,
+        time_embed_dim: int = 512,
+        supported_attention_backends: set[AttentionBackendEnum] | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        # 1. Attention
+        self.norm1 = GlmImageAdaLayerNormZero(time_embed_dim, dim)
+
+        self.attn1 = GlmImageAttention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            out_dim=dim,
+            bias=True,
+            qk_norm="layer_norm",
+            elementwise_affine=False,
+            eps=1e-5,
+            supported_attention_backends=supported_attention_backends,
+            prefix=f"{prefix}.attn1",
+        )
+
+        # 2. Feedforward with fused ScaleResidualLayerNorm
+        self.norm2 = ScaleResidualLayerNormScaleShift(
+            dim, norm_type="layer", eps=1e-5, elementwise_affine=False
+        )
+        self.norm2_context = ScaleResidualLayerNormScaleShift(
+            dim, norm_type="layer", eps=1e-5, elementwise_affine=False
+        )
+        self.ff = MLP(input_dim=dim, mlp_hidden_dim=dim * 4, output_dim=dim, act_type="gelu_pytorch_tanh")
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[
+            Union[
+                Tuple[torch.Tensor, torch.Tensor],
+                List[Tuple[torch.Tensor, torch.Tensor]],
+            ]
+        ] = None,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        kv_cache: Optional[GlmImageLayerKVCache] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # 1. Timestep conditioning
+        (
+            norm_hidden_states,
+            gate_msa,
+            shift_mlp,
+            scale_mlp,
+            gate_mlp,
+            norm_encoder_hidden_states,
+            c_gate_msa,
+            c_shift_mlp,
+            c_scale_mlp,
+            c_gate_mlp,
+        ) = self.norm1(hidden_states, encoder_hidden_states, temb)
+
+        # 2. Attention
+        if attention_kwargs is None:
+            attention_kwargs = {}
+
+        attn_hidden_states, attn_encoder_hidden_states = self.attn1(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+            kv_cache=kv_cache,
+            **attention_kwargs,
+        )
+
+        # 3. Feedforward (fused residual + norm + scale/shift)
+        norm_hidden_states, hidden_states = self.norm2(
+            hidden_states,
+            attn_hidden_states,
+            gate_msa.unsqueeze(1),
+            shift_mlp.unsqueeze(1),
+            scale_mlp.unsqueeze(1),
+        )
+        norm_encoder_hidden_states, encoder_hidden_states = self.norm2_context(
+            encoder_hidden_states,
+            attn_encoder_hidden_states,
+            c_gate_msa.unsqueeze(1),
+            c_shift_mlp.unsqueeze(1),
+            c_scale_mlp.unsqueeze(1),
+        )
+
+        ff_output = self.ff(norm_hidden_states)
+        ff_output_context = self.ff(norm_encoder_hidden_states)
+        hidden_states = hidden_states + ff_output * gate_mlp.unsqueeze(1)
+        encoder_hidden_states = (
+            encoder_hidden_states + ff_output_context * c_gate_mlp.unsqueeze(1)
+        )
+
+        return hidden_states, encoder_hidden_states
+
+
+# =============================================================================
+# Rotary Positional Embedding
+# =============================================================================
+class GlmImageRotaryPosEmbed(nn.Module):
+
+    def __init__(self, dim: int, patch_size: int, theta: float = 10000.0) -> None:
+        super().__init__()
+        self.dim = dim
+        self.patch_size = patch_size
+        self.theta = theta
+        self._cache_key: tuple | None = None
+        self._cache_value: tuple[torch.Tensor, torch.Tensor] | None = None
+
+    def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        _, _, raw_h, raw_w = hidden_states.shape
+        height = raw_h // self.patch_size
+        width = raw_w // self.patch_size
+        device = hidden_states.device
+
+        cache_key = (height, width, device.type,
+                     device.index if device.index is not None else -1)
+        if self._cache_key == cache_key and self._cache_value is not None:
+            return self._cache_value
+
+        dim_h, dim_w = self.dim // 2, self.dim // 2
+        h_inv_freq = 1.0 / (
+            self.theta
+            ** (
+                torch.arange(0, dim_h, 2, dtype=torch.float32, device=device)[
+                    : (dim_h // 2)
+                ]
+                / dim_h
+            )
+        )
+        w_inv_freq = 1.0 / (
+            self.theta
+            ** (
+                torch.arange(0, dim_w, 2, dtype=torch.float32, device=device)[
+                    : (dim_w // 2)
+                ]
+                / dim_w
+            )
+        )
+        h_seq = torch.arange(height, device=device)
+        w_seq = torch.arange(width, device=device)
+        freqs_h = torch.outer(h_seq, h_inv_freq).unsqueeze(1).expand(height, width, -1)
+        freqs_w = torch.outer(w_seq, w_inv_freq).unsqueeze(0).expand(height, width, -1)
+
+        freqs = torch.cat([freqs_h, freqs_w], dim=-1).reshape(height * width, -1)
+        result = (freqs.cos(), freqs.sin())
+        self._cache_key = cache_key
+        self._cache_value = result
+        return result
+
+
+# =============================================================================
+# Final AdaLayerNorm
+# =============================================================================
+class GlmImageAdaLayerNormContinuous(nn.Module):
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        conditioning_embedding_dim: int,
+        elementwise_affine: bool = True,
+        eps: float = 1e-5,
+        bias: bool = True,
+        norm_type: str = "layer_norm",
+    ):
+        super().__init__()
+        self.linear = nn.Linear(
+            conditioning_embedding_dim, embedding_dim * 2, bias=bias
+        )
+        if norm_type == "layer_norm":
+            self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias)
+        elif norm_type == "rms_norm":
+            self.norm = nn.RMSNorm(embedding_dim, eps, elementwise_affine)
+        else:
+            raise ValueError(f"unknown norm_type {norm_type}")
+
+    def forward(
+        self, x: torch.Tensor, conditioning_embedding: torch.Tensor
+    ) -> torch.Tensor:
+        emb = self.linear(conditioning_embedding.to(x.dtype))
+        scale, shift = torch.chunk(emb, 2, dim=1)
+        x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
+        return x
+
+
+# =============================================================================
+# Main Model
+# =============================================================================
+class GlmImageTransformer2DModel(BaseDiT):
+
+    _fsdp_shard_conditions = GlmImageDiTConfig().arch_config._fsdp_shard_conditions
+    _compile_conditions = GlmImageDiTConfig().arch_config._compile_conditions
+    _supported_attention_backends = (
+        AttentionBackendEnum.FLASH_ATTN,
+        AttentionBackendEnum.TORCH_SDPA,
+    )
+
+    param_names_mapping = GlmImageDiTConfig().arch_config.param_names_mapping
+    reverse_param_names_mapping = {}
+    lora_param_names_mapping = {}
+
+    def __init__(
+        self,
+        config: GlmImageDiTConfig,
+        hf_config: dict[str, Any],
+    ):
+        super().__init__(config=config, hf_config=hf_config)
+
+        arch_config = config.arch_config
+
+        self.in_channels = arch_config.in_channels
+        self.out_channels = arch_config.out_channels
+        self.patch_size = arch_config.patch_size
+        self.num_layers = arch_config.num_layers
+        self.attention_head_dim = arch_config.attention_head_dim
+        self.num_attention_heads = arch_config.num_attention_heads
+        self.text_embed_dim = arch_config.text_embed_dim
+        self.time_embed_dim = arch_config.time_embed_dim
+
+        # GlmImage uses 2 additional SDXL-like conditions - target_size, crop_coords
+        # Each of these are sincos embeddings of shape 2 * condition_dim
+        pooled_projection_dim = 2 * 2 * arch_config.condition_dim
+        inner_dim = arch_config.num_attention_heads * arch_config.attention_head_dim
+
+        self.hidden_size = inner_dim
+        self.num_channels_latents = arch_config.out_channels
+
+        # 1. RoPE
+        self.rotary_emb = GlmImageRotaryPosEmbed(
+            arch_config.attention_head_dim, arch_config.patch_size, theta=10000.0
+        )
+
+        # 2. Patch & Text-timestep embedding
+        self.image_projector = GlmImageImageProjector(
+            arch_config.in_channels, inner_dim, arch_config.patch_size
+        )
+        self.glyph_projector = MLP(
+            input_dim=arch_config.text_embed_dim,
+            mlp_hidden_dim=inner_dim,
+            output_dim=inner_dim,
+            act_type="gelu",
+        )
+        self.prior_token_embedding = nn.Embedding(
+            arch_config.prior_vq_quantizer_codebook_size, inner_dim
+        )
+        self.prior_projector = MLP(
+            input_dim=inner_dim,
+            mlp_hidden_dim=inner_dim,
+            output_dim=inner_dim,
+            act_type="silu",
+        )
+
+        self.time_condition_embed = GlmImageCombinedTimestepSizeEmbeddings(
+            embedding_dim=arch_config.time_embed_dim,
+            condition_dim=arch_config.condition_dim,
+            pooled_projection_dim=pooled_projection_dim,
+            timesteps_dim=arch_config.time_embed_dim,
+        )
+
+        # 3. Transformer blocks
+        self.transformer_blocks = nn.ModuleList(
+            [
+                GlmImageTransformerBlock(
+                    inner_dim,
+                    arch_config.num_attention_heads,
+                    arch_config.attention_head_dim,
+                    arch_config.time_embed_dim,
+                    supported_attention_backends=self._supported_attention_backends,
+                    prefix=f"transformer_blocks.{i}",
+                )
+                for i in range(arch_config.num_layers)
+            ]
+        )
+
+        # 4. Output projection
+        self.norm_out = GlmImageAdaLayerNormContinuous(
+            inner_dim, arch_config.time_embed_dim, elementwise_affine=False
+        )
+        self.proj_out = nn.Linear(
+            inner_dim,
+            arch_config.patch_size * arch_config.patch_size * arch_config.out_channels,
+            bias=True,
+        )
+
+        self.gradient_checkpointing = False
+
+        self.__post_init__()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        prior_token_id: torch.Tensor,
+        prior_token_drop: torch.Tensor,
+        timestep: torch.LongTensor,
+        target_size: torch.Tensor,
+        crop_coords: torch.Tensor,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        kv_caches: Optional[GlmImageKVCache] = None,
+        kv_caches_mode: Optional[str] = None,
+        freqs_cis: Optional[
+            Union[
+                Tuple[torch.Tensor, torch.Tensor],
+                List[Tuple[torch.Tensor, torch.Tensor]],
+            ]
+        ] = None,
+        guidance: torch.Tensor = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        if kv_caches is not None:
+            kv_caches.set_mode(kv_caches_mode)
+
+        batch_size, num_channels, height, width = hidden_states.shape
+
+        if isinstance(encoder_hidden_states, list):
+            encoder_hidden_states = encoder_hidden_states[0]
+
+        # 1. RoPE
+        image_rotary_emb = freqs_cis
+        if image_rotary_emb is None:
+            image_rotary_emb = self.rotary_emb(hidden_states)
+
+        # 2. Patch & Timestep embeddings
+        p = self.patch_size
+        post_patch_height = height // p
+        post_patch_width = width // p
+
+        hidden_states = self.image_projector(hidden_states)
+        encoder_hidden_states = self.glyph_projector(encoder_hidden_states)
+        prior_embedding = self.prior_token_embedding(prior_token_id)
+        # Zero dropped priors by multiply: boolean indexing + .any() syncs each step.
+        keep = (~prior_token_drop).to(device=prior_embedding.device, dtype=prior_embedding.dtype)
+        while keep.dim() < prior_embedding.dim():
+            keep = keep.unsqueeze(-1)
+        prior_embedding = prior_embedding * keep
+        prior_hidden_states = self.prior_projector(prior_embedding)
+        hidden_states = hidden_states + prior_hidden_states
+
+        temb = self.time_condition_embed(
+            timestep, target_size, crop_coords, hidden_states.dtype
+        )
+        temb = F.silu(temb)
+
+        # 3. Transformer blocks
+        for idx, block in enumerate(self.transformer_blocks):
+            hidden_states, encoder_hidden_states = block(
+                hidden_states,
+                encoder_hidden_states,
+                temb,
+                image_rotary_emb,
+                attention_kwargs,
+                kv_cache=kv_caches[idx] if kv_caches is not None else None,
+            )
+
+        # 4. Output norm & projection
+        hidden_states = self.norm_out(hidden_states, temb)
+        hidden_states = self.proj_out(hidden_states)
+
+        # 5. Unpatchify
+        hidden_states = hidden_states.reshape(
+            batch_size, post_patch_height, post_patch_width, -1, p, p
+        )
+        output = hidden_states.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3)
+
+        return output.float()
+
+
+EntryClass = GlmImageTransformer2DModel

fastvideo/models/dits/ltx2.py

@@ -72,8 +72,13 @@ def _rms_norm_dispatch(
     """Use QuACK RMSNorm only for LTX-2 refine stage, else Torch RMSNorm."""
     if _is_ltx2_refine_stage():
         return _quack_rmsnorm(x, weight=weight, eps=eps)
+    # Official LTX-2 runs the DiT natively in bf16 with no autocast, so its
+    # rms_norm outputs stay bf16 on every torch version. torch 2.12 added
+    # rms_norm to autocast's float32 cast policy, which under our
+    # autocast(bfloat16) forward returns float32 instead. Restore the input
+    # dtype to match the official numerics (and torch <= 2.11 behavior).
     return torch.nn.functional.rms_norm(
-        x, (x.shape[-1], ), weight=weight, eps=eps)
+        x, (x.shape[-1], ), weight=weight, eps=eps).to(x.dtype)
 
 
 class StageAwareRMSNorm(nn.RMSNorm):

fastvideo/models/encoders/glm_image_ar_loader.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Sole HF-import boundary for GLM-Image's AR encoder (lazy-wrapper exception E001)."""
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class GlmImageARLoader(nn.Module):
+
+    def __init__(self, model_path: str, processor_path: str | None = None,
+                 *, torch_dtype: torch.dtype = torch.bfloat16,
+                 trust_remote_code: bool = True) -> None:
+        super().__init__()
+        from transformers import (AutoProcessor,
+                                  GlmImageForConditionalGeneration)
+        logger.info("Loading GLM-Image AR encoder from %s", model_path)
+        self._model = GlmImageForConditionalGeneration.from_pretrained(
+            model_path,
+            torch_dtype=torch_dtype,
+            trust_remote_code=trust_remote_code,
+        )
+        if processor_path is not None:
+            logger.info("Loading GLM-Image processor from %s", processor_path)
+            self.processor = AutoProcessor.from_pretrained(
+                processor_path, trust_remote_code=trust_remote_code)
+        else:
+            self.processor = None
+
+    @torch.no_grad()
+    def generate(self, *args: Any, **kwargs: Any) -> torch.Tensor:
+        return self._model.generate(*args, **kwargs)
+
+    @torch.no_grad()
+    def get_image_features(self, pixel_values: torch.Tensor,
+                           image_grid_thw: torch.Tensor) -> Any:
+        return self._model.get_image_features(pixel_values, image_grid_thw)
+
+    @torch.no_grad()
+    def get_image_tokens(self, image_embeds: torch.Tensor,
+                         image_grid_thw: torch.Tensor) -> torch.Tensor:
+        return self._model.get_image_tokens(image_embeds, image_grid_thw)
+
+    @property
+    def config(self):  # type: ignore[no-untyped-def]
+        return self._model.config
+
+    @property
+    def generation_config(self):  # type: ignore[no-untyped-def]
+        return self._model.generation_config
+
+    def to(self, *args, **kwargs):  # type: ignore[override]
+        self._model = self._model.to(*args, **kwargs)
+        return super().to(*args, **kwargs)
+
+    def eval(self):  # type: ignore[override]
+        self._model = self._model.eval()
+        return super().eval()

fastvideo/models/loader/benchmarks/benchmark_weight_loading_comparison.py

@@ -42,7 +42,7 @@ def load_independent(files: list[str], device: str):
     """Before-PR behavior: every rank reads every tensor from disk to GPU."""
     for st_file in files:
         with safe_open(st_file, framework="pt", device=device) as f:
-            for name in f:
+            for name in f.keys():  # noqa: SIM118
                 param = f.get_tensor(name)
                 yield name, param
 
@@ -54,7 +54,7 @@ def load_broadcast(files: list[str], device: str, node_group,
     handles = []
     for st_file in files:
         with safe_open(st_file, framework="pt", device=device) as f:
-            for name in f:
+            for name in f.keys():  # noqa: SIM118
                 if local_rank == 0:
                     param = f.get_tensor(name)
                 else:

fastvideo/models/loader/component_loader.py

@@ -95,6 +95,8 @@ class ComponentLoader(ABC):
             "image_processor": (ImageProcessorLoader, "transformers"),
             "feature_extractor": (ImageProcessorLoader, "transformers"),
             "image_encoder": (ImageEncoderLoader, "transformers"),
+            "vision_language_encoder": (VisionLanguageEncoderLoader, "transformers"),
+            "processor": (ProcessorLoader, "transformers"),
             "upsampler": (UpsamplerLoader, "diffusers"),
             "upsampler_2": (UpsamplerLoader, "diffusers"),
             # Stable Audio's `StableAudioMultiConditioner` bundles T5 +
@@ -523,6 +525,39 @@ class ImageEncoderLoader(TextEncoderLoader):
         )
 
 
+class VisionLanguageEncoderLoader(ComponentLoader):
+    """Loader for vision-language autoregressive encoders."""
+
+    def load(self, model_path: str, fastvideo_args: FastVideoArgs):
+        from fastvideo.distributed.parallel_state import get_local_torch_device
+        from fastvideo.models.encoders.glm_image_ar_loader import (
+            GlmImageARLoader)
+
+        logger.info("Loading vision-language encoder from %s", model_path)
+        target_device = get_local_torch_device()
+        loader = GlmImageARLoader(
+            model_path,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=fastvideo_args.trust_remote_code,
+        ).to(target_device).eval()
+        return loader
+
+
+class ProcessorLoader(ComponentLoader):
+    """Loader for HF processors that pair with vision-language encoders."""
+
+    def load(self, model_path: str, fastvideo_args: FastVideoArgs):
+        from transformers import AutoProcessor
+
+        logger.info("Loading processor from %s", model_path)
+        processor = AutoProcessor.from_pretrained(
+            model_path,
+            trust_remote_code=fastvideo_args.trust_remote_code,
+        )
+        logger.info("Loaded processor: %s", processor.__class__.__name__)
+        return processor
+
+
 class ImageProcessorLoader(ComponentLoader):
     """Loader for image processor."""
 

fastvideo/models/registry.py

@@ -61,6 +61,11 @@ _IMAGE_TO_VIDEO_DIT_MODELS = {
     "MatrixGame3WanModel": ("dits", "matrixgame3", "MatrixGame3WanModel"),
 }
 
+# Text-to-image DiT models (2D image generation)
+_TEXT_TO_IMAGE_DIT_MODELS = {
+    "GlmImageTransformer2DModel": ("dits", "glm_image", "GlmImageTransformer2DModel"),
+}
+
 _TEXT_ENCODER_MODELS = {
     "CLIPTextModel": ("encoders", "clip", "CLIPTextModel"),
     "CLIPTextModelWithProjection":
@@ -134,6 +139,7 @@ _UPSAMPLERS = {
 _LEGACY_FAST_VIDEO_MODELS = {
     **_TEXT_TO_VIDEO_DIT_MODELS,
     **_IMAGE_TO_VIDEO_DIT_MODELS,
+    **_TEXT_TO_IMAGE_DIT_MODELS,
     **_TEXT_ENCODER_MODELS,
     **_IMAGE_ENCODER_MODELS,
     **_VAE_MODELS,

fastvideo/pipelines/basic/flux/__init__.py

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

fastvideo/pipelines/basic/flux/flux_pipeline.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.stages.flux_stages import (
+    FluxConditioningStage,
+    FluxDecodingStage,
+    FluxDenoisingStage,
+    FluxInputValidationStage,
+    FluxLatentPreparationStage,
+    FluxTimestepPreparationStage,
+)
+from fastvideo.pipelines.stages.text_encoding import TextEncodingStage
+
+
+class FluxPipeline(ComposedPipelineBase):
+    """FLUX.1-dev T2I (Diffusers module layout, packed latents, embedded guidance)."""
+
+    _required_config_modules = [
+        "scheduler",
+        "transformer",
+        "vae",
+        "text_encoder",
+        "text_encoder_2",
+        "tokenizer",
+        "tokenizer_2",
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        self.add_stage(stage_name="input_validation_stage", stage=FluxInputValidationStage())
+
+        self.add_stage(
+            stage_name="text_encoding_stage",
+            stage=TextEncodingStage(
+                text_encoders=[
+                    self.get_module("text_encoder"),
+                    self.get_module("text_encoder_2"),
+                ],
+                tokenizers=[
+                    self.get_module("tokenizer"),
+                    self.get_module("tokenizer_2"),
+                ],
+            ),
+        )
+
+        self.add_stage(stage_name="flux_conditioning_stage", stage=FluxConditioningStage())
+
+        self.add_stage(
+            stage_name="timestep_preparation_stage",
+            stage=FluxTimestepPreparationStage(scheduler=self.get_module("scheduler")),
+        )
+
+        self.add_stage(
+            stage_name="latent_preparation_stage",
+            stage=FluxLatentPreparationStage(scheduler=self.get_module("scheduler")),
+        )
+
+        self.add_stage(
+            stage_name="denoising_stage",
+            stage=FluxDenoisingStage(
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler"),
+            ),
+        )
+
+        self.add_stage(
+            stage_name="decoding_stage",
+            stage=FluxDecodingStage(vae=self.get_module("vae")),
+        )
+
+
+EntryClass = FluxPipeline

fastvideo/pipelines/basic/glm_image/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+"""GLM-Image pipeline package."""
+
+from fastvideo.pipelines.basic.glm_image.glm_image_pipeline import (
+    GlmImagePipeline, )
+
+__all__ = ["GlmImagePipeline"]

fastvideo/pipelines/basic/glm_image/glm_image_pipeline.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler, )
+from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
+from fastvideo.pipelines.basic.glm_image.stages import (
+    GlmImageBeforeDenoisingStage,
+    GlmImageConditionEncodingStage,
+    GlmImageDecodingStage,
+    GlmImageDenoisingStage,
+)
+from fastvideo.pipelines.stages import InputValidationStage
+
+logger = init_logger(__name__)
+
+
+class GlmImagePipeline(LoRAPipeline, ComposedPipelineBase):
+    pipeline_name = "GlmImagePipeline"
+
+    _required_config_modules = [
+        "text_encoder",
+        "tokenizer",
+        "vae",
+        "transformer",
+        "scheduler",
+        "vision_language_encoder",
+        "processor",
+    ]
+
+    _optional_config_modules: list[str] = []
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None:
+        self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(shift=1.0)
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        self.add_stage(
+            stage_name="input_validation_stage",
+            stage=InputValidationStage(),
+        )
+
+        self.add_stage(
+            stage_name="glm_image_before_denoising_stage",
+            stage=GlmImageBeforeDenoisingStage(
+                vae=self.get_module("vae"),
+                text_encoder=self.get_module("text_encoder"),
+                tokenizer=self.get_module("tokenizer"),
+                processor=self.get_module("processor"),
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler"),
+                vision_language_encoder=self.get_module("vision_language_encoder"),
+            ),
+        )
+
+        self.add_stage(
+            stage_name="glm_image_condition_encoding_stage",
+            stage=GlmImageConditionEncodingStage(
+                vae=self.get_module("vae"),
+                transformer=self.get_module("transformer"),
+            ),
+        )
+
+        self.add_stage(
+            stage_name="denoising_stage",
+            stage=GlmImageDenoisingStage(
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler"),
+                vae=self.get_module("vae"),
+                pipeline=self,
+            ),
+        )
+
+        self.add_stage(
+            stage_name="decoding_stage",
+            stage=GlmImageDecodingStage(
+                vae=self.get_module("vae"),
+                pipeline=self,
+            ),
+        )
+
+
+EntryClass = GlmImagePipeline

fastvideo/pipelines/basic/glm_image/stages/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0
+"""GLM-Image pipeline stages."""
+
+from fastvideo.pipelines.basic.glm_image.stages.before_denoising import (GlmImageBeforeDenoisingStage)
+from fastvideo.pipelines.basic.glm_image.stages.condition_encoding import (GlmImageConditionEncodingStage)
+from fastvideo.pipelines.basic.glm_image.stages.decoding import (GlmImageDecodingStage)
+from fastvideo.pipelines.basic.glm_image.stages.denoising import (GlmImageDenoisingStage)
+
+__all__ = [
+    "GlmImageBeforeDenoisingStage",
+    "GlmImageConditionEncodingStage",
+    "GlmImageDecodingStage",
+    "GlmImageDenoisingStage",
+]

fastvideo/pipelines/basic/glm_image/stages/before_denoising.py

@@ -0,0 +1,269 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import re
+from math import sqrt
+
+import numpy as np
+import torch
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+from fastvideo.pipelines.stages.validators import StageValidators as V
+from fastvideo.pipelines.stages.validators import VerificationResult
+
+logger = init_logger(__name__)
+
+
+def calculate_shift(
+    image_seq_len: int,
+    base_seq_len: int = 256,
+    base_shift: float = 0.25,
+    max_shift: float = 0.75,
+) -> float:
+    return (image_seq_len / base_seq_len)**0.5 * max_shift + base_shift
+
+
+def get_glyph_texts(prompt: str | list[str]) -> list[str] | list[list[str]]:
+    if isinstance(prompt, str):
+        prompts: list[str] = [prompt]
+        is_batch = False
+    else:
+        prompts = prompt
+        is_batch = True
+    out: list[list[str]] = []
+    for p in prompts:
+        out.append(
+            re.findall(r"'([^']*)'", p) + re.findall(r"“([^“”]*)”", p) + re.findall(r'"([^"]*)"', p) +
+            re.findall(r"「([^「」]*)」", p))
+    return out if is_batch else out[0]
+
+
+def compute_glyph_embeds(
+    prompts: list[str],
+    tokenizer,
+    text_encoder,
+    device: torch.device,
+    dtype: torch.dtype,
+    max_sequence_length: int = 2048,
+) -> torch.Tensor:
+    all_glyph_texts = get_glyph_texts(prompts)
+    all_glyph_embeds = []
+    for glyph_texts in all_glyph_texts:
+        if len(glyph_texts) == 0:
+            glyph_texts = [""]
+        input_ids = tokenizer(
+            glyph_texts,
+            max_length=max_sequence_length,
+            truncation=True,
+        ).input_ids
+        input_ids = [[tokenizer.pad_token_id] * ((len(input_ids) + 1) % 2) + ids for ids in input_ids]
+        max_length = max(len(ids) for ids in input_ids)
+        attention_mask = torch.tensor(
+            [[1] * len(ids) + [0] * (max_length - len(ids)) for ids in input_ids],
+            device=device,
+        )
+        input_ids_t = torch.tensor(
+            [ids + [tokenizer.pad_token_id] * (max_length - len(ids)) for ids in input_ids],
+            device=device,
+        )
+        outputs = text_encoder(input_ids_t, attention_mask=attention_mask)
+        glyph_embeds = outputs.last_hidden_state[attention_mask.bool()].unsqueeze(0)
+        all_glyph_embeds.append(glyph_embeds)
+
+    max_seq_len = max(emb.size(1) for emb in all_glyph_embeds)
+    padded = []
+    for emb in all_glyph_embeds:
+        if emb.size(1) < max_seq_len:
+            pad = torch.zeros(emb.size(0), max_seq_len - emb.size(1), emb.size(2), device=device, dtype=emb.dtype)
+            emb = torch.cat([pad, emb], dim=1)
+        padded.append(emb)
+    return torch.cat(padded, dim=0).to(device=device, dtype=dtype)
+
+
+def _grid_dims(height: int, width: int) -> tuple[int, int, int, int]:
+    th, tw = height // 32, width // 32
+    ratio = th / tw
+    pth = int(sqrt(ratio) * 16)
+    ptw = int(sqrt(1 / ratio) * 16)
+    return th, tw, pth, ptw
+
+
+def _upsample_d32_to_d16(tokens: torch.Tensor, th: int, tw: int) -> torch.Tensor:
+    tokens = tokens.view(1, 1, th, tw).float()
+    tokens = torch.nn.functional.interpolate(tokens, scale_factor=2, mode="nearest").long()
+    return tokens.view(1, -1)
+
+
+class GlmImageBeforeDenoisingStage(PipelineStage):
+
+    def __init__(self,
+                 vae,
+                 text_encoder,
+                 tokenizer,
+                 processor,
+                 transformer,
+                 scheduler,
+                 vision_language_encoder=None) -> None:
+        super().__init__()
+        self.vae = vae
+        self.text_encoder = text_encoder
+        self.tokenizer = tokenizer
+        self.processor = processor
+        self.transformer = transformer
+        self.scheduler = scheduler
+        if isinstance(vision_language_encoder, tuple):
+            self.vision_language_encoder, self.vl_processor = (vision_language_encoder[0], vision_language_encoder[1]
+                                                               or processor)
+        else:
+            self.vision_language_encoder = vision_language_encoder
+            self.vl_processor = processor
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        device = get_local_torch_device()
+        dtype = torch.bfloat16
+        th, tw, pth, ptw = _grid_dims(batch.height, batch.width)
+
+        if batch.seed is not None:
+            torch.manual_seed(batch.seed)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed(batch.seed)
+
+        # 1-3. AR token generation. I2I prepends the condition image and uses a
+        # single-scale target grid; T2I is multi-scale.
+        is_t2i = batch.pil_image is None
+        if self.vision_language_encoder is not None:
+            content = [{"type": "text", "text": batch.prompt}]
+            if not is_t2i:
+                content.insert(0, {"type": "image", "image": batch.pil_image})
+            messages = [{"role": "user", "content": content}]
+            inputs = self.vl_processor.apply_chat_template(messages,
+                                                           tokenize=True,
+                                                           target_h=batch.height,
+                                                           target_w=batch.width,
+                                                           return_dict=True,
+                                                           return_tensors="pt").to(device)
+
+            if is_t2i:
+                up_h, up_w = th, tw
+                large_start, large_count = pth * ptw, th * tw
+                max_new = large_count + (pth * ptw) + 1
+            else:
+                # Condition grid(s) first, target grid last.
+                _, t_h, t_w = inputs["image_grid_thw"][-1].tolist()
+                up_h, up_w = int(t_h), int(t_w)
+                large_start, large_count = 0, up_h * up_w
+                max_new = large_count + 1
+
+            outputs = self.vision_language_encoder.generate(**inputs, max_new_tokens=max_new, do_sample=True)
+            gen_tokens = outputs[0][inputs.input_ids.shape[-1]:]
+            if gen_tokens.shape[0] >= large_start + large_count:
+                large_tokens = gen_tokens[large_start:large_start + large_count]
+            else:
+                available = gen_tokens[large_start:]
+                large_tokens = torch.zeros(large_count, dtype=gen_tokens.dtype, device=gen_tokens.device)
+                if available.shape[0] > 0:
+                    large_tokens[:min(available.shape[0], large_count)] = available[:large_count]
+                logger.warning("AR generated %d tokens, expected %d. Padding with zeros.", gen_tokens.shape[0],
+                               large_start + large_count)
+            batch.prior_token_id = _upsample_d32_to_d16(large_tokens, up_h, up_w)
+            batch.prior_token_drop = torch.zeros(batch.prior_token_id.shape, dtype=torch.bool, device=device)
+
+            if not is_t2i:
+                self._compute_source_prior_tokens(batch, inputs)
+        else:
+            num_prior_tokens = 4 * th * tw
+            logger.warning("No vision_language_encoder provided; using random dropped priors.")
+            batch.prior_token_id = torch.randint(0, 16384, (1, num_prior_tokens), device=device)
+            batch.prior_token_drop = torch.ones(batch.prior_token_id.shape, dtype=torch.bool, device=device)
+
+        # 4. Glyph T5 encoding.
+        prompts = [batch.prompt] if isinstance(batch.prompt, str) else list(batch.prompt)
+        prompt_embeds = compute_glyph_embeds(prompts, self.tokenizer, self.text_encoder, device, dtype)
+
+        # 5. CFG-side negative encoding.
+        if batch.do_classifier_free_guidance:
+            neg_prompts = [batch.negative_prompt or ""] * len(prompts)
+            neg_embeds = compute_glyph_embeds(neg_prompts, self.tokenizer, self.text_encoder, device, dtype)
+            L_pos, L_neg = prompt_embeds.shape[1], neg_embeds.shape[1]
+            max_L = max(L_pos, L_neg)
+            if L_pos < max_L:
+                pad = torch.zeros(prompt_embeds.shape[0],
+                                  max_L - L_pos,
+                                  prompt_embeds.shape[2],
+                                  device=device,
+                                  dtype=dtype)
+                prompt_embeds = torch.cat([pad, prompt_embeds], dim=1)
+            if L_neg < max_L:
+                pad = torch.zeros(neg_embeds.shape[0], max_L - L_neg, neg_embeds.shape[2], device=device, dtype=dtype)
+                neg_embeds = torch.cat([pad, neg_embeds], dim=1)
+            # Row 0 conditional (positive), row 1 unconditional (negative).
+            prompt_embeds = torch.cat([prompt_embeds, neg_embeds], dim=0)
+            att_pos = torch.ones((1, max_L), device=device)
+            att_neg = torch.ones((1, max_L), device=device)
+            if L_pos < max_L:
+                att_pos[:, :max_L - L_pos] = 0
+            if L_neg < max_L:
+                att_neg[:, :max_L - L_neg] = 0
+            attention_mask = torch.cat([att_pos, att_neg], dim=0)
+        else:
+            attention_mask = torch.ones((1, prompt_embeds.shape[1]), device=device)
+
+        batch.prompt_embeds = [prompt_embeds]
+        batch.attention_mask = attention_mask
+
+        # 6. Latents + dynamic flow shift.
+        if batch.seed is not None:
+            torch.manual_seed(batch.seed)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed(batch.seed)
+        batch.latents = torch.randn((1, 16, 1, batch.height // 8, batch.width // 8), device=device, dtype=dtype)
+
+        # Integer-cast linspace timesteps with resolution-dependent shift applied to
+        # sigmas only; the DiT is conditioned on the unshifted integer timesteps.
+        ntt = self.scheduler.config.num_train_timesteps
+        patch_size = self.transformer.patch_size
+        image_seq_len = ((batch.height // 8) * (batch.width // 8)) // (patch_size**2)
+        sched_timesteps = np.linspace(ntt, 1.0, batch.num_inference_steps + 1)[:-1].astype(np.int64).astype(np.float32)
+        sched_sigmas = sched_timesteps / ntt
+        self.scheduler.set_shift(calculate_shift(image_seq_len))
+        self.scheduler.set_timesteps(batch.num_inference_steps,
+                                     device=device,
+                                     sigmas=sched_sigmas.tolist(),
+                                     timesteps=sched_timesteps.tolist())
+        batch.timesteps = self.scheduler.timesteps
+        return batch
+
+    @torch.no_grad()
+    def _compute_source_prior_tokens(self, batch: ForwardBatch, inputs) -> None:
+        image_grid_thw = inputs["image_grid_thw"]
+        num_condition_images = image_grid_thw.shape[0] - 1
+        source_grids = image_grid_thw[:num_condition_images]
+        image_features = self.vision_language_encoder.get_image_features(inputs["pixel_values"], source_grids)
+        image_feature_parts = getattr(image_features, "pooler_output", image_features)
+        embed = torch.cat(image_feature_parts, dim=0)
+        src_ids_d32 = self.vision_language_encoder.get_image_tokens(embed, source_grids)
+        split_sizes = source_grids.prod(dim=-1).tolist()
+        upsampled = [
+            _upsample_d32_to_d16(ids, int(grid[1]), int(grid[2])).squeeze(0)
+            for ids, grid in zip(torch.split(src_ids_d32, split_sizes), source_grids, strict=False)
+        ]
+        src_grids_up = source_grids.clone()
+        src_grids_up[:, 1] *= 2
+        src_grids_up[:, 2] *= 2
+        batch.extra["glm_prior_token_image_ids"] = torch.cat(upsampled, dim=0)
+        batch.extra["glm_source_image_grid_thw"] = src_grids_up
+
+    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("prompt", batch.prompt, V.string_not_empty)
+        return result
+
+    def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("prior_token_id", batch.prior_token_id, V.is_tensor)
+        return result

fastvideo/pipelines/basic/glm_image/stages/condition_encoding.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import torch
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.image_processor import ImageProcessor
+from fastvideo.models.dits.glm_image import GlmImageKVCache
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+
+_CONDITION_MULTIPLE_OF = 16  # vae_scale_factor (8) * DiT patch_size (2)
+
+
+class GlmImageConditionEncodingStage(PipelineStage):
+
+    def __init__(self, vae, transformer) -> None:
+        super().__init__()
+        self.vae = vae
+        self.transformer = transformer
+        self.image_processor = ImageProcessor(vae_scale_factor=_CONDITION_MULTIPLE_OF)
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if batch.pil_image is None:
+            return batch
+
+        device = get_local_torch_device()
+        dtype = torch.bfloat16
+        self.vae.to(device)
+
+        prior_ids = batch.extra["glm_prior_token_image_ids"].to(device)
+        if prior_ids.dim() == 1:
+            prior_ids = prior_ids.unsqueeze(0)
+
+        # Latent patch count must match the source prior tokens; mismatch is fatal.
+        src_grid = batch.extra["glm_source_image_grid_thw"][0]
+        cond_h = int(src_grid[1]) * _CONDITION_MULTIPLE_OF
+        cond_w = int(src_grid[2]) * _CONDITION_MULTIPLE_OF
+        cond_img = self.image_processor.preprocess(batch.pil_image, cond_h, cond_w).to(device=device,
+                                                                                       dtype=torch.float32)
+        latent = self.vae.encode(cond_img).latent_dist.mode()
+        # NOTE: at runtime self.vae.config is a diffusers FrozenDict with flat
+        # latents_mean/latents_std fields. Access the flat fields directly.
+        cfg = self.vae.config
+        mean = torch.tensor(cfg.latents_mean, device=device, dtype=torch.float32).view(1, -1, 1, 1)
+        std = torch.tensor(cfg.latents_std, device=device, dtype=torch.float32).view(1, -1, 1, 1)
+        latent = ((latent - mean) / std).to(dtype)
+
+        kv_caches = GlmImageKVCache(num_layers=self.transformer.num_layers)
+        empty_text = batch.prompt_embeds[0][:1, :0, :].to(device=device, dtype=dtype)
+        with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=batch):
+            self.transformer(
+                hidden_states=latent,
+                encoder_hidden_states=empty_text,
+                prior_token_id=prior_ids,
+                prior_token_drop=torch.zeros((prior_ids.shape[0], ), dtype=torch.bool, device=device),
+                timestep=torch.zeros((1, ), device=device),
+                target_size=torch.tensor([tuple(cond_img.shape[-2:])], device=device, dtype=torch.long),
+                crop_coords=torch.zeros((1, 2), device=device, dtype=torch.long),
+                kv_caches=kv_caches,
+                kv_caches_mode="write",
+            )
+        batch.extra["glm_kv_caches"] = kv_caches
+        return batch

fastvideo/pipelines/basic/glm_image/stages/decoding.py

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import torch
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.stages.decoding import DecodingStage
+from fastvideo.utils import PRECISION_TO_TYPE
+
+
+class GlmImageDecodingStage(DecodingStage):
+
+    @torch.no_grad()
+    def decode(self, latents: torch.Tensor, fastvideo_args: FastVideoArgs) -> torch.Tensor:
+        self.vae.to(get_local_torch_device())
+        latents = latents.to(get_local_torch_device())
+
+        vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
+        vae_autocast = (vae_dtype != torch.float32 and not fastvideo_args.disable_autocast)
+
+        latents = self._denormalize_latents(latents)
+        if latents.dim() == 5:
+            latents = latents.squeeze(2)
+
+        with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast):
+            if fastvideo_args.pipeline_config.vae_tiling:
+                self.vae.enable_tiling()
+            if not vae_autocast:
+                latents = latents.to(vae_dtype)
+            decoded = self.vae.decode(latents)
+
+        image = decoded.sample if hasattr(decoded, "sample") else decoded
+        image = (image / 2 + 0.5).clamp(0, 1)
+        return image.unsqueeze(2)

fastvideo/pipelines/basic/glm_image/stages/denoising.py

@@ -0,0 +1,192 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CFG convention (both denoise paths): row 0 conditional (positive), row 1 unconditional."""
+from __future__ import annotations
+
+import torch
+
+from fastvideo.attention.backends.sdpa import SDPAMetadata
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.denoising import DenoisingStage
+from fastvideo.pipelines.stages.validators import StageValidators as V
+from fastvideo.pipelines.stages.validators import VerificationResult
+from fastvideo.platforms import AttentionBackendEnum
+
+
+class GlmImageDenoisingStage(DenoisingStage):
+
+    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.min_dims(1)])
+        latents = getattr(batch, "latent", getattr(batch, "latents", None))
+        result.add_check("latents", latents, [V.is_tensor, V.with_dims(5)])
+        result.add_check("num_inference_steps", batch.num_inference_steps, V.positive_int)
+        result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty)
+        return result
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        device = get_local_torch_device()
+        dtype = torch.bfloat16
+        guidance_scale = batch.guidance_scale
+        do_cfg = guidance_scale > 1.0
+
+        latents = getattr(batch, "latent", getattr(batch, "latents", None))
+        if latents is None:
+            raise ValueError("No latents found in batch.")
+        if latents.dim() == 5:
+            latents = latents.squeeze(2)
+
+        prompt_embeds = batch.prompt_embeds[0]
+        text_attention_mask = getattr(batch, "attention_mask", None)
+        timesteps = batch.timesteps
+
+        patch_size = self.transformer.patch_size
+        _, _, h, w = latents.shape
+        image_seq_length = (h // patch_size) * (w // patch_size)
+        text_seq_length = prompt_embeds.shape[1] if prompt_embeds.dim() >= 2 else 0
+
+        first_block = self.transformer.transformer_blocks[0]
+        backend = getattr(first_block.attn1.attn, "backend", None)
+        sdpa = backend == AttentionBackendEnum.TORCH_SDPA and text_attention_mask is not None
+
+        kv_caches = batch.extra.get("glm_kv_caches")
+        if kv_caches is None:
+            self._denoise_t2i(batch, latents, prompt_embeds, text_attention_mask, timesteps, do_cfg, guidance_scale,
+                              text_seq_length, image_seq_length, sdpa, device, dtype)
+        else:
+            self._denoise_i2i(batch, latents, prompt_embeds, text_attention_mask, timesteps, do_cfg, guidance_scale,
+                              text_seq_length, image_seq_length, sdpa, kv_caches, device, dtype)
+        return batch
+
+    def _denoise_t2i(self, batch, latents, prompt_embeds, text_attention_mask, timesteps, do_cfg, guidance_scale,
+                     text_seq_length, image_seq_length, sdpa, device, dtype) -> None:
+        num_inference_steps = batch.num_inference_steps
+        bs = 2 if do_cfg else 1
+        target_size = torch.tensor([[batch.height, batch.width]], device=device, dtype=torch.long).repeat(bs, 1)
+        crop_coords = torch.zeros((bs, 2), device=device, dtype=torch.long)
+
+        prior_token_id = batch.prior_token_id
+        if do_cfg and prior_token_id.shape[0] == 1:
+            prior_token_id = prior_token_id.repeat(2, 1)
+        if do_cfg:
+            prior_token_drop = torch.tensor([False, True], device=device)
+        else:
+            prior_token_drop = getattr(batch, "prior_token_drop", torch.tensor([False], device=device))
+
+        attention_mask_kv = None
+        if sdpa:
+            if (text_attention_mask.shape[0] == 1 and bs > 1):
+                text_attention_mask = text_attention_mask.repeat(bs, 1)
+            mix_attn_mask = torch.ones((bs, text_seq_length + image_seq_length), device=device, dtype=torch.float32)
+            mix_attn_mask[:, :text_seq_length] = (text_attention_mask.float().to(device))
+            attention_mask_kv = (mix_attn_mask > 0).unsqueeze(1).unsqueeze(2)
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                latent_model_input = torch.cat([latents] * 2) if do_cfg else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t).to(dtype)
+                t_expand = t.expand(latent_model_input.shape[0]) - 1
+
+                attn_metadata = (SDPAMetadata(current_timestep=i, attn_mask=attention_mask_kv)
+                                 if attention_mask_kv is not None else None)
+
+                with torch.no_grad(), set_forward_context(current_timestep=i,
+                                                          attn_metadata=attn_metadata,
+                                                          forward_batch=batch):
+                    noise_pred = self.transformer(
+                        latent_model_input,
+                        prompt_embeds,
+                        prior_token_id,
+                        prior_token_drop,
+                        t_expand,
+                        target_size,
+                        crop_coords,
+                    )
+
+                    if do_cfg:
+                        noise_pred_cond, noise_pred_uncond = noise_pred.float().chunk(2)
+                        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
+
+                        guidance_rescale = getattr(batch, "guidance_rescale", 0.0)
+                        if guidance_rescale > 0.0:
+                            dims = list(range(1, noise_pred_cond.ndim))
+                            std_text = noise_pred_cond.std(dim=dims, keepdim=True)
+                            std_cfg = noise_pred.std(dim=dims, keepdim=True)
+                            rescaled = noise_pred * (std_text / std_cfg)
+                            noise_pred = (guidance_rescale * rescaled + (1 - guidance_rescale) * noise_pred)
+
+                    latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+                progress_bar.update()
+
+        batch.latents = latents.unsqueeze(2)
+
+    def _denoise_i2i(self, batch, latents, prompt_embeds, text_attention_mask, timesteps, do_cfg, guidance_scale,
+                     text_seq_length, image_seq_length, sdpa, kv_caches, device, dtype) -> None:
+        """Two separate transformer calls (cond reads the cache, uncond skips it):
+        the cache mode is one global flag with batch-1 k/v, so a 2-row CFG call
+        cannot express both."""
+        num_inference_steps = batch.num_inference_steps
+        target_size = torch.tensor([[batch.height, batch.width]], device=device, dtype=torch.long)
+        crop_coords = torch.zeros((1, 2), device=device, dtype=torch.long)
+        prior_token_id = batch.prior_token_id[:1]
+        drop_keep = torch.zeros((1, ), dtype=torch.bool, device=device)
+        drop_all = torch.ones((1, ), dtype=torch.bool, device=device)
+        cache_len = kv_caches[0].k_cache.shape[1] if kv_caches[0].k_cache is not None else 0
+
+        def _mask(row: int, with_cache: bool):
+            if not sdpa:
+                return None
+            prefix = cache_len if with_cache else 0
+            m = torch.ones((1, prefix + text_seq_length + image_seq_length), device=device, dtype=torch.float32)
+            m[:, prefix:prefix + text_seq_length] = text_attention_mask[row:row + 1].float().to(device)
+            return (m > 0).unsqueeze(1).unsqueeze(2)
+
+        cond_mask = _mask(0, with_cache=True)
+        uncond_mask = _mask(1, with_cache=False) if do_cfg else None
+        pos_embeds = prompt_embeds[:1]
+        neg_embeds = prompt_embeds[1:2] if do_cfg else None
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                latent_model_input = self.scheduler.scale_model_input(latents, t).to(dtype)
+                t_expand = t.expand(1) - 1
+
+                cond_meta = (SDPAMetadata(current_timestep=i, attn_mask=cond_mask) if cond_mask is not None else None)
+                with torch.no_grad(), set_forward_context(current_timestep=i,
+                                                          attn_metadata=cond_meta,
+                                                          forward_batch=batch):
+                    noise_pred = self.transformer(latent_model_input,
+                                                  pos_embeds,
+                                                  prior_token_id,
+                                                  drop_keep,
+                                                  t_expand,
+                                                  target_size,
+                                                  crop_coords,
+                                                  kv_caches=kv_caches,
+                                                  kv_caches_mode="read")
+
+                if do_cfg:
+                    uncond_meta = (SDPAMetadata(current_timestep=i, attn_mask=uncond_mask)
+                                   if uncond_mask is not None else None)
+                    with torch.no_grad(), set_forward_context(current_timestep=i,
+                                                              attn_metadata=uncond_meta,
+                                                              forward_batch=batch):
+                        noise_pred_uncond = self.transformer(latent_model_input,
+                                                             neg_embeds,
+                                                             prior_token_id,
+                                                             drop_all,
+                                                             t_expand,
+                                                             target_size,
+                                                             crop_coords,
+                                                             kv_caches=kv_caches,
+                                                             kv_caches_mode="skip")
+                    noise_pred = noise_pred_uncond.float() + guidance_scale * (noise_pred.float() -
+                                                                               noise_pred_uncond.float())
+
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+                progress_bar.update()
+
+        kv_caches.clear()
+        batch.latents = latents.unsqueeze(2)

fastvideo/pipelines/pipeline_batch_info.py

@@ -109,6 +109,10 @@ class ForwardBatch:
     max_sequence_length: int | None = None
     prompt_template: dict[str, Any] | None = None
     do_classifier_free_guidance: bool = False
+    # When True, ``guidance_scale`` is passed into models that use embedded guidance (e.g. FLUX)
+    # and must not imply classic dual-forward CFG. Use ``true_cfg_scale > 1`` for true CFG.
+    use_embedded_guidance: bool = False
+    true_cfg_scale: float = 1.0
 
     # Batch info
     batch_size: int | None = None
@@ -252,9 +256,12 @@ class ForwardBatch:
     def __post_init__(self):
         """Initialize dependent fields after dataclass initialization."""
 
-        # Enable CFG for standard guidance_scale and LTX-2 text CFG scales.
+        # LTX-2 text CFG scales; FLUX uses ``use_embedded_guidance`` so ``guidance_scale > 1`` alone
+        # does not enable classifier-free guidance.
         ltx2_text_cfg_enabled = (self.ltx2_cfg_scale_video != 1.0 or self.ltx2_cfg_scale_audio != 1.0)
-        if self.guidance_scale > 1.0 or ltx2_text_cfg_enabled:
+        if self.use_embedded_guidance:
+            self.do_classifier_free_guidance = (self.true_cfg_scale > 1.0) or ltx2_text_cfg_enabled
+        elif self.guidance_scale > 1.0 or ltx2_text_cfg_enabled:
             self.do_classifier_free_guidance = True
         if self.negative_prompt_embeds is None:
             self.negative_prompt_embeds = []

fastvideo/pipelines/stages/flux_stages.py

@@ -0,0 +1,424 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import inspect
+from typing import Any
+
+import torch
+from diffusers.utils.torch_utils import randn_tensor
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+from fastvideo.pipelines.stages.input_validation import InputValidationStage
+from fastvideo.pipelines.stages.timestep_preparation import TimestepPreparationStage
+from fastvideo.utils import PRECISION_TO_TYPE
+
+logger = init_logger(__name__)
+
+
+def _pack_latents(
+    latents: torch.Tensor,
+    batch_size: int,
+    num_channels_latents: int,
+    height: int,
+    width: int,
+) -> torch.Tensor:
+    """Diffusers ``_pack_latents`` for FLUX (2×2 spatial pack in latent space)."""
+    latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
+    latents = latents.permute(0, 2, 4, 1, 3, 5)
+    return latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4)
+
+
+def _unpack_latents(
+    latents: torch.Tensor,
+    batch_size: int,
+    num_channels_latents: int,
+    height: int,
+    width: int,
+) -> torch.Tensor:
+    """Inverse of ``_pack_latents``."""
+    latents = latents.reshape(batch_size, height // 2, width // 2, num_channels_latents, 2, 2)
+    latents = latents.permute(0, 3, 1, 4, 2, 5)
+    return latents.reshape(batch_size, num_channels_latents, height, width)
+
+
+def _prepare_latent_image_ids(
+    patch_height: int,
+    patch_width: int,
+    device: torch.device,
+    dtype: torch.dtype = torch.long,
+) -> torch.Tensor:
+    """Match Diffusers ``FluxPipeline._prepare_latent_image_ids`` (no batch dim)."""
+    latent_image_ids = torch.zeros(patch_height, patch_width, 3, device=device, dtype=torch.float32)
+    latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(patch_height, device=device)[:, None]
+    latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(patch_width, device=device)[None, :]
+    h, w, c = latent_image_ids.shape
+    return latent_image_ids.reshape(h * w, c).to(dtype=dtype)
+
+
+class FluxInputValidationStage(InputValidationStage):
+    """Require height/width divisible by 16 (VAE scale × 2 for FLUX packing)."""
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        if (batch.height is not None and batch.width is not None and (batch.height % 16 != 0 or batch.width % 16 != 0)):
+            raise ValueError("FLUX expects height and width divisible by 16 "
+                             f"(VAE latent grid × 2× packing); got {batch.height}×{batch.width}.")
+        return super().forward(batch, fastvideo_args)
+
+
+class FluxConditioningStage(PipelineStage):
+    """Build CLIP pooled + T5 sequence + ``text_ids`` (and optional negative for true CFG)."""
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if len(batch.prompt_embeds) < 2:
+            raise ValueError("FluxConditioningStage expects 2 prompt_embeds (CLIP pooled, T5 sequence), "
+                             f"got {len(batch.prompt_embeds)}")
+
+        device = get_local_torch_device()
+        target_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
+
+        pooled = batch.prompt_embeds[0].to(device=device, dtype=target_dtype)
+        enc = batch.prompt_embeds[1].to(device=device, dtype=target_dtype)
+        seq_len = enc.shape[1]
+        text_ids = torch.zeros(seq_len, 3, device=device, dtype=torch.long)
+
+        batch.extra["flux_pooled_projections"] = pooled
+        batch.extra["flux_encoder_hidden_states"] = enc
+        batch.extra["flux_text_ids"] = text_ids
+
+        if batch.do_classifier_free_guidance:
+            if not batch.negative_prompt_embeds or len(batch.negative_prompt_embeds) < 2:
+                raise ValueError("True CFG requires two negative_prompt_embeds (CLIP, T5).")
+            neg_pooled = batch.negative_prompt_embeds[0].to(device=device, dtype=target_dtype)
+            neg_enc = batch.negative_prompt_embeds[1].to(device=device, dtype=target_dtype)
+            batch.extra["flux_negative_pooled_projections"] = neg_pooled
+            batch.extra["flux_negative_encoder_hidden_states"] = neg_enc
+
+        return batch
+
+
+class FluxTimestepPreparationStage(TimestepPreparationStage):
+    """Flow Match with resolution-dependent ``mu`` from packed image sequence length."""
+
+    @staticmethod
+    def _calculate_mu(
+        image_seq_len: int,
+        base_seq_len: int = 256,
+        max_seq_len: int = 4096,
+        base_shift: float = 0.5,
+        max_shift: float = 1.15,
+    ) -> float:
+        m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+        b = base_shift - m * base_seq_len
+        return float(image_seq_len) * m + b
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        sig = inspect.signature(self.scheduler.set_timesteps)
+        if "mu" not in sig.parameters:
+            logger.warning(
+                "FLUX timestep prep: scheduler %s.set_timesteps does not accept 'mu'; falling back to the base "
+                "timestep schedule. FLUX expects a FlowMatchEulerDiscreteScheduler with resolution-dependent "
+                "dynamic shifting — output quality may degrade.",
+                type(self.scheduler).__name__)
+            return super().forward(batch, fastvideo_args)
+
+        cfg = getattr(self.scheduler, "config", None)
+        use_dynamic = bool(getattr(cfg, "use_dynamic_shifting", False)) if cfg is not None else False
+        if not use_dynamic:
+            logger.warning(
+                "FLUX timestep prep: scheduler has use_dynamic_shifting=False; falling back to the base timestep "
+                "schedule and skipping the resolution-dependent 'mu' shift. FLUX requires dynamic shifting for "
+                "correct timesteps — output quality may degrade.")
+            return super().forward(batch, fastvideo_args)
+
+        if batch.height is None or batch.width is None:
+            raise ValueError("height/width must be set before FluxTimestepPreparationStage")
+
+        vae_arch = fastvideo_args.pipeline_config.vae_config.arch_config
+        spatial_ratio = int(getattr(vae_arch, "spatial_compression_ratio", 8))
+        h_lat = batch.height // spatial_ratio
+        w_lat = batch.width // spatial_ratio
+        if h_lat % 2 != 0 or w_lat % 2 != 0:
+            raise ValueError(
+                f"Latent spatial dims must be even for FLUX packing; got {h_lat}×{w_lat} from {batch.height}×{batch.width}."
+            )
+        image_seq_len = (h_lat // 2) * (w_lat // 2)
+
+        base_seq_len = int(getattr(cfg, "base_image_seq_len", 256))
+        max_seq_len = int(getattr(cfg, "max_image_seq_len", 4096))
+        base_shift = float(getattr(cfg, "base_shift", 0.5))
+        max_shift = float(getattr(cfg, "max_shift", 1.15))
+
+        device = get_local_torch_device()
+        mu = self._calculate_mu(
+            image_seq_len=image_seq_len,
+            base_seq_len=base_seq_len,
+            max_seq_len=max_seq_len,
+            base_shift=base_shift,
+            max_shift=max_shift,
+        )
+        self.scheduler.set_timesteps(batch.num_inference_steps, device=device, mu=mu)
+        batch.timesteps = self.scheduler.timesteps
+        return batch
+
+
+class FluxLatentPreparationStage(PipelineStage):
+
+    def __init__(self, scheduler) -> None:
+        super().__init__()
+        self.scheduler = scheduler
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if batch.height is None or batch.width is None:
+            raise ValueError("height/width required for FluxLatentPreparationStage")
+
+        if isinstance(batch.prompt, list):
+            batch_size = len(batch.prompt)
+        elif batch.prompt is not None:
+            batch_size = 1
+        else:
+            if not batch.prompt_embeds:
+                raise ValueError("prompt or prompt_embeds must be provided")
+            batch_size = batch.prompt_embeds[0].shape[0]
+
+        batch_size *= batch.num_videos_per_prompt
+
+        if isinstance(batch.generator, list) and len(batch.generator) != batch_size:
+            raise ValueError(f"generator list length {len(batch.generator)} does not match batch_size {batch_size}")
+
+        arch = fastvideo_args.pipeline_config.dit_config.arch_config
+        in_channels = int(getattr(arch, "in_channels", 64))
+        num_channels_latents = in_channels // 4
+
+        vae_arch = fastvideo_args.pipeline_config.vae_config.arch_config
+        spatial_ratio = int(getattr(vae_arch, "spatial_compression_ratio", 8))
+
+        h_lat = batch.height // spatial_ratio
+        w_lat = batch.width // spatial_ratio
+
+        dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
+        device = get_local_torch_device()
+
+        shape = (batch_size, num_channels_latents, h_lat, w_lat)
+        latents = batch.latents
+        if latents is None:
+            latents = randn_tensor(shape, generator=batch.generator, device=device, dtype=dtype)
+            if hasattr(self.scheduler, "init_noise_sigma"):
+                latents = latents * self.scheduler.init_noise_sigma
+        else:
+            latents = latents.to(device=device, dtype=dtype)
+            if latents.shape != shape:
+                raise ValueError(f"Expected latents shape {shape}, got {tuple(latents.shape)}")
+            if hasattr(self.scheduler, "init_noise_sigma"):
+                latents = latents * self.scheduler.init_noise_sigma
+
+        packed = _pack_latents(latents, batch_size, num_channels_latents, h_lat, w_lat)
+
+        patch_h, patch_w = h_lat // 2, w_lat // 2
+        img_ids = _prepare_latent_image_ids(patch_h, patch_w, device, dtype=torch.long)
+
+        batch.latents = packed
+        batch.raw_latent_shape = shape
+        batch.extra["flux_h_lat"] = h_lat
+        batch.extra["flux_w_lat"] = w_lat
+        batch.extra["flux_num_channels_latents"] = num_channels_latents
+        batch.extra["flux_latent_image_ids"] = img_ids
+
+        return batch
+
+
+class FluxDenoisingStage(PipelineStage):
+
+    def __init__(self, transformer, scheduler) -> None:
+        super().__init__()
+        self.transformer = transformer
+        self.scheduler = scheduler
+
+    @staticmethod
+    def _step_kwargs(scheduler_step, batch: ForwardBatch) -> dict[str, Any]:
+        kwargs: dict[str, Any] = {}
+        sig = inspect.signature(scheduler_step)
+        if "generator" in sig.parameters:
+            gen = batch.generator[0] if isinstance(batch.generator, list) else batch.generator
+            kwargs["generator"] = gen
+        return kwargs
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if batch.timesteps is None:
+            raise ValueError("timesteps must be set before FluxDenoisingStage")
+        if batch.latents is None:
+            raise ValueError("latents must be set before FluxDenoisingStage")
+
+        packed = batch.latents
+        timesteps = batch.timesteps
+
+        pooled = batch.extra["flux_pooled_projections"]
+        enc = batch.extra["flux_encoder_hidden_states"]
+        txt_ids = batch.extra["flux_text_ids"]
+        img_ids = batch.extra["flux_latent_image_ids"]
+
+        neg_pooled = batch.extra.get("flux_negative_pooled_projections")
+        neg_enc = batch.extra.get("flux_negative_encoder_hidden_states")
+
+        true_cfg_scale = float(batch.true_cfg_scale)
+        use_true_cfg = batch.do_classifier_free_guidance and true_cfg_scale > 1.0
+
+        # Prefer the loaded transformer's arch (HF ``guidance_embeds``), not static pipeline defaults.
+        tr_arch = self.transformer.fastvideo_config.arch_config
+        guidance_embeds = bool(getattr(tr_arch, "guidance_embeds", False))
+
+        device = get_local_torch_device()
+        target_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
+        autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast
+
+        bs = packed.shape[0]
+        if guidance_embeds:
+            guidance = torch.full((bs, ), float(batch.guidance_scale), device=device, dtype=torch.float32)
+        else:
+            guidance = None
+
+        step_extras = self._step_kwargs(self.scheduler.step, batch)
+
+        for t in timesteps:
+            t_scalar = t
+            if not isinstance(t_scalar, torch.Tensor):
+                t_scalar = torch.tensor([t_scalar], device=device, dtype=torch.float32)
+            t_scalar = t_scalar.to(device=device, dtype=torch.float32)
+
+            timestep_model = t_scalar.expand(bs).float() / 1000.0
+            timestep_model = timestep_model.to(dtype=target_dtype)
+
+            ts_ctx = int(t_scalar.reshape(-1)[0].item())
+            with (
+                    torch.autocast(
+                        device_type="cuda",
+                        dtype=target_dtype,
+                        enabled=autocast_enabled and device.type == "cuda",
+                    ),
+                    set_forward_context(
+                        current_timestep=ts_ctx,
+                        attn_metadata=None,
+                        forward_batch=batch,
+                    ),
+            ):
+                if use_true_cfg:
+                    assert neg_enc is not None and neg_pooled is not None
+                    n_neg = self.transformer(
+                        hidden_states=packed,
+                        encoder_hidden_states=neg_enc,
+                        pooled_projections=neg_pooled,
+                        timestep=timestep_model,
+                        guidance=guidance,
+                        txt_ids=txt_ids,
+                        img_ids=img_ids,
+                        return_dict=False,
+                    )[0]
+                    n_pos = self.transformer(
+                        hidden_states=packed,
+                        encoder_hidden_states=enc,
+                        pooled_projections=pooled,
+                        timestep=timestep_model,
+                        guidance=guidance,
+                        txt_ids=txt_ids,
+                        img_ids=img_ids,
+                        return_dict=False,
+                    )[0]
+                    noise_pred = n_neg + true_cfg_scale * (n_pos - n_neg)
+                else:
+                    noise_pred = self.transformer(
+                        hidden_states=packed,
+                        encoder_hidden_states=enc,
+                        pooled_projections=pooled,
+                        timestep=timestep_model,
+                        guidance=guidance,
+                        txt_ids=txt_ids,
+                        img_ids=img_ids,
+                        return_dict=False,
+                    )[0]
+
+            packed = self.scheduler.step(
+                noise_pred,
+                t_scalar,
+                packed,
+                return_dict=False,
+                **step_extras,
+            )[0]
+
+        batch.latents = packed
+        return batch
+
+
+class FluxDecodingStage(PipelineStage):
+    """Unpack latents, apply VAE scaling/shift, decode to pixels (5D output ``B×3×1×H×W``)."""
+
+    def __init__(self, vae) -> None:
+        super().__init__()
+        self.vae = vae
+
+    @staticmethod
+    def _denormalize_latents(latents: torch.Tensor, vae: Any) -> torch.Tensor:
+        cfg = getattr(vae, "config", None)
+        sf = getattr(cfg, "scaling_factor", None) if cfg is not None else None
+        sh = getattr(cfg, "shift_factor", None) if cfg is not None else None
+        if sf is None and hasattr(vae, "scaling_factor"):
+            sf = vae.scaling_factor
+        if sh is None and hasattr(vae, "shift_factor"):
+            sh = vae.shift_factor
+        if sf is not None:
+            latents = latents / (sf.to(latents.device, latents.dtype) if isinstance(sf, torch.Tensor) else sf)
+            if sh is not None:
+                latents = latents + (sh.to(latents.device, latents.dtype) if isinstance(sh, torch.Tensor) else sh)
+        return latents
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        packed = batch.latents
+        if packed is None:
+            raise ValueError("latents must be set before FluxDecodingStage")
+
+        h_lat = int(batch.extra["flux_h_lat"])
+        w_lat = int(batch.extra["flux_w_lat"])
+        num_ch = int(batch.extra["flux_num_channels_latents"])
+        raw_shape = batch.raw_latent_shape
+        if raw_shape is None:
+            raise ValueError("raw_latent_shape missing; FluxLatentPreparationStage must run first.")
+        batch_size = int(raw_shape[0])
+
+        infer_device = get_local_torch_device()
+        packed = packed.to(infer_device)
+
+        latents_4d = _unpack_latents(packed, batch_size, num_ch, h_lat, w_lat)
+        latents_4d = self._denormalize_latents(latents_4d, self.vae)
+
+        vae_device = next(self.vae.parameters()).device
+        latents_4d = latents_4d.to(device=vae_device)
+
+        vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
+        autocast_enabled = (vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
+        use_cuda_autocast = autocast_enabled and vae_device.type == "cuda"
+
+        with torch.autocast(
+                device_type="cuda",
+                dtype=vae_dtype,
+                enabled=use_cuda_autocast,
+        ):
+            if not autocast_enabled:
+                latents_4d = latents_4d.to(dtype=vae_dtype)
+            dec = self.vae.decode(latents_4d)
+            image = dec.sample if hasattr(dec, "sample") else dec[0]
+
+        image = (image / 2 + 0.5).clamp(0, 1)
+        batch.output = image.unsqueeze(2).detach().float().cpu()
+        return batch

fastvideo/registry.py

@@ -58,11 +58,14 @@ from fastvideo.configs.pipelines.wan import (
     WanT2V480PConfig,
     WanT2V720PConfig,
 )
+from fastvideo.configs.pipelines.glm_image import GlmImageConfig
+from fastvideo.configs.pipelines.flux import FluxPipelineConfig
 from fastvideo.configs.pipelines.sd35 import SD35Config
 from fastvideo.configs.pipelines.stable_audio import (StableAudioOpenSmallConfig, StableAudioT2AConfig)
 from fastvideo.api.sampling_param import SamplingParam
 from fastvideo.api.matrixgame2 import MatrixGame2SamplingParam
 from fastvideo.api.matrixgame3 import MatrixGame3SamplingParam
+from fastvideo.api.flux import FluxSamplingParam
 
 from fastvideo.fastvideo_args import WorkloadType
 from fastvideo.logger import init_logger
@@ -1074,6 +1077,33 @@ def _register_configs() -> None:
         default_preset="sd35_medium",
     )
 
+    # GLM-Image
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=GlmImageConfig,
+        hf_model_paths=[
+            "zai-org/GLM-Image",
+        ],
+        model_detectors=[lambda path: "glmimage" in path.lower() or "glm-image" in path.lower()],
+        workload_types=(WorkloadType.T2I, ),
+        model_family="glm_image",
+    )
+
+    # FLUX.1-dev (Diffusers)
+    register_configs(
+        sampling_param_cls=FluxSamplingParam,
+        pipeline_config_cls=FluxPipelineConfig,
+        workload_types=(WorkloadType.T2I, ),
+        hf_model_paths=[
+            "black-forest-labs/FLUX.1-dev",
+        ],
+        model_detectors=[
+            lambda path: "fluxpipeline" in path,
+            lambda path: "flux.1-dev" in path or "flux_1_dev" in path,
+            lambda path: "/flux/" in path or path.endswith("/flux"),
+        ],
+    )
+
 
 # --- Part 3: Main Resolver ---
 

fastvideo/tests/api/test_parser.py

@@ -183,6 +183,7 @@ def test_load_run_config_supports_yaml_roundtrip(tmp_path) -> None:
                 "guidance_scale_2": None,
                 "guidance_rescale": 0.0,
                 "true_cfg_scale": None,
+                "use_embedded_guidance": None,
                 "boundary_ratio": None,
                 "sigmas": None,
             },

fastvideo/tests/attention/test_sdpa_head_sizes.py

@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU tests: SDPA must not restrict attention head sizes.
+
+``torch.nn.functional.scaled_dot_product_attention`` is head-size
+agnostic, but ``SDPABackend.get_supported_head_sizes()`` used to declare
+a list copied from FlashAttentionBackend that omitted e.g. 80 (the head
+size of CLIP-L vision encoders such as matrixgame2's). That made
+capability validation flag SDPA as incompatible with head-80 layers even
+though it handles them fine.
+
+CPU-only; no GPU or flash-attn dependency.
+"""
+
+import torch
+
+from fastvideo.attention.backends.sdpa import SDPABackend, SDPAImpl, SDPAMetadata
+from fastvideo.platforms.cuda import CudaPlatformBase
+from fastvideo.platforms.interface import AttentionBackendEnum
+
+
+def test_sdpa_head_sizes_unrestricted() -> None:
+    """SDPA declares no head-size restriction (None), so 80 is allowed."""
+    assert SDPABackend.get_supported_head_sizes() is None
+    # Once the capability-validation API (#1494) lands, head-80 must pass.
+    validate = getattr(SDPABackend, "validate_compatibility", None)
+    if validate is not None:
+        assert validate(head_size=80, dtype=torch.bfloat16) is None
+
+
+def test_explicit_torch_sdpa_pin_head_80_resolves() -> None:
+    """An explicit TORCH_SDPA pin on a head-80 layer resolves to SDPABackend."""
+    cls_str = CudaPlatformBase.get_attn_backend_cls(AttentionBackendEnum.TORCH_SDPA, head_size=80, dtype=torch.bfloat16)
+    assert cls_str == "fastvideo.attention.backends.sdpa.SDPABackend"
+
+
+def test_sdpa_impl_computes_head_80_attention() -> None:
+    """SDPAImpl produces correct output for head_size=80 (not in the old list)."""
+    bs, seq, heads, head_dim = 2, 8, 2, 80
+    softmax_scale = head_dim**-0.5
+    impl = SDPAImpl(num_heads=heads, head_size=head_dim, causal=False, softmax_scale=softmax_scale)
+    torch.manual_seed(0)
+    q, k, v = (torch.randn(bs, seq, heads, head_dim) for _ in range(3))
+
+    out = impl.forward(q, k, v, SDPAMetadata(current_timestep=0))
+
+    assert out.shape == (bs, seq, heads, head_dim)
+    qt, kt, vt = (t.transpose(1, 2) for t in (q, k, v))
+    ref = torch.softmax(qt @ kt.transpose(-2, -1) * softmax_scale, dim=-1) @ vt
+    torch.testing.assert_close(out, ref.transpose(1, 2), rtol=1e-4, atol=1e-4)

fastvideo/tests/ops/test_rotary_pos_embed_cache.py

@@ -0,0 +1,190 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Unit tests for the get_rotary_pos_embed memoization cache."""
+
+import pytest
+import torch
+
+from fastvideo.layers.rotary_embedding import (
+    _ROTARY_POS_EMBED_CACHE,
+    _ROTARY_POS_EMBED_CACHE_MAXSIZE,
+    get_rotary_pos_embed,
+)
+
+
+def _rope_dim_list(hidden_size: int, heads_num: int) -> list[int]:
+    """Return the default 3-axis rope_dim_list used by the video DiTs."""
+    d = hidden_size // heads_num
+    return [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
+
+
+def _call(
+    rope_sizes=(21, 30, 52),
+    hidden_size=1536,
+    heads_num=12,
+    rope_dim_list="default",
+    rope_theta=10000.0,
+    dtype=torch.float64,
+    start_frame=0,
+    use_real=True,
+    **kwargs,
+):
+    """Thin wrapper around get_rotary_pos_embed with DiT-like defaults."""
+    if rope_dim_list == "default":
+        rope_dim_list = _rope_dim_list(hidden_size, heads_num)
+    return get_rotary_pos_embed(
+        rope_sizes,
+        hidden_size,
+        heads_num,
+        rope_dim_list,
+        rope_theta,
+        dtype=dtype,
+        start_frame=start_frame,
+        use_real=use_real,
+        **kwargs,
+    )
+
+
+@pytest.fixture(autouse=True)
+def _clear_cache():
+    """Isolate every test by clearing the module-level cache around it."""
+    _ROTARY_POS_EMBED_CACHE.clear()
+    yield
+    _ROTARY_POS_EMBED_CACHE.clear()
+
+
+def test_repeated_call_hits_cache():
+    """A second identical call returns the exact same tensor objects."""
+    cos1, sin1 = _call()
+    cos2, sin2 = _call()
+    assert cos1 is cos2 and sin1 is sin2
+    assert len(_ROTARY_POS_EMBED_CACHE) == 1
+
+
+def test_many_identical_calls_keep_single_entry():
+    """Many identical calls never grow the cache beyond one entry."""
+    for _ in range(10):
+        _call()
+    assert len(_ROTARY_POS_EMBED_CACHE) == 1
+
+
+@pytest.mark.parametrize(
+    "rope_sizes,hidden_size,heads_num,dtype,use_real",
+    [
+        ((21, 30, 52), 1536, 12, torch.float64, True),
+        ((21, 45, 80), 5120, 40, torch.float64, True),
+        ((1, 16, 16), 1536, 12, torch.float32, True),
+        ((4, 8, 8), 1536, 12, torch.float64, False),
+    ],
+)
+def test_cached_matches_fresh_recompute(rope_sizes, hidden_size, heads_num,
+                                        dtype, use_real):
+    """Cached tensors are bitwise-equal to a fresh uncached recompute."""
+    cos_cached, sin_cached = _call(rope_sizes=rope_sizes,
+                                   hidden_size=hidden_size,
+                                   heads_num=heads_num,
+                                   dtype=dtype,
+                                   use_real=use_real)
+    _ROTARY_POS_EMBED_CACHE.clear()
+    cos_fresh, sin_fresh = _call(rope_sizes=rope_sizes,
+                                 hidden_size=hidden_size,
+                                 heads_num=heads_num,
+                                 dtype=dtype,
+                                 use_real=use_real)
+    assert torch.equal(cos_cached, cos_fresh)
+    assert torch.equal(sin_cached, sin_fresh)
+
+
+@pytest.mark.parametrize(
+    "kwargs_a,kwargs_b",
+    [
+        ({"rope_sizes": (21, 30, 52)}, {"rope_sizes": (21, 45, 80)}),
+        ({"dtype": torch.float64}, {"dtype": torch.float32}),
+        ({"use_real": True}, {"use_real": False}),
+        ({"start_frame": 0}, {"start_frame": 3}),
+        ({"rope_theta": 10000.0}, {"rope_theta": 5000.0}),
+        ({"shard_dim": 0}, {"shard_dim": 1}),
+    ],
+)
+def test_distinct_args_create_distinct_entries(kwargs_a, kwargs_b):
+    """Any output-affecting argument difference yields a separate cache entry."""
+    _call(**kwargs_a)
+    _call(**kwargs_b)
+    assert len(_ROTARY_POS_EMBED_CACHE) == 2
+
+
+def test_none_rope_dim_list_shares_key_with_equivalent_list():
+    """None rope_dim_list and its derived explicit list map to one entry."""
+    # head_dim must be divisible by 3 for the None branch to stay valid.
+    hidden_size, heads_num = 1536, 16  # head_dim == 96 -> [32, 32, 32]
+    _call(rope_dim_list=None, hidden_size=hidden_size, heads_num=heads_num)
+    before = len(_ROTARY_POS_EMBED_CACHE)
+    _call(rope_dim_list=[32, 32, 32], hidden_size=hidden_size, heads_num=heads_num)
+    assert len(_ROTARY_POS_EMBED_CACHE) == before == 1
+
+
+def test_use_real_controls_last_dim():
+    """use_real=True spans full head_dim; use_real=False spans half."""
+    cos_full, _ = _call(use_real=True)
+    cos_half, _ = _call(use_real=False)
+    assert cos_full.shape[-1] == 128
+    assert cos_half.shape[-1] == 64
+
+
+@pytest.mark.parametrize("rope_sizes", [(1, 1, 1), (1, 30, 52), (21, 1, 1)])
+def test_degenerate_grid_shapes(rope_sizes):
+    """Degenerate single-element axes still produce a correctly sized table."""
+    cos, sin = _call(rope_sizes=rope_sizes)
+    expected = rope_sizes[0] * rope_sizes[1] * rope_sizes[2]
+    assert cos.shape[0] == expected
+    assert sin.shape[0] == expected
+
+
+def test_scalar_and_list_factors_are_hashable_and_distinct():
+    """List-valued rescale factors are hashable and keyed apart from scalars."""
+    _call(theta_rescale_factor=1.0)
+    _call(theta_rescale_factor=[1.0, 1.0, 1.0])
+    assert len(_ROTARY_POS_EMBED_CACHE) == 2
+
+
+def test_caller_device_copy_does_not_corrupt_cache():
+    """The .to()/.float() copy callers perform must not mutate cached tensors."""
+    cos, _ = _call()
+    snapshot = cos.clone()
+    _ = cos.to("cpu").float()
+    cos_again, _ = _call()
+    assert torch.equal(cos_again, snapshot)
+
+
+def test_start_frame_offsets_values():
+    """A non-zero start_frame shifts the temporal positions, changing output."""
+    cos0, _ = _call(start_frame=0)
+    cos3, _ = _call(start_frame=3)
+    assert not torch.equal(cos0, cos3)
+    assert len(_ROTARY_POS_EMBED_CACHE) == 2
+
+
+def test_cache_is_bounded_and_evicts_oldest():
+    """The cache caps at the max size and evicts the oldest entry first."""
+    # Tiny grids keep this lightweight; each start_frame is a distinct key.
+    overshoot = _ROTARY_POS_EMBED_CACHE_MAXSIZE + 4
+    for frame in range(overshoot):
+        _call(rope_sizes=(2, 2, 2), start_frame=frame)
+        assert len(_ROTARY_POS_EMBED_CACHE) <= _ROTARY_POS_EMBED_CACHE_MAXSIZE
+    assert len(_ROTARY_POS_EMBED_CACHE) == _ROTARY_POS_EMBED_CACHE_MAXSIZE
+    # The earliest-inserted frames must have been evicted; the latest survive.
+    surviving = {key[-2] for key in _ROTARY_POS_EMBED_CACHE}  # start_frame slot
+    assert overshoot - 1 in surviving
+    assert 0 not in surviving
+
+
+def test_cache_hit_refreshes_recency():
+    """Re-accessing an entry protects it from eviction over an untouched one."""
+    _call(rope_sizes=(2, 2, 2), start_frame=0)  # entry we will keep hot
+    for frame in range(1, _ROTARY_POS_EMBED_CACHE_MAXSIZE):
+        _call(rope_sizes=(2, 2, 2), start_frame=frame)
+    assert len(_ROTARY_POS_EMBED_CACHE) == _ROTARY_POS_EMBED_CACHE_MAXSIZE
+    _call(rope_sizes=(2, 2, 2), start_frame=0)   # hit -> frame 0 becomes most recent
+    _call(rope_sizes=(2, 2, 2), start_frame=99)  # miss -> evicts now-oldest (frame 1)
+    surviving = {key[-2] for key in _ROTARY_POS_EMBED_CACHE}
+    assert 0 in surviving
+    assert 1 not in surviving

fastvideo/tests/performance/compare_baseline.py

@@ -89,6 +89,14 @@ COMPARISON_IDENTITY_KEYS = (
     "hardware_profile_id",
     "software_profile_id",
 )
+# Comparator verdicts (issue #1532), stored as record["comparator_status"].
+STATUS_PASS = "PASS"
+STATUS_REGRESSION = "REGRESSION"
+STATUS_CALIBRATION_NEEDED = "CALIBRATION_NEEDED"
+STATUS_RECIPE_MISMATCH = "RECIPE_MISMATCH"
+STATUS_INFRA_ERROR = "INFRA_ERROR"
+# Reserved for the promoted-baseline workflow; never emitted by this comparator.
+STATUS_QUALITY_BLOCKED = "QUALITY_BLOCKED"
 
 
 def _should_persist_tracking() -> bool:
@@ -343,6 +351,90 @@ def _check_regressions(
     return failures
 
 
+def _recipe_mismatch_fingerprints(
+    record: dict[str, Any],
+    identity_filters: dict[str, str],
+) -> list[str]:
+    """Return baseline recipe fingerprints for the same variant cohort.
+
+    Non-empty means the same (workload, variant, version, hardware, software)
+    cohort has baseline history under a DIFFERENT recipe fingerprint: the
+    recipe changed without being represented as a new variant.
+    """
+    if not identity_filters:
+        return []
+    variant_filters = {key: value for key, value in identity_filters.items() if key != "recipe_fingerprint"}
+    same_variant = load_records_for_model(
+        TRACKING_ROOT,
+        record["model_id"],
+        record["gpu_type"],
+        **variant_filters,
+        successful_only=True,
+        baseline_eligible_only=True,
+    )
+    fingerprints = {str(r.get("recipe_fingerprint")) for r in same_variant}
+    return sorted(fingerprints - {identity_filters["recipe_fingerprint"]})
+
+
+def _compare_record(
+    record: dict[str, Any],
+    metric_policies: tuple[MetricPolicy, ...],
+) -> tuple[list[str], list[dict[str, Any]]]:
+    """Compare one normalized record against its exact-identity cohort.
+
+    Sets ``record["comparator_status"]`` (and ``record["baseline_status"]``
+    for cohort bookkeeping) and returns ``(failures, baseline_records)``.
+    """
+    model = record.get("model_id", "unknown")
+    try:
+        identity_filters = _comparison_identity_filters(record)
+        baseline_records = load_records_for_model(
+            TRACKING_ROOT,
+            record["model_id"],
+            record["gpu_type"],
+            **identity_filters,
+            last_n=5,
+            successful_only=True,
+            baseline_eligible_only=True,
+        )
+        if baseline_records:
+            record["baseline_status"] = "compared"
+            failures = _check_regressions(record, baseline_records, metric_policies)
+            record["comparator_status"] = STATUS_REGRESSION if failures else STATUS_PASS
+            return failures, baseline_records
+
+        mismatched = _recipe_mismatch_fingerprints(record, identity_filters)
+    except Exception as exc:
+        record["comparator_status"] = STATUS_INFRA_ERROR
+        return [f"{model} baseline comparison hit an infra error: {exc}"], []
+
+    if mismatched:
+        record["comparator_status"] = STATUS_RECIPE_MISMATCH
+        return [
+            f"{model} recipe fingerprint {record.get('recipe_fingerprint')} does not "
+            f"match baseline fingerprint(s) {', '.join(mismatched)} for variant "
+            f"{record.get('variant_id')}. Represent recipe changes as a new "
+            "variant_id, or reseed the baseline."
+        ], []
+
+    # No comparable baseline anywhere: a brand-new cohort. Make that loud and
+    # machine-readable instead of an indistinguishable pass, so a cohort shift
+    # (intended or accidental, e.g. an identity-field change) never silently
+    # blinds the comparison — and never silently seeds a passing baseline.
+    record["baseline_status"] = "initialized_new_cohort"
+    record["comparator_status"] = STATUS_CALIBRATION_NEEDED
+    print("=" * 72)
+    print(f"CALIBRATION_NEEDED: no comparable baseline for {model} on "
+          f"{record.get('gpu_type', 'unknown')}"
+          f"{_format_identity_filters(identity_filters)}")
+    print("Regression gating is INACTIVE for this record and it will NOT "
+          "seed a baseline. Seed the cohort explicitly via the reseed "
+          "workflow. If this cohort shift is unexpected, check the identity "
+          "fields above.")
+    print("=" * 72)
+    return [], []
+
+
 def _compact_value(value: float | None, precision: int = 3) -> str:
     if value is None:
         return "n/a"
@@ -394,6 +486,7 @@ def _build_summary_row(
     return {
         "model_id": record["model_id"],
         "gpu_type": record["gpu_type"],
+        "comparator_status": record.get("comparator_status", STATUS_PASS),
         "baseline_n": len(baseline_records),
         "metrics": metric_values,
         "worst_regression_pct": worst_regression_pct,
@@ -435,7 +528,10 @@ def _build_markdown_summary(
             else "none"
         )
         failing_metrics = ", ".join(row["failing_metrics"]) if row["failing_metrics"] else "none"
-        status = "FAIL" if row["failed"] else "PASS"
+        status = row["comparator_status"]
+        if row["failed"] and status == STATUS_PASS:
+            # Baseline comparison passed but the fixed-threshold phase failed.
+            status = "FAIL"
 
         lines.append(f"| {row['model_id']} | {row['gpu_type']} | "
                      f"{row['baseline_n']} | "
@@ -501,46 +597,26 @@ def main() -> int:
         if identity_filters is None:
             # Records without the full v2 identity block skip rolling-baseline
             # comparison entirely: only the static thresholds gate them and
-            # they never become baseline eligible.
+            # they never become baseline eligible. Nothing was compared and
+            # nothing failed, so the comparator verdict is PASS;
+            # baseline_status keeps the skip machine-readable.
             record["baseline_status"] = "skipped_missing_identity"
+            record["comparator_status"] = STATUS_PASS
             failures: list[str] = []
         else:
-            baseline_records = load_records_for_model(
-                TRACKING_ROOT,
-                record["model_id"],
-                record["gpu_type"],
-                **identity_filters,
-                last_n=5,
-                successful_only=True,
-                baseline_eligible_only=True,
-            )
-
-            if not baseline_records:
-                # A brand-new cohort has NO regression gating until history
-                # accumulates — make that loud and machine-readable instead of
-                # an indistinguishable pass, so a cohort shift (intended or
-                # accidental, e.g. an identity-field change) never silently
-                # blinds the comparison.
-                record["baseline_status"] = "initialized_new_cohort"
-                print("=" * 72)
-                print(f"WARNING: NO BASELINE — initializing a NEW cohort for "
-                      f"{record['model_id']} on {record['gpu_type']}"
-                      f"{_format_identity_filters(identity_filters)}")
-                print("Regression gating is INACTIVE for this cohort until "
-                      "baseline history accumulates. If this cohort shift is "
-                      "unexpected, check the identity fields above.")
-                print("=" * 72)
-                failures = []
-            else:
-                record["baseline_status"] = "compared"
-                failures = _check_regressions(record, baseline_records, metric_policies)
+            failures, baseline_records = _compare_record(record, metric_policies)
         if static_threshold_failed:
             failures.append(f"{record['model_id']} fixed-threshold phase failed "
                             f"(PERF_PYTEST_RC={os.environ.get('PERF_PYTEST_RC')})")
 
         record["success"] = not failures
+        # Only compared-and-PASS records may advance the rolling baseline: a
+        # CALIBRATION_NEEDED record must never silently seed a new cohort, and
+        # records without the v2 identity block never become eligible.
         record["baseline_eligible"] = (
-            identity_filters is not None and _is_baseline_eligible(record["run_source"], record["success"])
+            identity_filters is not None
+            and record["comparator_status"] == STATUS_PASS
+            and _is_baseline_eligible(record["run_source"], record["success"])
         )
         all_failures.extend(failures)
 

fastvideo/tests/performance/test_compare_baseline_policy.py

@@ -421,3 +421,196 @@ def test_informational_metric_remains_visible_without_failing():
     assert row["metrics"]["throughput"]["regressed"] is False
     assert row["threshold_exceeded_metrics"] == ["throughput"]
     assert row["failing_metrics"] == []
+
+
+def _v2_raw_result(**overrides):
+    raw = _raw_result()
+    raw.update({
+        "workload_id": "wan-t2v",
+        "variant_id": "1.3b-sp2",
+        "benchmark_version": 2,
+        "recipe_fingerprint": "recipe-1",
+        "hardware_profile_id": "hw-1",
+        "software_profile_id": "sw-1",
+    })
+    raw.update(overrides)
+    return raw
+
+
+def _v2_baseline_record(**overrides):
+    record = {
+        "gpu_type": "NVIDIA L40S",
+        "workload_id": "wan-t2v",
+        "variant_id": "1.3b-sp2",
+        "benchmark_version": 2,
+        "recipe_fingerprint": "recipe-1",
+        "hardware_profile_id": "hw-1",
+        "software_profile_id": "sw-1",
+        "latency": 10.0,
+        "throughput": 4.5,
+        "memory": 10000.0,
+        "timestamp": "2026-06-15T00:00:00+00:00",
+        "success": True,
+        "run_source": "scheduled_main",
+        "baseline_eligible": True,
+    }
+    record.update(overrides)
+    return record
+
+
+def _run_compare(monkeypatch, tmp_path, raw_result, baseline_records):
+    """Run compare_baseline.main() against a local tracking root.
+
+    Returns (exit_code, normalized_record, markdown_report).
+    """
+    results_dir = tmp_path / "results"
+    results_dir.mkdir()
+    (results_dir / "perf_current.json").write_text(json.dumps(raw_result))
+
+    model_dir = tmp_path / "tracking" / raw_result["benchmark_id"]
+    model_dir.mkdir(parents=True)
+    for index, record in enumerate(baseline_records):
+        (model_dir / f"rec{index}.json").write_text(json.dumps(record))
+
+    reports_dir = tmp_path / "reports"
+    monkeypatch.setattr(compare_baseline, "RESULTS_DIR", str(results_dir))
+    monkeypatch.setattr(compare_baseline, "TRACKING_ROOT", str(tmp_path / "tracking"))
+    monkeypatch.setattr(compare_baseline, "PERF_REPORTS_DIR", str(reports_dir))
+    monkeypatch.setattr(compare_baseline, "UPLOAD_POLICY", "never")
+    monkeypatch.setattr(compare_baseline, "sync_from_hf", lambda *args, **kwargs: None)
+    monkeypatch.setenv("PERF_RUN_SOURCE", "scheduled_main")
+    monkeypatch.delenv("PERF_PYTEST_RC", raising=False)
+    monkeypatch.delenv("GITHUB_STEP_SUMMARY", raising=False)
+
+    exit_code = compare_baseline.main()
+
+    normalized_paths = list((reports_dir / "results").glob("normalized_perf_*.json"))
+    assert len(normalized_paths) == 1
+    markdown = "\n".join(path.read_text() for path in reports_dir.glob("perf_*.md"))
+    return exit_code, json.loads(normalized_paths[0].read_text()), markdown
+
+
+def test_main_pass_with_comparable_baseline(monkeypatch, tmp_path):
+    exit_code, record, markdown = _run_compare(
+        monkeypatch, tmp_path, _v2_raw_result(), [_v2_baseline_record()])
+
+    assert exit_code == 0
+    assert record["comparator_status"] == "PASS"
+    assert record["baseline_status"] == "compared"
+    assert record["success"] is True
+    # Scheduled-main PASS records keep advancing the rolling baseline.
+    assert record["baseline_eligible"] is True
+    assert "| PASS |" in markdown
+
+
+def test_main_regression_on_gated_metric_fails_ci(monkeypatch, tmp_path):
+    exit_code, record, markdown = _run_compare(
+        monkeypatch, tmp_path, _v2_raw_result(avg_generation_time_s=20.0),
+        [_v2_baseline_record()])
+
+    assert exit_code == 1
+    assert record["comparator_status"] == "REGRESSION"
+    assert record["success"] is False
+    assert record["baseline_eligible"] is False
+    assert "| REGRESSION |" in markdown
+
+
+def test_main_missing_baseline_is_calibration_needed_and_does_not_seed(monkeypatch, tmp_path):
+    exit_code, record, markdown = _run_compare(
+        monkeypatch, tmp_path, _v2_raw_result(), [])
+
+    assert exit_code == 0
+    assert record["comparator_status"] == "CALIBRATION_NEEDED"
+    assert record["baseline_status"] == "initialized_new_cohort"
+    assert record["success"] is True
+    # Visible, but never silently seeds a passing baseline.
+    assert record["baseline_eligible"] is False
+    assert "| CALIBRATION_NEEDED |" in markdown
+
+
+def test_main_recipe_change_without_new_variant_is_recipe_mismatch(monkeypatch, tmp_path):
+    exit_code, record, markdown = _run_compare(
+        monkeypatch, tmp_path, _v2_raw_result(recipe_fingerprint="recipe-2"),
+        [_v2_baseline_record()])
+
+    assert exit_code == 1
+    assert record["comparator_status"] == "RECIPE_MISMATCH"
+    assert record["success"] is False
+    assert record["baseline_eligible"] is False
+    assert "| RECIPE_MISMATCH |" in markdown
+
+
+def test_main_recipe_change_as_new_variant_is_calibration_needed(monkeypatch, tmp_path):
+    exit_code, record, _ = _run_compare(
+        monkeypatch, tmp_path,
+        _v2_raw_result(variant_id="1.3b-sp2-r2", recipe_fingerprint="recipe-2"),
+        [_v2_baseline_record()])
+
+    assert exit_code == 0
+    assert record["comparator_status"] == "CALIBRATION_NEEDED"
+
+
+def test_main_v2_record_never_compares_against_v1_baselines(monkeypatch, tmp_path):
+    legacy_v1_baseline = {
+        "gpu_type": "NVIDIA L40S",
+        # Would be a >100% latency regression if the comparator (wrongly)
+        # matched the v2 record against v1 history.
+        "latency": 1.0,
+        "timestamp": "2026-06-15T00:00:00+00:00",
+        "success": True,
+    }
+
+    exit_code, record, _ = _run_compare(
+        monkeypatch, tmp_path, _v2_raw_result(), [legacy_v1_baseline])
+
+    assert exit_code == 0
+    assert record["comparator_status"] == "CALIBRATION_NEEDED"
+
+
+def test_main_legacy_v1_record_skips_comparison_with_pass_verdict(monkeypatch, tmp_path):
+    legacy_v1_baseline = {
+        "gpu_type": "NVIDIA L40S",
+        # Would be a >100% latency regression if the comparator (wrongly)
+        # compared the identity-less v1 record against this history.
+        "latency": 1.0,
+        "timestamp": "2026-06-15T00:00:00+00:00",
+        "success": True,
+    }
+
+    exit_code, record, markdown = _run_compare(
+        monkeypatch, tmp_path, _raw_result(), [legacy_v1_baseline])
+
+    assert exit_code == 0
+    assert record["comparator_status"] == "PASS"
+    assert record["baseline_status"] == "skipped_missing_identity"
+    assert record["baseline_eligible"] is False
+    assert "| PASS |" in markdown
+
+
+def test_main_partial_identity_skips_comparison(monkeypatch, tmp_path):
+    raw = _v2_raw_result()
+    del raw["software_profile_id"]
+
+    exit_code, record, _ = _run_compare(monkeypatch, tmp_path, raw, [])
+
+    assert exit_code == 0
+    assert record["comparator_status"] == "PASS"
+    assert record["baseline_status"] == "skipped_missing_identity"
+    assert record["success"] is True
+    assert record["baseline_eligible"] is False
+
+
+def test_main_baseline_load_failure_is_infra_error_and_fails_ci(monkeypatch, tmp_path):
+    def _boom(*_args, **_kwargs):
+        raise OSError("hf store unavailable")
+
+    monkeypatch.setattr(compare_baseline, "load_records_for_model", _boom)
+
+    exit_code, record, markdown = _run_compare(
+        monkeypatch, tmp_path, _v2_raw_result(), [])
+
+    assert exit_code == 1
+    assert record["comparator_status"] == "INFRA_ERROR"
+    assert record["success"] is False
+    assert record["baseline_eligible"] is False
+    assert "| INFRA_ERROR |" in markdown

fastvideo/tests/ssim/inference_similarity_utils.py

@@ -62,12 +62,29 @@ def resolve_inference_device_reference_folder(logger: Logger) -> str:
     return device_reference_folder
 
 
-def _find_reference_video(reference_folder: str, prompt: str) -> str:
+def _find_reference_media(
+    reference_folder: str,
+    prompt: str,
+    *,
+    media_extension: str,
+) -> str:
+    """Pick a reference file whose basename contains the prompt prefix."""
     prompt_prefix = prompt[:100].strip()
+    allowed = (media_extension.lower(), ".mp4", ".png", ".jpg", ".jpeg")
+    matches: list[str] = []
     for filename in os.listdir(reference_folder):
-        if filename.endswith(".mp4") and prompt_prefix in filename:
-            return os.path.join(reference_folder, filename)
-    raise FileNotFoundError("Reference video missing")
+        low = filename.lower()
+        if not any(low.endswith(ext) for ext in allowed):
+            continue
+        if prompt_prefix in filename:
+            matches.append(filename)
+    if not matches:
+        raise FileNotFoundError("Reference media missing")
+    preferred = media_extension.lower().lstrip(".")
+    for name in matches:
+        if name.lower().endswith(f".{preferred}"):
+            return os.path.join(reference_folder, name)
+    return os.path.join(reference_folder, matches[0])
 
 
 def _remove_stale_generated_video(output_dir: str, output_video_name: str) -> None:
@@ -80,21 +97,23 @@ def _assert_similarity(
     *,
     logger: Logger,
     output_dir: str,
-    output_video_name: str,
+    output_media_name: str,
     reference_folder: str,
     prompt: str,
     num_inference_steps: int,
     min_acceptable_ssim: float,
     model_id: str,
     attention_backend_name: str,
+    media_extension: str,
 ) -> None:
-    generated_video_path = os.path.join(output_dir, output_video_name)
+    generated_media_path = os.path.join(output_dir, output_media_name)
+    artifact_kind = "image" if media_extension.lower() in (".png", ".jpg", ".jpeg") else "video"
     if not os.path.exists(reference_folder):
         logger.error("Reference folder missing: %s", reference_folder)
         xfail_missing_reference_in_bootstrap_mode(
-            generated_artifact_path=generated_video_path,
+            generated_artifact_path=generated_media_path,
             reference_folder=reference_folder,
-            artifact_kind="video",
+            artifact_kind=artifact_kind,
         )
         error_msg = (
             f"Reference video folder does not exist: {reference_folder}\n"
@@ -104,28 +123,32 @@ def _assert_similarity(
         raise FileNotFoundError(error_msg)
 
     try:
-        reference_video_path = _find_reference_video(reference_folder, prompt)
+        reference_media_path = _find_reference_media(
+            reference_folder,
+            prompt,
+            media_extension=media_extension,
+        )
     except FileNotFoundError as error:
         logger.error(
-            "Reference video not found for prompt: %s with backend: %s",
+            "Reference media not found for prompt: %s with backend: %s",
             prompt,
             attention_backend_name,
         )
         xfail_missing_reference_in_bootstrap_mode(
-            generated_artifact_path=generated_video_path,
+            generated_artifact_path=generated_media_path,
             reference_folder=reference_folder,
-            artifact_kind="video",
+            artifact_kind=artifact_kind,
         )
         raise error
 
     logger.info(
         "Computing SSIM between %s and %s",
-        reference_video_path,
-        generated_video_path,
+        reference_media_path,
+        generated_media_path,
     )
     ssim_values = compute_video_ssim_torchvision(
-        reference_video_path,
-        generated_video_path,
+        reference_media_path,
+        generated_media_path,
         use_ms_ssim=True,
     )
 
@@ -136,8 +159,8 @@ def _assert_similarity(
     success = write_ssim_results(
         output_dir,
         ssim_values,
-        reference_video_path,
-        generated_video_path,
+        reference_media_path,
+        generated_media_path,
         num_inference_steps,
         prompt,
     )
@@ -222,6 +245,7 @@ def run_text_to_video_similarity_test(
     min_acceptable_ssim: float,
     init_kwargs_override: dict[str, object] | None = None,
     generation_kwargs_override: dict[str, object] | None = None,
+    media_extension: str = ".mp4",
 ) -> None:
     with attention_backend(attention_backend_name):
         output_dir = build_generated_output_dir(
@@ -230,9 +254,9 @@ def run_text_to_video_similarity_test(
             model_id,
             attention_backend_name,
         )
-        output_video_name = f"{prompt[:100].strip()}.mp4"
+        output_media_name = f"{prompt[:100].strip()}{media_extension}"
         os.makedirs(output_dir, exist_ok=True)
-        _remove_stale_generated_video(output_dir, output_video_name)
+        _remove_stale_generated_video(output_dir, output_media_name)
 
         params_map = select_ssim_params(
             default_params_map,
@@ -274,13 +298,14 @@ def run_text_to_video_similarity_test(
     _assert_similarity(
         logger=logger,
         output_dir=output_dir,
-        output_video_name=output_video_name,
+        output_media_name=output_media_name,
         reference_folder=reference_folder,
         prompt=prompt,
         num_inference_steps=num_inference_steps,
         min_acceptable_ssim=min_acceptable_ssim,
         model_id=model_id,
         attention_backend_name=attention_backend_name,
+        media_extension=media_extension,
     )
 
 
@@ -306,9 +331,9 @@ def run_image_to_video_similarity_test(
             model_id,
             attention_backend_name,
         )
-        output_video_name = f"{prompt[:100].strip()}.mp4"
+        output_media_name = f"{prompt[:100].strip()}.mp4"
         os.makedirs(output_dir, exist_ok=True)
-        _remove_stale_generated_video(output_dir, output_video_name)
+        _remove_stale_generated_video(output_dir, output_media_name)
 
         params_map = select_ssim_params(
             default_params_map,
@@ -351,11 +376,12 @@ def run_image_to_video_similarity_test(
     _assert_similarity(
         logger=logger,
         output_dir=output_dir,
-        output_video_name=output_video_name,
+        output_media_name=output_media_name,
         reference_folder=reference_folder,
         prompt=prompt,
         num_inference_steps=num_inference_steps,
         min_acceptable_ssim=min_acceptable_ssim,
         model_id=model_id,
         attention_backend_name=attention_backend_name,
+        media_extension=".mp4",
     )

fastvideo/tests/ssim/test_flux_t2i_similarity.py

@@ -0,0 +1,134 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import os
+
+import pytest
+import torch
+
+from fastvideo.api.flux import FluxSamplingParam
+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
+
+# MS-SSIM gate (see module docstring).
+FLUX_T2I_MIN_SSIM = 0.98
+
+FLUX_MODEL_PATH = os.getenv(
+    "FLUX_T2I_MODEL_DIR",
+    "black-forest-labs/FLUX.1-dev",
+)
+
+device_reference_folder = resolve_device_reference_folder(
+    (
+        ("A40", "A40"),
+        ("L40S", "L40S"),
+        ("H100", "H100"),
+        ("H200", "H200"),
+        ("RTX 4090", "RTX4090"),
+        ("4090", "RTX4090"),
+    ),
+    device_name=get_cuda_device_name(),
+    fallback_device_prefix="L40S",
+    logger=logger,
+)
+
+# Folder token must match Hub path with slashes → double underscore (SD3.5
+# pattern in ``test_sd35_similarity.py``).
+MODEL_ID = "black-forest-labs__FLUX.1-dev"
+
+TEST_PROMPTS = [
+    "a photo of a cat",
+]
+
+FLUX_DEFAULT_PARAMS: dict[str, object] = {
+    "num_gpus": 1,
+    "model_path": FLUX_MODEL_PATH,
+    "sp_size": 1,
+    "tp_size": 1,
+    "height": 256,
+    "width": 256,
+    "num_frames": 1,
+    "fps": 1,
+    "num_inference_steps": 8,
+    "guidance_scale": 3.5,
+    "seed": 0,
+}
+
+_flux_full_defaults = FluxSamplingParam()
+FLUX_FULL_QUALITY_PARAMS: dict[str, object] = {
+    "num_gpus": 1,
+    "model_path": FLUX_MODEL_PATH,
+    "sp_size": 1,
+    "tp_size": 1,
+    "height": _flux_full_defaults.height,
+    "width": _flux_full_defaults.width,
+    "num_frames": 1,
+    "fps": _flux_full_defaults.fps,
+    "num_inference_steps": _flux_full_defaults.num_inference_steps,
+    "guidance_scale": _flux_full_defaults.guidance_scale,
+    "seed": _flux_full_defaults.seed,
+}
+
+FLUX_MODEL_TO_PARAMS = {
+    MODEL_ID: FLUX_DEFAULT_PARAMS,
+}
+FLUX_FULL_QUALITY_MODEL_TO_PARAMS = {
+    MODEL_ID: FLUX_FULL_QUALITY_PARAMS,
+}
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="FLUX T2I SSIM test requires CUDA",
+)
+@pytest.mark.parametrize("prompt", TEST_PROMPTS)
+@pytest.mark.parametrize("attention_backend_name", ["TORCH_SDPA"])
+@pytest.mark.parametrize("model_id", list(FLUX_MODEL_TO_PARAMS.keys()))
+def test_flux_t2i_similarity(
+    prompt: str,
+    attention_backend_name: str,
+    model_id: str,
+) -> None:
+    is_hf_repo = "/" in FLUX_MODEL_PATH and not FLUX_MODEL_PATH.startswith("/")
+    if not is_hf_repo and not os.path.isdir(FLUX_MODEL_PATH):
+        pytest.skip(
+            f"FLUX weights not found at {FLUX_MODEL_PATH} "
+            f"(set FLUX_T2I_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=FLUX_MODEL_TO_PARAMS,
+        full_quality_params_map=FLUX_FULL_QUALITY_MODEL_TO_PARAMS,
+        min_acceptable_ssim=FLUX_T2I_MIN_SSIM,
+        media_extension=".png",
+        init_kwargs_override={
+            "workload_type": "t2i",
+            "use_fsdp_inference": False,
+            "text_encoder_cpu_offload": False,
+            "vae_cpu_offload": False,
+            "image_encoder_cpu_offload": False,
+            "pin_cpu_memory": False,
+        },
+        generation_kwargs_override={
+            "save_video": True,
+            "use_embedded_guidance": True,
+            "true_cfg_scale": 1.0,
+        },
+    )

fastvideo/tests/ssim/test_glm_image_similarity.py

@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: Apache-2.0
+"""SSIM-based regression test for GLM-Image generation."""
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+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
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+LOCAL_WEIGHTS_DIR = Path(
+    os.getenv("GLM_IMAGE_LOCAL_WEIGHTS_DIR",
+              REPO_ROOT / "official_weights" / "glm_image"))
+GLM_IMAGE_MODEL_PATH = os.getenv("GLM_IMAGE_MODEL_DIR", str(LOCAL_WEIGHTS_DIR))
+
+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,
+)
+
+MODEL_ID = "zai-org__GLM-Image"
+
+TEST_PROMPTS = [
+    "A beautiful landscape photography with rolling hills, "
+    "a winding river, and a vibrant sunset in the background. "
+    "Warm golden light, photorealistic style.",
+]
+
+GLM_IMAGE_PARAMS = {
+    "num_gpus": 1,
+    "model_path": GLM_IMAGE_MODEL_PATH,
+    "sp_size": 1,
+    "tp_size": 1,
+    "height": 256,
+    "width": 256,
+    "num_frames": 1,
+    "fps": 1,
+    "num_inference_steps": 4,
+    "guidance_scale": 1.5,
+    "seed": 0,
+    "neg_prompt": "",
+}
+
+GLM_IMAGE_FULL_QUALITY_PARAMS = {
+    "num_gpus": 1,
+    "model_path": GLM_IMAGE_MODEL_PATH,
+    "sp_size": 1,
+    "tp_size": 1,
+    "height": 1024,
+    "width": 1024,
+    "num_frames": 1,
+    "fps": 1,
+    "num_inference_steps": 50,
+    "guidance_scale": 1.5,
+    "seed": 0,
+    "neg_prompt": "",
+}
+
+GLM_IMAGE_MODEL_TO_PARAMS = {
+    MODEL_ID: GLM_IMAGE_PARAMS,
+}
+GLM_IMAGE_FULL_QUALITY_MODEL_TO_PARAMS = {
+    MODEL_ID: GLM_IMAGE_FULL_QUALITY_PARAMS,
+}
+
+
+def _has_weights() -> bool:
+    required = ["transformer", "vae", "text_encoder",
+                "vision_language_encoder", "processor", "tokenizer",
+                "scheduler"]
+    return all((LOCAL_WEIGHTS_DIR / r).exists() for r in required)
+
+
+def _upstream_glm_image_available() -> bool:
+    try:
+        import transformers
+    except ImportError:
+        return False
+    return hasattr(transformers, "GlmImageForConditionalGeneration")
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="GLM-Image SSIM test requires CUDA",
+)
+@pytest.mark.skipif(
+    not _has_weights(),
+    reason=f"GLM-Image full weights not found at {LOCAL_WEIGHTS_DIR}.",
+)
+@pytest.mark.skipif(
+    not _upstream_glm_image_available(),
+    reason="GLM-Image needs transformers>=5.0.0rc0 (ships the AR encoder).",
+)
+@pytest.mark.parametrize("prompt", TEST_PROMPTS)
+@pytest.mark.parametrize("attention_backend_name", ["TORCH_SDPA"])
+@pytest.mark.parametrize("model_id", list(GLM_IMAGE_MODEL_TO_PARAMS.keys()))
+def test_glm_image_similarity(
+    prompt: str,
+    attention_backend_name: str,
+    model_id: str,
+) -> None:
+    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=GLM_IMAGE_MODEL_TO_PARAMS,
+        full_quality_params_map=GLM_IMAGE_FULL_QUALITY_MODEL_TO_PARAMS,
+        min_acceptable_ssim=0.98,
+        init_kwargs_override={
+            "trust_remote_code": True,
+            "use_fsdp_inference": False,
+        },
+        generation_kwargs_override={
+            "save_video": True,
+        },
+    )

fastvideo/tests/transformers/test_flux.py

@@ -0,0 +1,185 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import glob
+import os
+
+import pytest
+import torch
+from diffusers import FluxTransformer2DModel as HFFluxTransformer2DModel
+from torch.testing import assert_close
+
+from fastvideo.configs.models.dits.flux import FluxDiTConfig
+from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.models.loader.component_loader import TransformerLoader
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+os.environ.setdefault("MASTER_ADDR", "localhost")
+os.environ.setdefault("MASTER_PORT", "29517")
+
+_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
+_DEFAULT_FLUX_TRANSFORMER = os.path.join(
+    _REPO_ROOT,
+    "official_weights",
+    "FLUX.1-dev",
+    "transformer",
+)
+
+
+def _flux_transformer_path() -> str:
+    return os.environ.get("FLUX_TRANSFORMER_PATH", _DEFAULT_FLUX_TRANSFORMER)
+
+
+def _prepare_latent_image_ids(
+    height: int,
+    width: int,
+    device: torch.device,
+    dtype: torch.dtype = torch.long,
+) -> torch.Tensor:
+    """Match Diffusers ``FluxPipeline._prepare_latent_image_ids`` (batch omitted)."""
+    latent_image_ids = torch.zeros(height, width, 3)
+    latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height)[:, None]
+    latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width)[None, :]
+    h, w, c = latent_image_ids.shape
+    latent_image_ids = latent_image_ids.reshape(h * w, c)
+    return latent_image_ids.to(device=device, dtype=dtype)
+
+
+@pytest.fixture
+def torch_sdpa_attention_backend(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
+
+
+requires_cuda = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="FLUX DiT parity test requires CUDA",
+)
+
+requires_weights = pytest.mark.skipif(
+    not glob.glob(os.path.join(_flux_transformer_path(), "*.safetensors")),
+    reason=(
+        f"No safetensors under {_flux_transformer_path()} — download FLUX.1-dev "
+        "transformer or set FLUX_TRANSFORMER_PATH"
+    ),
+)
+
+
+@requires_cuda
+@requires_weights
+@pytest.mark.usefixtures("distributed_setup", "torch_sdpa_attention_backend")
+def test_flux_transformer_parity_vs_diffusers() -> None:
+    """Single forward: FastVideo DiT vs Diffusers ``FluxTransformer2DModel``."""
+    device = torch.device("cuda:0")
+    precision = torch.bfloat16
+    transformer_path = _flux_transformer_path()
+
+    args = FastVideoArgs(
+        model_path=transformer_path,
+        dit_cpu_offload=False,
+        dit_layerwise_offload=False,
+        pipeline_config=PipelineConfig(dit_config=FluxDiTConfig(), dit_precision="bf16"),
+    )
+    args.device = device
+
+    generator = torch.Generator(device=device).manual_seed(0)
+    torch.manual_seed(0)
+
+    batch_size = 1
+    latent_h, latent_w = 4, 4
+    img_seq = latent_h * latent_w
+    text_len = 32
+
+    hidden_states = torch.randn(
+        batch_size,
+        img_seq,
+        64,
+        device=device,
+        dtype=precision,
+        generator=generator,
+    )
+    encoder_hidden_states = torch.randn(
+        batch_size,
+        text_len,
+        4096,
+        device=device,
+        dtype=precision,
+        generator=generator,
+    )
+    pooled_projections = torch.randn(
+        batch_size,
+        768,
+        device=device,
+        dtype=precision,
+        generator=generator,
+    )
+
+    # Diffusers pipeline passes scheduler timesteps / 1000 (float, same dtype as latents).
+    timestep = torch.tensor([512.0], device=device, dtype=precision) / 1000.0
+    guidance = torch.full((batch_size,), 3.5, device=device, dtype=torch.float32)
+
+    txt_ids = torch.zeros(text_len, 3, device=device, dtype=torch.long)
+    img_ids = _prepare_latent_image_ids(latent_h, latent_w, device, dtype=torch.long)
+
+    forward_batch = ForwardBatch(data_type="dummy")
+
+    # One ~12B model at a time avoids peak VRAM from holding both checkpoints.
+    loader = TransformerLoader()
+    fv_model = loader.load(transformer_path, args).to(device=device, dtype=precision)
+    fv_model.eval()
+    with (
+        torch.no_grad(),
+        torch.amp.autocast("cuda", dtype=precision),
+        set_forward_context(
+            current_timestep=512,
+            attn_metadata=None,
+            forward_batch=forward_batch,
+        ),
+    ):
+        fv_out = fv_model(
+            hidden_states=hidden_states.clone(),
+            encoder_hidden_states=encoder_hidden_states.clone(),
+            pooled_projections=pooled_projections.clone(),
+            timestep=timestep.clone(),
+            guidance=guidance.clone(),
+            txt_ids=txt_ids,
+            img_ids=img_ids,
+            return_dict=False,
+        )[0]
+    fv_out_cpu = fv_out.detach().float().cpu()
+    del fv_model
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    hf_model = (
+        HFFluxTransformer2DModel.from_pretrained(
+            transformer_path,
+            torch_dtype=precision,
+        )
+        .to(device)
+        .eval()
+    )
+    with torch.no_grad(), torch.amp.autocast("cuda", dtype=precision):
+        hf_out = hf_model(
+            hidden_states=hidden_states.clone(),
+            encoder_hidden_states=encoder_hidden_states.clone(),
+            pooled_projections=pooled_projections.clone(),
+            timestep=timestep.clone(),
+            guidance=guidance.clone(),
+            txt_ids=txt_ids,
+            img_ids=img_ids,
+            return_dict=False,
+        )[0]
+
+    assert hf_out.shape == fv_out_cpu.shape
+    hf_cpu = hf_out.float().cpu()
+    abs_diff = (hf_cpu - fv_out_cpu).abs()
+    print(f"[FLUX DiT parity] max_diff={abs_diff.max():.4f}  mean_diff={abs_diff.mean():.4f}  "
+          f"median_diff={abs_diff.median():.4f}  p99_diff="
+          f"{abs_diff.flatten().kthvalue(int(0.99 * abs_diff.numel())).values:.4f}")
+    # bfloat16 accumulation over 57 transformer layers produces tail errors up to ~0.5
+    # on isolated elements (median=0, mean~0.04 on L40S).  atol=0.5 catches real bugs
+    # (wrong weights / missing layers) which produce mean_diff >> 0.1.
+    assert_close(hf_cpu, fv_out_cpu, atol=0.5, rtol=0.0)

fastvideo/tests/utils.py

@@ -77,13 +77,32 @@ def _read_video_frames(path: str) -> torch.Tensor:
     return torch.stack(frames)
 
 
+def _read_image_as_single_frame_video(path: str) -> torch.Tensor:
+    """Read one image as a single-frame ``(1, C, H, W)`` uint8 tensor."""
+    from torchvision.io import read_image
+
+    img = read_image(path)
+    return img.unsqueeze(0)
+
+
+def _read_visual_frames(path: str) -> torch.Tensor:
+    """Read a video or a single image as ``(T, C, H, W)`` uint8."""
+    ext = os.path.splitext(path)[1].lower()
+    if ext in {".png", ".jpg", ".jpeg", ".webp"}:
+        return _read_image_as_single_frame_video(path)
+    return _read_video_frames(path)
+
+
 def compute_video_ssim_torchvision(video1_path, video2_path, use_ms_ssim=True):
     """
-    Compute SSIM between two videos.
+    Compute SSIM between two videos or single-frame image files.
+
+    Image paths (``.png``, ``.jpg``, ``.jpeg``, ``.webp``) are treated as
+    one-frame clips so T2I SSIM can share the same MS-SSIM path as video.
 
     Args:
-        video1_path: Path to the first video.
-        video2_path: Path to the second video.
+        video1_path: Path to the first video or image.
+        video2_path: Path to the second video or image.
         use_ms_ssim: Whether to use Multi-Scale Structural Similarity(MS-SSIM) instead of SSIM.
     """
     from pytorch_msssim import ms_ssim, ssim
@@ -94,8 +113,8 @@ def compute_video_ssim_torchvision(video1_path, video2_path, use_ms_ssim=True):
     if not os.path.exists(video2_path):
         raise FileNotFoundError(f"Video2 not found: {video2_path}")
 
-    frames1 = _read_video_frames(video1_path)
-    frames2 = _read_video_frames(video2_path)
+    frames1 = _read_visual_frames(video1_path)
+    frames2 = _read_visual_frames(video2_path)
 
     # Ensure same number of frames
     min_frames = min(frames1.shape[0], frames2.shape[0])

pyproject.toml

@@ -21,7 +21,9 @@ dependencies = [
     "requests>=2.32.2",
 
     # Machine Learning & Transformers
-    "transformers>=4.57.3",
+    # GLM-Image's AR encoder (GlmImageForConditionalGeneration) first ships in
+    # transformers 5.0.0; floor bumped from >=4.57.3 to >=5.0.0 (stable, not rc).
+    "transformers>=5.0.0",
     # <0.23: tokenizers 0.23 renamed RobertaProcessing's binding args, so
     # transformers' CLIP-style tokenizer loading dies with
     # "RobertaProcessing.__new__() got an unexpected keyword argument 'cls'".

pyproject_other.toml

@@ -21,7 +21,9 @@ dependencies = [
     "requests>=2.32.2",
 
     # Machine Learning & Transformers
-    "transformers>=4.57.3",
+    # GLM-Image's AR encoder (GlmImageForConditionalGeneration) first ships in
+    # transformers 5.0.0; floor bumped from >=4.57.3 to >=5.0.0 (stable, not rc).
+    "transformers>=5.0.0",
     "tokenizers>=0.20.1",
     "sentencepiece>=0.2.0",
     "timm>=1.0.11",

tests/local_tests/flux/PORT_STATUS.md

@@ -0,0 +1,59 @@
+# FLUX.1-dev Port Status
+
+Model family: flux
+Official ref: black-forest-labs/FLUX.1-dev (Diffusers layout)
+Workload: T2I
+Last updated: 2026-06-20
+
+## Component Status
+
+| Component | Type | Parity test | Status | Notes |
+|---|---|---|---|---|
+| FluxTransformer2DModel | DiT (ported) | fastvideo/tests/transformers/test_flux.py | PASS (A40, 2026-05-11) | max_diff=0.5, mean_diff=0.04, median=0 — bf16 tail error, see notes |
+| AutoencoderKL | VAE (reused) | tests/local_tests/flux/test_flux_dev_component_loaders.py | PASS (A40, 2026-05-11) | Loader smoke: 54.14s |
+| CLIPTextModel | encoder (reused) | tests/local_tests/flux/test_flux_dev_component_loaders.py | PASS (A40, 2026-05-11) | Shared component |
+| T5EncoderModel | encoder (reused) | tests/local_tests/flux/test_flux_dev_component_loaders.py | PASS (A40, 2026-05-11) | Shared component |
+| FlowMatchEulerDiscreteScheduler | scheduler (reused) | tests/local_tests/flux/test_flux_dev_component_loaders.py | PASS (A40, 2026-05-11) | Shared component |
+
+## Conversion
+
+No conversion script needed — FLUX.1-dev uses native Diffusers checkpoint layout.
+Components load via production loaders with no key remapping.
+Strict-load evidence: PASS — component loader test runs strict load of all components (2026-05-11).
+
+## Pipeline
+
+| Test | Status | Notes |
+|---|---|---|
+| Pipeline smoke (2-step, 256×256) | PASS (A40, 2026-05-11) | 91.14s — output finite, non-null |
+| Pipeline parity vs Diffusers FluxPipeline | PASS (A40, 2026-05-11) | Output finite, in [0,1], correct shape — pixel parity not enforced (pipelines sample noise independently) |
+
+## Quality
+
+| Item | Status |
+|---|---|
+| SSIM test | written (fastvideo/tests/ssim/test_flux_t2i_similarity.py) |
+| Reference images committed | TORCH_SDPA committed (A40, 2026-05-11); FLASH_ATTN pending |
+
+## DiT Parity Notes
+
+DiT parity test compares FastVideo `FluxTransformer2DModel` vs Diffusers under identical inputs (bfloat16, L40S).
+- max_diff=0.5000, mean_diff=0.0382, median_diff=0.0000, p99_diff=0.2500
+- Median=0 and mean=0.04 confirm implementations are equivalent; tail errors up to 0.5 are expected
+  bfloat16 accumulation over 57 transformer layers (eps~7.8e-3, accumulated per-GEMM error ~0.4).
+- Test uses atol=0.5 which catches real bugs (wrong weights/layers produce mean_diff >> 0.1).
+
+## Known Blockers
+
+None. (Both prior blockers resolved — see below.)
+
+## Resolved
+
+1. SSIM reference image — RESOLVED. The TORCH_SDPA reference (A40, 256×256, 8 steps,
+   seed=0) is committed, so the SSIM gate runs against it. A FLASH_ATTN reference is
+   still pending (see Quality table) but is not required for the default backend.
+
+2. Registry registration — RESOLVED. FLUX is registered unconditionally in
+   `_register_configs()` (fastvideo/registry.py); there is no early-return guard that
+   could silently skip it. The earlier `_CONFIG_REGISTRY` pre-population concern no
+   longer applies after the registry refactor.

tests/local_tests/flux/README.md

@@ -0,0 +1,73 @@
+# FLUX.1-dev Local Tests
+
+## Prerequisites
+
+- CUDA GPU
+- `official_weights/FLUX.1-dev` (Diffusers layout) or set `FLUX_DEV_ROOT`
+- `FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA` for parity tests
+
+Download weights:
+```bash
+huggingface-cli download black-forest-labs/FLUX.1-dev --local-dir official_weights/FLUX.1-dev
+```
+
+## Component Loader Smoke Test
+
+Verifies CLIP/T5 encoders, tokenizers, VAE, and FlowMatch scheduler all load
+from the Diffusers checkpoint layout.
+
+```bash
+pytest tests/local_tests/flux/test_flux_dev_component_loaders.py -vs
+```
+
+Status: requires weights — not run in CI.
+
+## Pipeline Smoke Test
+
+Short denoise + decode (2 steps, 256×256) end-to-end.
+
+```bash
+FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA \
+pytest tests/local_tests/pipelines/test_flux_dev_pipeline_smoke.py -vs
+```
+
+Status: requires weights — not run in CI.
+
+## Pipeline Parity Test
+
+Compares FastVideo FluxPipeline decoded image output against Diffusers
+FluxPipeline under identical prompt, seed, and inference parameters (4
+steps, 256×256, seed=42). Tolerance: atol=5e-2.
+
+```bash
+FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA \
+pytest tests/local_tests/pipelines/test_flux_dev_pipeline_parity.py -vs
+```
+
+Status: PASS on A40 (2026-05-11) — requires weights; not run in CI.
+
+## DiT Parity Test
+
+Single-forward comparison of FastVideo `FluxTransformer2DModel` against
+Diffusers `FluxTransformer2DModel` under identical random inputs.
+
+```bash
+FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA \
+pytest fastvideo/tests/transformers/test_flux.py -vs
+```
+
+Status: requires weights — not run in CI.
+Pass evidence: PASS recorded in PORT_STATUS.md (A40, 2026-05-11).
+
+## SSIM Regression Test
+
+Full-quality image generation compared against per-device reference images.
+Reference images must be seeded first via `reference_videos_cli.py`.
+
+```bash
+FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA \
+FLUX_T2I_MODEL_DIR=official_weights/FLUX.1-dev \
+pytest fastvideo/tests/ssim/test_flux_t2i_similarity.py -vs
+```
+
+Status: seeded reference images committed (TORCH_SDPA, A40, 2026-05-11) — see PORT_STATUS.md.

tests/local_tests/flux/test_flux_dev_component_loaders.py

@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Loader smoke tests for FLUX.1-dev component loading from a local checkpoint.
+
+Verifies that tokenizers, CLIP/T5 text encoders, VAE, and FlowMatch scheduler
+all load correctly from a Diffusers-layout ``FLUX.1-dev`` directory.
+
+Requires ``official_weights/FLUX.1-dev`` (or set ``FLUX_DEV_ROOT``) and CUDA.
+
+Run from repo root::
+
+    pytest tests/local_tests/flux/test_flux_dev_component_loaders.py -vs
+"""
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import pytest
+import torch
+
+from fastvideo.configs.models.encoders import (
+    BaseEncoderOutput,
+    CLIPTextConfig,
+    T5LargeConfig,
+)
+from fastvideo.configs.models.vaes.autoencoder_kl import AutoencoderKLVAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig, preprocess_text
+
+os.environ.setdefault("MASTER_ADDR", "localhost")
+os.environ.setdefault("MASTER_PORT", "29517")
+
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+_FLUX_DEV_ROOT = os.environ.get(
+    "FLUX_DEV_ROOT",
+    str(_REPO_ROOT / "official_weights" / "FLUX.1-dev"),
+)
+
+
+def _flux_clip_post(outputs: BaseEncoderOutput) -> torch.Tensor:
+    if outputs.pooler_output is None:
+        raise RuntimeError("CLIP pooled output required")
+    return outputs.pooler_output
+
+
+def _flux_t5_post(outputs: BaseEncoderOutput) -> torch.Tensor:
+    if outputs.last_hidden_state is None:
+        raise RuntimeError("T5 last_hidden_state required")
+    return outputs.last_hidden_state
+
+
+@dataclass
+class _FluxDevLoaderPipelineConfig(PipelineConfig):
+    """Two text encoders (CLIP + T5) matching FLUX.1-dev model_index.json."""
+
+    vae_config: AutoencoderKLVAEConfig = field(
+        default_factory=AutoencoderKLVAEConfig
+    )
+    text_encoder_configs: tuple[CLIPTextConfig, T5LargeConfig] = field(
+        default_factory=lambda: (CLIPTextConfig(), T5LargeConfig())
+    )
+    text_encoder_precisions: tuple[str, ...] = ("fp32", "bf16")
+    preprocess_text_funcs: tuple = field(
+        default_factory=lambda: (preprocess_text, preprocess_text)
+    )
+    postprocess_text_funcs: tuple = field(
+        default_factory=lambda: (_flux_clip_post, _flux_t5_post)
+    )
+
+
+def test_flux_dev_model_index_components_load() -> None:
+    """CLIP/T5 encoders, tokenizers, VAE, and FlowMatch scheduler load."""
+    if not torch.cuda.is_available():
+        pytest.skip("FLUX component loader test requires CUDA")
+    if not Path(_FLUX_DEV_ROOT, "model_index.json").is_file():
+        pytest.skip(
+            "official_weights/FLUX.1-dev missing "
+            "(set FLUX_DEV_ROOT or download black-forest-labs/FLUX.1-dev)"
+        )
+
+    from fastvideo.distributed import (
+        cleanup_dist_env_and_memory,
+        maybe_init_distributed_environment_and_model_parallel,
+    )
+    from fastvideo.fastvideo_args import FastVideoArgs
+    from fastvideo.models.loader.component_loader import (
+        SchedulerLoader,
+        TextEncoderLoader,
+        TokenizerLoader,
+        VAELoader,
+    )
+
+    maybe_init_distributed_environment_and_model_parallel(1, 1)
+    try:
+        args = FastVideoArgs(
+            model_path=_FLUX_DEV_ROOT,
+            pipeline_config=_FluxDevLoaderPipelineConfig(),
+            hsdp_shard_dim=1,
+            pin_cpu_memory=False,
+        )
+
+        tok_clip = TokenizerLoader().load(
+            os.path.join(_FLUX_DEV_ROOT, "tokenizer"), args
+        )
+        tok_t5 = TokenizerLoader().load(
+            os.path.join(_FLUX_DEV_ROOT, "tokenizer_2"), args
+        )
+        assert "CLIP" in tok_clip.__class__.__name__
+        assert "T5" in tok_t5.__class__.__name__
+
+        te_clip = TextEncoderLoader().load(
+            os.path.join(_FLUX_DEV_ROOT, "text_encoder"), args
+        )
+        te_t5 = TextEncoderLoader().load(
+            os.path.join(_FLUX_DEV_ROOT, "text_encoder_2"), args
+        )
+        # Distributed init may wrap encoders in an FSDP shell.
+        assert te_clip.__class__.__name__.endswith("CLIPTextModel")
+        assert te_t5.__class__.__name__.endswith("T5EncoderModel")
+
+        vae = VAELoader().load(os.path.join(_FLUX_DEV_ROOT, "vae"), args)
+        assert vae.__class__.__name__ == "AutoencoderKL"
+
+        scheduler = SchedulerLoader().load(
+            os.path.join(_FLUX_DEV_ROOT, "scheduler"), args
+        )
+        assert scheduler.__class__.__name__ == "FlowMatchEulerDiscreteScheduler"
+        scheduler.set_timesteps(4, mu=0.7)
+        assert len(scheduler.timesteps) == 4
+    finally:
+        cleanup_dist_env_and_memory()

tests/local_tests/glm_image/README.md

@@ -0,0 +1,74 @@
+# GLM-Image Local Tests
+
+Local-only parity and smoke tests for the `glm_image` port: FastVideo vs the
+official HF `transformers` + `diffusers` GLM-Image implementation. These are
+**not run in CI** — they need the full weights, CUDA, and (for the reference
+side) the diffusers GLM-Image classes; each test skips with an actionable
+message when a prerequisite is missing.
+
+## What you need
+
+| | |
+|---|---|
+| Weights | `zai-org/GLM-Image` → `official_weights/glm_image` (~30 GB: AR encoder 4 shards, transformer 3 shards, vae + text_encoder) |
+| Runtime deps | `transformers>=5.0.0` (first release with the AR encoder `GlmImageForConditionalGeneration`) and `diffusers>=0.38.0` — both committed in `pyproject.toml` |
+| Test-only deps | the diffusers `GlmImageTransformer2DModel` / `GlmImagePipeline` reference classes need `diffusers>=0.37.0.dev0` (used **only** by the parity tests) |
+
+Download weights (from the repo root; cache via `HF_HOME` if `/` is small):
+
+```bash
+python ".agents/skills/add-model-01-prep/scripts/download_hf_weights.py" \
+    "zai-org/GLM-Image" "official_weights/glm_image"
+```
+
+Optionally install the diffusers reference classes to run the parity side:
+
+```bash
+uv pip install -U "git+https://github.com/huggingface/diffusers.git@ff3b86b4"
+```
+
+Parity was checked against transformers `a65bf6c9` + diffusers `ff3b86b4`.
+
+## Parity results (GB200)
+
+| Component | Test | Result |
+|---|---|---|
+| `transformer` (DiT) | `test_transformer_parity.py` | fp32 cosine `1.000000` (100% within 5e-3); bf16 cosine `0.999969` (element-wise rounding over 30 layers vs a different SDPA kernel — not a bug; RoPE proven exactly equivalent). fp32 + bf16 + strict-load |
+| `vae` (`AutoencoderKL`) | `test_vae_parity.py` | decode + encode + config, 3/3 |
+| `text_encoder` (T5) + **ByT5** tokenizer | `test_t5_parity.py` | tokenizer + glyph extract + embeds, 3/3 |
+| `vision_language_encoder` (GLM-4-9B AR) | `test_ar_parity.py` | surface + to/eval + greedy determinism, 3/3 |
+| `pipeline` (T2I) | `test_pipeline_parity.py` | injected-input determinism: latent cosine `0.999997`, decoded-image MAE `0.159/255`. validity + deterministic parity, 2/2 |
+| `pipeline` (I2I / edit) | `test_edit_pipeline_parity.py` | validity / wiring gate (condition image enters via a KV-cache write pass, not noise-the-latent). diffusers distribution parity deferred (stochastic AR) |
+
+`processor` (`GlmImageProcessor`) and `scheduler` (`FlowMatchEulerDiscreteScheduler`,
+shared flow-matching) are exercised via the AR and pipeline tests.
+
+Run:
+
+```bash
+pytest tests/local_tests/glm_image/test_transformer_parity.py -v -s
+pytest tests/local_tests/glm_image/test_vae_parity.py -v -s
+pytest tests/local_tests/glm_image/test_t5_parity.py -v -s
+pytest tests/local_tests/glm_image/test_ar_parity.py -v -s
+pytest tests/local_tests/glm_image/test_pipeline_parity.py -v -s
+pytest tests/local_tests/glm_image/test_edit_pipeline_parity.py -v -s
+```
+
+## Design notes (load-bearing — don't regress)
+
+- **Pipeline parity is by injection.** Full-pipeline pixel parity is ill-posed
+  (stochastic AR + independent latent RNG), so the pipeline tests inject matched
+  `(prompt_embeds, prior_token_ids, latents)` into both the diffusers pipeline
+  and the real `GlmImageDenoisingStage`, then compare denoised latents + decoded
+  image — not raw end-to-end pixels.
+- **HF imports are confined** to the lazy-wrapper loader
+  `fastvideo/models/encoders/glm_image_ar_loader.py`, never the production
+  loaders (production-boundary rule).
+- **Tokenizer is ByT5**, matching the diffusers reference (not plain T5);
+  `compute_glyph_embeds` mirrors diffusers `_get_glyph_embeds`.
+- **Model-specific stages** live under
+  `fastvideo/pipelines/basic/glm_image/stages/` per the `add-model` Files Map.
+- **Param mapping** is handled in-place at load via `param_names_mapping`
+  (`fastvideo/configs/models/dits/glm_image.py`) + the native VAE/encoder
+  configs — no conversion script; the loader raises on any unmatched param and
+  `test_transformer_strict_load.py` asserts completeness with `strict=True`.

tests/local_tests/glm_image/test_ar_parity.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import pytest
+import torch
+
+os.environ.setdefault("MASTER_ADDR", "localhost")
+os.environ.setdefault("MASTER_PORT", "29519")
+os.environ.setdefault("DISABLE_SP", "1")
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+FAMILY = "glm_image"
+LOCAL_WEIGHTS_DIR = Path(
+    os.getenv("GLM_IMAGE_LOCAL_WEIGHTS_DIR",
+              REPO_ROOT / "official_weights" / FAMILY))
+AR_DIR = LOCAL_WEIGHTS_DIR / "vision_language_encoder"
+PROCESSOR_DIR = LOCAL_WEIGHTS_DIR / "processor"
+
+
+def _has_weights() -> bool:
+    return AR_DIR.exists() and any(AR_DIR.glob("*.safetensors"))
+
+
+def _has_glm_image_transformers() -> bool:
+    import importlib.util
+    spec = importlib.util.find_spec("transformers")
+    if spec is None:
+        return False
+    import transformers
+    return hasattr(transformers, "GlmImageForConditionalGeneration")
+
+
+pytestmark = [
+    pytest.mark.skipif(
+        not _has_weights(),
+        reason=f"GLM-Image AR encoder weights not found at {AR_DIR}.",
+    ),
+    pytest.mark.skipif(
+        not _has_glm_image_transformers(),
+        reason=("transformers in this env lacks "
+                "`GlmImageForConditionalGeneration` (main pins 4.57.3; needs "
+                "5.0.0rc0+). Bump transformers locally to exercise the AR "
+                "lazy-wrapper."),
+    ),
+]
+
+
+SAMPLE_PROMPT = (
+    "A beautiful landscape photography with rolling hills, a winding river, "
+    "and a vibrant sunset in the background. Photorealistic style."
+)
+
+
+@pytest.fixture(scope="module")
+def device() -> torch.device:
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required for GLM-Image AR encoder parity.")
+    return torch.device("cuda")
+
+
+@pytest.fixture(scope="module")
+def fastvideo_ar(device):
+    pytest.importorskip("fastvideo")
+    try:
+        from fastvideo.models.encoders.glm_image_ar_loader import (
+            GlmImageARLoader)
+    except ImportError as e:
+        pytest.skip(f"FastVideo lazy-wrapper not yet at target path: {e}")
+    loader = GlmImageARLoader(str(AR_DIR), str(PROCESSOR_DIR),
+                              torch_dtype=torch.bfloat16)
+    loader.to(device).eval()
+    return loader
+
+
+def test_ar_wrapper_exposes_expected_surface(fastvideo_ar):
+    assert callable(getattr(fastvideo_ar, "generate", None))
+    assert getattr(fastvideo_ar, "processor", None) is not None
+    assert getattr(fastvideo_ar, "config", None) is not None
+    assert getattr(fastvideo_ar, "generation_config", None) is not None
+    from transformers import GlmImageForConditionalGeneration
+    assert isinstance(fastvideo_ar._model, GlmImageForConditionalGeneration)
+
+
+def test_ar_to_and_eval_propagate(fastvideo_ar, device):
+    assert next(fastvideo_ar._model.parameters()).device.type == device.type
+    assert not fastvideo_ar._model.training
+
+
+def test_ar_generate_is_deterministic_under_do_sample_false(fastvideo_ar,
+                                                            device):
+    processor = fastvideo_ar.processor
+    messages = [{
+        "role": "user",
+        "content": [{"type": "text", "text": SAMPLE_PROMPT}],
+    }]
+    inputs = processor.apply_chat_template(messages, tokenize=True,
+                                           target_h=1024, target_w=1024,
+                                           return_dict=True,
+                                           return_tensors="pt").to(device)
+    with torch.no_grad():
+        out_a = fastvideo_ar.generate(**inputs, max_new_tokens=32,
+                                      do_sample=False)
+        out_b = fastvideo_ar.generate(**inputs, max_new_tokens=32,
+                                      do_sample=False)
+    assert out_a.shape == out_b.shape
+    assert torch.equal(out_a, out_b)

tests/local_tests/glm_image/test_edit_pipeline_parity.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+"""End-to-end scaffold for the GLM-Image image-to-image (edit) pipeline.
+
+Drives FastVideo `VideoGenerator` in edit mode (a condition image is passed) and
+checks that the unified pipeline routes through the condition-encoding stage +
+KV-cache denoising path and emits a well-formed image. Like the T2I parity test,
+pixel-exact parity vs diffusers is NOT asserted: the AR prior is sampled
+(`do_sample=True`) and the diffusion latents draw from independent RNG streams,
+so the codebases produce different valid samples of the same edit. We gate on
+validity (and, when diffusers is present, a comparable brightness regime).
+
+Skips cleanly until weights are available; GPU-heavy.
+"""
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+
+os.environ.setdefault("MASTER_ADDR", "localhost")
+os.environ.setdefault("MASTER_PORT", "29520")
+os.environ.setdefault("DISABLE_SP", "1")
+os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+FAMILY = "glm_image"
+LOCAL_WEIGHTS_DIR = Path(
+    os.getenv("GLM_IMAGE_LOCAL_WEIGHTS_DIR",
+              REPO_ROOT / "official_weights" / FAMILY))
+CONDITION_IMAGE = REPO_ROOT / "assets" / "images" / "couple.jpg"
+
+
+def _has_weights() -> bool:
+    required = ["transformer", "vae", "text_encoder",
+                "vision_language_encoder", "processor", "tokenizer",
+                "scheduler"]
+    return all((LOCAL_WEIGHTS_DIR / r).exists() for r in required)
+
+
+def _upstream_glm_image_available() -> bool:
+    try:
+        import transformers
+    except ImportError:
+        return False
+    return hasattr(transformers, "GlmImageForConditionalGeneration")
+
+
+pytestmark = [
+    pytest.mark.skipif(
+        not _has_weights(),
+        reason=f"GLM-Image full weights not found at {LOCAL_WEIGHTS_DIR}.",
+    ),
+    pytest.mark.skipif(
+        not _upstream_glm_image_available(),
+        reason=("Edit pipeline needs transformers>=5.0.0rc0 "
+                "(GlmImageForConditionalGeneration); main pin predates it. "
+                "Bump locally to run."),
+    ),
+]
+
+
+EDIT_PROMPT = "Make the scene a snowy winter landscape."
+SEED = 0
+HEIGHT = 512
+WIDTH = 512
+STEPS = 8  # low step count keeps the test under 5 min on a single GPU
+
+
+@pytest.fixture(scope="module")
+def device():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required for GLM-Image edit pipeline.")
+    return torch.device("cuda")
+
+
+def _to_uint8_hwc(a) -> np.ndarray:
+    """Coerce a frame buffer (np or torch, [0,1] float or uint8, with optional
+    leading batch/time dims, HWC or CHW) to a single (H, W, 3) uint8 image."""
+    if torch.is_tensor(a):
+        a = a.float().cpu().numpy()
+    a = np.asarray(a)
+    while a.ndim > 3:
+        a = a[0]
+    if a.ndim == 3 and a.shape[0] in (1, 3) and a.shape[-1] not in (1, 3):
+        a = np.transpose(a, (1, 2, 0))
+    if a.dtype != np.uint8:
+        scale = 255.0 if float(a.max()) <= 1.5 else 1.0
+        a = np.clip(a * scale, 0, 255).astype(np.uint8)
+    return a
+
+
+def _fastvideo_edit_image(device) -> np.ndarray:
+    pytest.importorskip("fastvideo")
+    try:
+        from fastvideo import VideoGenerator
+    except ImportError as e:
+        pytest.skip(f"FastVideo VideoGenerator unavailable: {e}")
+    condition = Image.open(CONDITION_IMAGE).convert("RGB")
+    gen = VideoGenerator.from_pretrained(str(LOCAL_WEIGHTS_DIR), num_gpus=1,
+                                         trust_remote_code=True)
+    result = gen.generate_video(prompt=EDIT_PROMPT,
+                                pil_image=condition,
+                                save_video=False,
+                                return_frames=True,
+                                height=HEIGHT, width=WIDTH,
+                                num_inference_steps=STEPS,
+                                guidance_scale=1.5,
+                                seed=SEED)
+    gen.shutdown()
+    return _to_uint8_hwc(result["frames"][0])
+
+
+def test_edit_pipeline_produces_valid_image(device):
+    """Wiring gate for the edit path: passing a condition image routes through
+    the condition-encoding (KV-cache write) stage and the read/skip denoising
+    path, and the pipeline emits a well-formed, non-degenerate image of the
+    requested size. Component numerical correctness is covered by the (green)
+    DiT / VAE / T5 / AR component-parity tests."""
+    fv = _fastvideo_edit_image(device)
+    assert fv.shape == (HEIGHT, WIDTH, 3), f"unexpected shape {fv.shape}"
+    assert fv.dtype == np.uint8 and np.isfinite(fv).all()
+    assert fv.std() > 10.0, f"image is near-constant (std={fv.std():.2f})"

tests/local_tests/glm_image/test_pipeline_parity.py

@@ -0,0 +1,331 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+
+os.environ.setdefault("MASTER_ADDR", "localhost")
+os.environ.setdefault("MASTER_PORT", "29519")
+os.environ.setdefault("DISABLE_SP", "1")
+os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+FAMILY = "glm_image"
+LOCAL_WEIGHTS_DIR = Path(
+    os.getenv("GLM_IMAGE_LOCAL_WEIGHTS_DIR",
+              REPO_ROOT / "official_weights" / FAMILY))
+TRANSFORMER_DIR = LOCAL_WEIGHTS_DIR / "transformer"
+
+
+def _has_weights() -> bool:
+    required = ["transformer", "vae", "text_encoder",
+                "vision_language_encoder", "processor", "tokenizer",
+                "scheduler"]
+    return all((LOCAL_WEIGHTS_DIR / r).exists() for r in required)
+
+
+def _upstream_glm_image_available() -> bool:
+    try:
+        import transformers
+        import diffusers
+    except ImportError:
+        return False
+    return (hasattr(transformers, "GlmImageForConditionalGeneration")
+            and hasattr(diffusers, "GlmImagePipeline"))
+
+
+pytestmark = [
+    pytest.mark.skipif(
+        not _has_weights(),
+        reason=f"GLM-Image full weights not found at {LOCAL_WEIGHTS_DIR}.",
+    ),
+    pytest.mark.skipif(
+        not _upstream_glm_image_available(),
+        reason=("Pipeline parity needs transformers>=5.0.0rc0 and "
+                "diffusers>=0.37.0.dev0; main pins predate both. Bump locally "
+                "to run."),
+    ),
+]
+
+
+SAMPLE_PROMPT = (
+    "A landscape photo with rolling green hills under a clear blue sky.")
+SEED = 0
+HEIGHT = 512
+WIDTH = 512
+STEPS = 8
+
+# bf16 + a different SDPA kernel across 30 DiT layers x STEPS steps leaves a
+# small residual; a real wiring/schedule bug is far larger than these bounds.
+LATENT_COSINE_MIN = 0.995
+IMAGE_MAE_MAX = 5.0  # /255
+
+
+@pytest.fixture(scope="module")
+def device():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required for GLM-Image pipeline parity.")
+    return torch.device("cuda")
+
+
+def _to_uint8_hwc(a) -> np.ndarray:
+    if torch.is_tensor(a):
+        a = a.detach().float().cpu().numpy()
+    a = np.asarray(a)
+    while a.ndim > 3:
+        a = a[0]
+    if a.ndim == 3 and a.shape[0] in (1, 3) and a.shape[-1] not in (1, 3):
+        a = np.transpose(a, (1, 2, 0))
+    if a.dtype != np.uint8:
+        scale = 255.0 if float(a.max()) <= 1.5 else 1.0
+        a = np.clip(a * scale, 0, 255).astype(np.uint8)
+    return a
+
+
+# --------------------------------------------------------------------------- #
+# FastVideo native DiT loader (mirrors test_transformer_parity.py).
+# --------------------------------------------------------------------------- #
+def _load_state_dict(dir_path: Path) -> dict[str, torch.Tensor]:
+    safetensors = pytest.importorskip("safetensors.torch")
+    sd: dict[str, torch.Tensor] = {}
+    for shard in sorted(dir_path.glob("*.safetensors")):
+        sd.update(safetensors.load_file(str(shard)))
+    return sd
+
+
+def _apply_param_mapping(sd, mapping):
+    import re
+    out = {}
+    for k, v in sd.items():
+        new_k = k
+        for pat, repl in mapping.items():
+            if re.match(pat, k):
+                new_k = re.sub(pat, repl, k)
+                break
+        out[new_k] = v
+    return out
+
+
+def _ensure_distributed():
+    """The denoising stage calls get_local_torch_device(), which needs
+    FastVideo's world/TP groups. Initialize a single-rank (world_size=1) group
+    in-process (mirrors tests/local_tests/sd35/test_sd35_component_parity.py)."""
+    import torch.distributed as dist
+    from fastvideo.distributed.parallel_state import (
+        get_tp_group, init_distributed_environment, initialize_model_parallel)
+    try:
+        get_tp_group()
+        return
+    except Exception:
+        pass
+    if not dist.is_initialized():
+        os.environ.setdefault("RANK", "0")
+        os.environ.setdefault("WORLD_SIZE", "1")
+        os.environ.setdefault("LOCAL_RANK", "0")
+        torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+        store_path = f"/tmp/fastvideo_glm_pg_{os.getpid()}.store"
+        dist.init_process_group(backend="nccl",
+                                init_method=f"file://{store_path}",
+                                rank=0, world_size=1)
+    init_distributed_environment(world_size=1, rank=0, local_rank=0,
+                                 distributed_init_method="env://")
+    try:
+        get_tp_group()
+    except Exception:
+        initialize_model_parallel(tensor_model_parallel_size=1,
+                                  sequence_model_parallel_size=1,
+                                  data_parallel_size=1)
+
+
+def _load_fastvideo_transformer(device, dtype):
+    from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig
+    from fastvideo.models.dits.glm_image import GlmImageTransformer2DModel
+    cfg = GlmImageDiTConfig()
+    model = GlmImageTransformer2DModel(
+        cfg, {"_class_name": "GlmImageTransformer2DModel"})
+    sd = _load_state_dict(TRANSFORMER_DIR)
+    sd = _apply_param_mapping(sd, cfg.arch_config.param_names_mapping)
+    missing, unexpected = model.load_state_dict(sd, strict=False)
+    assert not missing, f"FastVideo DiT missing keys: {missing[:10]}"
+    assert not unexpected, f"FastVideo DiT unexpected keys: {unexpected[:10]}"
+    return model.to(device, dtype=dtype).eval()
+
+
+def _fastvideo_denoise_latents(device, dtype, *, prompt_embeds, prior_token_ids,
+                               init_latents):
+    """Drive the real GlmImageDenoisingStage with injected, matched inputs and
+    return the denoised latents (1, 16, H/8, W/8)."""
+    from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+        FlowMatchEulerDiscreteScheduler)
+    from fastvideo.pipelines.basic.glm_image.stages.before_denoising import (
+        calculate_shift)
+    from fastvideo.pipelines.basic.glm_image.stages.denoising import (
+        GlmImageDenoisingStage)
+    from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+    _ensure_distributed()
+    transformer = _load_fastvideo_transformer(device, dtype)
+    scheduler = FlowMatchEulerDiscreteScheduler(shift=1.0)
+    stage = GlmImageDenoisingStage(transformer=transformer, scheduler=scheduler)
+
+    # Reproduce the production timestep/sigma schedule (before_denoising.py):
+    # integer-cast linspace timesteps, resolution-dependent shift on sigmas.
+    ntt = scheduler.config.num_train_timesteps
+    patch = transformer.patch_size
+    image_seq_len = ((HEIGHT // 8) * (WIDTH // 8)) // (patch**2)
+    sched_t = np.linspace(ntt, 1.0, STEPS + 1)[:-1].astype(
+        np.int64).astype(np.float32)
+    scheduler.set_shift(calculate_shift(image_seq_len))
+    scheduler.set_timesteps(STEPS, device=device,
+                            sigmas=(sched_t / ntt).tolist(),
+                            timesteps=sched_t.tolist())
+
+    batch = ForwardBatch(data_type="image")
+    batch.prompt_embeds = [prompt_embeds.to(device, dtype)]
+    batch.attention_mask = None  # match diffusers (no text padding mask)
+    batch.prior_token_id = prior_token_ids.to(device)
+    batch.prior_token_drop = torch.zeros_like(prior_token_ids,
+                                              dtype=torch.bool, device=device)
+    batch.latents = init_latents.clone().unsqueeze(2).to(device, dtype)
+    batch.timesteps = scheduler.timesteps
+    batch.height, batch.width = HEIGHT, WIDTH
+    batch.num_inference_steps = STEPS
+    batch.guidance_scale = 1.5
+    batch.do_classifier_free_guidance = True
+    batch.seed = SEED
+    batch.extra = {}  # no glm_kv_caches -> T2I denoise path
+
+    out = stage.forward(batch, fastvideo_args=None)
+    latents = out.latents
+    if latents.dim() == 5:
+        latents = latents.squeeze(2)
+    latents = latents.detach().float().cpu()
+    del transformer, stage
+    torch.cuda.empty_cache()
+    import torch.distributed as dist
+    if dist.is_initialized():
+        dist.destroy_process_group()
+    return latents
+
+
+def _fastvideo_image(device) -> np.ndarray:
+    pytest.importorskip("fastvideo")
+    try:
+        from fastvideo import VideoGenerator
+    except ImportError as e:
+        pytest.skip(f"FastVideo VideoGenerator unavailable: {e}")
+    gen = VideoGenerator.from_pretrained(str(LOCAL_WEIGHTS_DIR), num_gpus=1,
+                                         trust_remote_code=True)
+    result = gen.generate_video(prompt=SAMPLE_PROMPT,
+                                save_video=False,
+                                return_frames=True,
+                                height=HEIGHT, width=WIDTH,
+                                num_inference_steps=STEPS,
+                                guidance_scale=1.5,
+                                seed=SEED)
+    gen.shutdown()
+    return _to_uint8_hwc(result["frames"][0])
+
+
+def test_fastvideo_pipeline_produces_valid_image(device):
+    fv = _fastvideo_image(device)
+    assert fv.shape == (HEIGHT, WIDTH, 3), f"unexpected shape {fv.shape}"
+    assert fv.dtype == np.uint8 and np.isfinite(fv).all()
+    assert fv.std() > 10.0, f"image is near-constant (std={fv.std():.2f})"
+
+
+def test_pipeline_denoise_parity_deterministic(device):
+    """Inject identical (prompt_embeds, prior_token_ids, init_latents) into the
+    diffusers pipeline and FastVideo's GlmImageDenoisingStage, then compare the
+    denoised latents and the decoded images. With the stochastic AR and the
+    independent latent RNG removed, the two implementations must agree to within
+    bf16/SDPA-kernel noise."""
+    # Inference only: the T5 glyph encoder leaves a grad graph on the embeds,
+    # which later breaks image_processor.postprocess()'s .numpy() on the decoded
+    # image. Disabling grad for the whole test also trims memory.
+    torch.set_grad_enabled(False)
+    diffusers = pytest.importorskip("diffusers")
+    from diffusers.utils.torch_utils import randn_tensor
+
+    dtype = torch.bfloat16
+    pipe = diffusers.GlmImagePipeline.from_pretrained(
+        str(LOCAL_WEIGHTS_DIR), torch_dtype=dtype).to(device)
+
+    gen = torch.Generator(device=device).manual_seed(SEED)
+
+    # --- shared, deterministic inputs (computed once, fed to both) ---------- #
+    pos, neg = pipe.encode_prompt(SAMPLE_PROMPT,
+                                  do_classifier_free_guidance=True,
+                                  device=device, dtype=dtype)
+    assert pos.shape[1] == neg.shape[1], (
+        f"quote-free prompt should give equal pos/neg glyph lengths, got "
+        f"{pos.shape[1]} vs {neg.shape[1]}")
+    prior_token_ids, _, _ = pipe.generate_prior_tokens(
+        SAMPLE_PROMPT, HEIGHT, WIDTH, image=None, device=device, generator=gen)
+    latent_ch = pipe.transformer.config.in_channels
+    init_latents = randn_tensor((1, latent_ch, HEIGHT // 8, WIDTH // 8),
+                                generator=gen, device=device, dtype=dtype)
+
+    # --- official denoised latents ----------------------------------------- #
+    # prompt=None: diffusers check_inputs rejects passing prompt + prompt_embeds
+    # together; the embeds (and prior tokens) fully specify the run.
+    official = pipe(prompt=None,
+                    prompt_embeds=pos, negative_prompt_embeds=neg,
+                    prior_token_ids=prior_token_ids,
+                    latents=init_latents.clone(),
+                    height=HEIGHT, width=WIDTH,
+                    num_inference_steps=STEPS, guidance_scale=1.5,
+                    output_type="latent").images.float().cpu()
+
+    # FastVideo packs CFG as a single [pos; neg] 2-row tensor (no padding here
+    # because L_pos == L_neg, asserted above).
+    packed = torch.cat([pos, neg], dim=0)
+
+    # --- FastVideo denoised latents (real GlmImageDenoisingStage) ---------- #
+    fv = _fastvideo_denoise_latents(device, dtype,
+                                    prompt_embeds=packed,
+                                    prior_token_ids=prior_token_ids,
+                                    init_latents=init_latents)
+
+    assert official.shape == fv.shape, (
+        f"latent shape mismatch: {official.shape} vs {fv.shape}")
+
+    # --- latent-level agreement -------------------------------------------- #
+    cos = torch.nn.functional.cosine_similarity(
+        official.flatten(), fv.flatten(), dim=0).item()
+    lat_mae = (official - fv).abs().mean().item()
+    lat_scale = official.abs().mean().item()
+    print(f"[glm-image denoise parity] latent cosine={cos:.6f} "
+          f"MAE={lat_mae:.5f} (|latent| mean={lat_scale:.5f}, "
+          f"rel={lat_mae / max(lat_scale, 1e-6):.4f})")
+
+    # --- decoded-image agreement (same VAE both sides isolates denoise) ----- #
+    def _decode(lat):
+        lat = lat.to(device, dtype)
+        mean = torch.tensor(pipe.vae.config.latents_mean).view(
+            1, -1, 1, 1).to(device, dtype)
+        std = torch.tensor(pipe.vae.config.latents_std).view(
+            1, -1, 1, 1).to(device, dtype)
+        img = pipe.vae.decode((lat * std + mean), return_dict=False)[0]
+        return _to_uint8_hwc(pipe.image_processor.postprocess(
+            img, output_type="np")[0])
+
+    off_img, fv_img = _decode(official), _decode(fv)
+    img_mae = np.abs(off_img.astype(np.float32)
+                     - fv_img.astype(np.float32)).mean()
+    print(f"[glm-image denoise parity] decoded image MAE={img_mae:.3f}/255 "
+          f"(diffusers mean={off_img.mean():.1f}, fv mean={fv_img.mean():.1f})")
+
+    del pipe
+    torch.cuda.empty_cache()
+
+    assert cos > LATENT_COSINE_MIN, (
+        f"denoised-latent cosine {cos:.6f} < {LATENT_COSINE_MIN}: the FastVideo "
+        "and official denoise diverge well beyond kernel noise (wiring bug).")
+    assert img_mae < IMAGE_MAE_MAX, (
+        f"decoded-image MAE {img_mae:.3f}/255 >= {IMAGE_MAE_MAX}: outputs are "
+        "not equivalent under matched inputs (wiring/schedule bug).")

tests/local_tests/glm_image/test_t5_parity.py

@@ -0,0 +1,180 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Parity scaffold for GLM-Image glyph T5 encoder.
+
+The text_encoder is a small byte-level T5 (d_model=1472, gated-gelu, ByT5
+tokenizer, vocab_size=384). This test exercises the glyph-only encoding path
+that the diffusers pipeline implements as `_get_glyph_embeds` and compares
+against the FastVideo glyph-encoding stage logic. Until the FastVideo stage is
+refactored to mirror diffusers (per-prompt loop, even-length pad-prefix,
+attention-mask flattening, left-padded batch), the parity test asserts the
+high-impact pieces individually so failures point to the exact path that
+diverged.
+"""
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import pytest
+import torch
+
+os.environ.setdefault("MASTER_ADDR", "localhost")
+os.environ.setdefault("MASTER_PORT", "29519")
+os.environ.setdefault("DISABLE_SP", "1")
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+FAMILY = "glm_image"
+LOCAL_WEIGHTS_DIR = Path(
+    os.getenv("GLM_IMAGE_LOCAL_WEIGHTS_DIR",
+              REPO_ROOT / "official_weights" / FAMILY))
+TEXT_ENCODER_DIR = LOCAL_WEIGHTS_DIR / "text_encoder"
+TOKENIZER_DIR = LOCAL_WEIGHTS_DIR / "tokenizer"
+
+
+def _has_weights() -> bool:
+    return (TEXT_ENCODER_DIR / "model.safetensors").exists() and TOKENIZER_DIR.exists()
+
+
+pytestmark = pytest.mark.skipif(
+    not _has_weights(),
+    reason=f"GLM-Image text_encoder/tokenizer not found under {LOCAL_WEIGHTS_DIR}.",
+)
+
+
+SAMPLE_PROMPT = (
+    'A photo of a coffee shop with a wooden sign reading "Daily Grind" '
+    "and a chalk menu next to it"
+)
+
+
+@pytest.fixture(scope="module")
+def device() -> torch.device:
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required for GLM-Image glyph encoder parity.")
+    return torch.device("cuda")
+
+
+@pytest.fixture(scope="module")
+def official_tokenizer():
+    transformers = pytest.importorskip("transformers")
+    return transformers.ByT5Tokenizer.from_pretrained(str(TOKENIZER_DIR))
+
+
+@pytest.fixture(scope="module")
+def official_text_encoder(device):
+    transformers = pytest.importorskip("transformers")
+    return transformers.T5EncoderModel.from_pretrained(
+        str(TEXT_ENCODER_DIR), torch_dtype=torch.float32).to(device).eval()
+
+
+def test_tokenizer_is_byt5(official_tokenizer):
+    """Glyph tokenizer must be ByT5 (byte-level), not generic T5."""
+    assert type(official_tokenizer).__name__ == "ByT5Tokenizer"
+    assert official_tokenizer.vocab_size <= 400, (
+        "ByT5 vocab is byte-level (~256+special)")
+
+
+def test_glyph_text_extraction():
+    """Glyph extraction must pull single+double+CJK quoted spans."""
+    pytest.importorskip("fastvideo")
+    try:
+        from fastvideo.pipelines.basic.glm_image.stages.before_denoising import (  # noqa
+            get_glyph_texts)
+    except ImportError as e:
+        pytest.skip(f"FastVideo glyph extractor not yet at target path: {e}")
+    out = get_glyph_texts(SAMPLE_PROMPT)
+    assert "Daily Grind" in out
+
+
+def _diffusers_reference_get_glyph_embeds(prompts, tokenizer, text_encoder,
+                                          device, dtype,
+                                          max_sequence_length=2048):
+    """Inlined copy of diffusers `_get_glyph_embeds` + `get_glyph_texts`.
+
+    Source: `reference/diffusers/src/diffusers/pipelines/glm_image/`
+    `pipeline_glm_image.py::GlmImagePipeline._get_glyph_embeds` at commit
+    ff3b86b4755b46a7b5656dfcf84d25bd25ad4740 (diffusers main, 0.37.0.dev0).
+    Inlined here because the installed diffusers package (`__version__`
+    typically <0.37.0) does not yet ship `GlmImagePipeline`. This function is
+    the reference oracle for FastVideo's `compute_glyph_embeds`.
+    """
+    import re
+
+    def get_glyph_texts(prompt):
+        if isinstance(prompt, str):
+            prompt = [prompt]
+        out = []
+        for p in prompt:
+            out.append(
+                re.findall(r"'([^']*)'", p)
+                + re.findall(r"“([^“”]*)”", p)
+                + re.findall(r'"([^"]*)"', p)
+                + re.findall(r"「([^「」]*)」", p))
+        return out
+
+    all_glyph_texts = get_glyph_texts(prompts)
+    all_glyph_embeds = []
+    for glyph_texts in all_glyph_texts:
+        if len(glyph_texts) == 0:
+            glyph_texts = [""]
+        input_ids = tokenizer(
+            glyph_texts,
+            max_length=max_sequence_length,
+            truncation=True,
+        ).input_ids
+        input_ids = [
+            [tokenizer.pad_token_id] * ((len(input_ids) + 1) % 2) + ids
+            for ids in input_ids
+        ]
+        max_length = max(len(ids) for ids in input_ids)
+        attention_mask = torch.tensor(
+            [[1] * len(ids) + [0] * (max_length - len(ids))
+             for ids in input_ids],
+            device=device,
+        )
+        input_ids_t = torch.tensor(
+            [
+                ids + [tokenizer.pad_token_id] * (max_length - len(ids))
+                for ids in input_ids
+            ],
+            device=device,
+        )
+        outputs = text_encoder(input_ids_t, attention_mask=attention_mask)
+        glyph_embeds = outputs.last_hidden_state[attention_mask.bool()].unsqueeze(0)
+        all_glyph_embeds.append(glyph_embeds)
+
+    max_seq_len = max(emb.size(1) for emb in all_glyph_embeds)
+    padded = []
+    for emb in all_glyph_embeds:
+        if emb.size(1) < max_seq_len:
+            pad = torch.zeros(emb.size(0), max_seq_len - emb.size(1),
+                              emb.size(2), device=device, dtype=emb.dtype)
+            emb = torch.cat([pad, emb], dim=1)
+        padded.append(emb)
+    return torch.cat(padded, dim=0).to(device=device, dtype=dtype)
+
+
+def test_get_glyph_embeds_matches_diffusers(official_tokenizer,
+                                            official_text_encoder, device):
+    """FastVideo `compute_glyph_embeds` must equal the diffusers reference."""
+    pytest.importorskip("fastvideo")
+    try:
+        from fastvideo.pipelines.basic.glm_image.stages.before_denoising import (  # noqa
+            compute_glyph_embeds)
+    except ImportError as e:
+        pytest.skip(f"FastVideo glyph encoder not yet at target path: {e}")
+
+    with torch.no_grad():
+        official_embeds = _diffusers_reference_get_glyph_embeds(
+            [SAMPLE_PROMPT], official_tokenizer, official_text_encoder, device,
+            torch.float32)
+        fv_embeds = compute_glyph_embeds([SAMPLE_PROMPT],
+                                         tokenizer=official_tokenizer,
+                                         text_encoder=official_text_encoder,
+                                         device=device,
+                                         dtype=torch.float32,
+                                         max_sequence_length=2048)
+    assert fv_embeds.shape == official_embeds.shape, (
+        f"shape mismatch: {fv_embeds.shape} vs {official_embeds.shape}")
+    torch.testing.assert_close(fv_embeds, official_embeds, atol=1e-4,
+                               rtol=1e-4)

tests/local_tests/glm_image/test_transformer_parity.py

@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Component parity scaffold for GLM-Image DiT (GlmImageTransformer2DModel).
+
+Compares the FastVideo-native port at `fastvideo.models.dits.glm_image` against
+the diffusers reference `GlmImageTransformer2DModel` from
+`diffusers.models.transformers.transformer_glm_image`.
+
+Skips cleanly until both the FastVideo class and the local weights exist; the
+test is real (loads weights, forwards a fixed input, compares output tensors).
+"""
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import pytest
+import torch
+
+os.environ.setdefault("MASTER_ADDR", "localhost")
+os.environ.setdefault("MASTER_PORT", "29519")
+os.environ.setdefault("DISABLE_SP", "1")
+os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+FAMILY = "glm_image"
+LOCAL_WEIGHTS_DIR = Path(
+    os.getenv("GLM_IMAGE_LOCAL_WEIGHTS_DIR",
+              REPO_ROOT / "official_weights" / FAMILY))
+TRANSFORMER_DIR = LOCAL_WEIGHTS_DIR / "transformer"
+
+def _has_weights() -> bool:
+    return TRANSFORMER_DIR.exists() and any(
+        TRANSFORMER_DIR.glob("*.safetensors"))
+
+
+def _has_glm_image_diffusers() -> bool:
+    """`GlmImageTransformer2DModel` lands in diffusers>=0.37.0.dev0.
+
+    FastVideo main pins `diffusers>=0.33.1`; the installed env is 0.36.0. The
+    numerical parity test needs the diffusers class as an oracle — until the
+    installed diffusers is upgraded, this test skips. The structural surface
+    is still verified by `test_glm_image_transformer_strict_load.py`.
+    """
+    import importlib.util
+    spec = importlib.util.find_spec("diffusers")
+    if spec is None:
+        return False
+    import diffusers
+    return hasattr(diffusers, "GlmImageTransformer2DModel")
+
+
+pytestmark = [
+    pytest.mark.skipif(
+        not _has_weights(),
+        reason=(
+            f"GLM-Image transformer weights not found at {TRANSFORMER_DIR}. "
+            "Download via "
+            "`python .agents/skills/add-model-01-prep/scripts/download_hf_weights.py "
+            "zai-org/GLM-Image official_weights/glm_image`."),
+    ),
+    pytest.mark.skipif(
+        not _has_glm_image_diffusers(),
+        reason=("diffusers in this env lacks `GlmImageTransformer2DModel` "
+                "(main allows >=0.33.1; installed 0.36.0; class lands in "
+                "0.37.0.dev0). Bump diffusers locally to run numerical parity. "
+                "Structural surface is covered by "
+                "`test_glm_image_transformer_strict_load.py`."),
+    ),
+]
+
+
+@pytest.fixture(scope="module")
+def device() -> torch.device:
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required for GLM-Image transformer parity.")
+    return torch.device("cuda")
+
+
+def _load_official(device, dtype):
+    diffusers = pytest.importorskip("diffusers")
+    cls = diffusers.GlmImageTransformer2DModel
+    return cls.from_pretrained(str(TRANSFORMER_DIR),
+                               torch_dtype=dtype).to(device).eval()
+
+
+def _load_fastvideo(device, dtype):
+    pytest.importorskip("fastvideo")
+    try:
+        from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig
+        from fastvideo.models.dits.glm_image import GlmImageTransformer2DModel
+    except ImportError as e:
+        pytest.skip(f"FastVideo GLM-Image DiT not yet ported: {e}")
+    cfg = GlmImageDiTConfig()
+    model = GlmImageTransformer2DModel(cfg,
+                                       {"_class_name":
+                                        "GlmImageTransformer2DModel"})
+    sd = _load_state_dict(TRANSFORMER_DIR)
+    sd = _apply_param_mapping(sd, cfg.arch_config.param_names_mapping)
+    missing, unexpected = model.load_state_dict(sd, strict=False)
+    assert not missing, f"FastVideo DiT missing keys: {missing[:10]}"
+    assert not unexpected, f"FastVideo DiT unexpected keys: {unexpected[:10]}"
+    return model.to(device, dtype=dtype).eval()
+
+
+def _forward_pair(device, dtype):
+    """Run both DiTs on identical synthetic inputs; return (fv_out, ref_out)."""
+    from fastvideo.forward_context import set_forward_context
+    ref = _load_official(device, dtype)
+    fv = _load_fastvideo(device, dtype)
+    inputs = _make_inputs(device, dtype=dtype)
+    with torch.no_grad():
+        ref_out = ref(**inputs, return_dict=False)[0]
+        with set_forward_context(current_timestep=0, attn_metadata=None,
+                                 forward_batch=None):
+            fv_out = fv(**inputs)
+        if isinstance(fv_out, tuple):
+            fv_out = fv_out[0]
+    assert ref_out.shape == fv_out.shape, (
+        f"shape mismatch: {ref_out.shape} vs {fv_out.shape}")
+    # Free the two 7B models before the next dtype runs (single-GPU budget).
+    del ref, fv
+    torch.cuda.empty_cache()
+    return fv_out.float(), ref_out.float()
+
+
+def _load_state_dict(dir_path: Path) -> dict[str, torch.Tensor]:
+    safetensors = pytest.importorskip("safetensors.torch")
+    sd: dict[str, torch.Tensor] = {}
+    for shard in sorted(dir_path.glob("*.safetensors")):
+        sd.update(safetensors.load_file(str(shard)))
+    return sd
+
+
+def _apply_param_mapping(sd, mapping):
+    import re
+    out = {}
+    for k, v in sd.items():
+        new_k = k
+        for pat, repl in mapping.items():
+            if re.match(pat, k):
+                new_k = re.sub(pat, repl, k)
+                break
+        out[new_k] = v
+    return out
+
+
+def _make_inputs(device, batch_size=1, height=64, width=64,
+                 text_seq_len=32, dtype=torch.bfloat16) -> dict[str, torch.Tensor]:
+    """Synthetic inputs matching the diffusers `GlmImageTransformer2DModel.forward`
+    contract. Latent shape is `(B, 16, H, W)`; post-patch grid is `(H/2, W/2)`
+    so `prior_token_id` has `(H/2)*(W/2)` entries per batch element."""
+    torch.manual_seed(0)
+    patch_size = 2
+    num_patches = (height // patch_size) * (width // patch_size)
+    return {
+        "hidden_states": torch.randn(batch_size, 16, height, width,
+                                     device=device, dtype=dtype),
+        "encoder_hidden_states": torch.randn(batch_size, text_seq_len, 1472,
+                                             device=device, dtype=dtype),
+        "prior_token_id": torch.randint(0, 16384, (batch_size, num_patches),
+                                        device=device, dtype=torch.long),
+        "prior_token_drop": torch.zeros(batch_size, device=device,
+                                        dtype=torch.bool),
+        "timestep": torch.tensor([500] * batch_size, device=device,
+                                 dtype=torch.long),
+        "target_size": torch.tensor([[height * 8, width * 8]] * batch_size,
+                                    device=device, dtype=torch.long),
+        "crop_coords": torch.zeros(batch_size, 2, device=device,
+                                   dtype=torch.long),
+    }
+
+
+# fp32 is the structural-correctness gate: with identical math the two
+# implementations agree to ~1e-3 (residual is SDPA reduction-order noise).
+ATOL_FP32 = 1e-2
+RTOL_FP32 = 1e-2
+# bf16 is the inference dtype. Rounding accumulates across 30 layers and the
+# attention kernels differ, so element-wise atol is meaningless; cosine
+# similarity is the right statistical bar for "same function, bf16 noise".
+COS_MIN_BF16 = 0.999
+
+
+def test_transformer_forward_matches_diffusers_fp32(device):
+    """Structural gate: in fp32 the FastVideo DiT is numerically faithful to
+    the diffusers reference (cosine ~1.0, element-wise close to ~1e-3)."""
+    fv_out, ref_out = _forward_pair(device, torch.float32)
+    cos = torch.nn.functional.cosine_similarity(
+        fv_out.flatten(), ref_out.flatten(), dim=0)
+    assert cos > 0.9999, f"fp32 cosine similarity too low: {cos:.6f}"
+    torch.testing.assert_close(fv_out, ref_out,
+                               atol=ATOL_FP32, rtol=RTOL_FP32)
+
+
+def test_transformer_forward_matches_diffusers_bf16(device):
+    """Inference-dtype sanity: bf16 output tracks the diffusers reference up to
+    precision accumulation (high cosine similarity). Element-wise atol is not a
+    meaningful bar for a 30-layer bf16 transformer vs a different attention
+    kernel (~28% of elements exceed atol=5e-3 purely from rounding) — see
+    PORT_STATUS I020."""
+    fv_out, ref_out = _forward_pair(device, torch.bfloat16)
+    cos = torch.nn.functional.cosine_similarity(
+        fv_out.flatten(), ref_out.flatten(), dim=0)
+    assert cos > COS_MIN_BF16, (
+        f"bf16 cosine similarity {cos:.6f} below {COS_MIN_BF16}; "
+        "indicates a structural divergence, not just precision.")

tests/local_tests/glm_image/test_transformer_strict_load.py

@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Strict-load smoke for GLM-Image DiT.
+
+Loads all three transformer shards from `official_weights/glm_image/transformer/`,
+applies the FastVideo `param_names_mapping`, and calls
+`GlmImageTransformer2DModel.load_state_dict(..., strict=True)`. Asserts the
+config-driven model instantiates and every checkpoint key has a matching
+FastVideo parameter with the same shape and dtype.
+
+This is weight-load verification only — it does not run a forward pass. Numerical
+parity against the diffusers reference lives in
+`test_glm_image_transformer_parity.py` and requires `diffusers >= 0.37.0.dev0`.
+"""
+from __future__ import annotations
+
+import os
+import re
+from pathlib import Path
+
+import pytest
+import torch
+
+os.environ.setdefault("MASTER_ADDR", "localhost")
+os.environ.setdefault("MASTER_PORT", "29519")
+os.environ.setdefault("DISABLE_SP", "1")
+os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+FAMILY = "glm_image"
+LOCAL_WEIGHTS_DIR = Path(
+    os.getenv("GLM_IMAGE_LOCAL_WEIGHTS_DIR",
+              REPO_ROOT / "official_weights" / FAMILY))
+TRANSFORMER_DIR = LOCAL_WEIGHTS_DIR / "transformer"
+
+
+def _has_shards() -> bool:
+    return TRANSFORMER_DIR.exists() and any(
+        TRANSFORMER_DIR.glob("*.safetensors"))
+
+
+pytestmark = pytest.mark.skipif(
+    not _has_shards(),
+    reason=f"GLM-Image transformer shards not at {TRANSFORMER_DIR}.",
+)
+
+
+def test_dit_strict_load_against_hf_checkpoint():
+    safetensors = pytest.importorskip("safetensors.torch")
+    from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig
+    from fastvideo.models.dits.glm_image import GlmImageTransformer2DModel
+
+    cfg = GlmImageDiTConfig()
+    mapping = cfg.arch_config.param_names_mapping
+
+    # 1. Build the FastVideo model on CPU/meta-free path.
+    torch.set_default_device("cpu")
+    hf_config = {"_class_name": "GlmImageTransformer2DModel"}
+    model = GlmImageTransformer2DModel(cfg, hf_config)
+
+    # 2. Load every shard and apply param_names_mapping.
+    raw_sd: dict[str, torch.Tensor] = {}
+    for shard in sorted(TRANSFORMER_DIR.glob("*.safetensors")):
+        raw_sd.update(safetensors.load_file(str(shard)))
+
+    def rename(k: str) -> str:
+        for pat, repl in mapping.items():
+            if re.match(pat, k):
+                return re.sub(pat, repl, k)
+        return k
+
+    renamed_sd = {rename(k): v for k, v in raw_sd.items()}
+
+    # 3. Compare key sets (sanity, redundant with strict=True but gives a
+    # clearer error message if it fails).
+    fv_keys = set(model.state_dict().keys())
+    ckpt_keys = set(renamed_sd.keys())
+    missing = fv_keys - ckpt_keys
+    unexpected = ckpt_keys - fv_keys
+    assert not missing, f"FastVideo DiT missing {len(missing)} keys, e.g. {sorted(missing)[:5]}"
+    assert not unexpected, f"checkpoint has {len(unexpected)} unexpected keys, e.g. {sorted(unexpected)[:5]}"
+
+    # 4. Shape and dtype sanity — strict_load with shape mismatch raises early.
+    fv_sd = model.state_dict()
+    shape_mismatches = []
+    for k, v in renamed_sd.items():
+        if v.shape != fv_sd[k].shape:
+            shape_mismatches.append((k, tuple(fv_sd[k].shape), tuple(v.shape)))
+    assert not shape_mismatches, (
+        f"shape mismatches: {shape_mismatches[:5]}")
+
+    # 5. Strict load.
+    incompatible = model.load_state_dict(renamed_sd, strict=True)
+    assert not incompatible.missing_keys
+    assert not incompatible.unexpected_keys

tests/local_tests/glm_image/test_vae_parity.py

@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Component parity scaffold for GLM-Image VAE (AutoencoderKL).
+
+Compares the FastVideo-native AutoencoderKL port against diffusers'
+`AutoencoderKL` loaded from `zai-org/GLM-Image/vae`. The published checkpoint
+uses `block_out_channels=[128, 512, 1024, 1024]`, `latent_channels=16`, and
+per-channel `latents_mean[16]` + `latents_std[16]` normalization (not a scalar
+`scaling_factor`).
+
+Skips cleanly until both the FastVideo class and the VAE shard exist.
+"""
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+import pytest
+import torch
+
+os.environ.setdefault("MASTER_ADDR", "localhost")
+os.environ.setdefault("MASTER_PORT", "29519")
+os.environ.setdefault("DISABLE_SP", "1")
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+FAMILY = "glm_image"
+LOCAL_WEIGHTS_DIR = Path(
+    os.getenv("GLM_IMAGE_LOCAL_WEIGHTS_DIR",
+              REPO_ROOT / "official_weights" / FAMILY))
+VAE_DIR = LOCAL_WEIGHTS_DIR / "vae"
+
+
+def _has_weights() -> bool:
+    return (VAE_DIR / "diffusion_pytorch_model.safetensors").exists()
+
+
+pytestmark = pytest.mark.skipif(
+    not _has_weights(),
+    reason=f"GLM-Image VAE weights not found at {VAE_DIR}.",
+)
+
+
+@pytest.fixture(scope="module")
+def device() -> torch.device:
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required for GLM-Image VAE parity.")
+    return torch.device("cuda")
+
+
+@pytest.fixture(scope="module")
+def vae_config():
+    with open(VAE_DIR / "config.json") as f:
+        return json.load(f)
+
+
+@pytest.fixture(scope="module")
+def official_vae(device):
+    diffusers = pytest.importorskip("diffusers")
+    vae = diffusers.AutoencoderKL.from_pretrained(
+        str(VAE_DIR), torch_dtype=torch.float32).to(device).eval()
+    return vae
+
+
+@pytest.fixture(scope="module")
+def fastvideo_vae(device):
+    pytest.importorskip("fastvideo")
+    try:
+        from fastvideo.configs.models.vaes.glm_image import GlmImageVAEConfig
+        from fastvideo.models.vaes.autoencoder_kl import AutoencoderKL
+    except ImportError as e:
+        pytest.skip(f"FastVideo native AutoencoderKL not yet ported: {e}")
+    cfg = GlmImageVAEConfig()
+    vae = AutoencoderKL(cfg)
+    safetensors = pytest.importorskip("safetensors.torch")
+    sd = safetensors.load_file(
+        str(VAE_DIR / "diffusion_pytorch_model.safetensors"))
+    missing, unexpected = vae.load_state_dict(sd, strict=False)
+    assert not missing, f"FastVideo VAE missing keys: {missing[:10]}"
+    assert not unexpected, f"FastVideo VAE unexpected keys: {unexpected[:10]}"
+    return vae.to(device, dtype=torch.float32).eval()
+
+
+def test_vae_config_matches_real_checkpoint(vae_config):
+    assert vae_config["block_out_channels"] == [128, 512, 1024, 1024]
+    assert vae_config["latent_channels"] == 16
+    assert "latents_mean" in vae_config and len(
+        vae_config["latents_mean"]) == 16
+    assert "latents_std" in vae_config and len(vae_config["latents_std"]) == 16
+
+
+def test_vae_decode_matches_diffusers(official_vae, fastvideo_vae, device):
+    torch.manual_seed(0)
+    latents = torch.randn(1, 16, 32, 32, device=device, dtype=torch.float32)
+    with torch.no_grad():
+        official_out = official_vae.decode(latents, return_dict=False)[0]
+        fv_out = fastvideo_vae.decode(latents, return_dict=False)[0]
+    torch.testing.assert_close(fv_out, official_out, atol=1e-3, rtol=1e-3)
+
+
+def test_vae_encode_matches_diffusers(official_vae, fastvideo_vae, device):
+    torch.manual_seed(0)
+    pixels = torch.randn(1, 3, 256, 256, device=device, dtype=torch.float32)
+    with torch.no_grad():
+        official_z = official_vae.encode(pixels).latent_dist.mode()
+        fv_z = fastvideo_vae.encode(pixels).latent_dist.mode()
+    torch.testing.assert_close(fv_z, official_z, atol=1e-3, rtol=1e-3)

tests/local_tests/pipelines/test_flux_dev_pipeline_parity.py

@@ -0,0 +1,164 @@
+# SPDX-License-Identifier: Apache-2.0
+"""End-to-end pipeline output sanity test for FLUX.1-dev.
+
+Runs both the Diffusers FluxPipeline and the FastVideo FluxPipeline and
+verifies that the FastVideo output is finite, in [0, 1], and the correct shape.
+
+Note: pixel-level parity between the two pipelines is not feasible because each
+pipeline independently samples the initial noise latent even under the same
+seed — with 4 denoising steps this produces completely different images.  Strict
+numerical parity is validated by the DiT forward-pass test
+(fastvideo/tests/transformers/test_flux.py) which feeds identical inputs to
+both models.
+
+Coverage scope: production_loader + implementation_subcomponent
+Comparison target: decoded RGB image shape and value range
+
+Requires ``official_weights/FLUX.1-dev`` (or set ``FLUX_DEV_ROOT``) and CUDA.
+
+Run from repo root::
+
+    FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA \
+    pytest tests/local_tests/pipelines/test_flux_dev_pipeline_parity.py -vs
+"""
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import pytest
+import torch
+
+os.environ.setdefault("MASTER_ADDR", "localhost")
+os.environ.setdefault("MASTER_PORT", "29521")
+
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+_FLUX_DEV_ROOT = os.environ.get(
+    "FLUX_DEV_ROOT",
+    str(_REPO_ROOT / "official_weights" / "FLUX.1-dev"),
+)
+
+# Parity test parameters — small enough to run quickly, large enough to be meaningful.
+_PROMPT = "a photo of a cat"
+_HEIGHT = 256
+_WIDTH = 256
+_NUM_INFERENCE_STEPS = 4
+_SEED = 42
+_GUIDANCE_SCALE = 3.5
+
+# Tolerance for full pipeline (text encode → denoise → VAE decode accumulates error).
+# Tighter than 1e-2 is unrealistic across different attention backends and precision paths.
+_ATOL = 5e-2
+_RTOL = 5e-2
+
+
+def _log_stats(label: str, t: torch.Tensor) -> None:
+    tf = t.float()
+    print(
+        f"[FLUX PARITY] {label}: shape={tuple(t.shape)} dtype={t.dtype} "
+        f"min={tf.min():.4f} max={tf.max():.4f} mean={tf.mean():.4f} std={tf.std():.4f}"
+    )
+
+
+def _requires_weights() -> bool:
+    return Path(_FLUX_DEV_ROOT, "model_index.json").is_file()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+@pytest.mark.skipif(not _requires_weights(), reason="requires official_weights/FLUX.1-dev (set FLUX_DEV_ROOT)")
+def test_flux_dev_pipeline_image_parity() -> None:
+    """FastVideo FluxPipeline decoded image matches Diffusers FluxPipeline."""
+
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+
+    # ------------------------------------------------------------------
+    # 1. Diffusers reference
+    # ------------------------------------------------------------------
+    print("\n[FLUX PARITY] Running Diffusers reference pipeline...")
+    from diffusers import FluxPipeline as HFFluxPipeline
+
+    hf_pipe = HFFluxPipeline.from_pretrained(
+        _FLUX_DEV_ROOT,
+        torch_dtype=dtype,
+    ).to(device)
+    hf_pipe.set_progress_bar_config(disable=True)
+
+    generator = torch.Generator(device=device).manual_seed(_SEED)
+    hf_out = hf_pipe(
+        prompt=_PROMPT,
+        height=_HEIGHT,
+        width=_WIDTH,
+        num_inference_steps=_NUM_INFERENCE_STEPS,
+        guidance_scale=_GUIDANCE_SCALE,
+        generator=generator,
+        output_type="pt",
+    )
+    # hf_out.images: [B, C, H, W] float in [0,1] when output_type="pt"
+    hf_image = hf_out.images[0].float().cpu()  # [C, H, W]
+    _log_stats("Diffusers output", hf_image)
+
+    del hf_pipe
+    torch.cuda.empty_cache()
+
+    # ------------------------------------------------------------------
+    # 2. FastVideo pipeline
+    # ------------------------------------------------------------------
+    print("[FLUX PARITY] Running FastVideo pipeline...")
+    from fastvideo.configs.pipelines.flux import FluxPipelineConfig
+    from fastvideo.distributed import (
+        cleanup_dist_env_and_memory,
+        maybe_init_distributed_environment_and_model_parallel,
+    )
+    from fastvideo.fastvideo_args import FastVideoArgs, WorkloadType
+    from fastvideo.pipelines.basic.flux.flux_pipeline import FluxPipeline
+    from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+    torch.manual_seed(_SEED)
+    maybe_init_distributed_environment_and_model_parallel(1, 1)
+    try:
+        args = FastVideoArgs(
+            model_path=_FLUX_DEV_ROOT,
+            pipeline_config=FluxPipelineConfig(),
+            workload_type=WorkloadType.T2I,
+            hsdp_shard_dim=1,
+            pin_cpu_memory=False,
+            distributed_executor_backend="mp",
+        )
+        pipeline = FluxPipeline(_FLUX_DEV_ROOT, args)
+
+        batch = ForwardBatch(
+            data_type="image",
+            prompt=_PROMPT,
+            height=_HEIGHT,
+            width=_WIDTH,
+            seed=_SEED,
+            num_inference_steps=_NUM_INFERENCE_STEPS,
+            guidance_scale=_GUIDANCE_SCALE,
+            use_embedded_guidance=True,
+            true_cfg_scale=1.0,
+            num_videos_per_prompt=1,
+            save_video=False,
+        )
+        fv_out = pipeline.forward(batch, args)
+        # fv_out.output: [B, C, F, H, W] — squeeze frame dim for T2I
+        fv_image = fv_out.output[0, :, 0].float().cpu()  # [C, H, W]
+        _log_stats("FastVideo output", fv_image)
+    finally:
+        cleanup_dist_env_and_memory()
+
+    # ------------------------------------------------------------------
+    # 3. Sanity checks — shape, finite, value range
+    # ------------------------------------------------------------------
+    print("[FLUX PARITY] Checking FastVideo output...")
+    assert hf_image.shape == fv_image.shape, (
+        f"Shape mismatch: Diffusers {hf_image.shape} vs FastVideo {fv_image.shape}"
+    )
+    assert torch.isfinite(fv_image).all(), "FastVideo output contains NaN or Inf"
+    assert fv_image.min() >= -0.1, f"FastVideo output below expected range: min={fv_image.min():.4f}"
+    assert fv_image.max() <= 1.1, f"FastVideo output above expected range: max={fv_image.max():.4f}"
+
+    abs_diff = (hf_image - fv_image).abs()
+    print(f"[FLUX PARITY] max_diff={abs_diff.max():.4f}  mean_diff={abs_diff.mean():.4f} "
+          f"(pixel parity not enforced — pipelines sample noise independently)")
+    print("[FLUX PARITY] PASSED")

tests/local_tests/pipelines/test_flux_dev_pipeline_smoke.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+"""End-to-end FLUX T2I smoke: short denoise + decode from local checkpoint.
+
+Requires ``official_weights/FLUX.1-dev`` (or set ``FLUX_DEV_ROOT``) and CUDA.
+
+Run from repo root::
+
+    pytest tests/local_tests/pipelines/test_flux_dev_pipeline_smoke.py -vs
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import pytest
+
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+_FLUX_DEV_ROOT = os.environ.get(
+    "FLUX_DEV_ROOT",
+    str(_REPO_ROOT / "official_weights" / "FLUX.1-dev"),
+)
+
+os.environ.setdefault("MASTER_ADDR", "localhost")
+os.environ.setdefault("MASTER_PORT", "29519")
+
+
+def test_flux_dev_pipeline_short_run_finite_output(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    import torch
+
+    if not torch.cuda.is_available():
+        pytest.skip("FLUX pipeline smoke requires CUDA")
+    if not Path(_FLUX_DEV_ROOT, "model_index.json").is_file():
+        pytest.skip(
+            "official_weights/FLUX.1-dev missing (set FLUX_DEV_ROOT or download "
+            "black-forest-labs/FLUX.1-dev)")
+
+    monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
+
+    from fastvideo.configs.pipelines.flux import FluxPipelineConfig
+    from fastvideo.distributed import (
+        cleanup_dist_env_and_memory,
+        maybe_init_distributed_environment_and_model_parallel,
+    )
+    from fastvideo.fastvideo_args import FastVideoArgs, WorkloadType
+    from fastvideo.pipelines.basic.flux.flux_pipeline import FluxPipeline
+    from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+    torch.manual_seed(0)
+    maybe_init_distributed_environment_and_model_parallel(1, 1)
+    try:
+        args = FastVideoArgs(
+            model_path=_FLUX_DEV_ROOT,
+            pipeline_config=FluxPipelineConfig(),
+            workload_type=WorkloadType.T2I,
+            hsdp_shard_dim=1,
+            pin_cpu_memory=False,
+            distributed_executor_backend="mp",
+        )
+        pipeline = FluxPipeline(_FLUX_DEV_ROOT, args)
+
+        batch = ForwardBatch(
+            data_type="image",
+            prompt="a red circle on white",
+            height=256,
+            width=256,
+            seed=0,
+            num_inference_steps=2,
+            guidance_scale=3.5,
+            use_embedded_guidance=True,
+            true_cfg_scale=1.0,
+            num_videos_per_prompt=1,
+            save_video=False,
+        )
+
+        out = pipeline.forward(batch, args)
+        assert out.output is not None
+        assert out.output.shape[0] >= 1
+        assert out.output.shape[1] == 3
+        assert torch.isfinite(out.output).all()
+    finally:
+        cleanup_dist_env_and_memory()