[misc]: simplify cosmos3 action (fold action split into _split_flat_latent)

Collapse the byte-identical _split_action_latent into the generic
_split_flat_latent; all three modality specs now share one splitter.
Behavior-preserving; 139 tests green.

.gitignore

@@ -34,8 +34,11 @@ env
 *.log
 weights/
 logs/
-official_weights/
-converted_weights/
+
+# Cosmos3 local parity assets (symlinked from main worktree)
+/official_weights/
+/converted_weights/
+/cosmos-framework
 
 # SSIM test outputs
 fastvideo/tests/ssim/generated_videos/

docs/design/inference_schema_parity_inventory.yaml

@@ -108,7 +108,6 @@ surfaces:
       vae_sp: generator.pipeline.preset_overrides.vae_sp
       dmd_denoising_steps: generator.pipeline.preset_overrides.dmd_denoising_steps
       ti2v_task: generator.pipeline.preset_overrides.ti2v_task
-      lucy_edit_task: generator.pipeline.preset_overrides.lucy_edit_task
       boundary_ratio: generator.pipeline.preset_overrides.boundary_ratio
     compatibility_only:
       model_path: "Redundant with generator.model_path."
@@ -126,18 +125,9 @@ surfaces:
       text_encoder_configs: "Legacy internal component config object."
       preprocess_text_funcs: "Internal text preprocessing hooks."
       postprocess_text_funcs: "Internal text postprocessing hooks."
-      scheduler_step_in_fp32: "Runtime scheduler precision toggle; not part of the public typed inference API."
 
   pipeline_config_extensions:
     preset_owned:
-      flux2_text_encoder_type:
-        sources:
-          - fastvideo.configs.pipelines.flux_2.Flux2PipelineConfig
-          - fastvideo.configs.pipelines.flux_2.Flux2KleinPipelineConfig
-      text_encoder_out_layers:
-        sources:
-          - fastvideo.configs.pipelines.flux_2.Flux2PipelineConfig
-          - fastvideo.configs.pipelines.flux_2.Flux2KleinPipelineConfig
       conditioning_strategy:
         sources:
           - fastvideo.configs.pipelines.cosmos.CosmosConfig
@@ -501,8 +491,6 @@ surfaces:
       inpaint_mask: request.extensions.stable_audio.inpaint_mask
     internal_only:
       data_type: "Derived from the request shape and not a public input."
-      latents: "Pre-generated diffusion latents supplied by parity/debug harnesses; not a public input."
-      max_sequence_length: "Model-specific text-encoder sequence cap; not part of the public typed inference API."
 
   sampling_param_extensions: {}
 

docs/inference/support_matrix.md

@@ -58,7 +58,6 @@ pipeline initialization and sampling.
 | FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ |
 | FastWan2.2 TI2V 5B Full Attn* | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ |
 | Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720P | ⭕ | ⭕ | ✅ | ⭕ | ⭕ |
-| Lucy Edit Dev 5B*** | `decart-ai/Lucy-Edit-Dev` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
 | Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480P<br>720P | ❌ | ❌ | ✅ | ⭕ | ⭕ |
 | Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480P<br>720P | ❌ | ❌ | ✅ | ⭕ | ⭕ |
 | HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720px1280p<br>544px960p | ❌ | ✅ | ✅ | ⭕ | ⭕ |
@@ -79,9 +78,6 @@ pipeline initialization and sampling.
 
 **Note**: Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue.
 
-***Lucy Edit Dev uses a non-commercial model license. FastVideo support is
-focused on inference integration for video editing workflows.
-
 `Sliding Tile Attn (Legacy Branch)` entries refer to the archived
 `sta_do_not_delete` branch workflow, not active `main` inference wiring.
 

examples/inference/basic/basic_cosmos3_i2v_new_api.py

@@ -0,0 +1,81 @@
+import os
+import time
+
+from fastvideo import VideoGenerator
+from fastvideo.api import (
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    InputConfig,
+    OffloadConfig,
+    OutputConfig,
+    SamplingConfig,
+)
+
+# NVIDIA Cosmos3-Nano omni world model — image-to-video (I2V) path through
+# FastVideo's native Cosmos3 pipeline. The input image conditions latent frame 0
+# (kept clean during denoising); the rest of the clip is generated to follow it.
+# Point COSMOS3_MODEL_PATH at a local diffusers checkpoint (e.g.
+# ``official_weights/cosmos3``) to skip the Hugging Face download.
+OUTPUT_PATH = "video_samples_cosmos3_i2v"
+
+
+def main():
+    model_name = os.environ.get("COSMOS3_MODEL_PATH", "nvidia/Cosmos3-Nano")
+    image_path = os.environ.get("COSMOS3_IMAGE_PATH", "assets/images/cyclist.jpg")
+
+    generator_config = GeneratorConfig(
+        model_path=model_name,
+        engine=EngineConfig(
+            num_gpus=1,
+            use_fsdp_inference=False,
+            offload=OffloadConfig(
+                text_encoder=True,
+                pin_cpu_memory=True,
+                dit=False,
+                vae=False,
+            ),
+        ),
+    )
+
+    load_start_time = time.perf_counter()
+    generator = VideoGenerator.from_config(generator_config)
+    load_time = time.perf_counter() - load_start_time
+
+    prompt = (
+        "A mountain biker rides forward along the sunlit forest trail, wheels "
+        "kicking up dust as trees and dappled light sweep past, smooth cinematic "
+        "tracking shot from behind."
+    )
+    request = GenerationRequest(
+        prompt=prompt,
+        inputs=InputConfig(image_path=image_path),
+        sampling=SamplingConfig(
+            # cosmos3_nano native defaults are 704x1280, 189 frames, 35 steps;
+            # overridable via env for quick smoke runs.
+            num_frames=int(os.environ.get("COSMOS3_NUM_FRAMES", "189")),
+            height=int(os.environ.get("COSMOS3_HEIGHT", "704")),
+            width=int(os.environ.get("COSMOS3_WIDTH", "1280")),
+            num_inference_steps=int(os.environ.get("COSMOS3_STEPS", "35")),
+            guidance_scale=6.0,
+            fps=24,
+            seed=1024,
+        ),
+        output=OutputConfig(
+            output_path=OUTPUT_PATH,
+            save_video=True,
+            return_frames=False,
+        ),
+    )
+
+    start_time = time.perf_counter()
+    result = generator.generate(request)
+    gen_time = time.perf_counter() - start_time
+
+    print(f"Time taken to load model: {load_time} seconds")
+    print(f"Time taken to generate video: {gen_time} seconds")
+    print(f"Output written to: {result.video_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_cosmos3_new_api.py

@@ -0,0 +1,77 @@
+import os
+import time
+
+from fastvideo import VideoGenerator
+from fastvideo.api import (
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    OffloadConfig,
+    OutputConfig,
+    SamplingConfig,
+)
+
+# NVIDIA Cosmos3-Nano omni world model — this example exercises the
+# text-to-video (T2V) path through FastVideo's native Cosmos3 pipeline.
+# Point COSMOS3_MODEL_PATH at a local diffusers checkpoint (e.g.
+# ``official_weights/cosmos3``) to skip the Hugging Face download.
+OUTPUT_PATH = "video_samples_cosmos3"
+
+
+def main():
+    model_name = os.environ.get("COSMOS3_MODEL_PATH", "nvidia/Cosmos3-Nano")
+
+    generator_config = GeneratorConfig(
+        model_path=model_name,
+        engine=EngineConfig(
+            num_gpus=1,
+            use_fsdp_inference=False,
+            offload=OffloadConfig(
+                text_encoder=True,
+                pin_cpu_memory=True,
+                dit=False,
+                vae=False,
+            ),
+        ),
+    )
+
+    load_start_time = time.perf_counter()
+    generator = VideoGenerator.from_config(generator_config)
+    load_time = time.perf_counter() - load_start_time
+
+    prompt = (
+        "A golden retriever puppy runs across a sunlit meadow toward the camera, "
+        "ears flopping and wildflowers swaying in the breeze. Shallow depth of "
+        "field, warm afternoon light, smooth cinematic tracking shot."
+    )
+    request = GenerationRequest(
+        prompt=prompt,
+        sampling=SamplingConfig(
+            # cosmos3_nano native defaults are 704x1280, 189 frames, 35 steps;
+            # overridable via env for quick smoke runs.
+            num_frames=int(os.environ.get("COSMOS3_NUM_FRAMES", "189")),
+            height=int(os.environ.get("COSMOS3_HEIGHT", "704")),
+            width=int(os.environ.get("COSMOS3_WIDTH", "1280")),
+            num_inference_steps=int(os.environ.get("COSMOS3_STEPS", "35")),
+            guidance_scale=6.0,
+            fps=24,
+            seed=1024,
+        ),
+        output=OutputConfig(
+            output_path=OUTPUT_PATH,
+            save_video=True,
+            return_frames=False,
+        ),
+    )
+
+    start_time = time.perf_counter()
+    result = generator.generate(request)
+    gen_time = time.perf_counter() - start_time
+
+    print(f"Time taken to load model: {load_time} seconds")
+    print(f"Time taken to generate video: {gen_time} seconds")
+    print(f"Output written to: {result.video_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_cosmos3_t2i_new_api.py

@@ -0,0 +1,78 @@
+import os
+import time
+
+from fastvideo import VideoGenerator
+from fastvideo.api import (
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    OffloadConfig,
+    OutputConfig,
+    SamplingConfig,
+)
+
+# NVIDIA Cosmos3-Nano omni world model — text-to-image (T2I) path through
+# FastVideo's native Cosmos3 pipeline. T2I is the single-frame case
+# (num_frames=1); the canonical Cosmos3 T2I resolution is 960x960 (the model's
+# "720" bucket, UniPC flow_shift=10.0). Point COSMOS3_MODEL_PATH at a local
+# diffusers checkpoint (e.g. ``official_weights/cosmos3``) to skip the download.
+OUTPUT_PATH = "video_samples_cosmos3_t2i"
+
+
+def main():
+    model_name = os.environ.get("COSMOS3_MODEL_PATH", "nvidia/Cosmos3-Nano")
+
+    generator_config = GeneratorConfig(
+        model_path=model_name,
+        engine=EngineConfig(
+            num_gpus=1,
+            use_fsdp_inference=False,
+            offload=OffloadConfig(
+                text_encoder=True,
+                pin_cpu_memory=True,
+                dit=False,
+                vae=False,
+            ),
+        ),
+    )
+
+    load_start_time = time.perf_counter()
+    generator = VideoGenerator.from_config(generator_config)
+    load_time = time.perf_counter() - load_start_time
+
+    prompt = (
+        "A photograph of a red panda sitting on a mossy log in a misty bamboo "
+        "forest, soft golden morning light filtering through the leaves, shallow "
+        "depth of field, crisp fur detail, serene atmosphere."
+    )
+    request = GenerationRequest(
+        prompt=prompt,
+        sampling=SamplingConfig(
+            # T2I is single-frame; canonical Cosmos3 T2I is 960x960. Overridable
+            # via env for quick smoke runs.
+            num_frames=1,
+            height=int(os.environ.get("COSMOS3_HEIGHT", "960")),
+            width=int(os.environ.get("COSMOS3_WIDTH", "960")),
+            num_inference_steps=int(os.environ.get("COSMOS3_STEPS", "35")),
+            guidance_scale=6.0,
+            fps=24,
+            seed=1024,
+        ),
+        output=OutputConfig(
+            output_path=OUTPUT_PATH,
+            save_video=True,
+            return_frames=False,
+        ),
+    )
+
+    start_time = time.perf_counter()
+    result = generator.generate(request)
+    gen_time = time.perf_counter() - start_time
+
+    print(f"Time taken to load model: {load_time} seconds")
+    print(f"Time taken to generate image: {gen_time} seconds")
+    print(f"Output written to: {result.video_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_cosmos3_t2vs_new_api.py

@@ -0,0 +1,67 @@
+import os
+import time
+
+# t2vs (text -> video + sound). The Cosmos3 denoise stage generates a joint
+# [vision | sound] latent and AVAE-decodes the sound to a waveform muxed into the
+# mp4. The joint-sound path is gated on COSMOS3_T2VS (set here for the example).
+os.environ.setdefault("COSMOS3_T2VS", "1")
+
+from fastvideo import VideoGenerator  # noqa: E402
+from fastvideo.api import (  # noqa: E402
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    OffloadConfig,
+    OutputConfig,
+    SamplingConfig,
+)
+
+OUTPUT_PATH = "video_samples_cosmos3_t2vs"
+
+
+def main():
+    model_name = os.environ.get("COSMOS3_MODEL_PATH", "nvidia/Cosmos3-Nano")
+
+    generator_config = GeneratorConfig(
+        model_path=model_name,
+        engine=EngineConfig(
+            num_gpus=1,
+            use_fsdp_inference=False,
+            offload=OffloadConfig(text_encoder=True, pin_cpu_memory=True, dit=False, vae=False),
+        ),
+    )
+
+    load_start = time.perf_counter()
+    generator = VideoGenerator.from_config(generator_config)
+    load_time = time.perf_counter() - load_start
+
+    prompt = (
+        "Ocean waves crash against a rocky shore at sunset, white foam spraying "
+        "into the air as seagulls wheel overhead. Golden light, cinematic wide "
+        "shot, the rhythmic roar of the surf."
+    )
+    request = GenerationRequest(
+        prompt=prompt,
+        sampling=SamplingConfig(
+            num_frames=int(os.environ.get("COSMOS3_NUM_FRAMES", "189")),
+            height=int(os.environ.get("COSMOS3_HEIGHT", "704")),
+            width=int(os.environ.get("COSMOS3_WIDTH", "1280")),
+            num_inference_steps=int(os.environ.get("COSMOS3_STEPS", "35")),
+            guidance_scale=6.0,
+            fps=24,
+            seed=1024,
+        ),
+        output=OutputConfig(output_path=OUTPUT_PATH, save_video=True, return_frames=False),
+    )
+
+    start = time.perf_counter()
+    result = generator.generate(request)
+    gen_time = time.perf_counter() - start
+
+    print(f"Time taken to load model: {load_time} seconds")
+    print(f"Time taken to generate video+sound: {gen_time} seconds")
+    print(f"Output written to: {result.video_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_flux2.py

@@ -1,124 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Run full Flux2 text-to-image generation through FastVideo.
-
-User story:
-    "I have a local or HF Diffusers-format full Flux2 checkpoint and want a
-    minimal text-to-image generation command that uses embedded guidance."
-"""
-import argparse
-import os
-from pathlib import Path
-
-from fastvideo import VideoGenerator
-from fastvideo.api import (
-    ComponentConfig,
-    EngineConfig,
-    GenerationRequest,
-    GeneratorConfig,
-    OffloadConfig,
-    OutputConfig,
-    ParallelismConfig,
-    PipelineSelection,
-    SamplingConfig,
-)
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Run full Flux2 text-to-image generation.")
-    parser.add_argument(
-        "--model-path",
-        default="black-forest-labs/FLUX.2-dev",
-        help="HF id or local diffusers-format full Flux2 weights directory.",
-    )
-    parser.add_argument(
-        "--output",
-        default="outputs/flux2/flux2.png",
-        help="Output PNG path.",
-    )
-    parser.add_argument(
-        "--prompt",
-        default="a photo of a banana on a wooden table, studio lighting",
-        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=4.0)
-    parser.add_argument("--max-sequence-length", type=int, default=None)
-    parser.add_argument("--seed", type=int, default=0)
-    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)
-    parser.add_argument(
-        "--backend",
-        default=None,
-        help="Set FASTVIDEO_ATTENTION_BACKEND, for example TORCH_SDPA.",
-    )
-    return parser.parse_args()
-
-
-def main() -> None:
-    args = parse_args()
-    if args.backend:
-        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = args.backend
-
-    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
-    )
-
-    generator_config = GeneratorConfig(
-        model_path=args.model_path,
-        engine=EngineConfig(
-            num_gpus=args.num_gpus,
-            parallelism=ParallelismConfig(tp_size=tp_size, sp_size=sp_size),
-            use_fsdp_inference=False,
-            offload=OffloadConfig(
-                dit=False,
-                vae=True,
-                text_encoder=True,
-                pin_cpu_memory=False,
-            ),
-        ),
-        pipeline=PipelineSelection(
-            workload_type="t2i",
-            components=ComponentConfig(override_pipeline_cls_name="Flux2Pipeline"),
-        ),
-    )
-
-    generator = VideoGenerator.from_config(generator_config)
-    try:
-        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,
-        )
-        extensions = {}
-        if args.max_sequence_length is not None:
-            extensions["max_sequence_length"] = args.max_sequence_length
-
-        request = GenerationRequest(
-            prompt=args.prompt,
-            sampling=sampling,
-            output=OutputConfig(
-                output_path=str(output),
-                save_video=True,
-                return_frames=False,
-            ),
-            extensions=extensions,
-        )
-        generator.generate(request)
-    finally:
-        generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()

examples/inference/basic/basic_flux2_klein.py

@@ -1,98 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Run Flux2 Klein text-to-image generation through FastVideo.
-
-User story:
-    "I need a short local smoke for the Flux2 Klein checkpoint before wiring it
-    into an image workflow. Use the model's distilled four-step defaults and
-    write a single PNG so I can compare the output against the reference."
-"""
-
-import argparse
-import os
-
-from fastvideo import VideoGenerator
-from fastvideo.api import (
-    EngineConfig,
-    GenerationRequest,
-    GeneratorConfig,
-    OffloadConfig,
-    OutputConfig,
-    PipelineSelection,
-    SamplingConfig,
-)
-
-
-DEFAULT_PROMPT = "a brushed steel espresso machine on a marble counter, morning window light"
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Run Flux2 Klein text-to-image generation.")
-    parser.add_argument(
-        "--model-path",
-        default="black-forest-labs/FLUX.2-klein-4B",
-        help="HF id or local diffusers-format Flux2 Klein weights directory.",
-    )
-    parser.add_argument(
-        "--output-path",
-        default="outputs/flux2/flux2_klein.png",
-        help="PNG output path or output directory.",
-    )
-    parser.add_argument("--prompt", default=DEFAULT_PROMPT, help="Prompt text.")
-    parser.add_argument("--seed", type=int, default=0, help="Generation seed.")
-    parser.add_argument("--height", type=int, default=1024, help="Output image height.")
-    parser.add_argument("--width", type=int, default=1024, help="Output image width.")
-    parser.add_argument("--steps", type=int, default=4, help="Number of denoising steps.")
-    parser.add_argument("--num-gpus", type=int, default=1, help="Number of GPUs to use.")
-    parser.add_argument(
-        "--backend",
-        default=None,
-        help="Set FASTVIDEO_ATTENTION_BACKEND, for example TORCH_SDPA.",
-    )
-    return parser.parse_args()
-
-
-def main() -> None:
-    args = parse_args()
-    if args.backend:
-        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = args.backend
-
-    generator_config = GeneratorConfig(
-        model_path=args.model_path,
-        engine=EngineConfig(
-            num_gpus=args.num_gpus,
-            use_fsdp_inference=False,
-            offload=OffloadConfig(
-                dit=False,
-                vae=True,
-                text_encoder=True,
-                pin_cpu_memory=False,
-            ),
-        ),
-        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=1.0,
-                seed=args.seed,
-            ),
-            output=OutputConfig(
-                output_path=args.output_path,
-                save_video=True,
-            ),
-        )
-        generator.generate(request)
-    finally:
-        generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()

examples/inference/basic/basic_ltx2_3_distilled_i2v.py

@@ -18,20 +18,20 @@ What the script does
 --------------------
 1. Loads FastVideo/LTX-2.3-Distilled-Diffusers (8 denoise + 3 refine steps,
    CFG=1, no refine LoRA — the distilled production recipe).
-2. Compiles the DiT, text encoder, and VAE (fullgraph, Inductor default
-   mode — autotune adds ~7 min cold-compile here with no measurable
-   e2e gain).
-3. Runs 2 warmup calls (untimed) + 2 measured calls. Two warmups are kept
-   as a safety net — the first call pays cold compile + first-shape guard
-   work, and a second warmup ensures any residual recompiles settle before
-   we measure.
+2. Compiles the DiT, text encoder, and VAE (fullgraph, max-autotune-no-
+   cudagraphs).
+3. Runs 2 warmup calls (untimed) + 2 measured calls. Two warmups are needed
+   even though distilled has no refine LoRA, because Inductor's per-shape
+   autotune for the first call still leaves a few cold guards on call 2;
+   the second warmup settles them. Skipping to a single warmup typically
+   inflates the first measured run by tens of seconds.
 4. Prints a per-stage breakdown and an average over the measured runs.
 
 Hardware notes
 --------------
 - Single-GPU example; for multi-GPU sequence-parallel see the gradio demo
   under `examples/inference/gradio/local/gradio_local_demo_ltx2_3/`.
-- First-time compile takes ~30-40 min on GB200 (~20 min on H100; cached
+- First-time compile + autotune takes ~10-30 min on H100 / GB200 (cached
   in `$TORCHINDUCTOR_CACHE_DIR` afterwards). Subsequent invocations only
   pay the one-time process load + a few seconds of dynamo trace.
 - On GB200 / Blackwell, run with `env -u LD_LIBRARY_PATH ...` to avoid a
@@ -151,13 +151,10 @@ def main() -> None:
     print(f"i2v image:       {I2V_IMAGE}")
     print(f"Output dir:      {OUTPUT_DIR.resolve()}")
 
-    # mode="default" — Inductor's default schedule matches max-autotune on
-    # this pipeline (denoise/refine/decode all within ~5 ms, n=2) while
-    # saving ~7 min of cold compile on a single GB200.
     torch_compile_kwargs = {
         "backend": "inductor",
         "fullgraph": True,
-        "mode": "default",
+        "mode": "max-autotune-no-cudagraphs",
         "dynamic": False,
     }
 

examples/inference/basic/basic_ltx2_3_distilled_i2v_typed.py

@@ -1,350 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""LTX-2.3 distilled image-to-video — typed API (``from_config`` / ``generate``).
-
-Identical generation behavior to ``basic_ltx2_3_distilled_i2v.py``, but
-expressed through the newer typed surface (``GeneratorConfig`` /
-``GenerationRequest``) instead of the ``from_pretrained(**legacy_kwargs)``
-bridge. The typed API is now the preferred entry point — the legacy
-example still works but emits a ``DeprecationWarning`` for the LTX-2.3
-specific knobs.
-
-Quick start
------------
-    export LTX23_I2V_IMAGE=/path/to/your/portrait_or_product.jpg
-    # optional overrides:
-    #   export LTX23_I2V_PROMPT="a fashion model walks toward camera..."
-    #   export LTX23_OUTPUT_DIR=outputs_video/ltx2_3_distilled_i2v_typed
-    python examples/inference/basic/basic_ltx2_3_distilled_i2v_typed.py
-
-What the script does
---------------------
-1. Loads FastVideo/LTX-2.3-Distilled-Diffusers (8 denoise + 3 refine
-   steps, CFG=1, no refine LoRA — the distilled production recipe).
-2. Compiles the DiT, text encoder, and VAE (fullgraph, Inductor default
-   mode — autotune adds ~7 min cold-compile here with no measurable
-   e2e gain).
-3. Runs 2 warmup calls (untimed) + 2 measured calls. Two warmups are
-   kept as a safety net — the first call pays cold compile + first-shape
-   guard work, and a second warmup ensures any residual recompiles
-   settle before we measure.
-4. Prints a per-stage breakdown and an average over the measured runs.
-
-Hardware notes
---------------
-- Single-GPU example; for multi-GPU sequence-parallel see the gradio
-  demo under ``examples/inference/gradio/local/gradio_local_demo_ltx2_3/``.
-- First-time compile takes ~30-40 min on GB200 (~20 min on H100;
-  cached in ``$TORCHINDUCTOR_CACHE_DIR`` afterwards). Subsequent
-  invocations only pay the one-time process load + a few seconds of
-  dynamo trace.
-- On GB200 / Blackwell, run with ``env -u LD_LIBRARY_PATH ...`` to
-  avoid a system-cuBLAS / torch-cuBLAS mismatch that fails every GEMM.
-  The ``_inductor.shape_padding = False`` line below also avoids a
-  ``pad_mm`` landmine on the same generation of cards.
-
-Typed-API mapping (legacy kwarg ↔ typed field)
-----------------------------------------------
-- ``num_gpus``                      ↔ ``engine.num_gpus``
-- ``enable_torch_compile``          ↔ ``engine.compile.enabled``
-- ``enable_torch_compile_text_encoder`` ↔ ``engine.compile.text_encoder_enabled``
-- ``enable_torch_compile_vae``      ↔ ``engine.compile.vae_enabled``
-- ``torch_compile_kwargs``          ↔ ``engine.compile.backend/fullgraph/mode/dynamic``
-- ``torch_compile_kwargs_vae``      ↔ empty ``compile.vae_kwargs`` (inherits master)
-- ``dit_cpu_offload``               ↔ ``engine.offload.dit``
-- ``text_encoder_cpu_offload``      ↔ ``engine.offload.text_encoder``
-- ``vae_cpu_offload``               ↔ ``engine.offload.vae``
-- ``ltx2_vae_tiling``               ↔ ``pipeline.vae_tiling``
-- ``ltx2_refine_enabled``           ↔ ``pipeline.preset_overrides["refine"]["enabled"]``
-- ``ltx2_refine_upsampler_path``    ↔ ``pipeline.components.upsampler_weights``
-- ``ltx2_refine_lora_path``         ↔ ``pipeline.components.lora_path``
-- ``ltx2_refine_num_inference_steps`` ↔ ``pipeline.preset_overrides["refine"]["num_inference_steps"]``
-- ``ltx2_refine_guidance_scale``    ↔ ``pipeline.preset_overrides["refine"]["guidance_scale"]``
-- ``ltx2_refine_add_noise``         ↔ ``pipeline.preset_overrides["refine"]["add_noise"]``
-- ``pipeline_config=PipelineConfig.from_pretrained(model_root)`` ↔ (no-op — ``PipelineConfig.from_kwargs`` already resolves the model-specific class from ``model_path``)
-- ``pipeline_config.dit_config.quant_config = None`` ↔ leave ``engine.quantization`` unset
-- ``ltx2_images`` / ``ltx2_image_crf`` ↔ ``request.extensions`` (LTX-2 specific, no
-  first-class typed field yet)
-"""
-from __future__ import annotations
-
-import os
-import time
-from collections import OrderedDict
-from pathlib import Path
-
-import torch._inductor.config as _inductor
-
-from fastvideo import VideoGenerator
-from fastvideo.api import (
-    CompileConfig,
-    ComponentConfig,
-    EngineConfig,
-    GenerationRequest,
-    GeneratorConfig,
-    OffloadConfig,
-    OutputConfig,
-    PipelineSelection,
-    SamplingConfig,
-)
-from fastvideo.utils import maybe_download_model
-
-os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN")
-os.environ.setdefault("FASTVIDEO_STAGE_LOGGING", "1")
-
-# Inductor knobs. ``shape_padding=False`` is mandatory on Blackwell to
-# avoid a cuBLAS INVALID_VALUE crash inside pad_mm during the refine
-# path. The rest are autotune-friendliness flags.
-_inductor.shape_padding = False
-_inductor.conv_1x1_as_mm = True
-_inductor.coordinate_descent_tuning = True
-_inductor.coordinate_descent_check_all_directions = True
-_inductor.epilogue_fusion = False
-
-MODEL_ID = os.path.expandvars(
-    os.path.expanduser(
-        os.getenv("LTX23_MODEL_PATH", "FastVideo/LTX-2.3-Distilled-Diffusers")
-    )
-)
-OUTPUT_DIR = Path(
-    os.getenv(
-        "LTX23_OUTPUT_DIR", "outputs_video/ltx2_3_distilled_i2v_typed"
-    )
-)
-I2V_IMAGE = os.getenv("LTX23_I2V_IMAGE", "")
-DEFAULT_PROMPT = (
-    "A fashion model takes a slow step forward and shifts her weight, "
-    "the soft fabric of her clothing swaying and rippling with the "
-    "motion, her hair shifting gently, soft even studio lighting on a "
-    "clean light background, elegant slow-motion runway feel."
-)
-PROMPT = os.getenv("LTX23_I2V_PROMPT", DEFAULT_PROMPT)
-
-
-def _print_stage_breakdown(result, label: str) -> float | None:
-    logging_info = getattr(result, "logging_info", None)
-    stages = getattr(logging_info, "stages", None) if logging_info else None
-    if not stages:
-        print(f"  [{label}] stage breakdown unavailable")
-        return None
-    print(f"  [{label}] stage breakdown:")
-    total = 0.0
-    for name, metrics in stages.items():
-        exec_s = float(metrics.get("execution_time", 0.0))
-        total += exec_s
-        print(f"    - {name}: {exec_s:.3f}s")
-    print(f"    - stage_sum: {total:.3f}s")
-    return total
-
-
-def _collect_stage_times(
-    result,
-    stage_times: dict[str, list[float]],
-    stage_order: OrderedDict[str, None],
-) -> None:
-    logging_info = getattr(result, "logging_info", None)
-    stages = getattr(logging_info, "stages", None) if logging_info else None
-    if not stages:
-        return
-    for name, metrics in stages.items():
-        stage_order.setdefault(name, None)
-        stage_times.setdefault(name, []).append(
-            float(metrics.get("execution_time", 0.0))
-        )
-
-
-def _resolve_refine_upsampler(model_root: str) -> Path:
-    for name in ("spatial_upscaler", "spatial_upsampler"):
-        cand = Path(model_root) / name
-        if (cand / "config.json").is_file():
-            return cand
-    raise FileNotFoundError(
-        f"No refine upsampler directory under {model_root}. "
-        f"Expected `{model_root}/spatial_upscaler/config.json`."
-    )
-
-
-def main() -> None:
-    if not I2V_IMAGE:
-        raise SystemExit(
-            "LTX23_I2V_IMAGE is required for i2v. Example:\n"
-            "  export LTX23_I2V_IMAGE=/path/to/portrait_or_product.jpg\n"
-            "  python examples/inference/basic/"
-            "basic_ltx2_3_distilled_i2v_typed.py"
-        )
-    if not Path(I2V_IMAGE).is_file():
-        raise SystemExit(f"LTX23_I2V_IMAGE not found: {I2V_IMAGE}")
-
-    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
-    model_root = maybe_download_model(MODEL_ID)
-    refine_upsampler_path = _resolve_refine_upsampler(model_root)
-    print(f"Model:           {model_root}")
-    print(f"Refine upsampler: {refine_upsampler_path}")
-    print(f"i2v image:       {I2V_IMAGE}")
-    print(f"Output dir:      {OUTPUT_DIR.resolve()}")
-
-    # mode="default" — Inductor's default schedule matches max-autotune on
-    # this pipeline (denoise/refine/decode all within ~5 ms, n=2) while
-    # saving ~7 min of cold compile on a single GB200.
-    generator_config = GeneratorConfig(
-        model_path=model_root,
-        engine=EngineConfig(
-            num_gpus=1,
-            # Keep DiT / text encoder / VAE resident on GPU — no CPU offload
-            # for serving-style runs. ``image_encoder`` and
-            # ``pin_cpu_memory`` are left at their schema defaults
-            # (matches the legacy example, which only set these three).
-            offload=OffloadConfig(
-                dit=False,
-                text_encoder=False,
-                vae=False,
-            ),
-            compile=CompileConfig(
-                enabled=True,
-                text_encoder_enabled=True,
-                # ``vae_enabled`` triggers ``_compile_with_conditions`` on
-                # ``LTX2CausalVideoAutoencoder``, which compiles just the
-                # encoder/decoder submodules and leaves the surrounding
-                # tiling control flow eager (required for ``fullgraph``).
-                # Empty ``vae_kwargs`` → inherits the master kwargs below.
-                vae_enabled=True,
-                backend="inductor",
-                fullgraph=True,
-                mode="default",
-                dynamic=False,
-            ),
-        ),
-        pipeline=PipelineSelection(
-            # ``PipelineConfig.from_kwargs`` resolves the model-specific
-            # pipeline-config class from ``model_path`` automatically, so we
-            # don't need to set ``components.pipeline_config_path`` — the
-            # model-specific VAE precision / decoder defaults are picked up
-            # the same way the legacy example's
-            # ``PipelineConfig.from_pretrained(model_root)`` did them.
-            components=ComponentConfig(
-                upsampler_weights=str(refine_upsampler_path),
-                # Distilled has no refine LoRA — omit ``lora_path``.
-            ),
-            vae_tiling=False,
-            preset_overrides={
-                "refine": {
-                    "enabled": True,
-                    "num_inference_steps": 3,
-                    "guidance_scale": 1.0,
-                    "add_noise": True,
-                },
-            },
-        ),
-    )
-
-    generator = VideoGenerator.from_config(generator_config)
-
-    def build_request(out_path: Path, seed: int) -> GenerationRequest:
-        return GenerationRequest(
-            prompt=PROMPT,
-            # distilled is CFG-free; no negative prompt
-            negative_prompt="",
-            sampling=SamplingConfig(
-                num_videos_per_prompt=1,
-                seed=seed,
-                height=1280,
-                width=832,
-                num_frames=121,
-                fps=24,
-                num_inference_steps=8,
-                guidance_scale=1.0,
-            ),
-            output=OutputConfig(
-                output_path=str(out_path),
-                save_video=True,
-                return_frames=False,
-            ),
-            # LTX-2.3 i2v fields don't have first-class typed slots yet;
-            # extensions is the documented bridge. ``ltx2_image_crf=0.0``
-            # skips an extra JPEG re-encode of an already JPEG image.
-            extensions={
-                "ltx2_images": [(I2V_IMAGE, 0, 1.0)],
-                "ltx2_image_crf": 0.0,
-            },
-        )
-
-    warmup_runs = 2
-    measured_runs = 2
-    warmup_secs: list[float] = []
-    measured_secs: list[float] = []
-    stage_times: dict[str, list[float]] = {}
-    stage_order: OrderedDict[str, None] = OrderedDict()
-
-    try:
-        for w in range(warmup_runs):
-            print(f"\n[warmup {w + 1}/{warmup_runs}] compiling + generating…")
-            t0 = time.perf_counter()
-            generator.generate(
-                build_request(
-                    OUTPUT_DIR / f"_warmup_{w + 1}.mp4", seed=7
-                )
-            )
-            dt = time.perf_counter() - t0
-            warmup_secs.append(dt)
-            print(f"[warmup {w + 1}/{warmup_runs}] wall={dt:.1f}s")
-
-        for w in range(warmup_runs):
-            (OUTPUT_DIR / f"_warmup_{w + 1}.mp4").unlink(missing_ok=True)
-
-        for m in range(measured_runs):
-            out_path = (
-                OUTPUT_DIR
-                / f"output_ltx2_3_distilled_i2v_typed_run_{m + 1}.mp4"
-            )
-            print(
-                f"\n[measured {m + 1}/{measured_runs}] generating: {out_path}"
-            )
-            t0 = time.perf_counter()
-            result = generator.generate(
-                build_request(out_path, seed=2002 + m)
-            )
-            wall = time.perf_counter() - t0
-            # ``e2e_latency`` is currently surfaced via ``result.extra``;
-            # ``GenerationResult`` exposes ``generation_time`` as a
-            # first-class field but the LTX-2 pipeline only fills the
-            # legacy ``e2e_latency`` key. Prefer the explicit one, fall
-            # back to wall-clock.
-            e2e = (
-                result.extra.get("e2e_latency")
-                if hasattr(result, "extra") else None
-            ) or wall
-            measured_secs.append(e2e)
-            print(
-                f"[measured {m + 1}/{measured_runs}] "
-                f"e2e={e2e:.2f}s wall={wall:.2f}s"
-            )
-            _print_stage_breakdown(result, f"measured {m + 1}")
-            _collect_stage_times(result, stage_times, stage_order)
-
-        print("\n=== summary ===")
-        print(
-            f"warmup wall-times:      "
-            f"{[round(x, 1) for x in warmup_secs]}"
-        )
-        if measured_secs:
-            avg = sum(measured_secs) / len(measured_secs)
-            print(
-                f"measured e2e (n={len(measured_secs)}): "
-                f"{[round(x, 2) for x in measured_secs]} -> avg {avg:.2f}s"
-            )
-        if stage_times:
-            print(f"average stage times over {measured_runs} measured runs:")
-            avg_total = 0.0
-            for name in stage_order:
-                vals = stage_times.get(name) or []
-                if not vals:
-                    continue
-                avg_v = sum(vals) / len(vals)
-                avg_total += avg_v
-                print(f"  - {name}: {avg_v:.3f}s")
-            print(f"  - stage_sum_avg: {avg_total:.3f}s")
-    finally:
-        generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()

examples/inference/basic/basic_lucy_edit.py

@@ -1,38 +0,0 @@
-from fastvideo import VideoGenerator
-
-OUTPUT_PATH = "video_samples_lucy_edit"
-
-
-def main():
-    generator = VideoGenerator.from_pretrained(
-        "decart-ai/Lucy-Edit-Dev",
-        num_gpus=1,
-        use_fsdp_inference=False,
-        dit_cpu_offload=True,
-        vae_cpu_offload=False,
-        text_encoder_cpu_offload=True,
-        pin_cpu_memory=True,
-    )
-
-    prompt = ("Change the apron and blouse to a classic clown costume: satin "
-              "polka-dot jumpsuit in bright primary colors, ruffled white collar, "
-              "oversized pom-pom buttons, white gloves, oversized red shoes, red "
-              "foam nose; soft window light from left, eye-level medium shot.")
-    video_path = "https://d2drjpuinn46lb.cloudfront.net/painter_original_edit.mp4"
-
-    generator.generate_video(
-        prompt,
-        negative_prompt="",
-        video_path=video_path,
-        output_path=OUTPUT_PATH,
-        save_video=True,
-        height=480,
-        width=832,
-        num_frames=81,
-        fps=24,
-        guidance_scale=5.0,
-    )
-
-
-if __name__ == "__main__":
-    main()

fastvideo/api/sampling_param.py

@@ -29,10 +29,6 @@ class SamplingParam:
     # Video inputs
     video_path: str | None = None
 
-    # Optional pre-generated diffusion latents. Used by parity/debug harnesses
-    # and advanced callers that need deterministic latent reuse.
-    latents: Any | None = None
-
     # Action control inputs (Matrix-Game)
     mouse_cond: Any | None = None  # Shape: (B, T, 2)
     keyboard_cond: Any | None = None  # Shape: (B, T, K)
@@ -68,7 +64,6 @@ class SamplingParam:
     # Text inputs
     prompt: str | list[str] | None = None
     negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
-    max_sequence_length: int | None = None
     prompt_path: str | None = None
     output_path: str = "outputs/"
     output_video_name: str | None = None

fastvideo/attention/backends/video_sparse_attn.py

@@ -150,11 +150,14 @@ class VideoSparseAttentionMetadata(AttentionMetadata):
     # in postprocess_output().  Avoids materializing the intermediate
     # ``[B, len(non_pad_index), H, D]`` tensor on every layer.
     untile_combined_index: torch.LongTensor
-    # Per-step shared padded buffer used by tile().  Inference can reuse this
-    # across VSA layers, but training disables it so activation checkpointing
-    # can release the large tiled QKVG scratch tensor after each attention call.
+    # Per-step shared padded buffer used by tile().  Lazily populated on
+    # the first layer's call and reused by every subsequent VSA layer in
+    # the same denoising step.  Scoping to metadata (not class/instance)
+    # makes the reuse thread-safe across concurrent requests and keeps
+    # the "pad positions are zero" invariant trivially true (the buffer
+    # is freshly zeroed alongside ``non_pad_index`` so the index set
+    # cannot drift between calls).
     tile_buf: torch.Tensor | None = None
-    cache_tile_buf: bool = True
 
 
 class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
@@ -172,7 +175,6 @@ class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
         patch_size: tuple[int, int, int],
         VSA_sparsity: float,
         device: torch.device,
-        cache_tile_buf: bool = True,
         **kwargs: dict[str, Any],
     ) -> VideoSparseAttentionMetadata:
         patch_size = patch_size
@@ -199,8 +201,7 @@ class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
             reverse_tile_partition_indices=reverse_tile_partition_indices,
             variable_block_sizes=variable_block_sizes,
             non_pad_index=non_pad_index,
-            untile_combined_index=untile_combined_index,
-            cache_tile_buf=cache_tile_buf)
+            untile_combined_index=untile_combined_index)
 
 
 class VideoSparseAttentionImpl(AttentionImpl):
@@ -236,11 +237,6 @@ class VideoSparseAttentionImpl(AttentionImpl):
         w_padded_size = num_tiles[2] * VSA_TILE_SIZE[2]
         target_shape = (x.shape[0], t_padded_size * h_padded_size * w_padded_size, x.shape[-2], x.shape[-1])
 
-        if not attn_metadata.cache_tile_buf:
-            buf = torch.zeros(target_shape, device=x.device, dtype=x.dtype)
-            buf[:, attn_metadata.non_pad_index] = x[:, attn_metadata.tile_partition_indices]
-            return buf
-
         # Reuse the per-step buffer stashed on metadata (lazily allocated
         # on the first VSA layer's call within a denoising step).  Pad
         # positions are zero from the initial torch.zeros and never

fastvideo/configs/models/dits/__init__.py

@@ -1,6 +1,5 @@
 from fastvideo.configs.models.dits.cosmos import CosmosVideoConfig
 from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig
-from fastvideo.configs.models.dits.flux_2 import Flux2Config
 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,5 +14,5 @@ from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig
 __all__ = [
     "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "CosmosVideoConfig",
     "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig",
-    "MagiHumanVideoConfig", "StableAudioConfig", "Flux2Config"
+    "MagiHumanVideoConfig", "StableAudioConfig"
 ]

fastvideo/configs/models/dits/base.py

@@ -14,11 +14,6 @@ class DiTArchConfig(ArchConfig):
     param_names_mapping: dict = field(default_factory=dict)
     reverse_param_names_mapping: dict = field(default_factory=dict)
     lora_param_names_mapping: dict = field(default_factory=dict)
-    # When True, the denoising stage casts text/prompt embeddings to the DiT's
-    # working dtype before the diffusion loop. Flux2 requires this (BFL casts ctx
-    # to bf16 before denoising); models with fp32 text encoders (Wan, Hunyuan15,
-    # SD3.5) leave it False to preserve full-precision embeddings.
-    cast_prompt_embeds_to_dit_dtype: bool = False
     _supported_attention_backends: tuple[AttentionBackendEnum,
                                          ...] = (AttentionBackendEnum.SAGE_ATTN, AttentionBackendEnum.FLASH_ATTN,
                                                  AttentionBackendEnum.TORCH_SDPA,

fastvideo/configs/models/dits/cosmos3.py

@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 VFM Transformer FastVideo dataclass configs.
+
+Architecture is 1:1 with the published ``nvidia/Cosmos3-Nano`` checkpoint
+(``transformer/config.json``; class ``Cosmos3OmniTransformer`` / framework
+``Cosmos3VFMNetwork``). Field values match that config so the FastVideo native
+DiT builds a parameter tree matching the checkpoint's state-dict surface
+(814 tensors / 44 patterns, validated 2026-06-06).
+
+Reference of record: ``cosmos-framework`` (NVIDIA). The checkpoint is a single
+``layers`` ModuleList of dual-pathway (understanding/text + generation/vision)
+decoder blocks; per layer: ``self_attn`` with und (``to_{q,k,v}``/``to_out``)
+and gen (``add_{q,k,v}_proj``/``to_add_out``) projections + QK-norms, plus
+``mlp`` (und) and ``mlp_moe_gen`` (gen), and four RMSNorms. Top level adds
+``embed_tokens``/``norm``/``norm_moe_gen``/``lm_head``/``proj_in``/``proj_out``/
+``time_embedder`` and dormant ``action_*``/``audio_*`` heads. The checkpoint
+remap lives in ``scripts/checkpoint_conversion/cosmos3_convert.py``.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+def is_cosmos3_transformer_block(name: str, module) -> bool:
+    """FSDP shard boundary: the dual-pathway decoder blocks ``layers.{i}``."""
+    del module
+    parts = name.split(".")
+    return "layers" in parts and parts[-1].isdigit()
+
+
+@dataclass
+class Cosmos3ArchConfig(DiTArchConfig):
+    """Architecture config for the Cosmos3 omni DiT (Cosmos3-Nano).
+
+    1:1 with ``transformer/config.json``. The action/sound heads ship in the
+    checkpoint, so they are constructed for strict-load parity even though the
+    PR1 video path (T2V/I2V/T2I) leaves them dormant.
+    """
+
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_cosmos3_transformer_block])
+
+    # Conversion is owned by scripts/checkpoint_conversion/cosmos3_convert.py;
+    # the native module tree is the source of truth for parameter names.
+    param_names_mapping: dict = field(default_factory=dict)
+
+    # ---- Backbone (Qwen3-VL-text) ----
+    hidden_size: int = 4096
+    num_hidden_layers: int = 36
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 8  # GQA (4 query groups)
+    head_dim: int = 128
+    intermediate_size: int = 12288
+    hidden_act: str = "silu"
+    vocab_size: int = 151936
+    rms_norm_eps: float = 1e-6
+    attention_bias: bool = False
+    qk_norm_for_diffusion: bool = True
+    qk_norm_for_text: bool = True
+    use_moe: bool = True  # dual-pathway weights; sparse routing unused
+    joint_attn_implementation: str = "two_way"
+    freeze_und: bool = False
+
+    # ---- Position embedding (unified 3D MRoPE) ----
+    position_embedding_type: str = "unified_3d_mrope"
+    rope_theta: float = 5_000_000.0
+    max_position_embeddings: int = 262144
+    mrope_section: list[int] = field(default_factory=lambda: [24, 20, 20])
+    mrope_interleaved: bool = True
+    unified_3d_mrope_reset_spatial_ids: bool = True
+    temporal_modality_margin: int = 15000  # unified_3d_mrope_temporal_modality_margin
+
+    # ---- VAE / patch geometry ----
+    latent_patch_size: int = 2
+    latent_channel: int = 48
+    patch_latent_dim: int = 192  # latent_patch_size**2 * latent_channel
+
+    # ---- Diffusion conditioning ----
+    timestep_scale: float = 0.001
+
+    # ---- Temporal / FPS modulation ----
+    base_fps: float = 24.0
+    temporal_compression_factor: int = 4
+    enable_fps_modulation: bool = True
+    video_temporal_causal: bool = False
+
+    # ---- Action generation head (dormant in PR1 video path) ----
+    action_gen: bool = True
+    action_dim: int = 64
+    max_action_dim: int = 64
+    num_embodiment_domains: int = 32
+
+    # ---- Sound generation head (dormant in PR1 video path) ----
+    sound_gen: bool = True
+    sound_dim: int = 64
+    sound_latent_fps: float = 25.0
+    temporal_compression_factor_sound: int = 1
+
+    # ---- BaseDiT bookkeeping ----
+    in_channels: int = 48
+    out_channels: int = 48
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        # Video DiT contract: latent channels == VAE z_dim.
+        self.num_channels_latents = self.latent_channel
+        if not self.out_channels:
+            self.out_channels = self.in_channels
+        # Derived: patchify packs latent_patch_size**2 spatial patches * channels.
+        self.patch_latent_dim = self.latent_patch_size**2 * self.latent_channel
+
+
+@dataclass
+class Cosmos3VideoConfig(DiTConfig):
+    """Pipeline-level Cosmos3 DiT config (T2V / I2V / T2I share this surface)."""
+
+    arch_config: DiTArchConfig = field(default_factory=Cosmos3ArchConfig)
+    prefix: str = "Cosmos3"

fastvideo/configs/models/dits/flux_2.py

@@ -1,77 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Copied and adapted from: https://github.com/sglang-ai/sglang
-from dataclasses import dataclass, field
-
-from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
-from fastvideo.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-@dataclass
-class Flux2ArchConfig(DiTArchConfig):
-    """Architecture configuration for Flux2 transformer model."""
-
-    cast_prompt_embeds_to_dit_dtype: bool = True
-
-    # Flux2-specific architecture parameters
-    patch_size: int = 1
-    in_channels: int = 64
-    out_channels: int | None = None
-    num_layers: int = 19  # Number of double-stream transformer blocks
-    num_single_layers: int = 38  # Number of single-stream transformer blocks
-    attention_head_dim: int = 128
-    num_attention_heads: int = 24
-    joint_attention_dim: int = 4096  # Dimension for text encoder output
-    timestep_guidance_channels: int = 256  # Dimension for timestep embedding
-    mlp_ratio: float = 3.0
-    axes_dims_rope: tuple[int, ...] = (32, 32, 32, 32)  # RoPE dimensions per axis (match diffusers Flux2)
-    rope_theta: int = 2000  # Base frequency for RoPE (match diffusers Flux2)
-    eps: float = 1e-6
-    guidance_embeds: bool = True  # Whether to use guidance embeddings
-    # When True, compute SwiGLU in fp32 inside ``ff_context`` only (bf16 noise mitigation).
-    ff_context_swiglu_fp32: bool = False
-
-    # Parameter name mapping for loading HuggingFace checkpoints
-    param_names_mapping: dict = field(default_factory=lambda: {
-        r"transformer\.(\w*)\.(.*)$": r"\1.\2",
-    })
-
-    def __post_init__(self) -> None:
-        super().__post_init__()
-        self.out_channels = self.out_channels or self.in_channels
-        self.hidden_size = self.num_attention_heads * self.attention_head_dim
-        self.num_channels_latents = self.out_channels
-
-    def update_from_weight_keys(self, all_keys: set[str]) -> None:
-        """Infer num_layers and num_single_layers from checkpoint weight keys so the model is built with the same number of blocks as the weights."""
-        if not all_keys:
-            return
-        num_layers = 0
-        num_single_layers = 0
-        for k in all_keys:
-            if "single_transformer_blocks." not in k and "transformer_blocks." in k:
-                parts = k.split("transformer_blocks.")[-1].split(".")
-                if parts[0].isdigit():
-                    num_layers = max(num_layers, int(parts[0]) + 1)
-            if "single_transformer_blocks." in k:
-                parts = k.split("single_transformer_blocks.")[-1].split(".")
-                if parts[0].isdigit():
-                    num_single_layers = max(num_single_layers, int(parts[0]) + 1)
-        if num_layers > 0:
-            self.num_layers = num_layers
-            logger.info("Inferred num_layers=%s from checkpoint keys", num_layers)
-        if num_single_layers > 0:
-            self.num_single_layers = num_single_layers
-            logger.info("Inferred num_single_layers=%s from checkpoint keys", num_single_layers)
-        if num_layers > 0 or num_single_layers > 0:
-            self.__post_init__()
-
-
-@dataclass
-class Flux2Config(DiTConfig):
-    """Configuration for Flux2 transformer model."""
-
-    arch_config: DiTArchConfig = field(default_factory=Flux2ArchConfig)
-
-    prefix: str = "Flux"

fastvideo/configs/models/encoders/__init__.py

@@ -7,8 +7,6 @@ from fastvideo.configs.models.encoders.qwen2_5 import Qwen2_5_VLConfig
 from fastvideo.configs.models.encoders.siglip import SiglipVisionConfig
 from fastvideo.configs.models.encoders.reason1 import Reason1ArchConfig, Reason1Config
 from fastvideo.configs.models.encoders.gemma import LTX2GemmaConfig
-from fastvideo.configs.models.encoders.mistral3 import Mistral3TextConfig
-from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig
 from fastvideo.configs.models.encoders.stable_audio_conditioner import (StableAudioConditionerArchConfig,
                                                                         StableAudioConditionerConfig)
 from fastvideo.configs.models.encoders.t5gemma import T5GemmaEncoderConfig
@@ -17,5 +15,5 @@ __all__ = [
     "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig", "BaseEncoderOutput", "CLIPTextConfig",
     "CLIPVisionConfig", "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig", "Qwen2_5_VLConfig",
     "Reason1ArchConfig", "Reason1Config", "LTX2GemmaConfig", "SiglipVisionConfig", "StableAudioConditionerArchConfig",
-    "StableAudioConditionerConfig", "T5GemmaEncoderConfig", "Qwen3TextConfig", "Mistral3TextConfig"
+    "StableAudioConditionerConfig", "T5GemmaEncoderConfig"
 ]

fastvideo/configs/models/encoders/base.py

@@ -36,11 +36,6 @@ class TextEncoderArchConfig(EncoderArchConfig):
         default_factory=list)  # mapping from huggingface weight names to custom names
     tokenizer_kwargs: dict[str, Any] = field(default_factory=dict)
     _fsdp_shard_conditions: list = field(default_factory=lambda: [])
-    # When True, the tokenizer loader prefers AutoProcessor over AutoTokenizer
-    # for encoders whose tokenizer dir ships a processor_config.json (e.g. Flux2
-    # full's Mistral3 multimodal processor). Default False keeps every existing
-    # encoder on the historical AutoTokenizer path.
-    require_processor: bool = False
 
     def __post_init__(self) -> None:
         self.tokenizer_kwargs = {

fastvideo/configs/models/encoders/mistral3.py

@@ -1,38 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Mistral3 text encoder configuration for full Flux2."""
-from dataclasses import dataclass, field
-
-from fastvideo.configs.models.encoders.base import (
-    TextEncoderArchConfig,
-    TextEncoderConfig,
-)
-
-
-@dataclass
-class Mistral3TextArchConfig(TextEncoderArchConfig):
-    """Architecture config for the Mistral3 text encoder used by full Flux2."""
-
-    architectures: list[str] = field(default_factory=lambda: ["Mistral3ForConditionalGeneration"])
-    hidden_size: int = 5120
-    num_hidden_layers: int = 40
-    text_len: int = 512
-    output_hidden_states: bool = True
-    # Mistral3 (full Flux2) ships a multimodal processor; load via AutoProcessor.
-    require_processor: bool = True
-
-    def __post_init__(self) -> None:
-        self.tokenizer_kwargs = {
-            "padding": "max_length",
-            "truncation": True,
-            "max_length": self.text_len,
-            "return_tensors": "pt",
-        }
-
-
-@dataclass
-class Mistral3TextConfig(TextEncoderConfig):
-    """Top-level config for the Mistral3 full Flux2 text encoder."""
-
-    arch_config: TextEncoderArchConfig = field(default_factory=Mistral3TextArchConfig)
-    prefix: str = "mistral3"
-    is_chat_model: bool = True

fastvideo/configs/models/encoders/qwen3.py

@@ -1,82 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Ported from SGLang: python/sglang/multimodal_gen/configs/models/encoders/qwen3.py
-"""Qwen3 text encoder configuration for FastVideo diffusion models (e.g. Flux2 Klein)."""
-from dataclasses import dataclass, field
-from typing import Any
-
-from fastvideo.configs.models.encoders.base import (
-    TextEncoderArchConfig,
-    TextEncoderConfig,
-)
-
-
-def _is_transformer_layer(n: str, m: Any) -> bool:
-    return "layers" in n and str.isdigit(n.split(".")[-1])
-
-
-def _is_embeddings(n: str, m: Any) -> bool:
-    return n.endswith("embed_tokens")
-
-
-def _is_final_norm(n: str, m: Any) -> bool:
-    return n.endswith("norm")
-
-
-@dataclass
-class Qwen3TextArchConfig(TextEncoderArchConfig):
-    """Architecture config for Qwen3 text encoder.
-
-    Qwen3 is similar to LLaMA but with QK-Norm (RMSNorm on Q and K before attention).
-    Used by Flux2 Klein.
-    """
-
-    vocab_size: int = 151936
-    hidden_size: int = 2560
-    intermediate_size: int = 9728
-    num_hidden_layers: int = 36
-    num_attention_heads: int = 32
-    num_key_value_heads: int = 8
-    hidden_act: str = "silu"
-    max_position_embeddings: int = 40960
-    initializer_range: float = 0.02
-    rms_norm_eps: float = 1e-6
-    use_cache: bool = True
-    pad_token_id: int = 151643
-    bos_token_id: int = 151643
-    eos_token_id: int = 151645
-    tie_word_embeddings: bool = True
-    rope_theta: float = 1000000.0
-    rope_scaling: dict | None = None
-    attention_bias: bool = False
-    attention_dropout: float = 0.0
-    mlp_bias: bool = False
-    head_dim: int = 128
-    text_len: int = 512
-    output_hidden_states: bool = True  # Klein needs hidden states from layers 9, 18, 27
-
-    stacked_params_mapping: list[tuple[str, str, str | int]] = field(default_factory=lambda: [
-        (".qkv_proj", ".q_proj", "q"),
-        (".qkv_proj", ".k_proj", "k"),
-        (".qkv_proj", ".v_proj", "v"),
-        (".gate_up_proj", ".gate_proj", 0),
-        (".gate_up_proj", ".up_proj", 1),
-    ])
-    _fsdp_shard_conditions: list = field(
-        default_factory=lambda: [_is_transformer_layer, _is_embeddings, _is_final_norm])
-
-    def __post_init__(self) -> None:
-        self.tokenizer_kwargs = {
-            "padding": "max_length",
-            "truncation": True,
-            "max_length": self.text_len,
-            "return_tensors": "pt",
-        }
-
-
-@dataclass
-class Qwen3TextConfig(TextEncoderConfig):
-    """Top-level config for Qwen3 text encoder."""
-
-    arch_config: TextEncoderArchConfig = field(default_factory=Qwen3TextArchConfig)
-    prefix: str = "qwen3"
-    is_chat_model: bool = True

fastvideo/configs/models/vaes/__init__.py

@@ -1,12 +1,12 @@
 from fastvideo.configs.models.vaes.cosmosvae import CosmosVAEConfig
 from fastvideo.configs.models.vaes.cosmos2_5vae import Cosmos25VAEConfig
+from fastvideo.configs.models.vaes.cosmos3vae import Cosmos3VAEConfig
 from fastvideo.configs.models.vaes.gamecraftvae import GameCraftVAEConfig
 from fastvideo.configs.models.vaes.gen3cvae import Gen3CVAEConfig
 from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
 from fastvideo.configs.models.vaes.hunyuan15vae import Hunyuan15VAEConfig
 from fastvideo.configs.models.vaes.ltx2vae import LTX2VAEConfig
 from fastvideo.configs.models.vaes.oobleck import OobleckVAEArchConfig, OobleckVAEConfig
-from fastvideo.configs.models.vaes.flux2vae import Flux2VAEConfig
 from fastvideo.configs.models.vaes.wanvae import WanVAEConfig
 
 __all__ = [
@@ -15,10 +15,10 @@ __all__ = [
     "WanVAEConfig",
     "CosmosVAEConfig",
     "Cosmos25VAEConfig",
+    "Cosmos3VAEConfig",
     "Gen3CVAEConfig",
     "Hunyuan15VAEConfig",
     "LTX2VAEConfig",
     "OobleckVAEArchConfig",
     "OobleckVAEConfig",
-    "Flux2VAEConfig",
 ]

fastvideo/configs/models/vaes/cosmos3vae.py

@@ -0,0 +1,277 @@
+"""Cosmos3 (Wan2.2-TI2V-5B) VAE config and checkpoint-key mapping.
+
+The Cosmos3 checkpoint VAE is literally ``Wan-AI/Wan2.2-TI2V-5B-Diffusers``
+(diffusers ``AutoencoderKLWan``), so this config locks the Wan2.2 geometry:
+residual down/up blocks, ``patch_size=2``, ``z_dim=48``, ``base_dim=160``,
+``decoder_base_dim=256``, and ``scale_factor_spatial=16``. The 48-dim
+``latents_mean``/``latents_std`` are taken verbatim from the Cosmos3
+checkpoint's ``vae/config.json`` (identical to the canonical Wan2.2-TI2V-5B
+statistics).
+
+Mirrors the :class:`Cosmos25VAEArchConfig` pattern. ``param_names_mapping`` /
+``map_official_key`` translate the *official* Wan2.2 VAE state-dict keys
+(nested-residual naming, e.g. ``encoder.downsamples.{b}.downsamples.{j}`` and
+``decoder.upsamples.{b}.upsamples.{j}``) into FastVideo's ``AutoencoderKLWan``
+key space. The standard diffusers checkpoint already ships native FastVideo
+keys, so these helpers exist for parity tooling and official ``.pth`` loading.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.vaes.wanvae import WanVAEArchConfig, WanVAEConfig
+
+
+@dataclass
+class Cosmos3VAEArchConfig(WanVAEArchConfig):
+    # Wan2.2-TI2V-5B geometry (differs from the Wan2.1 WanVAEArchConfig
+    # defaults: residual blocks, patch_size=2, z_dim=48, base_dim=160,
+    # decoder_base_dim=256, scale_factor_spatial=16, 12 patch channels).
+    _name_or_path: str = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+    base_dim: int = 160
+    decoder_base_dim: int | None = 256
+    z_dim: int = 48
+    dim_mult: tuple[int, ...] = (1, 2, 4, 4)
+    num_res_blocks: int = 2
+    attn_scales: tuple[float, ...] = ()
+    temperal_downsample: tuple[bool, ...] = (False, True, True)
+    dropout: float = 0.0
+    is_residual: bool = True
+    in_channels: int = 12
+    out_channels: int = 12
+    patch_size: int | None = 2
+    scale_factor_temporal: int = 4
+    scale_factor_spatial: int = 16
+    clip_output: bool = False
+
+    # 48-dim statistics copied verbatim from the Cosmos3 checkpoint
+    # (official_weights/cosmos3/vae/config.json).
+    latents_mean: tuple[float, ...] = (
+        -0.2289,
+        -0.0052,
+        -0.1323,
+        -0.2339,
+        -0.2799,
+        0.0174,
+        0.1838,
+        0.1557,
+        -0.1382,
+        0.0542,
+        0.2813,
+        0.0891,
+        0.157,
+        -0.0098,
+        0.0375,
+        -0.1825,
+        -0.2246,
+        -0.1207,
+        -0.0698,
+        0.5109,
+        0.2665,
+        -0.2108,
+        -0.2158,
+        0.2502,
+        -0.2055,
+        -0.0322,
+        0.1109,
+        0.1567,
+        -0.0729,
+        0.0899,
+        -0.2799,
+        -0.123,
+        -0.0313,
+        -0.1649,
+        0.0117,
+        0.0723,
+        -0.2839,
+        -0.2083,
+        -0.052,
+        0.3748,
+        0.0152,
+        0.1957,
+        0.1433,
+        -0.2944,
+        0.3573,
+        -0.0548,
+        -0.1681,
+        -0.0667,
+    )
+    latents_std: tuple[float, ...] = (
+        0.4765,
+        1.0364,
+        0.4514,
+        1.1677,
+        0.5313,
+        0.499,
+        0.4818,
+        0.5013,
+        0.8158,
+        1.0344,
+        0.5894,
+        1.0901,
+        0.6885,
+        0.6165,
+        0.8454,
+        0.4978,
+        0.5759,
+        0.3523,
+        0.7135,
+        0.6804,
+        0.5833,
+        1.4146,
+        0.8986,
+        0.5659,
+        0.7069,
+        0.5338,
+        0.4889,
+        0.4917,
+        0.4069,
+        0.4999,
+        0.6866,
+        0.4093,
+        0.5709,
+        0.6065,
+        0.6415,
+        0.4944,
+        0.5726,
+        1.2042,
+        0.5458,
+        1.6887,
+        0.3971,
+        1.06,
+        0.3943,
+        0.5537,
+        0.5444,
+        0.4089,
+        0.7468,
+        0.7744,
+    )
+
+    # Simple 1:1 renames. The nested-residual block remapping (encoder
+    # downsamples / decoder upsamples / middle / head) is handled by
+    # ``map_official_key()``.
+    param_names_mapping: dict[str, str] = field(
+        default_factory=lambda: {
+            r"^conv1\.(.*)$": r"quant_conv.\1",
+            r"^conv2\.(.*)$": r"post_quant_conv.\1",
+            r"^encoder\.conv1\.(.*)$": r"encoder.conv_in.\1",
+            r"^decoder\.conv1\.(.*)$": r"decoder.conv_in.\1",
+            r"^encoder\.head\.0\.gamma$": r"encoder.norm_out.gamma",
+            r"^encoder\.head\.2\.(.*)$": r"encoder.conv_out.\1",
+            r"^decoder\.head\.0\.gamma$": r"decoder.norm_out.gamma",
+            r"^decoder\.head\.2\.(.*)$": r"decoder.conv_out.\1",
+        })
+
+    @staticmethod
+    def map_official_key(key: str) -> str | None:
+        """Map a single official Wan2.2 VAE key into FastVideo key space.
+
+        Handles the residual (Wan2.2) module layout where each down/up block
+        is a nested ``Sequential`` (``downsamples.{b}.downsamples.{j}`` /
+        ``upsamples.{b}.upsamples.{j}``) rather than the flat Wan2.1 indexing.
+        Returns ``None`` for keys with no FastVideo counterpart.
+        """
+
+        def map_residual_subkey(prefix: str, sub: str) -> str | None:
+            if re.match(r"^residual\.0\.gamma$", sub):
+                return f"{prefix}.norm1.gamma"
+            m = re.match(r"^residual\.2\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.conv1.{m.group(1)}"
+            if re.match(r"^residual\.3\.gamma$", sub):
+                return f"{prefix}.norm2.gamma"
+            m = re.match(r"^residual\.6\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.conv2.{m.group(1)}"
+            m = re.match(r"^shortcut\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.conv_shortcut.{m.group(1)}"
+            return None
+
+        def map_attn_subkey(prefix: str, sub: str) -> str | None:
+            if re.match(r"^norm\.gamma$", sub):
+                return f"{prefix}.norm.gamma"
+            m = re.match(r"^to_qkv\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.to_qkv.{m.group(1)}"
+            m = re.match(r"^proj\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.proj.{m.group(1)}"
+            return None
+
+        def map_resample_subkey(prefix: str, sub: str) -> str | None:
+            m = re.match(r"^resample\.1\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.resample.1.{m.group(1)}"
+            m = re.match(r"^time_conv\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.time_conv.{m.group(1)}"
+            return None
+
+        m = re.match(r"^conv1\.(weight|bias)$", key)
+        if m:
+            return f"quant_conv.{m.group(1)}"
+        m = re.match(r"^conv2\.(weight|bias)$", key)
+        if m:
+            return f"post_quant_conv.{m.group(1)}"
+        m = re.match(r"^(encoder|decoder)\.conv1\.(weight|bias)$", key)
+        if m:
+            return f"{m.group(1)}.conv_in.{m.group(2)}"
+        m = re.match(r"^(encoder|decoder)\.head\.0\.gamma$", key)
+        if m:
+            return f"{m.group(1)}.norm_out.gamma"
+        m = re.match(r"^(encoder|decoder)\.head\.2\.(weight|bias)$", key)
+        if m:
+            return f"{m.group(1)}.conv_out.{m.group(2)}"
+
+        m = re.match(r"^(encoder|decoder)\.middle\.0\.(.*)$", key)
+        if m:
+            return map_residual_subkey(f"{m.group(1)}.mid_block.resnets.0", m.group(2))
+        m = re.match(r"^(encoder|decoder)\.middle\.1\.(.*)$", key)
+        if m:
+            return map_attn_subkey(f"{m.group(1)}.mid_block.attentions.0", m.group(2))
+        m = re.match(r"^(encoder|decoder)\.middle\.2\.(.*)$", key)
+        if m:
+            return map_residual_subkey(f"{m.group(1)}.mid_block.resnets.1", m.group(2))
+
+        # Encoder: downsamples.{block}.downsamples.{j}.* (nested residual layout)
+        m = re.match(r"^encoder\.downsamples\.(\d+)\.downsamples\.(\d+)\.(.*)$", key)
+        if m:
+            block_i, res_i, sub = int(m.group(1)), int(m.group(2)), m.group(3)
+            if sub.startswith("resample.") or sub.startswith("time_conv."):
+                return map_resample_subkey(f"encoder.down_blocks.{block_i}.downsampler", sub)
+            return map_residual_subkey(f"encoder.down_blocks.{block_i}.resnets.{res_i}", sub)
+
+        # Decoder: upsamples.{block}.upsamples.{j}.* (nested residual layout)
+        m = re.match(r"^decoder\.upsamples\.(\d+)\.upsamples\.(\d+)\.(.*)$", key)
+        if m:
+            block_i, res_i, sub = int(m.group(1)), int(m.group(2)), m.group(3)
+            if sub.startswith("resample.") or sub.startswith("time_conv."):
+                return map_resample_subkey(f"decoder.up_blocks.{block_i}.upsampler", sub)
+            return map_residual_subkey(f"decoder.up_blocks.{block_i}.resnets.{res_i}", sub)
+
+        return None
+
+    # ``__post_init__`` (scaling_factor / shift_factor / compression ratios) is
+    # inherited unchanged from ``WanVAEArchConfig``.
+
+
+@dataclass
+class Cosmos3VAEConfig(WanVAEConfig):
+    """Cosmos3 VAE config (reuses FastVideo's Wan2.2 ``AutoencoderKLWan``).
+
+    Subclasses :class:`WanVAEConfig` so the model reads the same runtime flags
+    (``use_feature_cache``, ``use_light_vae``, tiling) and only swaps in the
+    Cosmos3 = Wan2.2 ``arch_config``.
+    """
+
+    arch_config: Cosmos3VAEArchConfig = field(default_factory=Cosmos3VAEArchConfig)
+
+    use_feature_cache: bool = True
+    use_tiling: bool = False
+    use_temporal_tiling: bool = False
+    use_parallel_tiling: bool = False
+
+    # ``__post_init__`` (blend_num_frames) is inherited from ``WanVAEConfig``.

fastvideo/configs/models/vaes/flux2vae.py

@@ -1,58 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Copied and adapted from: https://github.com/sglang-ai/sglang
-from dataclasses import dataclass, field
-
-from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig
-
-
-@dataclass
-class Flux2VAEArchConfig(VAEArchConfig):
-    """Architecture configuration for Flux2 VAE model."""
-
-    # Flux2 VAE-specific architecture parameters
-    in_channels: int = 3
-    out_channels: int = 3
-    down_block_types: tuple[str, ...] = (
-        "DownEncoderBlock2D",
-        "DownEncoderBlock2D",
-        "DownEncoderBlock2D",
-        "AttnDownEncoderBlock2D",
-    )
-    up_block_types: tuple[str, ...] = (
-        "AttnUpDecoderBlock2D",
-        "UpDecoderBlock2D",
-        "UpDecoderBlock2D",
-        "UpDecoderBlock2D",
-    )
-    block_out_channels: tuple[int, ...] = (128, 256, 512, 512)
-    layers_per_block: int = 2
-    act_fn: str = "silu"
-    latent_channels: int = 16
-    norm_num_groups: int = 32
-    sample_size: int = 512
-    force_upcast: bool = False
-    use_quant_conv: bool = True
-    use_post_quant_conv: bool = True
-    mid_block_add_attention: bool = True
-    batch_norm_eps: float = 1e-5
-    batch_norm_momentum: float = 0.1
-    patch_size: tuple[int, int] = (1, 1)
-
-    # Latent scaling for decode: avoid division-by-zero; match Flux/Flux2 convention (e.g. 0.13025)
-    scaling_factor: float = 0.13025
-
-    # Spatial compression (for images, this is typically 8)
-    spatial_compression_ratio: int = 8
-    temporal_compression_ratio: int = 1  # Images don't have temporal dimension
-
-
-@dataclass
-class Flux2VAEConfig(VAEConfig):
-    """Configuration for Flux2 VAE model."""
-
-    arch_config: Flux2VAEArchConfig = field(default_factory=Flux2VAEArchConfig)
-
-    # Flux2 is an image model, so disable temporal tiling
-    use_tiling: bool = False
-    use_temporal_tiling: bool = False
-    use_parallel_tiling: bool = False

fastvideo/configs/pipelines/__init__.py

@@ -9,12 +9,12 @@ from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig
 from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig
 from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
 from fastvideo.registry import get_pipeline_config_cls_from_name
-from fastvideo.configs.pipelines.wan import (LucyEditDevConfig, SelfForcingWanT2V480PConfig, WanI2V480PConfig,
-                                             WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig)
+from fastvideo.configs.pipelines.wan import (SelfForcingWanT2V480PConfig, WanI2V480PConfig, WanI2V720PConfig,
+                                             WanT2V480PConfig, WanT2V720PConfig)
 
 __all__ = [
     "HunyuanConfig", "FastHunyuanConfig", "HunyuanGameCraftPipelineConfig", "PipelineConfig", "Hunyuan15T2V480PConfig",
     "Hunyuan15T2V720PConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig",
-    "SelfForcingWanT2V480PConfig", "LucyEditDevConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig",
-    "HYWorldConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
+    "SelfForcingWanT2V480PConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig", "HYWorldConfig",
+    "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
 ]

fastvideo/configs/pipelines/base.py

@@ -35,11 +35,6 @@ class PipelineConfig:
     flow_shift: float | None = None
     flow_shift_sr: float | None = None
     disable_autocast: bool = False
-    # When True, the scheduler's Euler update runs in fp32 outside the autocast
-    # block (Diffusers-style; avoids BF16 drift over multiple steps). Flux2 sets
-    # this True for reference parity; other models keep the legacy in-autocast
-    # behavior to preserve existing SSIM references.
-    scheduler_step_in_fp32: bool = False
     is_causal: bool = False
 
     # Model configuration
@@ -69,9 +64,8 @@ class PipelineConfig:
     # DMD parameters
     dmd_denoising_steps: list[int] | None = field(default=None)
 
-    # Wan2.2 task modifiers
+    # Wan2.2 TI2V parameters
     ti2v_task: bool = False
-    lucy_edit_task: bool = False
     boundary_ratio: float | None = None
 
     # Compilation

fastvideo/configs/pipelines/cosmos3.py

@@ -0,0 +1,69 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 pipeline configuration.
+
+Reference of record: the official ``cosmos-framework`` / ``nvidia/Cosmos3-Nano``
+checkpoint (``model_index.json``). Cosmos3 is structurally different from
+Cosmos 2.5:
+
+- Dual-pathway (UND + GEN) DiT lives entirely inside ``Cosmos3VFMTransformer``
+  (``Cosmos3VideoConfig``).
+- No separate text encoder — the Qwen3-VL-text backbone is inside the DiT, so
+  ``text_encoder_configs`` is the empty tuple. The Qwen2 tokenizer is loaded as
+  the ``text_tokenizer`` checkpoint module by the component loader.
+- VAE is Wan2.2 ``AutoencoderKLWan`` (z_dim=48, scale_factor_spatial=16),
+  configured by ``Cosmos3VAEConfig`` (the checkpoint's exact latents_mean/std).
+- Scheduler is FastVideo-native ``UniPCMultistepScheduler`` configured for
+  pure flow matching (flow_prediction, use_flow_sigmas), equivalent to the
+  framework's ``FlowUniPCMultistepScheduler``. The checkpoint's diffusers-style
+  scheduler config (karras/sigma_min/max) is coerced to the flow setup in
+  ``Cosmos3OmniDiffusersPipeline.initialize_pipeline``.
+- T2I default ``flow_shift`` is 3.0 (set per-request by ``_set_flow_shift``);
+  T2V/I2V use the engine-init default of 1.0 baked into this config.
+"""
+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.cosmos3 import (Cosmos3ArchConfig, Cosmos3VideoConfig)
+from fastvideo.configs.models.encoders import BaseEncoderOutput
+from fastvideo.configs.models.vaes import Cosmos3VAEConfig  # Wan2.2 AutoencoderKLWan
+from fastvideo.configs.pipelines.base import PipelineConfig
+
+
+@dataclass
+class Cosmos3Config(PipelineConfig):
+    """Configuration for the Cosmos3 video generation pipeline (T2V/I2V/T2I).
+
+    Wires the framework-parity-verified Cosmos3 components: the native
+    ``Cosmos3VideoConfig`` DiT, the Wan2.2 ``Cosmos3VAEConfig`` VAE, the Qwen2
+    tokenizer (loaded as ``text_tokenizer``), and the UniPC scheduler.
+    """
+
+    dit_config: DiTConfig = field(default_factory=lambda: Cosmos3VideoConfig(arch_config=Cosmos3ArchConfig()))
+
+    vae_config: VAEConfig = field(default_factory=Cosmos3VAEConfig)
+
+    # No separate text encoder: the Qwen3-VL-text backbone lives inside the DiT
+    # and the pipeline tokenizes in Cosmos3DenoisingStage, so all three
+    # text-encoder lists are empty (the generic text-encode stage is not used).
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=tuple)
+    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(default_factory=tuple)
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], ...] = field(default_factory=tuple)
+
+    dit_precision: str = "bf16"
+    vae_precision: str = "bf16"
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=tuple)
+
+    embedded_cfg_scale: float = 0.0
+    # T2V/I2V engine-init flow_shift (framework text2video/image2video default);
+    # T2I overrides to 3.0 per request via Cosmos3DenoisingStage._set_flow_shift.
+    flow_shift: float = 10.0
+
+    vae_tiling: bool = False
+    vae_sp: bool = False
+
+    def __post_init__(self):
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True

fastvideo/configs/pipelines/flux_2.py

@@ -1,86 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Copied and adapted from: https://github.com/sglang-ai/sglang
-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.flux_2 import Flux2Config
-from fastvideo.configs.models.encoders import BaseEncoderOutput
-from fastvideo.configs.models.encoders.base import EncoderArchConfig
-from fastvideo.configs.models.encoders.mistral3 import Mistral3TextConfig
-from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig
-from fastvideo.configs.models.vaes.flux2vae import Flux2VAEConfig
-from fastvideo.configs.pipelines.base import PipelineConfig, preprocess_text
-
-
-@dataclass
-class Flux2PipelineConfig(PipelineConfig):
-    """Configuration for Flux2 image generation pipeline."""
-
-    # Flux2-specific parameters
-    embedded_cfg_scale: float | None = 4.0
-    scheduler_step_in_fp32: bool = True
-    flux2_text_encoder_type: str = "mistral3"
-    text_encoder_out_layers: tuple[int, ...] = (10, 20, 30)
-
-    # DiT configuration
-    dit_config: DiTConfig = field(default_factory=Flux2Config)
-    dit_precision: str = "bf16"
-
-    # VAE configuration
-    vae_config: VAEConfig = field(default_factory=Flux2VAEConfig)
-    vae_precision: str = "fp32"
-    vae_tiling: bool = False  # Flux2 is image model, disable tiling by default
-    vae_sp: bool = False
-
-    # Text encoder configuration (full Flux2 uses Mistral3)
-    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (Mistral3TextConfig(), ))
-    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", ))
-
-    # Default postprocess function (can be overridden)
-    @staticmethod
-    def default_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
-        """Default text postprocessing for Flux2."""
-        return outputs.last_hidden_state
-
-    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
-                                  ...] = field(default_factory=lambda: (Flux2PipelineConfig.default_postprocess_text, ))
-
-
-def flux2_klein_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
-    """Klein postprocess: hidden states from layers 9, 18, 27 (Qwen3)."""
-    hidden_states_layers: list[int] = [9, 18, 27]
-    if outputs.hidden_states is None:
-        raise ValueError("Flux2 Klein requires output_hidden_states=True from text encoder")
-    out = torch.stack([outputs.hidden_states[k] for k in hidden_states_layers], dim=1)
-    batch_size, num_channels, seq_len, hidden_dim = out.shape
-    prompt_embeds = out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, num_channels * hidden_dim)
-    return prompt_embeds
-
-
-@dataclass
-class Flux2KleinEncoderArchConfig(EncoderArchConfig):
-    """Encoder arch config for Flux2 Klein (Qwen3); needs hidden states for layers 9, 18, 27."""
-    output_hidden_states: bool = True
-
-
-@dataclass
-class Flux2KleinTextEncoderConfig(EncoderConfig):
-    """Text encoder config for Flux2 Klein (Qwen3)."""
-    arch_config: EncoderArchConfig = field(default_factory=Flux2KleinEncoderArchConfig)
-
-
-@dataclass
-class Flux2KleinPipelineConfig(Flux2PipelineConfig):
-    """Configuration for Flux2 Klein (distilled, 4-step, no guidance)."""
-    embedded_cfg_scale: float | None = None  # Klein distilled: no guidance embedding (matches Diffusers)
-    scheduler_step_in_fp32: bool = True
-    flux2_text_encoder_type: str = "qwen3"
-    text_encoder_out_layers: tuple[int, ...] = (9, 18, 27)
-    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (Qwen3TextConfig(), ))
-    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", ))
-    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(default_factory=lambda: (preprocess_text, ))
-    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
-                                  ...] = field(default_factory=lambda: (flux2_klein_postprocess_text, ))

fastvideo/configs/pipelines/wan.py

@@ -6,11 +6,9 @@ import torch
 
 from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
 from fastvideo.configs.models.dits import WanVideoConfig
-from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig
 from fastvideo.configs.models.encoders import (BaseEncoderOutput, CLIPVisionConfig, T5Config,
                                                WAN2_1ControlCLIPVisionConfig)
 from fastvideo.configs.models.vaes import WanVAEConfig
-from fastvideo.configs.models.vaes.wanvae import WanVAEArchConfig
 from fastvideo.configs.pipelines.base import PipelineConfig
 
 
@@ -122,142 +120,6 @@ class Wan2_2_TI2V_5B_Config(WanT2V480PConfig):
     expand_timesteps: bool = True
 
     def __post_init__(self) -> None:
-        assert not (self.ti2v_task and self.lucy_edit_task)
-        self.vae_config.load_encoder = True
-        self.vae_config.load_decoder = True
-        self.dit_config.expand_timesteps = self.expand_timesteps
-
-
-@dataclass
-class LucyEditDevConfig(Wan2_2_TI2V_5B_Config):
-    """Configuration for Decart Lucy Edit Dev video editing."""
-
-    dit_config: DiTConfig = field(default_factory=lambda: WanVideoConfig(arch_config=WanVideoArchConfig(
-        num_attention_heads=24,
-        in_channels=96,
-        out_channels=48,
-        ffn_dim=14336,
-        num_layers=30,
-    )))
-    vae_config: VAEConfig = field(default_factory=lambda: WanVAEConfig(arch_config=WanVAEArchConfig(
-        base_dim=160,
-        decoder_base_dim=256,
-        z_dim=48,
-        in_channels=12,
-        out_channels=12,
-        scale_factor_spatial=16,
-        patch_size=2,
-        is_residual=True,
-        clip_output=False,
-        latents_mean=(
-            -0.2289,
-            -0.0052,
-            -0.1323,
-            -0.2339,
-            -0.2799,
-            0.0174,
-            0.1838,
-            0.1557,
-            -0.1382,
-            0.0542,
-            0.2813,
-            0.0891,
-            0.1570,
-            -0.0098,
-            0.0375,
-            -0.1825,
-            -0.2246,
-            -0.1207,
-            -0.0698,
-            0.5109,
-            0.2665,
-            -0.2108,
-            -0.2158,
-            0.2502,
-            -0.2055,
-            -0.0322,
-            0.1109,
-            0.1567,
-            -0.0729,
-            0.0899,
-            -0.2799,
-            -0.1230,
-            -0.0313,
-            -0.1649,
-            0.0117,
-            0.0723,
-            -0.2839,
-            -0.2083,
-            -0.0520,
-            0.3748,
-            0.0152,
-            0.1957,
-            0.1433,
-            -0.2944,
-            0.3573,
-            -0.0548,
-            -0.1681,
-            -0.0667,
-        ),
-        latents_std=(
-            0.4765,
-            1.0364,
-            0.4514,
-            1.1677,
-            0.5313,
-            0.4990,
-            0.4818,
-            0.5013,
-            0.8158,
-            1.0344,
-            0.5894,
-            1.0901,
-            0.6885,
-            0.6165,
-            0.8454,
-            0.4978,
-            0.5759,
-            0.3523,
-            0.7135,
-            0.6804,
-            0.5833,
-            1.4146,
-            0.8986,
-            0.5659,
-            0.7069,
-            0.5338,
-            0.4889,
-            0.4917,
-            0.4069,
-            0.4999,
-            0.6866,
-            0.4093,
-            0.5709,
-            0.6065,
-            0.6415,
-            0.4944,
-            0.5726,
-            1.2042,
-            0.5458,
-            1.6887,
-            0.3971,
-            1.0600,
-            0.3943,
-            0.5537,
-            0.5444,
-            0.4089,
-            0.7468,
-            0.7744,
-        ),
-    )))
-    ti2v_task: bool = False
-    lucy_edit_task: bool = True
-
-    def __post_init__(self) -> None:
-        assert not (self.ti2v_task and self.lucy_edit_task)
-        # Lucy uses Wan2.2's enhanced 48-channel VAE latents. Denoising
-        # concatenates noise + video latents, matching the 96-channel
-        # transformer input declared above.
         self.vae_config.load_encoder = True
         self.vae_config.load_decoder = True
         self.dit_config.expand_timesteps = self.expand_timesteps

fastvideo/entrypoints/video_generator.py

@@ -517,7 +517,6 @@ class VideoGenerator:
             if _ek in kwargs:
                 extra_overrides[_ek] = kwargs.pop(_ek)
 
-        prompt_embeds = kwargs.pop("prompt_embeds", None)
         sampling_param.update(kwargs)
         kwargs["_extra_overrides"] = extra_overrides
 
@@ -568,8 +567,6 @@ class VideoGenerator:
             raise ValueError("Either prompt or prompt_txt must be provided")
         output_path = self._prepare_output_path(sampling_param.output_path, prompt)
         kwargs["output_path"] = output_path
-        if prompt_embeds is not None:
-            kwargs["prompt_embeds"] = prompt_embeds
         return self._generate_single_video(
             prompt=prompt,
             sampling_param=sampling_param,
@@ -672,7 +669,6 @@ class VideoGenerator:
         prompt = prompt.strip()
         sampling_param = deepcopy(sampling_param)
         output_path = kwargs["output_path"]
-        prompt_embeds = kwargs.get("prompt_embeds")
         sampling_param.prompt = prompt
         # Process negative prompt
         if sampling_param.negative_prompt is not None:
@@ -719,9 +715,6 @@ class VideoGenerator:
             n_tokens=n_tokens,
             VSA_sparsity=fastvideo_args.VSA_sparsity,
         )
-        # Allow precomputed prompt_embeds (e.g. from diffusers) to skip text encoding
-        if prompt_embeds is not None:
-            batch.prompt_embeds = (list(prompt_embeds) if isinstance(prompt_embeds, list | tuple) else [prompt_embeds])
 
         extra_overrides = kwargs.pop("_extra_overrides", {})
         for _ek, _ev in extra_overrides.items():

fastvideo/layers/linear.py

@@ -219,15 +219,6 @@ class ReplicatedLinear(LinearBase):
                         (e.g. model.layers.0.qkv_proj)
     """
 
-    # Opt-in instrumentation: when ``enable_shape_tracking`` is set to True,
-    # ``forward`` records every unique ``(input_shape, output_shape)`` pair
-    # observed across all ``ReplicatedLinear`` instances, along with the
-    # subclass name that produced it. Used by upcoming QAT-aware backends
-    # to discover which GEMM shapes need quantized kernels. Defaults to
-    # False; default forward path is bit-identical to pre-slice behavior.
-    enable_shape_tracking = False
-    _shape_to_layer_types: dict[tuple[torch.Size, torch.Size], set[str]] = {}
-
     def __init__(
         self,
         input_size: int,
@@ -294,8 +285,6 @@ class ReplicatedLinear(LinearBase):
         bias = self.bias if not self.skip_bias_add else None
         assert self.quant_method is not None
         output = self.quant_method.apply(self, x, bias)
-        if self.enable_shape_tracking:
-            self._track_shape(x.shape, output.shape)
         output_bias = self.bias if self.skip_bias_add else None
         return output, output_bias
 
@@ -305,41 +294,6 @@ class ReplicatedLinear(LinearBase):
         s += f", bias={self.bias is not None}"
         return s
 
-    @classmethod
-    def get_shape_mapping(cls) -> dict:
-        """Get the mapping from (input_shape, output_shape) to layer types."""
-        return cls._shape_to_layer_types.copy()
-
-    @classmethod
-    def reset_shape_tracking(cls) -> None:
-        """Clear tracked shapes and layer type mappings."""
-        cls._shape_to_layer_types.clear()
-
-    def _track_shape(self, input_shape: torch.Size, output_shape: torch.Size) -> None:
-        shape_key = (input_shape, output_shape)
-        if shape_key not in self._shape_to_layer_types:
-            self._shape_to_layer_types[shape_key] = set()
-            logger.debug("Layer: %s | input shape: %s --> output shape: %s, Quant Method: %s", self.prefix, input_shape,
-                         output_shape, self.quant_method.__class__.__name__)
-        self._shape_to_layer_types[shape_key].add(self.__class__.__name__)
-
-    @classmethod
-    def print_shape_summary(cls) -> None:
-        """Log a summary of all unique shapes and their layer types."""
-        if not cls._shape_to_layer_types:
-            logger.info("No shapes have been processed yet.")
-            return
-
-        lines = [
-            "=== Matrix Multiplication Shape Summary ===",
-            f"Total unique shapes: {len(cls._shape_to_layer_types)}",
-        ]
-        for i, (shape_key, layer_types) in enumerate(cls._shape_to_layer_types.items(), 1):
-            input_shape, output_shape = shape_key
-            lines.append(f"{i}. Input: {input_shape} → Output: {output_shape}")
-            lines.append(f"   Layer types: {', '.join(sorted(layer_types))}")
-        logger.info("\n".join(lines))
-
 
 class ColumnParallelLinear(LinearBase):
     """Linear layer with column parallelism.

fastvideo/layers/mlp.py

@@ -5,7 +5,6 @@ import torch.nn as nn
 
 from fastvideo.layers.activation import get_act_fn
 from fastvideo.layers.linear import ReplicatedLinear
-from fastvideo.layers.quantization import QuantizationConfig
 
 
 class MLP(nn.Module):
@@ -22,27 +21,18 @@ class MLP(nn.Module):
         act_type: str = "gelu_pytorch_tanh",
         dtype: torch.dtype | None = None,
         prefix: str = "",
-        quant_config: QuantizationConfig | None = None,
     ):
         super().__init__()
         self.fc_in = ReplicatedLinear(
             input_dim,
             mlp_hidden_dim,  # For activation func like SiLU that need 2x width
             bias=bias,
-            params_dtype=dtype,
-            quant_config=quant_config,
-            prefix=f"{prefix}.fc_in",
-        )
+            params_dtype=dtype)
 
         self.act = get_act_fn(act_type)
         if output_dim is None:
             output_dim = input_dim
-        self.fc_out = ReplicatedLinear(mlp_hidden_dim,
-                                       output_dim,
-                                       bias=bias,
-                                       params_dtype=dtype,
-                                       quant_config=quant_config,
-                                       prefix=f"{prefix}.fc_out")
+        self.fc_out = ReplicatedLinear(mlp_hidden_dim, output_dim, bias=bias, params_dtype=dtype)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x, _ = self.fc_in(x)

fastvideo/layers/rotary_embedding.py

@@ -52,7 +52,6 @@ def apply_rotary_emb(
     freqs_cis: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
     use_real: bool = True,
     use_real_unbind_dim: int = -1,
-    sequence_dim: int = 2,
 ) -> torch.Tensor:
     """
     Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
@@ -61,23 +60,17 @@ def apply_rotary_emb(
     tensors contain rotary embeddings and are returned as real tensors.
     Args:
         x (`torch.Tensor`):
-            Query or key tensor to apply rotary embeddings. [B, H, S, D] if sequence_dim=2 else [B, S, H, D].
+            Query or key tensor to apply rotary embeddings. [B, H, S, D] xk (torch.Tensor): Key tensor to apply
         freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
-        sequence_dim: 1 = sequence at dim 1 (cos [1,S,1,D], x [B,S,H,D]); 2 = sequence at dim 2 (cos [1,1,S,D], x [B,H,S,D]).
     Returns:
         Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
     """
     if use_real:
         cos, sin = freqs_cis  # [S, D]
+        # Match Diffusers broadcasting (sequence_dim=2 case)
+        cos = cos[None, None, :, :]
+        sin = sin[None, None, :, :]
         cos, sin = cos.to(x.device), sin.to(x.device)
-        if sequence_dim == 2:
-            cos = cos[None, None, :, :]
-            sin = sin[None, None, :, :]
-        elif sequence_dim == 1:
-            cos = cos[None, :, None, :]
-            sin = sin[None, :, None, :]
-        else:
-            raise ValueError(f"sequence_dim must be 1 or 2, got {sequence_dim}")
 
         if use_real_unbind_dim == -1:
             # Used for flux, cogvideox, hunyuan-dit

fastvideo/models/audio/cosmos3_avae.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 sound tokenizer (AVAE) — decode path.
+
+The Cosmos3 ``sound_tokenizer`` is an AVAE (audio VAE). Its shipped diffusers
+checkpoint is **decoder-only** (``decoder.*``) in ``AutoencoderOobleck`` naming
+with SnakeBeta activations and ``weight_g``/``weight_v`` weight-norm — exactly
+FastVideo's native :class:`~fastvideo.models.vaes.oobleck.OobleckDecoder`
+(verified bit-exact vs the framework in ``test_cosmos3_avae_parity``). Text-to-
+video+sound (t2vs) only needs DECODE: the DiT generates the sound latent and
+this module decodes it to a waveform, so only the decoder is ported (the
+SpectrogramConvNeXt encoder is not exported in the checkpoint).
+
+Mirrors the framework ``AVAEModel.decode``: run the Oobleck decoder, then clamp
+to [-1, 1]. The VAE bottleneck's decode is the identity (the DiT already emits
+the post-bottleneck latent), so there is no bottleneck step here.
+"""
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass, field
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from fastvideo.logger import init_logger
+from fastvideo.models.vaes.oobleck import OobleckDecoder
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class Cosmos3SoundVAEArchConfig:
+    """Cosmos3 AVAE decoder constants (from ``sound_tokenizer/config.json``)."""
+
+    dec_dim: int = 320  # decoder base channels
+    vocoder_input_dim: int = 64  # latent channels in
+    dec_c_mults: list[int] = field(default_factory=lambda: [1, 2, 4, 8, 16])
+    dec_strides: list[int] = field(default_factory=lambda: [2, 4, 5, 6, 8])
+    audio_channels: int = 2  # stereo
+    sampling_rate: int = 48000
+
+    @property
+    def hop_size(self) -> int:
+        return int(np.prod(self.dec_strides))  # 1920
+
+
+class Cosmos3SoundVAE(nn.Module):
+    """Decoder-only Cosmos3 AVAE: latent ``[B, z, T]`` -> waveform ``[B, C, N]``."""
+
+    def __init__(self, arch: Cosmos3SoundVAEArchConfig | None = None) -> None:
+        super().__init__()
+        self.arch = arch or Cosmos3SoundVAEArchConfig()
+        self.decoder = OobleckDecoder(
+            channels=self.arch.dec_dim,
+            input_channels=self.arch.vocoder_input_dim,
+            audio_channels=self.arch.audio_channels,
+            # The framework builds decoder blocks from ``reversed(dec_strides)``
+            # (deepest first), so block strides are e.g. [8,6,5,4,2].
+            upsampling_ratios=list(reversed(self.arch.dec_strides)),
+            channel_multiples=list(self.arch.dec_c_mults),
+        )
+
+    @property
+    def sample_rate(self) -> int:
+        return self.arch.sampling_rate
+
+    @property
+    def audio_channels(self) -> int:
+        return self.arch.audio_channels
+
+    @property
+    def hop_size(self) -> int:
+        return self.arch.hop_size
+
+    def get_latent_num_samples(self, num_audio_samples: int) -> int:
+        """Latent length for a given audio length (``AVAEInterface``: ``N // hop``)."""
+        return int(num_audio_samples) // self.arch.hop_size
+
+    @torch.no_grad()
+    def decode(self, latent: torch.Tensor) -> torch.Tensor:
+        """Decode normalized latent ``[B, z, T]`` to waveform ``[B, C, N]`` in [-1, 1].
+
+        Matches ``AVAEModel.decode``: Oobleck decoder then clamp to [-1, 1] (the
+        VAE bottleneck decode is identity).
+        """
+        audio = self.decoder(latent)  # [B, C, N]
+        return audio.clamp(-1.0, 1.0)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_path: str,
+        *,
+        torch_dtype: torch.dtype | None = None,
+    ) -> "Cosmos3SoundVAE":
+        """Build + load the decoder from a ``sound_tokenizer`` directory.
+
+        Reads ``config.json`` (``dec_dim`` / ``vocoder_input_dim`` /
+        ``dec_c_mults`` / ``dec_strides`` / ``sampling_rate`` / ``stereo``) and
+        loads the ``decoder.*`` weights (the checkpoint is decoder-only).
+        """
+        from safetensors.torch import load_file
+
+        cfg_path = os.path.join(model_path, "config.json")
+        with open(cfg_path) as f:
+            cfg = json.load(f)
+        arch = Cosmos3SoundVAEArchConfig(
+            dec_dim=int(cfg["dec_dim"]),
+            vocoder_input_dim=int(cfg["vocoder_input_dim"]),
+            dec_c_mults=list(cfg["dec_c_mults"]),
+            dec_strides=list(cfg["dec_strides"]),
+            audio_channels=2 if cfg.get("stereo", True) else 1,
+            sampling_rate=int(cfg.get("sampling_rate", 48000)),
+        )
+        model = cls(arch)
+
+        weights_path = os.path.join(model_path, "diffusion_pytorch_model.safetensors")
+        state = load_file(weights_path)
+        # Decoder-only checkpoint: strip the ``decoder.`` prefix.
+        dec_state = {k[len("decoder."):]: v for k, v in state.items() if k.startswith("decoder.")}
+        model.decoder.load_state_dict(dec_state, strict=True)
+        logger.info("Loaded Cosmos3 sound AVAE decoder (%d params) from %s",
+                    sum(p.numel() for p in model.parameters()), model_path)
+
+        if torch_dtype is not None:
+            model = model.to(dtype=torch_dtype)
+        model.eval()
+        return model
+
+
+EntryClass = Cosmos3SoundVAE

fastvideo/models/dits/cosmos3.py

@@ -0,0 +1,991 @@
+# SPDX-License-Identifier: Apache-2.0
+"""FastVideo-native Cosmos3 omni DiT (``Cosmos3VFMTransformer``).
+
+Numerical-parity port of the official ``cosmos_framework`` ``Cosmos3VFMNetwork``
+(the ``Qwen3-VL-text`` MoT backbone + the VFM vision / action / sound heads).
+The module tree mirrors the published *diffusers* checkpoint layout so a
+converter can strict-load with a near-identity ``param_names_mapping``:
+
+* top level: ``embed_tokens`` / ``norm`` / ``norm_moe_gen`` / ``lm_head`` /
+  ``proj_in`` / ``proj_out`` / ``time_embedder.linear_{1,2}`` plus the dormant
+  ``action_*`` / ``audio_*`` heads (constructed for strict-load parity);
+* per layer ``layers.{i}``: dual-pathway ``self_attn`` with understanding
+  (``to_{q,k,v}`` / ``to_out``) and generation (``add_{q,k,v}_proj`` /
+  ``to_add_out``) projections + per-head QK-norms (``norm_{q,k}`` for und,
+  ``norm_added_{q,k}`` for gen), ``mlp`` (und) and ``mlp_moe_gen`` (gen) SwiGLU
+  blocks, and four RMSNorms.
+
+The forward replicates the framework contract for the video path:
+``patchify -> proj_in -> (+ 3d-rope additive latent pos emb) -> scatter-add
+timestep embeds onto noisy patches -> dual-pathway decoder (causal text +
+full vision two-way attention, GQA, per-head QK-norm before RoPE, unified
+3D-MRoPE) -> und/gen final norms -> proj_out on noisy vision patches ->
+unpatchify``.
+
+RoPE is applied with the framework's exact ``rotate_half`` math (contiguous
+split-half: ``q * cos + rotate_half(q) * sin``) rather than the interleaved
+``apply_rotary_emb`` helper, and attention uses plain SDPA so the whole module
+runs on CPU / float32 for parity testing. No diffusers / transformers
+model-class imports happen at runtime — the DiT is fully native.
+
+This module also keeps two pure, standalone math utilities used by the Tier-A
+scaffold parity tests: the unified-3D mRoPE position-ID generators
+(``compute_mrope_position_ids_text`` / ``compute_mrope_position_ids_vision``)
+and the batched ``patchify`` / ``unpatchify`` ``[B,C,T,H,W] <-> [B,N,p*p*C]``
+roundtrip helpers.
+"""
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from fastvideo.configs.models.dits.cosmos3 import Cosmos3VideoConfig
+from fastvideo.layers.layernorm import RMSNorm
+from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.layers.visual_embedding import timestep_embedding
+from fastvideo.models.dits.base import BaseDiT
+
+EntryClass = ["Cosmos3VFMTransformer"]
+
+
+# ===========================================================================
+# Standalone mRoPE position-ID generators (unified 3D mRoPE)
+# ===========================================================================
+def compute_mrope_position_ids_text(
+    num_tokens: int,
+    temporal_offset: int,
+) -> tuple[torch.Tensor, int]:
+    """Generate 3D mRoPE position IDs for text tokens.
+
+    Text tokens broadcast a single monotonically-increasing position-ID
+    sequence across all three (t, h, w) axes.
+    """
+    ids = torch.arange(num_tokens, dtype=torch.long) + temporal_offset
+    mrope_ids = ids.unsqueeze(0).expand(3, -1).contiguous()
+    return mrope_ids, temporal_offset + num_tokens
+
+
+def compute_mrope_position_ids_vision(
+    grid_t: int,
+    grid_h: int,
+    grid_w: int,
+    temporal_offset: int | float,
+    fps: float | None = None,
+    base_fps: float = 24.0,
+    temporal_compression_factor: int = 4,
+    base_temporal_compression_factor: int | None = None,
+    enable_fps_modulation: bool = True,
+    start_frame_offset: int = 0,
+) -> tuple[torch.Tensor, int]:
+    """Generate 3D mRoPE position IDs for vision tokens.
+
+    Builds a ``(t, h, w)`` position grid (Qwen3-VL style, spatial indices reset
+    per temporal segment) flattened in t-major order. Optionally modulates the
+    temporal axis by ``base_fps / tcf * (1 / (fps / tcf))`` so two clips at
+    different FPS retain wall-clock-aligned temporal positions.
+    """
+    fps_modulation = enable_fps_modulation and fps is not None
+
+    if fps_modulation:
+        assert fps is not None
+        tps = fps / temporal_compression_factor
+        effective_base_tcf = (base_temporal_compression_factor
+                              if base_temporal_compression_factor is not None else temporal_compression_factor)
+        base_tps = base_fps / effective_base_tcf
+        frame_indices = torch.arange(grid_t, dtype=torch.float32)
+        t_index = (((frame_indices + start_frame_offset) / tps * base_tps + temporal_offset).view(-1, 1).expand(
+            -1, grid_h * grid_w).flatten())
+    else:
+        t_index = (torch.arange(grid_t, dtype=torch.long).view(-1, 1).expand(-1, grid_h * grid_w).flatten() +
+                   int(temporal_offset) + start_frame_offset)
+
+    h_index = (torch.arange(grid_h, dtype=torch.long).view(1, -1, 1).expand(grid_t, -1, grid_w).flatten())
+    w_index = (torch.arange(grid_w, dtype=torch.long).view(1, 1, -1).expand(grid_t, grid_h, -1).flatten())
+
+    if fps_modulation:
+        mrope_ids = torch.stack([t_index, h_index.to(torch.float32), w_index.to(torch.float32)], dim=0)
+    else:
+        mrope_ids = torch.stack([t_index, h_index, w_index], dim=0)
+
+    next_offset = math.floor(mrope_ids.max().item()) + 1
+    return mrope_ids, next_offset
+
+
+# ===========================================================================
+# RoPE helpers (match cosmos_framework qwen3_vl rotate_half + apply_rotary_pos_emb)
+# ===========================================================================
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    """Contiguous split-half rotation, identical to qwen3_vl.rotate_half."""
+    x1 = x[..., :x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2:]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def _apply_rotary_pos_emb(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    unsqueeze_dim: int = 1,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Apply RoPE exactly as cosmos_framework qwen3_vl.apply_rotary_pos_emb.
+
+    ``q`` / ``k`` are ``[N, heads, head_dim]`` and ``cos`` / ``sin`` are
+    ``[N, head_dim]`` (per-token); ``unsqueeze_dim=1`` broadcasts cos/sin over
+    the head axis.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (_rotate_half(q) * sin)
+    k_embed = (k * cos) + (_rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+class Cosmos3TextRotaryEmbedding(nn.Module):
+    """Unified 3D-MRoPE rotary embedding (qwen3_vl Cosmos3 flavor).
+
+    Reproduces ``Qwen3VLTextRotaryEmbedding``: ``inv_freq`` from the default
+    RoPE init (``1 / theta**(arange(0, head_dim, 2) / head_dim)``), interleaved
+    MRoPE mixing of the T/H/W frequency bands per ``mrope_section``, and a final
+    ``cat([freqs, freqs])`` so cos/sin are ``[N, head_dim]``.
+    """
+
+    def __init__(
+        self,
+        head_dim: int,
+        rope_theta: float,
+        mrope_section: list[int],
+    ) -> None:
+        super().__init__()
+        self.head_dim = head_dim
+        self.rope_theta = float(rope_theta)
+        self.mrope_section = list(mrope_section)
+        self.register_buffer("inv_freq", self._compute_inv_freq(), persistent=False)
+        self.attention_scaling = 1.0  # default RoPE has unit attention scaling
+
+    def _compute_inv_freq(self, device: torch.device | None = None) -> torch.Tensor:
+        exponent = torch.arange(0, self.head_dim, 2, dtype=torch.int64, device=device).float() / self.head_dim
+        return 1.0 / (self.rope_theta**exponent)
+
+    def reset_inv_freq(self, device: torch.device) -> None:
+        """Recompute the non-persistent ``inv_freq`` on ``device`` after a
+        meta-device weight load (it is derived from ``rope_theta`` and is not
+        part of the checkpoint)."""
+        self.register_buffer("inv_freq", self._compute_inv_freq(device), persistent=False)
+
+    def _apply_interleaved_mrope(self, freqs: torch.Tensor) -> torch.Tensor:
+        """freqs: [3, N, head_dim//2] -> [N, head_dim//2] (interleaved T/H/W)."""
+        freqs_t = freqs[0].clone()
+        for dim, offset in enumerate((1, 2), start=1):  # H, W
+            length = self.mrope_section[dim] * 3
+            idx = slice(offset, length, 3)
+            freqs_t[..., idx] = freqs[dim, ..., idx]
+        return freqs_t
+
+    @torch.no_grad()
+    def forward(self, position_ids: torch.Tensor, device: torch.device,
+                dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]:
+        """position_ids: [N] (1D) or [3, N] (mrope) -> (cos, sin) each [N, head_dim]."""
+        if position_ids.ndim == 1:
+            position_ids = position_ids[None, :].expand(3, -1)  # [3, N]
+        position_ids = position_ids.to(device)
+        inv_freq = self.inv_freq.to(device).float()  # [head_dim//2]
+        inv_freq_expanded = inv_freq[None, :, None].expand(3, -1, 1)  # [3, head_dim//2, 1]
+        position_ids_expanded = position_ids[:, None, :].float()  # [3, 1, N]
+        freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)  # [3, N, head_dim//2]
+        freqs = self._apply_interleaved_mrope(freqs)  # [N, head_dim//2]
+        emb = torch.cat((freqs, freqs), dim=-1)  # [N, head_dim]
+        cos = (emb.cos() * self.attention_scaling).to(dtype)
+        sin = (emb.sin() * self.attention_scaling).to(dtype)
+        return cos, sin
+
+
+# ===========================================================================
+# 3D-RoPE additive latent position embedding (VideoRopePosition3DEmb)
+# ===========================================================================
+class Cosmos3VideoRopePosition3DEmb(nn.Module):
+    """Additive 3D-RoPE-style latent position embedding.
+
+    Mirrors ``cosmos_framework`` ``VideoRopePosition3DEmb`` (used when
+    ``position_embedding_type == "3d_rope"``). Produces an additive
+    ``[N_vision, head_dim]`` tensor concatenated over per-latent (t, h, w)
+    grids. ``enable_fps_modulation`` is supported for completeness but the
+    video path keeps it disabled by default.
+    """
+
+    def __init__(
+        self,
+        head_dim: int,
+        len_h: int,
+        len_w: int,
+        len_t: int,
+        base_fps: int = 24,
+        base_temporal_compression_factor: int = 4,
+        temporal_compression_factor: int = 4,
+        h_extrapolation_ratio: float = 1.0,
+        w_extrapolation_ratio: float = 1.0,
+        t_extrapolation_ratio: float = 1.0,
+        enable_fps_modulation: bool = False,
+    ) -> None:
+        super().__init__()
+        self.base_tps = base_fps / base_temporal_compression_factor
+        self.temporal_compression_factor = temporal_compression_factor
+        self.max_h = len_h
+        self.max_w = len_w
+        self.max_t = len_t
+        self.enable_fps_modulation = enable_fps_modulation
+
+        dim = head_dim
+        dim_h = dim // 6 * 2
+        dim_w = dim_h
+        dim_t = dim - 2 * dim_h
+        assert dim == dim_h + dim_w + dim_t, f"bad dim: {dim} != {dim_h} + {dim_w} + {dim_t}"
+        self._dim_h = dim_h
+        self._dim_t = dim_t
+
+        self.register_buffer(
+            "dim_spatial_range",
+            torch.arange(0, dim_h, 2)[:(dim_h // 2)].float() / dim_h,
+            persistent=True,
+        )
+        self.register_buffer(
+            "dim_temporal_range",
+            torch.arange(0, dim_t, 2)[:(dim_t // 2)].float() / dim_t,
+            persistent=True,
+        )
+        self.h_ntk_factor = h_extrapolation_ratio**(dim_h / (dim_h - 2))
+        self.w_ntk_factor = w_extrapolation_ratio**(dim_w / (dim_w - 2))
+        self.t_ntk_factor = t_extrapolation_ratio**(dim_t / (dim_t - 2))
+
+    def _generate(self, t: int, h: int, w: int, device: torch.device, fps: torch.Tensor | None,
+                  start_frame_offset: int) -> torch.Tensor:
+        tps = (fps / self.temporal_compression_factor) if fps is not None else None
+
+        h_theta = 10000.0 * self.h_ntk_factor
+        w_theta = 10000.0 * self.w_ntk_factor
+        t_theta = 10000.0 * self.t_ntk_factor
+
+        spatial = self.dim_spatial_range.to(device).float()
+        temporal = self.dim_temporal_range.to(device).float()
+        h_spatial_freqs = 1.0 / (h_theta**spatial)
+        w_spatial_freqs = 1.0 / (w_theta**spatial)
+        temporal_freqs = 1.0 / (t_theta**temporal)
+
+        max_needed = max(t, h, w)
+        seq = torch.arange(max_needed, device=device, dtype=torch.float32)
+        half_emb_h = torch.outer(seq[:h], h_spatial_freqs)  # [h, dim_h/2]
+        half_emb_w = torch.outer(seq[:w], w_spatial_freqs)  # [w, dim_w/2]
+        frame_indices = seq[:t]  # [t]
+
+        if self.enable_fps_modulation and tps is not None:
+            scaled_time = (frame_indices + start_frame_offset) / tps[:1] * self.base_tps
+            half_emb_t = torch.outer(scaled_time, temporal_freqs)
+        else:
+            half_emb_t = torch.outer(frame_indices, temporal_freqs)  # [t, dim_t/2]
+
+        emb_t = half_emb_t[:, None, None, :].expand(t, h, w, -1)
+        emb_h = half_emb_h[None, :, None, :].expand(t, h, w, -1)
+        emb_w = half_emb_w[None, None, :, :].expand(t, h, w, -1)
+        rope = torch.cat([emb_t, emb_h, emb_w] * 2, dim=-1)  # [t, h, w, head_dim]
+        return rope.reshape(t * h * w, -1).float()
+
+    def forward(
+        self,
+        token_shapes: list[tuple[int, int, int]],
+        fps: torch.Tensor | None = None,
+        start_frame_offset: int = 0,
+    ) -> torch.Tensor:
+        device = self.dim_spatial_range.device
+        out = []
+        for i, (t, h, w) in enumerate(token_shapes):
+            video_fps = fps[i:i + 1] if fps is not None else None
+            out.append(self._generate(t, h, w, device, video_fps, start_frame_offset))
+        return torch.cat(out, dim=0)  # [N_vision, head_dim]
+
+
+# ===========================================================================
+# Timestep embedder (DiT-style; matches modeling_utils.TimestepEmbedder)
+# ===========================================================================
+class Cosmos3TimestepEmbedder(nn.Module):
+    """Sinusoidal timestep -> MLP embedder. Checkpoint keys: ``linear_{1,2}``."""
+
+    def __init__(self, hidden_size: int, frequency_embedding_size: int = 256) -> None:
+        super().__init__()
+        self.frequency_embedding_size = frequency_embedding_size
+        self.hidden_size = hidden_size
+        self.linear_1 = ReplicatedLinear(frequency_embedding_size, hidden_size, bias=True)
+        self.act = nn.SiLU()
+        self.linear_2 = ReplicatedLinear(hidden_size, hidden_size, bias=True)
+
+    def forward(self, t: torch.Tensor) -> torch.Tensor:
+        t_freq = timestep_embedding(t, self.frequency_embedding_size)
+        # Sinusoidal runs in fp32 (timestep_scale makes inputs tiny); cast to the
+        # MLP weight dtype (bf16 at inference; no-op in the fp32 parity tests).
+        t_freq = t_freq.to(self.linear_1.weight.dtype)
+        h, _ = self.linear_1(t_freq)
+        h = self.act(h)
+        out, _ = self.linear_2(h)
+        return out
+
+
+# ===========================================================================
+# SwiGLU MLP (Qwen3-VL-text dense MLP)
+# ===========================================================================
+class Cosmos3MLP(nn.Module):
+    """SwiGLU MLP: ``down(act(gate(x)) * up(x))``. Keys: gate/up/down_proj."""
+
+    def __init__(self, hidden_size: int, intermediate_size: int) -> None:
+        super().__init__()
+        self.gate_proj = ReplicatedLinear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = ReplicatedLinear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = ReplicatedLinear(intermediate_size, hidden_size, bias=False)
+        self.act_fn = nn.SiLU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate, _ = self.gate_proj(x)
+        up, _ = self.up_proj(x)
+        out, _ = self.down_proj(self.act_fn(gate) * up)
+        return out
+
+
+# ===========================================================================
+# Dual-pathway packed two-way attention
+# ===========================================================================
+def _sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, is_causal: bool, scale: float) -> torch.Tensor:
+    """SDPA on ``[S, heads, head_dim]`` with GQA broadcast. Returns same layout."""
+    n_heads = q.shape[1]
+    n_kv = k.shape[1]
+    q_ = q.permute(1, 0, 2).unsqueeze(0)  # [1, heads, S, head_dim]
+    k_ = k.permute(1, 0, 2).unsqueeze(0)
+    v_ = v.permute(1, 0, 2).unsqueeze(0)
+    if n_kv != n_heads:
+        k_ = k_.repeat_interleave(n_heads // n_kv, dim=1)
+        v_ = v_.repeat_interleave(n_heads // n_kv, dim=1)
+    out = F.scaled_dot_product_attention(q_, k_, v_, is_causal=is_causal, scale=scale)
+    return out.squeeze(0).permute(1, 0, 2)  # [S, heads, head_dim]
+
+
+class Cosmos3DualAttention(nn.Module):
+    """Dual-pathway packed attention (understanding + generation).
+
+    Understanding (text) tokens use causal self-attention; generation (vision)
+    tokens use full attention where the gen query attends to ALL tokens
+    (und ++ gen). GQA with per-head QK-norm applied *before* RoPE. Checkpoint
+    keys: und ``to_{q,k,v}`` / ``to_out`` / ``norm_{q,k}``; gen
+    ``add_{q,k,v}_proj`` / ``to_add_out`` / ``norm_added_{q,k}``.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_attention_heads: int,
+        num_key_value_heads: int,
+        head_dim: int,
+        eps: float,
+        attention_bias: bool,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.scaling = head_dim**-0.5
+
+        q_dim = num_attention_heads * head_dim
+        kv_dim = num_key_value_heads * head_dim
+
+        # Understanding pathway
+        self.to_q = ReplicatedLinear(hidden_size, q_dim, bias=attention_bias)
+        self.to_k = ReplicatedLinear(hidden_size, kv_dim, bias=attention_bias)
+        self.to_v = ReplicatedLinear(hidden_size, kv_dim, bias=attention_bias)
+        self.to_out = ReplicatedLinear(q_dim, hidden_size, bias=attention_bias)
+        self.norm_q = RMSNorm(head_dim, eps=eps)
+        self.norm_k = RMSNorm(head_dim, eps=eps)
+
+        # Generation pathway
+        self.add_q_proj = ReplicatedLinear(hidden_size, q_dim, bias=attention_bias)
+        self.add_k_proj = ReplicatedLinear(hidden_size, kv_dim, bias=attention_bias)
+        self.add_v_proj = ReplicatedLinear(hidden_size, kv_dim, bias=attention_bias)
+        self.to_add_out = ReplicatedLinear(q_dim, hidden_size, bias=attention_bias)
+        self.norm_added_q = RMSNorm(head_dim, eps=eps)
+        self.norm_added_k = RMSNorm(head_dim, eps=eps)
+
+    def forward(
+        self,
+        und_seq: torch.Tensor,  # [N_und, hidden]
+        gen_seq: torch.Tensor,  # [N_gen, hidden]
+        cos_und: torch.Tensor,
+        sin_und: torch.Tensor,
+        cos_gen: torch.Tensor,
+        sin_gen: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        n_heads = self.num_attention_heads
+        n_kv = self.num_key_value_heads
+        d = self.head_dim
+
+        q_und, _ = self.to_q(und_seq)
+        k_und, _ = self.to_k(und_seq)
+        v_und, _ = self.to_v(und_seq)
+        q_gen, _ = self.add_q_proj(gen_seq)
+        k_gen, _ = self.add_k_proj(gen_seq)
+        v_gen, _ = self.add_v_proj(gen_seq)
+
+        q_und = q_und.view(-1, n_heads, d)
+        k_und = k_und.view(-1, n_kv, d)
+        v_und = v_und.view(-1, n_kv, d)
+        q_gen = q_gen.view(-1, n_heads, d)
+        k_gen = k_gen.view(-1, n_kv, d)
+        v_gen = v_gen.view(-1, n_kv, d)
+
+        # Per-head QK-norm BEFORE RoPE
+        q_und = self.norm_q(q_und)
+        k_und = self.norm_k(k_und)
+        q_gen = self.norm_added_q(q_gen)
+        k_gen = self.norm_added_k(k_gen)
+
+        # RoPE (heads-second layout; unsqueeze cos/sin over the head axis at dim=1)
+        q_und, k_und = _apply_rotary_pos_emb(q_und, k_und, cos_und, sin_und, unsqueeze_dim=1)
+        q_gen, k_gen = _apply_rotary_pos_emb(q_gen, k_gen, cos_gen, sin_gen, unsqueeze_dim=1)
+
+        # Causal self-attention over understanding (text) tokens
+        und_out = _sdpa(q_und, k_und, v_und, is_causal=True, scale=self.scaling)
+        und_out = und_out.reshape(und_out.shape[0], n_heads * d)
+
+        # Full attention: gen query attends to ALL tokens (und ++ gen)
+        all_k = torch.cat([k_und, k_gen], dim=0)
+        all_v = torch.cat([v_und, v_gen], dim=0)
+        gen_out = _sdpa(q_gen, all_k, all_v, is_causal=False, scale=self.scaling)
+        gen_out = gen_out.reshape(gen_out.shape[0], n_heads * d)
+
+        und_out, _ = self.to_out(und_out)
+        gen_out, _ = self.to_add_out(gen_out)
+        return und_out, gen_out
+
+
+# ===========================================================================
+# Dual-pathway decoder layer
+# ===========================================================================
+class Cosmos3DecoderLayer(nn.Module):
+    """MoT decoder layer: dual-pathway attention + dual SwiGLU MLP, 4 RMSNorms."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_attention_heads: int,
+        num_key_value_heads: int,
+        head_dim: int,
+        intermediate_size: int,
+        eps: float,
+        attention_bias: bool,
+    ) -> None:
+        super().__init__()
+        self.self_attn = Cosmos3DualAttention(
+            hidden_size=hidden_size,
+            num_attention_heads=num_attention_heads,
+            num_key_value_heads=num_key_value_heads,
+            head_dim=head_dim,
+            eps=eps,
+            attention_bias=attention_bias,
+        )
+        self.mlp = Cosmos3MLP(hidden_size, intermediate_size)
+        self.mlp_moe_gen = Cosmos3MLP(hidden_size, intermediate_size)
+        self.input_layernorm = RMSNorm(hidden_size, eps=eps)
+        self.input_layernorm_moe_gen = RMSNorm(hidden_size, eps=eps)
+        self.post_attention_layernorm = RMSNorm(hidden_size, eps=eps)
+        self.post_attention_layernorm_moe_gen = RMSNorm(hidden_size, eps=eps)
+
+    def forward(
+        self,
+        und_seq: torch.Tensor,
+        gen_seq: torch.Tensor,
+        cos_und: torch.Tensor,
+        sin_und: torch.Tensor,
+        cos_gen: torch.Tensor,
+        sin_gen: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # Pre-attention norm
+        und_norm = self.input_layernorm(und_seq)
+        gen_norm = self.input_layernorm_moe_gen(gen_seq)
+
+        und_attn, gen_attn = self.self_attn(und_norm, gen_norm, cos_und, sin_und, cos_gen, sin_gen)
+        und_res = und_seq + und_attn
+        gen_res = gen_seq + gen_attn
+
+        # Pre-MLP norm + dual SwiGLU
+        und_ln = self.post_attention_layernorm(und_res)
+        gen_ln = self.post_attention_layernorm_moe_gen(gen_res)
+        und_out = und_res + self.mlp(und_ln)
+        gen_out = gen_res + self.mlp_moe_gen(gen_ln)
+        return und_out, gen_out
+
+
+# ===========================================================================
+# Domain-aware linear (dormant action head; per-domain weight/bias embeddings)
+# ===========================================================================
+class _DomainAwareLinear(nn.Module):
+    """One weight/bias pair per embodiment domain (matches framework keys)."""
+
+    def __init__(self, input_size: int, output_size: int, num_domains: int) -> None:
+        super().__init__()
+        self.input_size = int(input_size)
+        self.output_size = int(output_size)
+        self.num_domains = int(num_domains)
+        self.fc = nn.Embedding(self.num_domains, self.output_size * self.input_size)
+        self.bias = nn.Embedding(self.num_domains, self.output_size)
+
+    def forward(self, x: torch.Tensor, domain_id: torch.Tensor) -> torch.Tensor:
+        if domain_id.ndim == 0:
+            domain_id = domain_id.unsqueeze(0)
+        domain_id = domain_id.to(device=x.device, dtype=torch.long).reshape(-1)
+        weight = self.fc(domain_id).view(domain_id.shape[0], self.input_size, self.output_size)
+        bias = self.bias(domain_id).view(domain_id.shape[0], self.output_size)
+        if x.ndim == 2:
+            return torch.bmm(x.unsqueeze(1), weight).squeeze(1) + bias
+        return torch.bmm(x, weight) + bias.unsqueeze(1)
+
+
+# ===========================================================================
+# Full DiT
+# ===========================================================================
+class Cosmos3VFMTransformer(BaseDiT):
+    """FastVideo-native Cosmos3 omni DiT.
+
+    Mirrors the published checkpoint's transformer key surface and the
+    ``cosmos_framework`` ``Cosmos3VFMNetwork`` forward math for the video path.
+    Runs on CPU / float32 (plain SDPA + native-math RoPE).
+    """
+
+    _fsdp_shard_conditions = Cosmos3VideoConfig().arch_config._fsdp_shard_conditions
+    _compile_conditions = Cosmos3VideoConfig().arch_config._compile_conditions
+    param_names_mapping = Cosmos3VideoConfig().arch_config.param_names_mapping
+    reverse_param_names_mapping: dict = {}
+
+    def __init__(self, config: Cosmos3VideoConfig, hf_config: dict[str, Any]) -> None:
+        super().__init__(config=config, hf_config=hf_config)
+        arch = config.arch_config
+
+        self.hidden_size = arch.hidden_size
+        self.num_attention_heads = arch.num_attention_heads
+        self.num_key_value_heads = arch.num_key_value_heads
+        self.head_dim = arch.head_dim
+        self.num_hidden_layers = arch.num_hidden_layers
+        self.intermediate_size = arch.intermediate_size
+        self.vocab_size = arch.vocab_size
+        self.rms_norm_eps = arch.rms_norm_eps
+        self.attention_bias = arch.attention_bias
+
+        # VAE / patch geometry
+        self.latent_patch_size = arch.latent_patch_size
+        self.latent_channel = arch.latent_channel
+        # Alias kept for the standalone patchify/unpatchify helpers + scaffold tests.
+        self.latent_channel_size = arch.latent_channel
+        self.patch_latent_dim = arch.patch_latent_dim
+        self.num_channels_latents = arch.latent_channel
+        self.timestep_scale = arch.timestep_scale
+        self.position_embedding_type = arch.position_embedding_type
+
+        # ---- Backbone ----
+        self.embed_tokens = nn.Embedding(self.vocab_size, self.hidden_size)
+        self.layers = nn.ModuleList([
+            Cosmos3DecoderLayer(
+                hidden_size=self.hidden_size,
+                num_attention_heads=self.num_attention_heads,
+                num_key_value_heads=self.num_key_value_heads,
+                head_dim=self.head_dim,
+                intermediate_size=self.intermediate_size,
+                eps=self.rms_norm_eps,
+                attention_bias=self.attention_bias,
+            ) for _ in range(self.num_hidden_layers)
+        ])
+        self.norm = RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
+        self.norm_moe_gen = RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
+        self.lm_head = nn.Linear(self.hidden_size, self.vocab_size, bias=False)
+
+        # Backbone RoPE (unified 3D-MRoPE)
+        self.rotary_emb = Cosmos3TextRotaryEmbedding(
+            head_dim=self.head_dim,
+            rope_theta=arch.rope_theta,
+            mrope_section=arch.mrope_section,
+        )
+
+        # ---- Vision head ----
+        self.proj_in = ReplicatedLinear(self.patch_latent_dim, self.hidden_size, bias=True)
+        self.proj_out = ReplicatedLinear(self.hidden_size, self.patch_latent_dim, bias=True)
+        self.time_embedder = Cosmos3TimestepEmbedder(self.hidden_size)
+
+        # Additive latent position embedding only for the legacy 3d_rope variant.
+        self.latent_pos_embed: Cosmos3VideoRopePosition3DEmb | None = None
+        if self.position_embedding_type == "3d_rope":
+            self.latent_pos_embed = Cosmos3VideoRopePosition3DEmb(
+                head_dim=self.hidden_size,
+                len_h=getattr(arch, "max_latent_h", 32),
+                len_w=getattr(arch, "max_latent_w", 32),
+                len_t=getattr(arch, "max_latent_t", 32),
+                base_fps=int(arch.base_fps),
+                base_temporal_compression_factor=arch.temporal_compression_factor,
+                temporal_compression_factor=arch.temporal_compression_factor,
+                enable_fps_modulation=arch.enable_fps_modulation,
+            )
+
+        # ---- Dormant action head (constructed for strict-load parity) ----
+        if getattr(arch, "action_gen", False):
+            self.action_dim = arch.action_dim
+            self.num_embodiment_domains = arch.num_embodiment_domains
+            self.action_proj_in = _DomainAwareLinear(self.action_dim, self.hidden_size, self.num_embodiment_domains)
+            self.action_proj_out = _DomainAwareLinear(self.hidden_size, self.action_dim, self.num_embodiment_domains)
+            self.action_modality_embed = nn.Parameter(torch.zeros(self.hidden_size))
+
+        # ---- Dormant audio / sound head (constructed for strict-load parity) ----
+        if getattr(arch, "sound_gen", False):
+            self.sound_dim = arch.sound_dim
+            self.audio_proj_in = ReplicatedLinear(self.sound_dim, self.hidden_size, bias=True)
+            self.audio_proj_out = ReplicatedLinear(self.hidden_size, self.sound_dim, bias=True)
+            self.audio_modality_embed = nn.Parameter(torch.zeros(self.hidden_size))
+
+        self.__post_init__()
+
+    # ------------------------------------------------------------------
+    # Standalone batched patchify / unpatchify (scaffold-test contract)
+    # ------------------------------------------------------------------
+    def _pad_to_patch_size(self, h: int, w: int) -> tuple[int, int, int, int]:
+        """Return ``(hp, wp, H_padded, W_padded)`` for ``latent_patch_size`` padding."""
+        p = self.latent_patch_size
+        h_padded = ((h + p - 1) // p) * p
+        w_padded = ((w + p - 1) // p) * p
+        return h_padded // p, w_padded // p, h_padded, w_padded
+
+    def patchify(self, latents: torch.Tensor, t: int, h: int, w: int) -> torch.Tensor:
+        """``[B, C, t, h, w] -> [B, t*hp*wp, p*p*C]`` (zero-pad h/w to patch multiples)."""
+        batch_size = latents.shape[0]
+        p = self.latent_patch_size
+        c = self.latent_channel_size
+        hp, wp, h_padded, w_padded = self._pad_to_patch_size(h, w)
+        if h_padded != h or w_padded != w:
+            latents = F.pad(latents, (0, w_padded - w, 0, h_padded - h))
+        x = latents.reshape(batch_size, c, t, hp, p, wp, p)
+        x = x.permute(0, 2, 3, 5, 4, 6, 1)
+        return x.reshape(batch_size, t * hp * wp, p * p * c)
+
+    def unpatchify(self, tokens: torch.Tensor, t: int, h: int, w: int) -> torch.Tensor:
+        """``[B, t*hp*wp, p*p*C] -> [B, C, t, h, w]`` (crop h/w padding)."""
+        batch_size = tokens.shape[0]
+        p = self.latent_patch_size
+        c = self.latent_channel_size
+        hp, wp, h_padded, w_padded = self._pad_to_patch_size(h, w)
+        x = tokens.reshape(batch_size, t, hp, wp, p, p, c)
+        x = x.permute(0, 6, 1, 2, 4, 3, 5)
+        x = x.reshape(batch_size, c, t, h_padded, w_padded)
+        if h_padded != h or w_padded != w:
+            x = x[:, :, :, :h, :w]
+        return x
+
+    # ------------------------------------------------------------------
+    # Framework-faithful packed patchify / unpatchify (einsum ordering)
+    # ------------------------------------------------------------------
+    def _patchify_and_pack(
+        self,
+        tokens_vision: list[torch.Tensor],
+        token_shapes: list[tuple[int, int, int]],
+    ) -> tuple[torch.Tensor, list[tuple[int, int, int]]]:
+        p = self.latent_patch_size
+        packed = []
+        original_shapes: list[tuple[int, int, int]] = []
+        for latent, (_t, _h, _w) in zip(tokens_vision, token_shapes):
+            latent = latent.squeeze(0) if latent.dim() == 5 else latent  # [C, T, H, W]
+            _, t_actual, h_actual, w_actual = latent.shape
+            original_shapes.append((t_actual, h_actual, w_actual))
+            h_padded = ((h_actual + p - 1) // p) * p
+            w_padded = ((w_actual + p - 1) // p) * p
+            if h_padded != h_actual or w_padded != w_actual:
+                padded = torch.zeros((self.latent_channel, t_actual, h_padded, w_padded),
+                                     device=latent.device,
+                                     dtype=latent.dtype)
+                padded[:, :, :h_actual, :w_actual] = latent
+                latent = padded
+            h_patches = h_padded // p
+            w_patches = w_padded // p
+            latent = latent.reshape(self.latent_channel, t_actual, h_patches, p, w_patches, p)
+            latent = torch.einsum("cthpwq->thwpqc", latent).reshape(-1, p * p * self.latent_channel)
+            packed.append(latent)
+        return torch.cat(packed, dim=0), original_shapes
+
+    def _unpatchify_and_unpack(
+        self,
+        packed_preds: torch.Tensor,
+        token_shapes: list[tuple[int, int, int]],
+        noisy_frame_indexes: list[torch.Tensor],
+        original_shapes: list[tuple[int, int, int]] | None,
+    ) -> list[torch.Tensor]:
+        p = self.latent_patch_size
+        outputs = []
+        start_idx = 0
+        for i, (t_c, _h_c, _w_c) in enumerate(token_shapes):
+            if original_shapes is not None:
+                _t_orig, h_orig, w_orig = original_shapes[i]
+                h_padded = ((h_orig + p - 1) // p) * p
+                w_padded = ((w_orig + p - 1) // p) * p
+                h_patches = h_padded // p
+                w_patches = w_padded // p
+            else:
+                h_orig, w_orig = _h_c * p, _w_c * p
+                h_patches, w_patches = _h_c, _w_c
+
+            nfi = noisy_frame_indexes[i]
+            t_n = len(nfi)
+            out = torch.zeros((self.latent_channel, t_c, h_orig, w_orig),
+                              device=packed_preds.device,
+                              dtype=packed_preds.dtype)
+            num_patches = t_n * h_patches * w_patches
+            if num_patches > 0:
+                end_idx = start_idx + num_patches
+                patches = packed_preds[start_idx:end_idx]
+                patches = patches.reshape(t_n, h_patches, w_patches, p, p, self.latent_channel)
+                latent = torch.einsum("thwpqc->cthpwq", patches)
+                latent = latent.reshape(self.latent_channel, t_n, h_patches * p, w_patches * p)
+                latent = latent[:, :, :h_orig, :w_orig]
+                out[:, nfi] = latent
+                start_idx = end_idx
+            outputs.append(out.unsqueeze(0))  # [1, C, T, H, W]
+        return outputs
+
+    def _scatter_timestep_embeds(
+        self,
+        packed_tokens: torch.Tensor,
+        packed_timestep_embeds: torch.Tensor,
+        noisy_frame_indexes: list[torch.Tensor],
+        token_shapes: list[tuple[int, ...]],
+    ) -> torch.Tensor:
+        """Add timestep embeds onto noisy patches (matches framework scatter_add)."""
+        start_noisy_index = 0
+        flat_idx = []
+        for noisy_i, shape_i in zip(noisy_frame_indexes, token_shapes):
+            spatial = math.prod(shape_i[1:])
+            spatial_idx = torch.arange(spatial, device=packed_tokens.device)
+            ni = (noisy_i * spatial).unsqueeze(-1).expand(-1, spatial)
+            ni = ni.clone() + spatial_idx + start_noisy_index
+            flat_idx.append(ni.flatten())
+            start_noisy_index += math.prod(shape_i)
+        flat = torch.cat(flat_idx, dim=0)
+        flat = flat.unsqueeze(-1).expand(-1, packed_tokens.shape[1])
+        return packed_tokens.scatter_add(0, flat, packed_timestep_embeds)
+
+    def materialize_non_persistent_buffers(self, device: torch.device, dtype: torch.dtype | None = None) -> None:
+        """Recompute non-persistent buffers lost by the meta-device FSDP load.
+
+        Only ``rotary_emb.inv_freq`` is non-persistent (derived from
+        ``rope_theta``, absent from the checkpoint). The FastVideo loader calls
+        this after ``load_model_from_full_model_state_dict``.
+        """
+        del dtype  # inv_freq stays float32 regardless of compute dtype
+        self.rotary_emb.reset_inv_freq(device)
+
+    # ------------------------------------------------------------------
+    # forward
+    # ------------------------------------------------------------------
+    def forward(  # type: ignore[override]
+        self,
+        text_ids: torch.Tensor,
+        text_indexes: torch.Tensor,
+        position_ids: torch.Tensor,
+        sequence_length: int,
+        split_lens: list[int],
+        attn_modes: list[str],
+        vision_tokens: list[torch.Tensor],
+        vision_token_shapes: list[tuple[int, int, int]],
+        vision_sequence_indexes: torch.Tensor,
+        vision_timesteps: torch.Tensor,
+        vision_mse_loss_indexes: torch.Tensor,
+        vision_noisy_frame_indexes: list[torch.Tensor],
+        fps_vision: torch.Tensor | None = None,
+        sound_tokens: list[torch.Tensor] | None = None,
+        sound_token_shapes: list[tuple[int, int, int]] | None = None,
+        sound_sequence_indexes: torch.Tensor | None = None,
+        sound_timesteps: torch.Tensor | None = None,
+        sound_mse_loss_indexes: torch.Tensor | None = None,
+        sound_noisy_frame_indexes: list[torch.Tensor] | None = None,
+        fps_sound: torch.Tensor | None = None,
+        action_tokens: list[torch.Tensor] | None = None,
+        action_token_shapes: list[tuple[int, ...]] | None = None,
+        action_sequence_indexes: torch.Tensor | None = None,
+        action_timesteps: torch.Tensor | None = None,
+        action_mse_loss_indexes: torch.Tensor | None = None,
+        action_noisy_frame_indexes: list[torch.Tensor] | None = None,
+        action_domain_id: list[torch.Tensor] | None = None,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        """Video-path forward returning ``{"last_hidden_state", "preds_vision"}``
+        (plus ``"preds_sound"`` when sound tokens are provided for t2vs).
+
+        Inputs mirror the framework ``PackedSequence`` fields for a single
+        text(causal)+vision(full) sample. This is the surface the parity test
+        and the converter exercise; a pipeline-facing wrapper composes these
+        from a ``PackedSequence`` builder.
+        """
+        device = self.embed_tokens.weight.device
+
+        # 1. Text embedding scattered into the packed sequence buffer
+        text_embed = self.embed_tokens(text_ids)  # [N_text, hidden]
+        target_dtype = text_embed.dtype
+        packed_sequence = text_embed.new_zeros((sequence_length, self.hidden_size))
+        packed_sequence[text_indexes] = text_embed
+
+        # 2. Vision: patchify -> proj_in -> (+ additive 3d-rope pos emb) -> timestep scatter
+        original_shapes: list[tuple[int, int, int]] | None = None
+        if vision_tokens:
+            packed_vision, original_shapes = self._patchify_and_pack(vision_tokens, vision_token_shapes)
+            # Vision latents arrive in fp32 (noise/VAE); cast to the model's
+            # compute dtype (bf16 at inference; no-op in the fp32 parity tests).
+            packed_vision, _ = self.proj_in(packed_vision.to(target_dtype))
+
+            if self.latent_pos_embed is not None:
+                pos_emb = self.latent_pos_embed(vision_token_shapes, fps=fps_vision).to(target_dtype)
+                packed_vision = packed_vision + pos_emb
+
+            if vision_mse_loss_indexes.numel() > 0:
+                timesteps = vision_timesteps * self.timestep_scale
+                ts_embeds = self.time_embedder(timesteps).to(target_dtype)
+                packed_vision = self._scatter_timestep_embeds(packed_vision, ts_embeds, vision_noisy_frame_indexes,
+                                                              vision_token_shapes)
+
+            packed_sequence[vision_sequence_indexes] = packed_vision
+
+        # 2b. Sound (t2vs): pack [C, T] -> audio_proj_in -> + modality embed ->
+        #     timestep scatter -> scatter into the (full) gen split. Mirrors the
+        #     framework ``_encode_sound``; sound shares the vision "full" split.
+        if sound_tokens:
+            packed_sound = torch.cat(
+                [s[:, :shp[0]].permute(1, 0) for s, shp in zip(sound_tokens, sound_token_shapes)],
+                dim=0,
+            )  # [total_sound_tokens, sound_dim]
+            packed_sound, _ = self.audio_proj_in(packed_sound.to(target_dtype))
+            packed_sound = packed_sound + self.audio_modality_embed
+            if sound_mse_loss_indexes is not None and sound_mse_loss_indexes.numel() > 0:
+                s_ts = sound_timesteps * self.timestep_scale
+                s_embeds = self.time_embedder(s_ts).to(target_dtype)
+                packed_sound = self._scatter_timestep_embeds(packed_sound, s_embeds, sound_noisy_frame_indexes,
+                                                             sound_token_shapes)
+            packed_sequence[sound_sequence_indexes] = packed_sound
+
+        # 2c. Action (DomainAwareLinear): pack [T, D] per sample with a per-token
+        #     domain id -> action_proj_in(domain) + action_modality_embed ->
+        #     timestep scatter -> scatter into the gen split (framework
+        #     ``_encode_action``). Action is domain-conditioned (per embodiment).
+        if action_tokens:
+            packed_action = torch.cat(
+                [a[:shp[0]] for a, shp in zip(action_tokens, action_token_shapes)], dim=0,
+            )  # [total_action_tokens, action_dim]
+            per_token_domain = torch.cat(
+                [d.reshape(1).expand(shp[0]) for d, shp in zip(action_domain_id, action_token_shapes)], dim=0,
+            )  # [total_action_tokens]
+            packed_action = self.action_proj_in(packed_action.to(target_dtype), per_token_domain)
+            packed_action = packed_action + self.action_modality_embed.view(1, -1)
+            if action_mse_loss_indexes is not None and action_mse_loss_indexes.numel() > 0:
+                a_ts = action_timesteps * self.timestep_scale
+                a_embeds = self.time_embedder(a_ts).to(target_dtype)
+                packed_action = self._scatter_timestep_embeds(packed_action, a_embeds, action_noisy_frame_indexes,
+                                                              action_token_shapes)
+            packed_sequence[action_sequence_indexes] = packed_action
+
+        # 3. RoPE for the full packed sequence (cos/sin per token); split und/gen
+        cos, sin = self.rotary_emb(position_ids, device=device, dtype=target_dtype)  # [N, head_dim]
+        und_idx, gen_idx = self._mode_indices(split_lens, attn_modes, device)
+        cos_und, sin_und = cos[und_idx], sin[und_idx]
+        cos_gen, sin_gen = cos[gen_idx], sin[gen_idx]
+
+        # 4. Dual-pathway decoder over (und = text causal, gen = vision full)
+        und_seq = packed_sequence[und_idx]
+        gen_seq = packed_sequence[gen_idx]
+        for layer in self.layers:
+            und_seq, gen_seq = layer(und_seq, gen_seq, cos_und, sin_und, cos_gen, sin_gen)
+
+        # 5. Final norms (und vs gen) and re-scatter into a joint buffer
+        und_seq = self.norm(und_seq)
+        gen_seq = self.norm_moe_gen(gen_seq)
+        last_hidden_state = packed_sequence.new_zeros((sequence_length, self.hidden_size))
+        last_hidden_state[und_idx] = und_seq
+        last_hidden_state[gen_idx] = gen_seq
+
+        # 6. Vision prediction: proj_out on noisy patches -> unpatchify
+        output: dict[str, Any] = {}
+        if vision_tokens and vision_mse_loss_indexes.numel() > 0:
+            preds, _ = self.proj_out(last_hidden_state[vision_mse_loss_indexes])
+            output["preds_vision"] = self._unpatchify_and_unpack(preds, vision_token_shapes,
+                                                                 vision_noisy_frame_indexes, original_shapes)
+
+        # 6b. Sound prediction (t2vs): audio_proj_out on noisy sound hidden
+        #     states -> unpack to per-sample [C, T] (framework ``_decode_sound``).
+        if sound_tokens and sound_mse_loss_indexes is not None and sound_mse_loss_indexes.numel() > 0:
+            preds_sound, _ = self.audio_proj_out(last_hidden_state[sound_mse_loss_indexes])
+            output["preds_sound"] = self._unpack_sound(preds_sound, sound_token_shapes, sound_noisy_frame_indexes)
+
+        # 6c. Action prediction: action_proj_out(per-token domain) on noisy hidden
+        #     states -> unpack to per-sample [T, D] (framework ``_decode_action``).
+        if action_tokens and action_mse_loss_indexes is not None and action_mse_loss_indexes.numel() > 0:
+            noisy_domain = torch.cat(
+                [d.reshape(1).expand(len(nfi)) for d, nfi in zip(action_domain_id, action_noisy_frame_indexes)], dim=0,
+            )
+            preds_action = self.action_proj_out(last_hidden_state[action_mse_loss_indexes], noisy_domain)
+            output["preds_action"] = self._unpack_action(preds_action, action_token_shapes, action_noisy_frame_indexes)
+
+        output["last_hidden_state"] = last_hidden_state
+        return output
+
+    def _unpack_action(
+        self,
+        packed_preds: torch.Tensor,
+        token_shapes: list[tuple[int, ...]],
+        noisy_frame_indexes: list[torch.Tensor],
+    ) -> list[torch.Tensor]:
+        """Scatter packed noisy action preds back to per-sample ``[T, D]`` (clean
+        frames left zero). Mirrors framework ``unpack_action``."""
+        outputs: list[torch.Tensor] = []
+        start_idx = 0
+        for shape, nfi in zip(token_shapes, noisy_frame_indexes):
+            t = shape[0]
+            out = torch.zeros((t, self.action_dim), device=packed_preds.device, dtype=packed_preds.dtype)
+            t_n = len(nfi)
+            if t_n > 0:
+                out[nfi] = packed_preds[start_idx:start_idx + t_n]
+                start_idx += t_n
+            outputs.append(out)
+        return outputs
+
+    def _unpack_sound(
+        self,
+        packed_preds: torch.Tensor,
+        token_shapes: list[tuple[int, int, int]],
+        noisy_frame_indexes: list[torch.Tensor],
+    ) -> list[torch.Tensor]:
+        """Scatter packed noisy sound preds back to per-sample ``[C, T]`` (clean
+        frames left zero). Mirrors framework ``unpack_sound_latents``."""
+        outputs: list[torch.Tensor] = []
+        start_idx = 0
+        for shape, nfi in zip(token_shapes, noisy_frame_indexes):
+            t = shape[0]
+            out = torch.zeros((self.sound_dim, t), device=packed_preds.device, dtype=packed_preds.dtype)
+            t_n = len(nfi)
+            if t_n > 0:
+                out[:, nfi] = packed_preds[start_idx:start_idx + t_n].T
+                start_idx += t_n
+            outputs.append(out)
+        return outputs
+
+    @staticmethod
+    def _mode_indices(split_lens: list[int], attn_modes: list[str],
+                      device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
+        """Token indexes for causal (und) and full (gen) splits, in pack order."""
+        und, gen = [], []
+        start = 0
+        for split_len, mode in zip(split_lens, attn_modes):
+            rng = range(start, start + split_len)
+            if mode == "causal":
+                und.extend(rng)
+            elif mode == "full":
+                gen.extend(rng)
+            start += split_len
+        return (torch.tensor(und, dtype=torch.long, device=device),
+                torch.tensor(gen, dtype=torch.long, device=device))

fastvideo/models/dits/wanvideo.py

@@ -26,7 +26,6 @@ from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed,
 from fastvideo.logger import init_logger
 from fastvideo.models.dits.base import BaseDiT
 from fastvideo.platforms import AttentionBackendEnum, current_platform
-from fastvideo.layers.quantization import QuantizationConfig
 
 from fastvideo.distributed.parallel_state import get_sp_world_size
 
@@ -107,9 +106,7 @@ class WanSelfAttention(nn.Module):
                  window_size=(-1, -1),
                  qk_norm=True,
                  eps=1e-6,
-                 parallel_attention=False,
-                 quant_config: QuantizationConfig | None = None,
-                 prefix: str = "") -> None:
+                 parallel_attention=False) -> None:
         assert dim % num_heads == 0
         super().__init__()
         self.dim = dim
@@ -121,10 +118,10 @@ class WanSelfAttention(nn.Module):
         self.parallel_attention = parallel_attention
 
         # layers
-        self.to_q = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_q")
-        self.to_k = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_k")
-        self.to_v = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_v")
-        self.to_out = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_out")
+        self.to_q = ReplicatedLinear(dim, dim)
+        self.to_k = ReplicatedLinear(dim, dim)
+        self.to_v = ReplicatedLinear(dim, dim)
+        self.to_out = ReplicatedLinear(dim, dim)
         self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
         self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
 
@@ -197,15 +194,13 @@ class WanI2VCrossAttention(WanSelfAttention):
         qk_norm=True,
         eps=1e-6,
         supported_attention_backends: tuple[AttentionBackendEnum, ...]
-        | None = None,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
+        | None = None
     ) -> None:
         super().__init__(dim, num_heads, window_size, qk_norm, eps,
-                         supported_attention_backends, quant_config=quant_config, prefix=prefix)
+                         supported_attention_backends)
 
-        self.add_k_proj = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.add_k_proj")
-        self.add_v_proj = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.add_v_proj")
+        self.add_k_proj = ReplicatedLinear(dim, dim)
+        self.add_v_proj = ReplicatedLinear(dim, dim)
         self.norm_added_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
         self.norm_added_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
 
@@ -251,17 +246,16 @@ class WanTransformerBlock(nn.Module):
                  added_kv_proj_dim: int | None = None,
                  supported_attention_backends: tuple[AttentionBackendEnum, ...]
                  | None = None,
-                 quant_config: QuantizationConfig | None = None,
                  prefix: str = ""):
         super().__init__()
 
         # 1. Self-attention
         self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
-        self.to_q = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_q")
-        self.to_k = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_k")
-        self.to_v = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_v")
+        self.to_q = ReplicatedLinear(dim, dim, bias=True)
+        self.to_k = ReplicatedLinear(dim, dim, bias=True)
+        self.to_v = ReplicatedLinear(dim, dim, bias=True)
 
-        self.to_out = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_out")
+        self.to_out = ReplicatedLinear(dim, dim, bias=True)
         self.attn1 = DistributedAttention(
             num_heads=num_heads,
             head_size=dim // num_heads,
@@ -296,17 +290,13 @@ class WanTransformerBlock(nn.Module):
             self.attn2 = WanI2VCrossAttention(dim,
                                               num_heads,
                                               qk_norm=qk_norm,
-                                              eps=eps,
-                                              quant_config=quant_config,
-                                              prefix=f"{prefix}.attn2")
+                                              eps=eps)
         else:
             # T2V
             self.attn2 = WanT2VCrossAttention(dim,
                                               num_heads,
                                               qk_norm=qk_norm,
-                                              eps=eps,
-                                              quant_config=quant_config,
-                                              prefix=f"{prefix}.attn2")
+                                              eps=eps)
         self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
             dim,
             norm_type="layer",
@@ -316,7 +306,7 @@ class WanTransformerBlock(nn.Module):
             compute_dtype=torch.float32)
 
         # 3. Feed-forward
-        self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config, prefix=f"{prefix}.ffn")
+        self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
         self.mlp_residual = ScaleResidual()
 
         self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
@@ -416,17 +406,17 @@ class WanTransformerBlock_VSA(nn.Module):
                  added_kv_proj_dim: int | None = None,
                  supported_attention_backends: tuple[AttentionBackendEnum, ...]
                  | None = None,
-                 quant_config: QuantizationConfig | None = None,
                  prefix: str = ""):
         super().__init__()
 
         # 1. Self-attention
         self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
-        self.to_q = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_q")
-        self.to_k = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_k")
-        self.to_v = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_v")
-        self.to_gate_compress = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_gate_compress")
-        self.to_out = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_out")
+        self.to_q = ReplicatedLinear(dim, dim, bias=True)
+        self.to_k = ReplicatedLinear(dim, dim, bias=True)
+        self.to_v = ReplicatedLinear(dim, dim, bias=True)
+        self.to_gate_compress = ReplicatedLinear(dim, dim, bias=True)
+
+        self.to_out = ReplicatedLinear(dim, dim, bias=True)
         self.attn1 = DistributedAttention_VSA(
             num_heads=num_heads,
             head_size=dim // num_heads,
@@ -461,17 +451,13 @@ class WanTransformerBlock_VSA(nn.Module):
             self.attn2 = WanI2VCrossAttention(dim,
                                               num_heads,
                                               qk_norm=qk_norm,
-                                              eps=eps,
-                                              quant_config=quant_config,
-                                              prefix=f"{prefix}.attn2")
+                                              eps=eps)
         else:
             # T2V
             self.attn2 = WanT2VCrossAttention(dim,
                                               num_heads,
                                               qk_norm=qk_norm,
-                                              eps=eps,
-                                              quant_config=quant_config,
-                                              prefix=f"{prefix}.attn2")
+                                              eps=eps)
         self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
             dim,
             norm_type="layer",
@@ -481,7 +467,7 @@ class WanTransformerBlock_VSA(nn.Module):
             compute_dtype=torch.float32)
 
         # 3. Feed-forward
-        self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config, prefix=f"{prefix}.ffn")
+        self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
         self.mlp_residual = ScaleResidual()
 
         self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
@@ -570,7 +556,6 @@ class WanTransformer3DModel(BaseDiT):
     def __init__(self, config: WanVideoConfig, hf_config: dict[str,
                                                                Any]) -> None:
         super().__init__(config=config, hf_config=hf_config)
-        self.quant_config = config.quant_config
 
         inner_dim = config.num_attention_heads * config.attention_head_dim
         self.hidden_size = config.hidden_size
@@ -609,7 +594,6 @@ class WanTransformer3DModel(BaseDiT):
                               config.eps,
                               config.added_kv_proj_dim,
                               self._supported_attention_backends,
-                              quant_config=config.quant_config,
                               prefix=f"{config.prefix}.blocks.{i}")
             for i in range(config.num_layers)
         ])

fastvideo/models/encoders/mistral3.py

@@ -1,48 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""HF-backed Mistral3 text encoder wrapper for full Flux2."""
-from typing import Any
-
-import torch
-from torch import nn
-
-from fastvideo.configs.models.encoders.mistral3 import Mistral3TextConfig
-from fastvideo.models.encoders.base import TextEncoder
-
-
-class Mistral3ForConditionalGeneration(TextEncoder):
-    """Loads the Transformers Mistral3 implementation for Flux2 text encoding."""
-
-    supports_hf_from_pretrained = True
-
-    def __init__(self, config: Mistral3TextConfig) -> None:
-        super().__init__(config)
-        self.config = config
-
-    @classmethod
-    def from_pretrained_local(
-        cls,
-        model_path: str,
-        model_config: Mistral3TextConfig,
-        dtype: torch.dtype,
-        device: torch.device,
-    ) -> nn.Module:
-        from transformers import AutoModelForImageTextToText
-
-        model = AutoModelForImageTextToText.from_pretrained(
-            model_path,
-            local_files_only=True,
-            torch_dtype=dtype,
-            low_cpu_mem_usage=True,
-        ).eval()
-        if device.type != "cpu":
-            model = model.to(device)
-        return model
-
-    def forward(self, *args: Any, **kwargs: Any) -> Any:
-        raise NotImplementedError(
-            "Mistral3ForConditionalGeneration is loaded through Transformers "
-            "via from_pretrained_local()."
-        )
-
-
-EntryClass = Mistral3ForConditionalGeneration

fastvideo/models/encoders/qwen3.py

@@ -1,461 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Ported from SGLang: python/sglang/multimodal_gen/runtime/models/encoders/qwen3.py
-"""Qwen3 causal LM text encoder for FastVideo diffusion models (e.g. Flux2 Klein)."""
-from collections.abc import Iterable
-from typing import Any
-
-import torch
-from torch import nn
-
-from fastvideo.attention import LocalAttention
-from fastvideo.configs.models.encoders import BaseEncoderOutput
-from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig
-from fastvideo.distributed import get_tp_world_size
-from fastvideo.layers.activation import SiluAndMul
-from fastvideo.layers.layernorm import RMSNorm
-from fastvideo.layers.linear import (
-    MergedColumnParallelLinear,
-    QKVParallelLinear,
-    RowParallelLinear,
-)
-from fastvideo.layers.quantization import QuantizationConfig
-from fastvideo.layers.rotary_embedding import get_rope
-from fastvideo.layers.vocab_parallel_embedding import VocabParallelEmbedding
-from fastvideo.models.encoders.base import TextEncoder
-from fastvideo.models.loader.weight_utils import (
-    default_weight_loader,
-    maybe_remap_kv_scale_name,
-)
-
-
-class Qwen3MLP(nn.Module):
-    """Qwen3 MLP with SwiGLU activation and tensor parallelism."""
-
-    def __init__(
-        self,
-        hidden_size: int,
-        intermediate_size: int,
-        hidden_act: str,
-        quant_config: QuantizationConfig | None = None,
-        bias: bool = False,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-        self.gate_up_proj = MergedColumnParallelLinear(
-            input_size=hidden_size,
-            output_sizes=[intermediate_size] * 2,
-            bias=bias,
-            quant_config=quant_config,
-            prefix=f"{prefix}.gate_up_proj",
-        )
-        self.down_proj = RowParallelLinear(
-            input_size=intermediate_size,
-            output_size=hidden_size,
-            bias=bias,
-            quant_config=quant_config,
-            prefix=f"{prefix}.down_proj",
-        )
-        if hidden_act != "silu":
-            raise ValueError(
-                f"Unsupported activation: {hidden_act}. Only silu is supported."
-            )
-        self.act_fn = SiluAndMul()
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x, _ = self.gate_up_proj(x)
-        x = self.act_fn(x)
-        x, _ = self.down_proj(x)
-        return x
-
-
-class Qwen3Attention(nn.Module):
-    """Qwen3 attention with QK-Norm and tensor parallelism.
-
-    Key difference from LLaMA: RMSNorm is applied to Q and K before attention.
-    """
-
-    def __init__(
-        self,
-        config: Qwen3TextConfig,
-        hidden_size: int,
-        num_heads: int,
-        num_kv_heads: int,
-        rope_theta: float = 1000000.0,
-        rope_scaling: dict[str, Any] | None = None,
-        max_position_embeddings: int = 40960,
-        quant_config: QuantizationConfig | None = None,
-        bias: bool = False,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-        self.hidden_size = hidden_size
-        tp_size = get_tp_world_size()
-        self.total_num_heads = num_heads
-        assert self.total_num_heads % tp_size == 0
-        self.num_heads = self.total_num_heads // tp_size
-        self.total_num_kv_heads = num_kv_heads
-        if self.total_num_kv_heads >= tp_size:
-            assert self.total_num_kv_heads % tp_size == 0
-        else:
-            assert tp_size % self.total_num_kv_heads == 0
-        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
-
-        self.head_dim = getattr(
-            config, "head_dim", self.hidden_size // self.total_num_heads
-        )
-        self.rotary_dim = self.head_dim
-        self.q_size = self.num_heads * self.head_dim
-        self.kv_size = self.num_kv_heads * self.head_dim
-        self.scaling = self.head_dim**-0.5
-        self.rope_theta = rope_theta
-        self.max_position_embeddings = max_position_embeddings
-
-        self.qkv_proj = QKVParallelLinear(
-            hidden_size=hidden_size,
-            head_size=self.head_dim,
-            total_num_heads=self.total_num_heads,
-            total_num_kv_heads=self.total_num_kv_heads,
-            bias=bias,
-            quant_config=quant_config,
-            prefix=f"{prefix}.qkv_proj",
-        )
-        self.o_proj = RowParallelLinear(
-            input_size=self.total_num_heads * self.head_dim,
-            output_size=hidden_size,
-            bias=bias,
-            quant_config=quant_config,
-            prefix=f"{prefix}.o_proj",
-        )
-
-        rms_norm_eps = getattr(config, "rms_norm_eps", 1e-6)
-        self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
-        self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
-
-        self.rotary_emb = get_rope(
-            self.head_dim,
-            rotary_dim=self.rotary_dim,
-            max_position=max_position_embeddings,
-            base=int(rope_theta),
-            rope_scaling=rope_scaling,
-            is_neox_style=True,
-        )
-
-        self.attn = LocalAttention(
-            self.num_heads,
-            self.head_dim,
-            self.num_kv_heads,
-            softmax_scale=self.scaling,
-            causal=True,
-            supported_attention_backends=config._supported_attention_backends,
-        )
-
-    def forward(
-        self,
-        positions: torch.Tensor,
-        hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        qkv, _ = self.qkv_proj(hidden_states)
-        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
-
-        batch_size, seq_len = q.shape[0], q.shape[1]
-        q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
-        k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
-        v = v.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
-
-        q = self.q_norm(q)
-        k = self.k_norm(k)
-
-        q = q.reshape(batch_size, seq_len, -1)
-        k = k.reshape(batch_size, seq_len, -1)
-
-        q, k = self.rotary_emb(positions, q, k)
-
-        q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
-        k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
-
-        if attention_mask is None:
-            attn_output = self.attn(q, k, v)
-        else:
-            q_sdpa = q.transpose(1, 2)
-            k_sdpa = k.transpose(1, 2)
-            v_sdpa = v.transpose(1, 2)
-            causal_mask = torch.ones(
-                seq_len,
-                seq_len,
-                device=q.device,
-                dtype=torch.bool,
-            ).tril()
-            key_mask = attention_mask.to(device=q.device, dtype=torch.bool)
-            attn_mask = causal_mask[None, None, :, :] & key_mask[:, None, None, :]
-            attn_output = torch.nn.functional.scaled_dot_product_attention(
-                q_sdpa,
-                k_sdpa,
-                v_sdpa,
-                attn_mask=attn_mask,
-                dropout_p=0.0,
-                is_causal=False,
-                scale=self.scaling,
-                enable_gqa=self.num_heads != self.num_kv_heads,
-            ).transpose(1, 2)
-
-        attn_output = attn_output.reshape(batch_size, seq_len, -1)
-
-        output, _ = self.o_proj(attn_output)
-        return output
-
-
-class Qwen3DecoderLayer(nn.Module):
-    """Qwen3 transformer decoder layer."""
-
-    def __init__(
-        self,
-        config: Qwen3TextConfig,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-        self.hidden_size = config.hidden_size
-        rope_theta = getattr(config, "rope_theta", 1000000.0)
-        rope_scaling = getattr(config, "rope_scaling", None)
-        max_position_embeddings = getattr(config, "max_position_embeddings", 40960)
-        attention_bias = getattr(config, "attention_bias", False)
-
-        self.self_attn = Qwen3Attention(
-            config=config,
-            hidden_size=self.hidden_size,
-            num_heads=config.num_attention_heads,
-            num_kv_heads=getattr(
-                config, "num_key_value_heads", config.num_attention_heads
-            ),
-            rope_theta=rope_theta,
-            rope_scaling=rope_scaling,
-            max_position_embeddings=max_position_embeddings,
-            quant_config=quant_config,
-            bias=attention_bias,
-            prefix=f"{prefix}.self_attn",
-        )
-        self.mlp = Qwen3MLP(
-            hidden_size=self.hidden_size,
-            intermediate_size=config.intermediate_size,
-            hidden_act=config.hidden_act,
-            quant_config=quant_config,
-            bias=getattr(config, "mlp_bias", False),
-            prefix=f"{prefix}.mlp",
-        )
-        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = RMSNorm(
-            config.hidden_size, eps=config.rms_norm_eps
-        )
-
-    def forward(
-        self,
-        positions: torch.Tensor,
-        hidden_states: torch.Tensor,
-        residual: torch.Tensor | None,
-        attention_mask: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        if residual is None:
-            residual = hidden_states
-            hidden_states = self.input_layernorm(hidden_states)
-        else:
-            hidden_states, residual = self.input_layernorm(hidden_states, residual)
-
-        hidden_states = self.self_attn(
-            positions=positions,
-            hidden_states=hidden_states,
-            attention_mask=attention_mask,
-        )
-
-        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
-        hidden_states = self.mlp(hidden_states)
-        return hidden_states, residual
-
-
-class Qwen3ForCausalLM(TextEncoder):
-    """Qwen3 causal language model for text encoding in diffusion models (e.g. Flux2 Klein).
-
-    Features:
-    - Tensor parallelism support
-    - FlashAttention/SDPA support via LocalAttention
-    - QK-Norm for better training stability
-    - output_hidden_states for Klein (layers 9, 18, 27)
-    """
-
-    supports_hf_from_pretrained = True
-
-    def __init__(self, config: Qwen3TextConfig) -> None:
-        super().__init__(config)
-
-        self.config = config
-        self.quant_config = getattr(config, "quant_config", None)
-
-        if getattr(config, "lora_config", None) is not None:
-            max_loras = getattr(config.lora_config, "max_loras", 1)
-            lora_vocab_size = getattr(config.lora_config, "lora_extra_vocab_size", 1)
-            lora_vocab = lora_vocab_size * max_loras
-        else:
-            lora_vocab = 0
-        self.vocab_size = config.vocab_size + lora_vocab
-        self.org_vocab_size = config.vocab_size
-
-        self.embed_tokens = VocabParallelEmbedding(
-            self.vocab_size,
-            config.hidden_size,
-            org_num_embeddings=config.vocab_size,
-            quant_config=self.quant_config,
-        )
-
-        self.layers = nn.ModuleList(
-            [
-                Qwen3DecoderLayer(
-                    config=config,
-                    quant_config=self.quant_config,
-                    prefix=f"{config.prefix}.layers.{i}",
-                )
-                for i in range(config.num_hidden_layers)
-            ]
-        )
-
-        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-
-    @classmethod
-    def from_pretrained_local(
-        cls,
-        model_path: str,
-        model_config: Qwen3TextConfig,
-        dtype: torch.dtype,
-        device: torch.device,
-    ) -> nn.Module:
-        from transformers import AutoModelForCausalLM
-
-        if device.type == "cpu" and torch.cuda.is_available():
-            from fastvideo.distributed import get_local_torch_device
-
-            device = get_local_torch_device()
-
-        return AutoModelForCausalLM.from_pretrained(
-            model_path,
-            local_files_only=True,
-            torch_dtype=dtype,
-            low_cpu_mem_usage=True,
-        ).eval().to(device)
-
-    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
-        return self.embed_tokens(input_ids)
-
-    def forward(
-        self,
-        input_ids: torch.Tensor | None = None,
-        position_ids: torch.Tensor | None = None,
-        attention_mask: torch.Tensor | None = None,
-        inputs_embeds: torch.Tensor | None = None,
-        output_hidden_states: bool | None = None,
-        **kwargs: Any,
-    ) -> BaseEncoderOutput:
-        output_hidden_states = (
-            output_hidden_states
-            if output_hidden_states is not None
-            else self.config.output_hidden_states
-        )
-
-        if inputs_embeds is not None:
-            hidden_states = inputs_embeds
-        else:
-            assert input_ids is not None
-            hidden_states = self.get_input_embeddings(input_ids)
-
-        residual = None
-
-        if position_ids is None:
-            position_ids = torch.arange(
-                0, hidden_states.shape[1], device=hidden_states.device
-            ).unsqueeze(0)
-
-        all_hidden_states: tuple[Any, ...] | None = (
-            () if output_hidden_states else None
-        )
-
-        for layer in self.layers:
-            if all_hidden_states is not None:
-                all_hidden_states += (
-                    (hidden_states,)
-                    if residual is None
-                    else (hidden_states + residual,)
-                )
-            hidden_states, residual = layer(
-                position_ids,
-                hidden_states,
-                residual,
-                attention_mask=attention_mask,
-            )
-
-        hidden_states, _ = self.norm(hidden_states, residual)
-
-        if all_hidden_states is not None:
-            all_hidden_states += (hidden_states,)
-
-        return BaseEncoderOutput(
-            last_hidden_state=hidden_states,
-            hidden_states=all_hidden_states,
-        )
-
-    def load_weights(
-        self, weights: Iterable[tuple[str, torch.Tensor]]
-    ) -> set[str]:
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-
-        for name, loaded_weight in weights:
-            if name.startswith("model."):
-                name = name[6:]
-
-            if "rotary_emb.inv_freq" in name:
-                continue
-            if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
-                continue
-
-            if "scale" in name:
-                kv_scale_name: str | None = maybe_remap_kv_scale_name(
-                    name, params_dict
-                )
-                if kv_scale_name is None:
-                    continue
-                name = kv_scale_name
-
-            for (
-                param_name,
-                weight_name,
-                shard_id,
-            ) in self.config.arch_config.stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-
-                if name not in params_dict:
-                    continue
-
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-
-                if name not in params_dict:
-                    continue
-
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-
-            loaded_params.add(name)
-
-        return loaded_params
-
-
-EntryClass = Qwen3ForCausalLM

fastvideo/models/loader/component_loader.py

@@ -13,7 +13,7 @@ from typing import cast
 import torch
 import torch.distributed as dist
 import torch.nn as nn
-from safetensors.torch import load_file as safetensors_load_file, safe_open
+from safetensors.torch import load_file as safetensors_load_file
 from torch.distributed import init_device_mesh
 from transformers import AutoImageProcessor, AutoProcessor, AutoTokenizer
 from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
@@ -92,6 +92,8 @@ class ComponentLoader(ABC):
             "tokenizer": (TokenizerLoader, "transformers"),
             "tokenizer_2": (TokenizerLoader, "transformers"),
             "tokenizer_3": (TokenizerLoader, "transformers"),
+            # Cosmos3's model_index names its Qwen2 tokenizer "text_tokenizer".
+            "text_tokenizer": (TokenizerLoader, "transformers"),
             "image_processor": (ImageProcessorLoader, "transformers"),
             "feature_extractor": (ImageProcessorLoader, "transformers"),
             "image_encoder": (ImageEncoderLoader, "transformers"),
@@ -573,53 +575,13 @@ class TokenizerLoader(ComponentLoader):
                 # If parsing fails, fall through to AutoTokenizer below.
                 pass
 
-        # Only Flux2 full's Mistral3 (require_processor=True) must load via
-        # AutoProcessor. Gate the processor_config.json shortcut on that flag so
-        # existing encoders (e.g. HunyuanVideo 1.5 / Qwen2.5-VL) stay on the
-        # historical AutoTokenizer path below even if their tokenizer dir happens
-        # to ship a processor_config.json.
-        require_processor = False
-        if hasattr(fastvideo_args.pipeline_config, "text_encoder_configs"):
-            try:
-                require_processor = any(
-                    getattr(getattr(cfg, "arch_config", None), "require_processor", False)
-                    for cfg in fastvideo_args.pipeline_config.text_encoder_configs)
-            except Exception:
-                require_processor = False
-
-        if require_processor and os.path.exists(os.path.join(resolved_model_path, "processor_config.json")):
-            processor = AutoProcessor.from_pretrained(
-                resolved_model_path,
-                local_files_only=os.path.isdir(resolved_model_path),
-                trust_remote_code=fastvideo_args.trust_remote_code,
-            )
-            logger.info(
-                "Loaded tokenizer/processor from %s: %s",
-                resolved_model_path,
-                processor.__class__.__name__,
-            )
-            return processor
-
-        try:
-            tokenizer = AutoTokenizer.from_pretrained(
-                resolved_model_path,  # "<path to model>/tokenizer"
-                # in v0, this was same string as encoder_name "ClipTextModel"
-                # TODO(will): pass these tokenizer kwargs from inference args? Maybe
-                # other method of config?
-                local_files_only=os.path.isdir(resolved_model_path),
-            )
-        except (OSError, ValueError):
-            tokenizer = AutoProcessor.from_pretrained(
-                resolved_model_path,
-                local_files_only=os.path.isdir(resolved_model_path),
-                trust_remote_code=fastvideo_args.trust_remote_code,
-            )
-            logger.info(
-                "Loaded tokenizer/processor from %s: %s",
-                resolved_model_path,
-                tokenizer.__class__.__name__,
-            )
-            return tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(
+            resolved_model_path,  # "<path to model>/tokenizer"
+            # in v0, this was same string as encoder_name "ClipTextModel"
+            # TODO(will): pass these tokenizer kwargs from inference args? Maybe
+            # other method of config?
+            local_files_only=os.path.isdir(resolved_model_path),
+        )
         padding_side = None
         if hasattr(fastvideo_args.pipeline_config, "text_encoder_configs"):
             try:
@@ -904,18 +866,6 @@ class VocoderLoader(ComponentLoader):
         return vocoder.eval()
 
 
-def _collect_safetensors_keys(safetensors_list: list) -> set:
-    """Collect all weight keys from safetensors files."""
-    all_keys: set[str] = set()
-    for path in safetensors_list:
-        try:
-            with safe_open(path, framework="pt") as f:
-                all_keys.update(f.keys())
-        except Exception as e:
-            logger.warning("Could not read keys from %s: %s", path, e)
-    return all_keys
-
-
 class TransformerLoader(ComponentLoader):
     """Loader for transformer."""
 
@@ -951,12 +901,6 @@ class TransformerLoader(ComponentLoader):
         if not safetensors_list:
             raise ValueError(f"No safetensors files found in {model_path}")
 
-        # arch_config can infer architecture from weight keys (e.g. Flux2 layer counts)
-        update_fn = getattr(dit_config.arch_config, "update_from_weight_keys", None)
-        if callable(update_fn):
-            weight_keys = _collect_safetensors_keys(safetensors_list)
-            update_fn(weight_keys)
-
         # Check if we should use custom initialization weights
         custom_weights_path = getattr(
             fastvideo_args, "init_weights_from_safetensors", None
@@ -1067,7 +1011,19 @@ class SchedulerLoader(ComponentLoader):
 
         scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name)
 
-        scheduler = scheduler_cls(**config)
+        # Diffusers checkpoints can carry newer scheduler config keys than the
+        # vendored scheduler accepts (e.g. shift_terminal / sigma_min / sigma_max
+        # from a newer diffusers release). Filter to the class's __init__ params,
+        # mirroring diffusers' ``from_config``, so loading is robust to schema
+        # drift instead of crashing on an unexpected kwarg.
+        import inspect
+        valid_params = set(inspect.signature(scheduler_cls.__init__).parameters)
+        filtered_config = {k: v for k, v in config.items() if k in valid_params}
+        dropped = sorted(set(config) - set(filtered_config))
+        if dropped:
+            logger.warning("Scheduler %s: dropping unsupported config keys %s", class_name, dropped)
+
+        scheduler = scheduler_cls(**filtered_config)
         if fastvideo_args.pipeline_config.flow_shift is not None:
             scheduler.set_shift(fastvideo_args.pipeline_config.flow_shift)
         return scheduler

fastvideo/models/loader/fsdp_load.py

@@ -332,7 +332,6 @@ def load_model_from_full_model_state_dict(
         NotImplementedError: If got FSDP with more than 1D.
     """
     meta_sd = model.state_dict()
-    named_parameters = dict(model.named_parameters())
     sharded_sd = {}
     custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict(
         full_sd_iterator, param_names_mapping)  # type: ignore
@@ -364,25 +363,8 @@ def load_model_from_full_model_state_dict(
             )
         if not hasattr(meta_sharded_param, "device_mesh"):
             full_tensor = full_tensor.to(device=device, dtype=param_dtype)
-            target_param = named_parameters.get(target_param_name)
-            weight_loader = getattr(target_param, "weight_loader", None)
-            # Gated on a shape mismatch: only fused/stacked params with a custom
-            # weight_loader (e.g. Qwen3's merged QKV/gate-up) take this path.
-            # Existing models whose unsharded params match the checkpoint shape
-            # fall through to the original `sharded_tensor = full_tensor` below.
-            if target_param is not None and callable(weight_loader) and tuple(target_param.shape) != tuple(
-                    full_tensor.shape):
-                loaded_param = nn.Parameter(torch.empty(tuple(target_param.shape),
-                                                        device=device,
-                                                        dtype=param_dtype),
-                                            requires_grad=False)
-                for attr_name, attr_value in vars(target_param).items():
-                    setattr(loaded_param, attr_name, attr_value)
-                weight_loader(loaded_param, full_tensor)
-                sharded_tensor = loaded_param.data
-            else:
-                # In cases where parts of the model aren't sharded, some parameters will be plain tensors.
-                sharded_tensor = full_tensor
+            # In cases where parts of the model aren't sharded, some parameters will be plain tensors
+            sharded_tensor = full_tensor
         else:
             full_tensor = full_tensor.to(device=device, dtype=param_dtype)
             sharded_tensor = distribute_tensor(

fastvideo/models/registry.py

@@ -35,6 +35,9 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
     "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
     "CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel"),
     "Cosmos25Transformer3DModel": ("dits", "cosmos2_5", "Cosmos25Transformer3DModel"),
+    # Cosmos3-Nano's checkpoint model_index names the DiT "Cosmos3OmniTransformer";
+    # map that HF class name to FastVideo's native Cosmos3VFMTransformer.
+    "Cosmos3OmniTransformer": ("dits", "cosmos3", "Cosmos3VFMTransformer"),
     "LongCatVideoTransformer3DModel": ("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"),  # Wrapper (Phase 1)
     "LongCatTransformer3DModel": ("dits", "longcat", "LongCatTransformer3DModel"),  # Native (Phase 2)
     "LTX2Transformer3DModel": ("dits", "ltx2", "LTX2Transformer3DModel"),
@@ -42,7 +45,6 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
     "LingBotWorldTransformer3DModel": ("dits", "lingbotworld", "LingBotWorldTransformer3DModel"),
     "Gen3CTransformer3DModel": ("dits", "gen3c", "Gen3CTransformer3DModel"),
     "Kandinsky5Transformer3DModel": ("dits", "kandinsky5", "Kandinsky5Transformer3DModel"),
-    "Flux2Transformer2DModel": ("dits", "flux_2", "Flux2Transformer2DModel"),
 }
 
 _IMAGE_TO_VIDEO_DIT_MODELS = {
@@ -70,9 +72,6 @@ _TEXT_ENCODER_MODELS = {
     "Qwen2_5_VLForConditionalGeneration":
     ("encoders", "reason1", "Reason1TextEncoder"),
     "LTX2GemmaTextEncoderModel": ("encoders", "gemma", "LTX2GemmaTextEncoderModel"),
-    "Qwen3ForCausalLM": ("encoders", "qwen3", "Qwen3ForCausalLM"),
-    "Mistral3ForConditionalGeneration":
-    ("encoders", "mistral3", "Mistral3ForConditionalGeneration"),
 }
 
 _IMAGE_ENCODER_MODELS: dict[str, tuple] = {
@@ -94,7 +93,6 @@ _VAE_MODELS = {
     ("vaes", "gen3c_tokenizer_vae", "AutoencoderKLGen3CTokenizer"),
     "AutoencoderKLStepvideo": ("vaes", "stepvideovae", "AutoencoderKLStepvideo"),
     "CausalVideoAutoencoder": ("vaes", "ltx2vae", "LTX2CausalVideoAutoencoder"),
-    "AutoencoderKLFlux2": ("vaes", "flux2vae", "AutoencoderKLFlux2"),
     # `stable-audio-open-1.0/vae/config.json` ships `_class_name="AutoencoderOobleck"`
     # (Diffusers' name); FastVideo's class is `OobleckVAE`.
     "AutoencoderOobleck": ("vaes", "oobleck", "OobleckVAE"),

fastvideo/models/vaes/flux2_components.py

@@ -1,532 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Copyright 2025 The HuggingFace Team. All rights reserved.
-# Adapted from: huggingface/diffusers `Encoder`/`Decoder` VAE components
-# at the installed 0.36.0 source surface used by Flux2 Klein.
-
-from dataclasses import dataclass
-from typing import Optional, Tuple
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from fastvideo.models.vaes.common import DiagonalGaussianDistribution
-
-
-@dataclass
-class AutoencoderKLOutput:
-    latent_dist: DiagonalGaussianDistribution
-
-    def __getitem__(self, idx: int):
-        return (self.latent_dist,)[idx]
-
-    def __getattr__(self, name: str):
-        # Existing local Flux2 parity tests used `vae.encode(x).mean` while
-        # diffusers-style callers use `vae.encode(x).latent_dist.mean`.
-        return getattr(self.latent_dist, name)
-
-
-@dataclass
-class DecoderOutput:
-    sample: torch.Tensor
-    commit_loss: Optional[torch.Tensor] = None
-
-    def __getitem__(self, idx: int):
-        return (self.sample, self.commit_loss)[idx]
-
-
-def get_activation(act_fn: str) -> nn.Module:
-    if act_fn in ("swish", "silu"):
-        return nn.SiLU()
-    if act_fn == "mish":
-        return nn.Mish()
-    if act_fn == "gelu":
-        return nn.GELU()
-    if act_fn == "relu":
-        return nn.ReLU()
-    raise ValueError(f"Unsupported activation function: {act_fn}")
-
-
-class AttnProcessor:
-    def __call__(self, attn: "Attention", hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
-        return attn._forward(hidden_states, temb=temb)
-
-
-class AttnAddedKVProcessor(AttnProcessor):
-    pass
-
-
-ADDED_KV_ATTENTION_PROCESSORS = frozenset({AttnAddedKVProcessor})
-CROSS_ATTENTION_PROCESSORS = frozenset({AttnProcessor})
-
-
-class Attention(nn.Module):
-    def __init__(
-        self,
-        query_dim: int,
-        heads: int = 8,
-        dim_head: int = 64,
-        dropout: float = 0.0,
-        bias: bool = False,
-        upcast_softmax: bool = False,
-        norm_num_groups: Optional[int] = None,
-        spatial_norm_dim: Optional[int] = None,
-        out_bias: bool = True,
-        eps: float = 1e-5,
-        rescale_output_factor: float = 1.0,
-        residual_connection: bool = False,
-        _from_deprecated_attn_block: bool = False,
-        **_: object,
-    ):
-        super().__init__()
-        self.inner_dim = dim_head * heads
-        self.query_dim = query_dim
-        self.heads = heads
-        self.dim_head = dim_head
-        self.scale = dim_head**-0.5
-        self.upcast_softmax = upcast_softmax
-        self.rescale_output_factor = rescale_output_factor
-        self.residual_connection = residual_connection
-        self._from_deprecated_attn_block = _from_deprecated_attn_block
-        self.spatial_norm = None
-        if spatial_norm_dim is not None:
-            raise ValueError("Flux2 VAE does not use spatial attention norm in this port")
-        self.group_norm = (
-            nn.GroupNorm(num_channels=query_dim, num_groups=norm_num_groups, eps=eps, affine=True)
-            if norm_num_groups is not None
-            else None
-        )
-        self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias)
-        self.to_k = nn.Linear(query_dim, self.inner_dim, bias=bias)
-        self.to_v = nn.Linear(query_dim, self.inner_dim, bias=bias)
-        self.to_out = nn.ModuleList([nn.Linear(self.inner_dim, query_dim, bias=out_bias), nn.Dropout(dropout)])
-        self.processor = AttnProcessor()
-
-    def set_processor(self, processor: AttnProcessor) -> None:
-        self.processor = processor
-
-    def get_processor(self) -> AttnProcessor:
-        return self.processor
-
-    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
-        return self.processor(self, hidden_states, temb=temb)
-
-    def _forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
-        residual = hidden_states
-        batch_size, channel, height, width = hidden_states.shape
-        if self.group_norm is not None:
-            hidden_states = self.group_norm(hidden_states)
-        hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
-
-        query = self.to_q(hidden_states)
-        key = self.to_k(hidden_states)
-        value = self.to_v(hidden_states)
-
-        query = query.view(batch_size, -1, self.heads, self.dim_head).transpose(1, 2)
-        key = key.view(batch_size, -1, self.heads, self.dim_head).transpose(1, 2)
-        value = value.view(batch_size, -1, self.heads, self.dim_head).transpose(1, 2)
-
-        if self.upcast_softmax:
-            query = query.float()
-            key = key.float()
-            value = value.float()
-        hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, scale=self.scale)
-        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, height * width, self.inner_dim)
-        hidden_states = hidden_states.to(self.to_out[0].weight.dtype)
-        hidden_states = self.to_out[0](hidden_states)
-        hidden_states = self.to_out[1](hidden_states)
-        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, channel, height, width)
-        if self.residual_connection:
-            hidden_states = hidden_states + residual
-        return hidden_states / self.rescale_output_factor
-
-
-class Downsample2D(nn.Module):
-    def __init__(self, channels: int, use_conv: bool = False, out_channels: Optional[int] = None, padding: int = 1, name: str = "conv"):
-        super().__init__()
-        self.channels = channels
-        self.out_channels = out_channels or channels
-        self.use_conv = use_conv
-        self.padding = padding
-        self.name = name
-        if use_conv:
-            conv = nn.Conv2d(self.channels, self.out_channels, kernel_size=3, stride=2, padding=padding)
-        else:
-            assert self.channels == self.out_channels
-            conv = nn.AvgPool2d(kernel_size=2, stride=2)
-        if name == "conv":
-            self.Conv2d_0 = conv
-            self.conv = conv
-        elif name == "Conv2d_0":
-            self.conv = conv
-        else:
-            self.conv = conv
-
-    def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
-        assert hidden_states.shape[1] == self.channels
-        if self.use_conv and self.padding == 0:
-            hidden_states = F.pad(hidden_states, (0, 1, 0, 1), mode="constant", value=0)
-        return self.conv(hidden_states)
-
-
-class Upsample2D(nn.Module):
-    def __init__(self, channels: int, use_conv: bool = False, out_channels: Optional[int] = None, name: str = "conv"):
-        super().__init__()
-        self.channels = channels
-        self.out_channels = out_channels or channels
-        self.use_conv = use_conv
-        self.use_conv_transpose = False
-        self.name = name
-        self.interpolate = True
-        conv = nn.Conv2d(self.channels, self.out_channels, kernel_size=3, padding=1) if use_conv else None
-        if name == "conv":
-            self.conv = conv
-        else:
-            self.Conv2d_0 = conv
-
-    def forward(self, hidden_states: torch.Tensor, output_size: Optional[int] = None, *args, **kwargs) -> torch.Tensor:
-        assert hidden_states.shape[1] == self.channels
-        dtype = hidden_states.dtype
-        if dtype == torch.bfloat16:
-            hidden_states = hidden_states.float()
-        if output_size is None:
-            hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest")
-        else:
-            hidden_states = F.interpolate(hidden_states, size=output_size, mode="nearest")
-        if dtype == torch.bfloat16:
-            hidden_states = hidden_states.to(dtype)
-        if self.use_conv:
-            hidden_states = self.conv(hidden_states) if self.name == "conv" else self.Conv2d_0(hidden_states)
-        return hidden_states
-
-
-class ResnetBlock2D(nn.Module):
-    def __init__(
-        self,
-        *,
-        in_channels: int,
-        out_channels: Optional[int] = None,
-        dropout: float = 0.0,
-        temb_channels: Optional[int] = 512,
-        groups: int = 32,
-        groups_out: Optional[int] = None,
-        eps: float = 1e-6,
-        non_linearity: str = "swish",
-        time_embedding_norm: str = "default",
-        output_scale_factor: float = 1.0,
-        use_in_shortcut: Optional[bool] = None,
-        conv_shortcut_bias: bool = True,
-        conv_2d_out_channels: Optional[int] = None,
-        **_: object,
-    ):
-        super().__init__()
-        if time_embedding_norm not in ("default", "scale_shift"):
-            raise ValueError(f"unknown time_embedding_norm: {time_embedding_norm}")
-        self.in_channels = in_channels
-        out_channels = in_channels if out_channels is None else out_channels
-        self.out_channels = out_channels
-        self.output_scale_factor = output_scale_factor
-        self.time_embedding_norm = time_embedding_norm
-        if groups_out is None:
-            groups_out = groups
-        self.norm1 = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
-        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
-        if temb_channels is not None:
-            self.time_emb_proj = nn.Linear(temb_channels, out_channels if time_embedding_norm == "default" else 2 * out_channels)
-        else:
-            self.time_emb_proj = None
-        self.norm2 = nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True)
-        self.dropout = nn.Dropout(dropout)
-        conv_2d_out_channels = conv_2d_out_channels or out_channels
-        self.conv2 = nn.Conv2d(out_channels, conv_2d_out_channels, kernel_size=3, stride=1, padding=1)
-        self.nonlinearity = get_activation(non_linearity)
-        self.use_in_shortcut = self.in_channels != conv_2d_out_channels if use_in_shortcut is None else use_in_shortcut
-        self.conv_shortcut = None
-        if self.use_in_shortcut:
-            self.conv_shortcut = nn.Conv2d(in_channels, conv_2d_out_channels, kernel_size=1, stride=1, padding=0, bias=conv_shortcut_bias)
-
-    def forward(self, input_tensor: torch.Tensor, temb: Optional[torch.Tensor] = None, *args, **kwargs) -> torch.Tensor:
-        hidden_states = self.norm1(input_tensor)
-        hidden_states = self.nonlinearity(hidden_states)
-        hidden_states = self.conv1(hidden_states)
-        if self.time_emb_proj is not None and temb is not None:
-            temb = self.nonlinearity(temb)
-            temb = self.time_emb_proj(temb)[:, :, None, None]
-        if self.time_embedding_norm == "default":
-            if temb is not None:
-                hidden_states = hidden_states + temb
-            hidden_states = self.norm2(hidden_states)
-        else:
-            if temb is None:
-                raise ValueError("temb cannot be None for scale_shift")
-            time_scale, time_shift = torch.chunk(temb, 2, dim=1)
-            hidden_states = self.norm2(hidden_states)
-            hidden_states = hidden_states * (1 + time_scale) + time_shift
-        hidden_states = self.nonlinearity(hidden_states)
-        hidden_states = self.dropout(hidden_states)
-        hidden_states = self.conv2(hidden_states)
-        if self.conv_shortcut is not None:
-            input_tensor = self.conv_shortcut(input_tensor.contiguous() if self.training else input_tensor)
-        return (input_tensor + hidden_states) / self.output_scale_factor
-
-
-class UNetMidBlock2D(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        temb_channels: Optional[int],
-        dropout: float = 0.0,
-        num_layers: int = 1,
-        resnet_eps: float = 1e-6,
-        resnet_time_scale_shift: str = "default",
-        resnet_act_fn: str = "swish",
-        resnet_groups: int = 32,
-        attn_groups: Optional[int] = None,
-        resnet_pre_norm: bool = True,
-        add_attention: bool = True,
-        attention_head_dim: int = 1,
-        output_scale_factor: float = 1.0,
-    ):
-        super().__init__()
-        if resnet_time_scale_shift == "spatial":
-            raise ValueError("Flux2 VAE does not use spatial resnet conditioning in this port")
-        resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
-        self.add_attention = add_attention
-        if attn_groups is None:
-            attn_groups = resnet_groups
-        resnets = [
-            ResnetBlock2D(
-                in_channels=in_channels,
-                out_channels=in_channels,
-                temb_channels=temb_channels,
-                eps=resnet_eps,
-                groups=resnet_groups,
-                dropout=dropout,
-                time_embedding_norm=resnet_time_scale_shift,
-                non_linearity=resnet_act_fn,
-                output_scale_factor=output_scale_factor,
-                pre_norm=resnet_pre_norm,
-            )
-        ]
-        attentions = []
-        if attention_head_dim is None:
-            attention_head_dim = in_channels
-        for _ in range(num_layers):
-            attentions.append(
-                Attention(
-                    in_channels,
-                    heads=in_channels // attention_head_dim,
-                    dim_head=attention_head_dim,
-                    rescale_output_factor=output_scale_factor,
-                    eps=resnet_eps,
-                    norm_num_groups=attn_groups,
-                    residual_connection=True,
-                    bias=True,
-                    upcast_softmax=True,
-                    _from_deprecated_attn_block=True,
-                )
-                if self.add_attention
-                else None
-            )
-            resnets.append(
-                ResnetBlock2D(
-                    in_channels=in_channels,
-                    out_channels=in_channels,
-                    temb_channels=temb_channels,
-                    eps=resnet_eps,
-                    groups=resnet_groups,
-                    dropout=dropout,
-                    time_embedding_norm=resnet_time_scale_shift,
-                    non_linearity=resnet_act_fn,
-                    output_scale_factor=output_scale_factor,
-                    pre_norm=resnet_pre_norm,
-                )
-            )
-        self.attentions = nn.ModuleList(attentions)
-        self.resnets = nn.ModuleList(resnets)
-        self.gradient_checkpointing = False
-
-    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
-        hidden_states = self.resnets[0](hidden_states, temb)
-        for attn, resnet in zip(self.attentions, self.resnets[1:]):
-            if attn is not None:
-                hidden_states = attn(hidden_states, temb=temb)
-            hidden_states = resnet(hidden_states, temb)
-        return hidden_states
-
-
-class DownEncoderBlock2D(nn.Module):
-    def __init__(self, in_channels: int, out_channels: int, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, resnet_act_fn: str = "swish", resnet_groups: int = 32, add_downsample: bool = True, downsample_padding: int = 1, **_: object):
-        super().__init__()
-        self.resnets = nn.ModuleList([
-            ResnetBlock2D(
-                in_channels=in_channels if i == 0 else out_channels,
-                out_channels=out_channels,
-                temb_channels=None,
-                eps=resnet_eps,
-                groups=resnet_groups,
-                dropout=dropout,
-                non_linearity=resnet_act_fn,
-            )
-            for i in range(num_layers)
-        ])
-        self.downsamplers = nn.ModuleList([Downsample2D(out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op")]) if add_downsample else None
-
-    def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
-        for resnet in self.resnets:
-            hidden_states = resnet(hidden_states, temb=None)
-        if self.downsamplers is not None:
-            for downsampler in self.downsamplers:
-                hidden_states = downsampler(hidden_states)
-        return hidden_states
-
-
-class AttnDownEncoderBlock2D(DownEncoderBlock2D):
-    def __init__(self, in_channels: int, out_channels: int, attention_head_dim: int = 1, **kwargs: object):
-        super().__init__(in_channels=in_channels, out_channels=out_channels, **kwargs)
-        resnet_groups = int(kwargs.get("resnet_groups", 32))
-        resnet_eps = float(kwargs.get("resnet_eps", 1e-6))
-        output_scale_factor = float(kwargs.get("output_scale_factor", 1.0))
-        if attention_head_dim is None:
-            attention_head_dim = out_channels
-        self.attentions = nn.ModuleList([
-            Attention(out_channels, heads=out_channels // attention_head_dim, dim_head=attention_head_dim, rescale_output_factor=output_scale_factor, eps=resnet_eps, norm_num_groups=resnet_groups, residual_connection=True, bias=True, upcast_softmax=True, _from_deprecated_attn_block=True)
-            for _ in self.resnets
-        ])
-
-    def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
-        for resnet, attn in zip(self.resnets, self.attentions):
-            hidden_states = resnet(hidden_states, temb=None)
-            hidden_states = attn(hidden_states)
-        if self.downsamplers is not None:
-            for downsampler in self.downsamplers:
-                hidden_states = downsampler(hidden_states)
-        return hidden_states
-
-
-class UpDecoderBlock2D(nn.Module):
-    def __init__(self, in_channels: int, out_channels: int, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, resnet_act_fn: str = "swish", resnet_groups: int = 32, add_upsample: bool = True, temb_channels: Optional[int] = None, **_: object):
-        super().__init__()
-        self.resnets = nn.ModuleList([
-            ResnetBlock2D(
-                in_channels=in_channels if i == 0 else out_channels,
-                out_channels=out_channels,
-                temb_channels=temb_channels,
-                eps=resnet_eps,
-                groups=resnet_groups,
-                dropout=dropout,
-                non_linearity=resnet_act_fn,
-            )
-            for i in range(num_layers)
-        ])
-        self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) if add_upsample else None
-        self.resolution_idx = None
-
-    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
-        for resnet in self.resnets:
-            hidden_states = resnet(hidden_states, temb=temb)
-        if self.upsamplers is not None:
-            for upsampler in self.upsamplers:
-                hidden_states = upsampler(hidden_states)
-        return hidden_states
-
-
-class AttnUpDecoderBlock2D(UpDecoderBlock2D):
-    def __init__(self, in_channels: int, out_channels: int, attention_head_dim: int = 1, **kwargs: object):
-        super().__init__(in_channels=in_channels, out_channels=out_channels, **kwargs)
-        resnet_groups = int(kwargs.get("resnet_groups", 32))
-        resnet_eps = float(kwargs.get("resnet_eps", 1e-6))
-        output_scale_factor = float(kwargs.get("output_scale_factor", 1.0))
-        if attention_head_dim is None:
-            attention_head_dim = out_channels
-        self.attentions = nn.ModuleList([
-            Attention(out_channels, heads=out_channels // attention_head_dim, dim_head=attention_head_dim, rescale_output_factor=output_scale_factor, eps=resnet_eps, norm_num_groups=resnet_groups, residual_connection=True, bias=True, upcast_softmax=True, _from_deprecated_attn_block=True)
-            for _ in self.resnets
-        ])
-
-    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
-        for resnet, attn in zip(self.resnets, self.attentions):
-            hidden_states = resnet(hidden_states, temb=temb)
-            hidden_states = attn(hidden_states, temb=temb)
-        if self.upsamplers is not None:
-            for upsampler in self.upsamplers:
-                hidden_states = upsampler(hidden_states)
-        return hidden_states
-
-
-def get_down_block(down_block_type: str, **kwargs: object) -> nn.Module:
-    if down_block_type == "DownEncoderBlock2D":
-        return DownEncoderBlock2D(**kwargs)
-    if down_block_type == "AttnDownEncoderBlock2D":
-        return AttnDownEncoderBlock2D(**kwargs)
-    raise ValueError(f"Unsupported Flux2 VAE down block type: {down_block_type}")
-
-
-def get_up_block(up_block_type: str, **kwargs: object) -> nn.Module:
-    kwargs.pop("prev_output_channel", None)
-    kwargs.pop("resolution_idx", None)
-    if up_block_type == "UpDecoderBlock2D":
-        return UpDecoderBlock2D(**kwargs)
-    if up_block_type == "AttnUpDecoderBlock2D":
-        return AttnUpDecoderBlock2D(**kwargs)
-    raise ValueError(f"Unsupported Flux2 VAE up block type: {up_block_type}")
-
-
-class Encoder(nn.Module):
-    def __init__(self, in_channels: int = 3, out_channels: int = 3, down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",), block_out_channels: Tuple[int, ...] = (64,), layers_per_block: int = 2, norm_num_groups: int = 32, act_fn: str = "silu", double_z: bool = True, mid_block_add_attention: bool = True):
-        super().__init__()
-        self.layers_per_block = layers_per_block
-        self.conv_in = nn.Conv2d(in_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1)
-        self.down_blocks = nn.ModuleList([])
-        output_channel = block_out_channels[0]
-        for i, down_block_type in enumerate(down_block_types):
-            input_channel = output_channel
-            output_channel = block_out_channels[i]
-            is_final_block = i == len(block_out_channels) - 1
-            self.down_blocks.append(get_down_block(down_block_type, num_layers=self.layers_per_block, in_channels=input_channel, out_channels=output_channel, add_downsample=not is_final_block, resnet_eps=1e-6, downsample_padding=0, resnet_act_fn=act_fn, resnet_groups=norm_num_groups, attention_head_dim=output_channel, temb_channels=None))
-        self.mid_block = UNetMidBlock2D(in_channels=block_out_channels[-1], resnet_eps=1e-6, resnet_act_fn=act_fn, output_scale_factor=1, resnet_time_scale_shift="default", attention_head_dim=block_out_channels[-1], resnet_groups=norm_num_groups, temb_channels=None, add_attention=mid_block_add_attention)
-        self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6)
-        self.conv_act = nn.SiLU()
-        conv_out_channels = 2 * out_channels if double_z else out_channels
-        self.conv_out = nn.Conv2d(block_out_channels[-1], conv_out_channels, 3, padding=1)
-        self.gradient_checkpointing = False
-
-    def forward(self, sample: torch.Tensor) -> torch.Tensor:
-        sample = self.conv_in(sample)
-        for down_block in self.down_blocks:
-            sample = down_block(sample)
-        sample = self.mid_block(sample)
-        sample = self.conv_norm_out(sample)
-        sample = self.conv_act(sample)
-        return self.conv_out(sample)
-
-
-class Decoder(nn.Module):
-    def __init__(self, in_channels: int = 3, out_channels: int = 3, up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",), block_out_channels: Tuple[int, ...] = (64,), layers_per_block: int = 2, norm_num_groups: int = 32, act_fn: str = "silu", norm_type: str = "group", mid_block_add_attention: bool = True):
-        super().__init__()
-        if norm_type != "group":
-            raise ValueError("Flux2 VAE Decoder only supports group norm in this port")
-        self.layers_per_block = layers_per_block
-        self.conv_in = nn.Conv2d(in_channels, block_out_channels[-1], kernel_size=3, stride=1, padding=1)
-        self.up_blocks = nn.ModuleList([])
-        self.mid_block = UNetMidBlock2D(in_channels=block_out_channels[-1], resnet_eps=1e-6, resnet_act_fn=act_fn, output_scale_factor=1, resnet_time_scale_shift="default", attention_head_dim=block_out_channels[-1], resnet_groups=norm_num_groups, temb_channels=None, add_attention=mid_block_add_attention)
-        reversed_block_out_channels = list(reversed(block_out_channels))
-        output_channel = reversed_block_out_channels[0]
-        for i, up_block_type in enumerate(up_block_types):
-            prev_output_channel = output_channel
-            output_channel = reversed_block_out_channels[i]
-            is_final_block = i == len(block_out_channels) - 1
-            self.up_blocks.append(get_up_block(up_block_type, num_layers=self.layers_per_block + 1, in_channels=prev_output_channel, out_channels=output_channel, prev_output_channel=prev_output_channel, add_upsample=not is_final_block, resnet_eps=1e-6, resnet_act_fn=act_fn, resnet_groups=norm_num_groups, attention_head_dim=output_channel, temb_channels=None, resnet_time_scale_shift=norm_type))
-        self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6)
-        self.conv_act = nn.SiLU()
-        self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, 3, padding=1)
-        self.gradient_checkpointing = False
-
-    def forward(self, sample: torch.Tensor, latent_embeds: Optional[torch.Tensor] = None) -> torch.Tensor:
-        sample = self.conv_in(sample)
-        sample = self.mid_block(sample, latent_embeds)
-        for up_block in self.up_blocks:
-            sample = up_block(sample, latent_embeds)
-        sample = self.conv_norm_out(sample)
-        sample = self.conv_act(sample)
-        return self.conv_out(sample)

fastvideo/models/vaes/flux2vae.py

@@ -1,533 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Copied and adapted from: https://github.com/sglang-ai/sglang
-import math
-from typing import Dict, Optional, Tuple, Union
-
-import torch
-import torch.nn as nn
-from fastvideo.configs.models.vaes.flux2vae import Flux2VAEConfig
-from fastvideo.models.vaes.common import (
-    DiagonalGaussianDistribution,
-    ParallelTiledVAE,
-)
-from fastvideo.models.vaes.flux2_components import (
-    ADDED_KV_ATTENTION_PROCESSORS,
-    CROSS_ATTENTION_PROCESSORS,
-    Attention,
-    AttnAddedKVProcessor,
-    AttnProcessor,
-    AutoencoderKLOutput,
-    Decoder,
-    DecoderOutput,
-    Encoder,
-)
-
-AttentionProcessor = AttnProcessor
-
-
-class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE):
-    r"""
-    A VAE model with KL loss for encoding images into latents and decoding latent representations into images.
-    
-    This model inherits from [`ParallelTiledVAE`] for tiling support and uses standard diffusers
-    Encoder/Decoder components for Flux2 image generation.
-    """
-
-    _supports_gradient_checkpointing = True
-    _no_split_modules = ["Attention", "ResnetBlock2D"]
-
-    def __init__(
-        self,
-        config: Flux2VAEConfig,
-    ):
-        nn.Module.__init__(self)
-        ParallelTiledVAE.__init__(self, config=config)
-
-        self.config = config
-        arch_config = config.arch_config
-
-        in_channels: int = arch_config.in_channels
-        out_channels: int = arch_config.out_channels
-        down_block_types: Tuple[str, ...] = arch_config.down_block_types
-        up_block_types: Tuple[str, ...] = arch_config.up_block_types
-        block_out_channels: Tuple[int, ...] = arch_config.block_out_channels
-        layers_per_block: int = arch_config.layers_per_block
-        act_fn: str = arch_config.act_fn
-        latent_channels: int = arch_config.latent_channels
-        norm_num_groups: int = arch_config.norm_num_groups
-        sample_size: int = arch_config.sample_size
-        force_upcast: bool = arch_config.force_upcast
-        use_quant_conv: bool = arch_config.use_quant_conv
-        use_post_quant_conv: bool = arch_config.use_post_quant_conv
-        mid_block_add_attention: bool = arch_config.mid_block_add_attention
-        batch_norm_eps: float = arch_config.batch_norm_eps
-        batch_norm_momentum: float = arch_config.batch_norm_momentum
-        patch_size: Tuple[int, int] = arch_config.patch_size
-        
-        # pass init params to Encoder
-        self.encoder = Encoder(
-            in_channels=in_channels,
-            out_channels=latent_channels,
-            down_block_types=down_block_types,
-            block_out_channels=block_out_channels,
-            layers_per_block=layers_per_block,
-            act_fn=act_fn,
-            norm_num_groups=norm_num_groups,
-            double_z=True,
-            mid_block_add_attention=mid_block_add_attention,
-        )
-
-        # pass init params to Decoder
-        self.decoder = Decoder(
-            in_channels=latent_channels,
-            out_channels=out_channels,
-            up_block_types=up_block_types,
-            block_out_channels=block_out_channels,
-            layers_per_block=layers_per_block,
-            norm_num_groups=norm_num_groups,
-            act_fn=act_fn,
-            mid_block_add_attention=mid_block_add_attention,
-        )
-
-        self.quant_conv = (
-            nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1)
-            if use_quant_conv
-            else None
-        )
-        self.post_quant_conv = (
-            nn.Conv2d(latent_channels, latent_channels, 1)
-            if use_post_quant_conv
-            else None
-        )
-
-        self.bn = nn.BatchNorm2d(
-            math.prod(patch_size) * latent_channels,
-            eps=batch_norm_eps,
-            momentum=batch_norm_momentum,
-            affine=False,
-            track_running_stats=True,
-        )
-
-        self.use_slicing = False
-        self.use_tiling = False
-
-        # only relevant if vae tiling is enabled
-        self.tile_sample_min_size = sample_size
-        sample_size_val = (
-            sample_size[0]
-            if isinstance(sample_size, (list, tuple))
-            else sample_size
-        )
-        self.tile_latent_min_size = int(
-            sample_size_val / (2 ** (len(block_out_channels) - 1))
-        )
-        self.tile_overlap_factor = 0.25
-
-    @property
-    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
-    def attn_processors(self) -> Dict[str, AttentionProcessor]:
-        r"""
-        Returns:
-            `dict` of attention processors: A dictionary containing all attention processors used in the model with
-            indexed by its weight name.
-        """
-        # set recursively
-        processors = {}
-
-        def fn_recursive_add_processors(
-            name: str,
-            module: torch.nn.Module,
-            processors: Dict[str, AttentionProcessor],
-        ):
-            if hasattr(module, "get_processor"):
-                processors[f"{name}.processor"] = module.get_processor()
-
-            for sub_name, child in module.named_children():
-                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
-
-            return processors
-
-        for name, module in self.named_children():
-            fn_recursive_add_processors(name, module, processors)
-
-        return processors
-
-    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
-    def set_attn_processor(
-        self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]
-    ):
-        r"""
-        Sets the attention processor to use to compute attention.
-
-        Parameters:
-            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
-                The instantiated processor class or a dictionary of processor classes that will be set as the processor
-                for **all** `Attention` layers.
-
-                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
-                processor. This is strongly recommended when setting trainable attention processors.
-
-        """
-        count = len(self.attn_processors.keys())
-
-        if isinstance(processor, dict) and len(processor) != count:
-            raise ValueError(
-                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
-                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
-            )
-
-        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
-            if hasattr(module, "set_processor"):
-                if not isinstance(processor, dict):
-                    module.set_processor(processor)
-                else:
-                    module.set_processor(processor.pop(f"{name}.processor"))
-
-            for sub_name, child in module.named_children():
-                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
-
-        for name, module in self.named_children():
-            fn_recursive_attn_processor(name, module, processor)
-
-    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor
-    def set_default_attn_processor(self):
-        """
-        Disables custom attention processors and sets the default attention implementation.
-        """
-        if all(
-            proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS
-            for proc in self.attn_processors.values()
-        ):
-            processor = AttnAddedKVProcessor()
-        elif all(
-            proc.__class__ in CROSS_ATTENTION_PROCESSORS
-            for proc in self.attn_processors.values()
-        ):
-            processor = AttnProcessor()
-        else:
-            raise ValueError(
-                f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
-            )
-
-        self.set_attn_processor(processor)
-
-    def _encode(self, x: torch.Tensor) -> torch.Tensor:
-        batch_size, num_channels, height, width = x.shape
-
-        if self.use_tiling and (
-            width > self.tile_sample_min_size or height > self.tile_sample_min_size
-        ):
-            return self._tiled_encode(x)
-
-        enc = self.encoder(x)
-        if self.quant_conv is not None:
-            enc = self.quant_conv(enc)
-
-        return enc
-
-    def encode(
-        self, x: torch.Tensor, return_dict: bool = True
-    ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
-        """
-        Encode a batch of images into latents.
-
-        Args:
-            x (`torch.Tensor`): Input batch of images.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
-
-        Returns:
-                The latent representations of the encoded images. If `return_dict` is True, a
-                [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned.
-        """
-
-        if x.ndim == 5:
-            assert x.shape[2] == 1
-            x = x.squeeze(2)
-
-        if self.use_slicing and x.shape[0] > 1:
-            encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)]
-            h = torch.cat(encoded_slices)
-        else:
-            h = self._encode(x)
-
-        posterior = DiagonalGaussianDistribution(h)
-        if not return_dict:
-            return (posterior,)
-
-        return AutoencoderKLOutput(latent_dist=posterior)
-
-    def _decode(
-        self, z: torch.Tensor, return_dict: bool = True
-    ) -> Union[DecoderOutput, torch.Tensor]:
-        if self.use_tiling and (
-            z.shape[-1] > self.tile_latent_min_size
-            or z.shape[-2] > self.tile_latent_min_size
-        ):
-            return self.tiled_decode(z, return_dict=return_dict)
-
-        if self.post_quant_conv is not None:
-            z = self.post_quant_conv(z)
-
-        dec = self.decoder(z)
-
-        if not return_dict:
-            return (dec,)
-
-        return DecoderOutput(sample=dec)
-
-    def decode(
-        self, z: torch.FloatTensor, return_dict: bool = True, generator=None
-    ) -> Union[DecoderOutput, torch.FloatTensor]:
-        """
-        Decode a batch of images.
-
-        Args:
-            z (`torch.Tensor`): Input batch of latent vectors.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
-
-        Returns:
-            [`~models.vae.DecoderOutput`] or `tuple`:
-                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
-                returned.
-
-        """
-        if self.use_slicing and z.shape[0] > 1:
-            decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
-            decoded = torch.cat(decoded_slices)
-        else:
-            decoded = self._decode(z).sample
-
-        if not return_dict:
-            return (decoded,)
-
-        return DecoderOutput(sample=decoded)
-
-    def blend_v(
-        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
-    ) -> torch.Tensor:
-        blend_extent = min(a.shape[2], b.shape[2], blend_extent)
-        for y in range(blend_extent):
-            b[:, :, y, :] = a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[
-                :, :, y, :
-            ] * (y / blend_extent)
-        return b
-
-    def blend_h(
-        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
-    ) -> torch.Tensor:
-        blend_extent = min(a.shape[3], b.shape[3], blend_extent)
-        for x in range(blend_extent):
-            b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[
-                :, :, :, x
-            ] * (x / blend_extent)
-        return b
-
-    def _tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
-        r"""Encode a batch of images using a tiled encoder.
-
-        When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several
-        steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is
-        different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the
-        tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the
-        output, but they should be much less noticeable.
-
-        Args:
-            x (`torch.Tensor`): Input batch of images.
-
-        Returns:
-            `torch.Tensor`:
-                The latent representation of the encoded videos.
-        """
-
-        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
-        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
-        row_limit = self.tile_latent_min_size - blend_extent
-
-        # Split the image into 512x512 tiles and encode them separately.
-        rows = []
-        for i in range(0, x.shape[2], overlap_size):
-            row = []
-            for j in range(0, x.shape[3], overlap_size):
-                tile = x[
-                    :,
-                    :,
-                    i : i + self.tile_sample_min_size,
-                    j : j + self.tile_sample_min_size,
-                ]
-                tile = self.encoder(tile)
-                if self.quant_conv is not None:
-                    tile = self.quant_conv(tile)
-                row.append(tile)
-            rows.append(row)
-        result_rows = []
-        for i, row in enumerate(rows):
-            result_row = []
-            for j, tile in enumerate(row):
-                # blend the above tile and the left tile
-                # to the current tile and add the current tile to the result row
-                if i > 0:
-                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
-                if j > 0:
-                    tile = self.blend_h(row[j - 1], tile, blend_extent)
-                result_row.append(tile[:, :, :row_limit, :row_limit])
-            result_rows.append(torch.cat(result_row, dim=3))
-
-        enc = torch.cat(result_rows, dim=2)
-        return enc
-
-    def tiled_encode(
-        self, x: torch.Tensor, return_dict: bool = True
-    ) -> AutoencoderKLOutput:
-        r"""Encode a batch of images using a tiled encoder.
-
-        When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several
-        steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is
-        different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the
-        tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the
-        output, but they should be much less noticeable.
-
-        Args:
-            x (`torch.Tensor`): Input batch of images.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
-
-        Returns:
-            [`~models.autoencoder_kl.AutoencoderKLOutput`] or `tuple`:
-                If return_dict is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain
-                `tuple` is returned.
-        """
-        deprecation_message = (
-            "The tiled_encode implementation supporting the `return_dict` parameter is deprecated. In the future, the "
-            "implementation of this method will be replaced with that of `_tiled_encode` and you will no longer be able "
-            "to pass `return_dict`. You will also have to create a `DiagonalGaussianDistribution()` from the returned value."
-        )
-
-        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
-        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
-        row_limit = self.tile_latent_min_size - blend_extent
-
-        # Split the image into 512x512 tiles and encode them separately.
-        rows = []
-        for i in range(0, x.shape[2], overlap_size):
-            row = []
-            for j in range(0, x.shape[3], overlap_size):
-                tile = x[
-                    :,
-                    :,
-                    i : i + self.tile_sample_min_size,
-                    j : j + self.tile_sample_min_size,
-                ]
-                tile = self.encoder(tile)
-                if self.quant_conv is not None:
-                    tile = self.quant_conv(tile)
-                row.append(tile)
-            rows.append(row)
-        result_rows = []
-        for i, row in enumerate(rows):
-            result_row = []
-            for j, tile in enumerate(row):
-                # blend the above tile and the left tile
-                # to the current tile and add the current tile to the result row
-                if i > 0:
-                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
-                if j > 0:
-                    tile = self.blend_h(row[j - 1], tile, blend_extent)
-                result_row.append(tile[:, :, :row_limit, :row_limit])
-            result_rows.append(torch.cat(result_row, dim=3))
-
-        moments = torch.cat(result_rows, dim=2)
-        posterior = DiagonalGaussianDistribution(moments)
-
-        if not return_dict:
-            return (posterior,)
-
-        return AutoencoderKLOutput(latent_dist=posterior)
-
-    def tiled_decode(
-        self, z: torch.Tensor, return_dict: bool = True
-    ) -> Union[DecoderOutput, torch.Tensor]:
-        r"""
-        Decode a batch of images using a tiled decoder.
-
-        Args:
-            z (`torch.Tensor`): Input batch of latent vectors.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
-
-        Returns:
-            [`~models.vae.DecoderOutput`] or `tuple`:
-                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
-                returned.
-        """
-        overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
-        blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
-        row_limit = self.tile_sample_min_size - blend_extent
-
-        # Split z into overlapping 64x64 tiles and decode them separately.
-        # The tiles have an overlap to avoid seams between tiles.
-        rows = []
-        for i in range(0, z.shape[2], overlap_size):
-            row = []
-            for j in range(0, z.shape[3], overlap_size):
-                tile = z[
-                    :,
-                    :,
-                    i : i + self.tile_latent_min_size,
-                    j : j + self.tile_latent_min_size,
-                ]
-                if self.post_quant_conv is not None:
-                    tile = self.post_quant_conv(tile)
-                decoded = self.decoder(tile)
-                row.append(decoded)
-            rows.append(row)
-        result_rows = []
-        for i, row in enumerate(rows):
-            result_row = []
-            for j, tile in enumerate(row):
-                # blend the above tile and the left tile
-                # to the current tile and add the current tile to the result row
-                if i > 0:
-                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
-                if j > 0:
-                    tile = self.blend_h(row[j - 1], tile, blend_extent)
-                result_row.append(tile[:, :, :row_limit, :row_limit])
-            result_rows.append(torch.cat(result_row, dim=3))
-
-        dec = torch.cat(result_rows, dim=2)
-        if not return_dict:
-            return (dec,)
-
-        return DecoderOutput(sample=dec)
-
-    def forward(
-        self,
-        sample: torch.Tensor,
-        sample_posterior: bool = False,
-        return_dict: bool = True,
-        generator: Optional[torch.Generator] = None,
-    ) -> Union[DecoderOutput, torch.Tensor]:
-        r"""
-        Args:
-            sample (`torch.Tensor`): Input sample.
-            sample_posterior (`bool`, *optional*, defaults to `False`):
-                Whether to sample from the posterior.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
-        """
-        x = sample
-        posterior = self.encode(x).latent_dist
-        if sample_posterior:
-            z = posterior.sample(generator=generator)
-        else:
-            z = posterior.mode()
-        dec = self.decode(z).sample
-
-        if not return_dict:
-            return (dec,)
-
-        return DecoderOutput(sample=dec)
-
-
-EntryClass = AutoencoderKLFlux2

fastvideo/models/vaes/oobleck.py

@@ -94,6 +94,9 @@ class OobleckDecoderBlock(nn.Module):
             input_dim, output_dim,
             kernel_size=2 * stride, stride=stride,
             padding=math.ceil(stride / 2),
+            # Clean L*stride upsample for both parities; a no-op (0) for even
+            # strides (Stable Audio), needed for odd strides (Cosmos3: 5).
+            output_padding=stride % 2,
         ))
         self.res_unit1 = OobleckResidualUnit(output_dim, dilation=1)
         self.res_unit2 = OobleckResidualUnit(output_dim, dilation=3)

fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py

@@ -0,0 +1,654 @@
+# SPDX-License-Identifier: Apache-2.0
+"""FastVideo-native Cosmos3 video pipeline (T2V / I2V / T2I).
+
+This replaces the earlier vllm-omni-derived skeleton with a native, stage-based
+:class:`ComposedPipelineBase` pipeline that wires the framework-parity-verified
+Cosmos3 components:
+
+* tokenizer: Qwen2 ``Qwen2TokenizerFast`` + chat template (the only allowed
+  third-party model-adjacent dependency; tokenizers are explicitly permitted),
+* VAE: FastVideo-native ``AutoencoderKLWan`` (Wan2.2) via ``Cosmos3VAEConfig``;
+  encode normalizes ``(mu - mean) * inv_std`` and decode denormalizes + clamps,
+* sequence-packing: :func:`pack_cosmos3_video_sequence` (native, parity-tested),
+* DiT: ``Cosmos3VFMTransformer`` (native, bit-identical to the framework),
+* scheduler: FastVideo-native ``UniPCMultistepScheduler`` configured for pure
+  flow matching (``flow_prediction`` + ``use_flow_sigmas``), numerically
+  equivalent to the framework's ``FlowUniPCMultistepScheduler`` (parity-tested
+  in ``test_cosmos3_scheduler_parity``).
+
+The denoise/CFG glue is a faithful port of the framework's
+``Cosmos3OmniDiffusersPipeline`` math (mirrored in the framework-equivalent
+``diffusers_cosmos3.pipeline``): per UniPC timestep, run a SEQUENTIAL conditional
+then unconditional pass (each repacks the sequence with the prompt / negative
+prompt token ids, forwards the DiT, and zeros the prediction on conditioning
+frames), then combine ``v = uncond + guidance * (cond - uncond)`` and take one
+``scheduler.step(model_output=v, timestep, sample=latent)``. ``timestep_scale``
+is applied to the per-token timesteps *inside* the DiT (its ``forward`` already
+multiplies ``vision_timesteps * timestep_scale`` before the time embedder), so
+the loop passes raw scheduler timesteps to the packer.
+
+The pure denoise math lives in :class:`Cosmos3DenoiseEngine` and the free
+function :func:`cosmos3_get_cfg_velocity` so it can be unit-/parity-tested
+directly against the framework oracle without constructing the full pipeline.
+
+No diffusers/transformers *model* classes are imported at runtime here; only the
+Qwen2 tokenizer (loaded by the component loader) and the UniPC scheduler are
+third-party, both explicitly allowed.
+"""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_unipc_multistep import (
+    UniPCMultistepScheduler, )
+from fastvideo.pipelines.basic.cosmos3.sequence_packing import (
+    Cosmos3ActionItem,
+    Cosmos3SampleInputs,
+    Cosmos3SoundItem,
+    Cosmos3VisionItem,
+    pack_cosmos3_video_sequence,
+)
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+
+logger = init_logger(__name__)
+
+# System prompts, verbatim from the framework (diffusers_cosmos3.pipeline).
+_SYSTEM_PROMPT_IMAGE = "You are a helpful assistant who will generate images from a give prompt."
+_SYSTEM_PROMPT_VIDEO = "You are a helpful assistant who will generate videos from a give prompt."
+
+
+# ===========================================================================
+# Special-token resolution (Qwen2 chat tokenizer)
+# ===========================================================================
+def cosmos3_special_tokens(tokenizer: Any) -> dict[str, int]:
+    """Resolve the Cosmos3 generation special tokens from a Qwen2 tokenizer.
+
+    Mirrors the framework's ``llm_special_tokens``:
+    ``start_of_generation=<|vision_start|>``, ``end_of_generation=<|vision_end|>``,
+    ``eos_token_id=tokenizer.eos_token_id``.
+    """
+    return {
+        "start_of_generation": int(tokenizer.convert_tokens_to_ids("<|vision_start|>")),
+        "end_of_generation": int(tokenizer.convert_tokens_to_ids("<|vision_end|>")),
+        "eos_token_id": int(tokenizer.eos_token_id),
+    }
+
+
+def cosmos3_tokenize_caption(
+    tokenizer: Any,
+    caption: str,
+    *,
+    is_video: bool = False,
+    use_system_prompt: bool = False,
+) -> list[int]:
+    """Tokenize a caption with the Qwen2 chat template (framework-faithful).
+
+    Optionally prepends an image/video system prompt; always adds the
+    generation prompt and disables ``add_vision_id`` (matching the framework's
+    ``tokenize_caption``).
+    """
+    conversations: list[dict[str, str]] = []
+    if use_system_prompt:
+        conversations.append({
+            "role": "system",
+            "content": _SYSTEM_PROMPT_VIDEO if is_video else _SYSTEM_PROMPT_IMAGE,
+        })
+    conversations.append({"role": "user", "content": caption})
+    token_ids = tokenizer.apply_chat_template(
+        conversations,
+        tokenize=True,
+        add_generation_prompt=True,
+        add_vision_id=False,
+    )
+    return list(token_ids)
+
+
+# ===========================================================================
+# VAE encode/decode bridge (normalize / denormalize, matching the framework)
+# ===========================================================================
+@dataclass
+class _VaeNorm:
+    """Cached ``mean`` / ``inv_std`` for VAE (de)normalization."""
+
+    mean: torch.Tensor  # [z_dim]
+    inv_std: torch.Tensor  # [z_dim]
+
+    @classmethod
+    def from_vae(cls, vae: Any, dtype: torch.dtype) -> _VaeNorm:
+        mean = torch.tensor(list(vae.config.latents_mean), dtype=dtype)
+        std = torch.tensor(list(vae.config.latents_std), dtype=dtype)
+        return cls(mean=mean, inv_std=1.0 / std)
+
+
+def cosmos3_vae_encode(vae: Any, video: torch.Tensor, norm: _VaeNorm) -> torch.Tensor:
+    """Encode ``[B, 3, T, H, W]`` pixels in [-1, 1] to NORMALIZED latents.
+
+    Matches the framework ``DiffusersWan22VAE.encode``: take the posterior mode
+    and apply ``(mu - mean) * inv_std``. FastVideo's ``AutoencoderKLWan.encode``
+    returns a ``DiagonalGaussianDistribution``; we read ``.mode()``.
+    """
+    in_dtype = video.dtype
+    device = video.device
+    mean = norm.mean.to(device=device, dtype=in_dtype).view(1, -1, 1, 1, 1)
+    inv_std = norm.inv_std.to(device=device, dtype=in_dtype).view(1, -1, 1, 1, 1)
+    raw_mu = vae.encode(video).mode()
+    return ((raw_mu - mean) * inv_std).to(in_dtype)
+
+
+def cosmos3_vae_decode(vae: Any, latents: torch.Tensor, norm: _VaeNorm) -> torch.Tensor:
+    """Decode NORMALIZED latents ``[B, z, T, H, W]`` to pixels ``[B, 3, T, H, W]``.
+
+    Inverts the normalization (``z / inv_std + mean``) then calls
+    ``vae.decode`` (which already clamps to [-1, 1]).
+    """
+    in_dtype = latents.dtype
+    device = latents.device
+    mean = norm.mean.to(device=device, dtype=in_dtype).view(1, -1, 1, 1, 1)
+    inv_std = norm.inv_std.to(device=device, dtype=in_dtype).view(1, -1, 1, 1, 1)
+    z_raw = latents / inv_std + mean
+    out = vae.decode(z_raw)
+    if isinstance(out, tuple):
+        out = out[0]
+    if hasattr(out, "sample"):
+        out = out.sample
+    return out.to(in_dtype)
+
+
+# ===========================================================================
+# Per-vision-item packing geometry
+# ===========================================================================
+@dataclass
+class Cosmos3VisionSpec:
+    """Geometry + conditioning for one vision item in a denoise run.
+
+    Args:
+        condition_frame_indexes: Latent-frame indices kept clean.
+        shape: ``(C, T, H, W)`` of the latent for this item.
+    """
+
+    shape: tuple[int, int, int, int]
+    condition_frame_indexes: list[int]
+
+    @property
+    def numel(self) -> int:
+        return int(math.prod(self.shape))
+
+
+# ===========================================================================
+# Pure denoise/CFG math (parity oracle target)
+# ===========================================================================
+def _split_flat_latent(flat: torch.Tensor, specs: list[Any]) -> list[torch.Tensor]:
+    """Split a flat vector into per-item tensors via each spec's ``numel``/``shape``.
+
+    Shared by vision (``[C, T, H, W]``), sound (``[C, T]``), and action
+    (``[T, D]``) specs — every spec exposes ``numel`` and ``shape``.
+    """
+    out: list[torch.Tensor] = []
+    offset = 0
+    for spec in specs:
+        out.append(flat[offset:offset + spec.numel].reshape(spec.shape))
+        offset += spec.numel
+    return out
+
+
+@dataclass
+class Cosmos3SoundSpec:
+    """Geometry + conditioning for one sound item in a denoise run.
+
+    Args:
+        shape: ``(C, T)`` of the sound latent (channels, temporal frames).
+        condition_frame_indexes: Latent-frame indices kept clean (``[]`` for t2vs).
+        fps: Sound latent FPS (``sound_latent_fps``); used iff fps modulation is on.
+    """
+
+    shape: tuple[int, int]
+    condition_frame_indexes: list[int] = field(default_factory=list)
+    fps: float | None = None
+
+    @property
+    def numel(self) -> int:
+        return int(math.prod(self.shape))
+
+
+@dataclass
+class Cosmos3ActionSpec:
+    """Geometry + conditioning for one action item in a denoise run.
+
+    Args:
+        shape: ``(T, action_dim)`` of the action latent.
+        condition_frame_indexes: Frame indices kept clean (conditioning actions).
+        domain_id: Embodiment domain id for the domain-aware action projection.
+        fps: Action FPS; used iff fps modulation is on.
+    """
+
+    shape: tuple[int, int]
+    condition_frame_indexes: list[int] = field(default_factory=list)
+    domain_id: int = 0
+    fps: float | None = None
+
+    @property
+    def numel(self) -> int:
+        return int(math.prod(self.shape))
+
+
+def cosmos3_get_cfg_velocity(
+    *,
+    transformer: Any,
+    flat_latent: torch.Tensor,
+    timestep: torch.Tensor,
+    guidance: float,
+    specs: list[Cosmos3VisionSpec],
+    cond_token_ids: list[int],
+    uncond_token_ids: list[int],
+    special_tokens: dict[str, int],
+    latent_patch_size: int,
+    temporal_modality_margin: int,
+    reset_spatial_ids: bool,
+    enable_fps_modulation: bool,
+    base_fps: float,
+    temporal_compression_factor: int,
+    include_end_of_generation_token: bool = False,
+    fps_per_item: list[float] | None = None,
+    normalize_cfg: bool = False,
+    sound_specs: list[Cosmos3SoundSpec] | None = None,
+    sound_fps_per_item: list[float] | None = None,
+    action_specs: list[Cosmos3ActionSpec] | None = None,
+    action_fps_per_item: list[float] | None = None,
+) -> torch.Tensor:
+    """Sequential-CFG velocity for one denoise step (framework math).
+
+    Replicates the framework ``get_cfg_velocity``:
+
+    1. split ``flat_latent`` into per-vision-item ``[C, T, H, W]`` latents,
+    2. run a conditional pass (prompt tokens) and an unconditional pass
+       (negative-prompt tokens); each repacks via
+       :func:`pack_cosmos3_video_sequence`, forwards the DiT to obtain
+       ``preds_vision`` (a list of ``[1, C, T, H, W]`` unpatchified noisy-frame
+       predictions), and zeros the prediction on conditioning frames
+       (``pred * (1 - condition_mask)``),
+    3. combine ``v = uncond + guidance * (cond - uncond)`` (optionally
+       norm-rescaled), returned flattened to match ``flat_latent``.
+
+    ``timestep`` is a scalar tensor (raw scheduler timestep); ``timestep_scale``
+    is applied inside the DiT, so it is passed through unscaled here.
+    """
+    assert timestep.numel() == 1, "timestep must be a scalar"
+    timestep_value = float(timestep.reshape(()).item())
+
+    # Combined flat layout: [all vision | all action | all sound], matching the
+    # framework per-sample concat order ([vision_i | action_i | sound_i]); single
+    # sample here.
+    vision_total = sum(spec.numel for spec in specs)
+    action_total = sum(spec.numel for spec in action_specs) if action_specs else 0
+    noise_x_vision = _split_flat_latent(flat_latent[:vision_total], specs)
+    noise_x_action = (_split_flat_latent(flat_latent[vision_total:vision_total +
+                                                     action_total], action_specs) if action_specs else None)
+    noise_x_sound = (_split_flat_latent(flat_latent[vision_total +
+                                                    action_total:], sound_specs) if sound_specs else None)
+    device = next(transformer.parameters()).device
+
+    def _run(token_ids: list[int]) -> torch.Tensor:
+        sound_items: list[Cosmos3SoundItem] = []
+        if sound_specs is not None and noise_x_sound is not None:
+            sound_items = [
+                Cosmos3SoundItem(
+                    latent=noise_x_sound[i],
+                    condition_frame_indexes=list(ss.condition_frame_indexes),
+                    fps=(sound_fps_per_item[i] if sound_fps_per_item is not None else None),
+                ) for i, ss in enumerate(sound_specs)
+            ]
+        action_items: list[Cosmos3ActionItem] = []
+        if action_specs is not None and noise_x_action is not None:
+            action_items = [
+                Cosmos3ActionItem(
+                    latent=noise_x_action[i],
+                    condition_frame_indexes=list(asp.condition_frame_indexes),
+                    domain_id=asp.domain_id,
+                    fps=(action_fps_per_item[i] if action_fps_per_item is not None else None),
+                ) for i, asp in enumerate(action_specs)
+            ]
+        samples = [
+            Cosmos3SampleInputs(
+                text_ids=list(token_ids),
+                vision=Cosmos3VisionItem(
+                    latent=latent,
+                    condition_frame_indexes=list(spec.condition_frame_indexes),
+                    fps=(fps_per_item[i] if fps_per_item is not None else None),
+                ),
+                sound=(sound_items[i] if i < len(sound_items) else None),
+                action=(action_items[i] if i < len(action_items) else None),
+                timestep=timestep_value,
+            ) for i, (latent, spec) in enumerate(zip(noise_x_vision, specs, strict=False))
+        ]
+        packed = pack_cosmos3_video_sequence(
+            samples,
+            special_tokens,
+            latent_patch_size=latent_patch_size,
+            include_end_of_generation_token=include_end_of_generation_token,
+            temporal_modality_margin=temporal_modality_margin,
+            reset_spatial_ids=reset_spatial_ids,
+            enable_fps_modulation=enable_fps_modulation,
+            base_fps=base_fps,
+            temporal_compression_factor=temporal_compression_factor,
+        )
+        out = transformer(**packed.to_dit_kwargs(device=device))
+
+        # Vision velocity: zero on conditioning frames, per item, flattened.
+        vision_vel = torch.zeros(vision_total, device=flat_latent.device, dtype=flat_latent.dtype)
+        preds = out.get("preds_vision")
+        if preds is not None:
+            items: list[torch.Tensor] = []
+            for pred, cond_mask in zip(preds, packed.vision_condition_mask, strict=False):
+                pred = pred.squeeze(0) if pred.dim() == 5 else pred  # [C, T, H, W]
+                keep = (1.0 - cond_mask).to(dtype=pred.dtype, device=pred.device)  # [T,1,1]
+                items.append(pred * keep if keep.sum() > 0 else torch.zeros_like(pred))
+            vision_vel = torch.cat([v.reshape(-1) for v in items]).to(flat_latent.dtype)
+
+        parts = [vision_vel]
+
+        if action_specs:
+            # Action velocity: preds_action are per-item [T, D], already zero on
+            # clean frames; zero on cond frames defensively.
+            action_vel = torch.zeros(action_total, device=flat_latent.device, dtype=flat_latent.dtype)
+            preds_a = out.get("preds_action")
+            if preds_a is not None:
+                a_items: list[torch.Tensor] = []
+                for pred, cond_mask in zip(preds_a, packed.action_condition_mask, strict=False):
+                    pred = pred.squeeze(0) if pred.dim() == 3 else pred  # [T, D]
+                    keep = (1.0 - cond_mask).reshape(-1, 1).to(dtype=pred.dtype, device=pred.device)  # [T, 1]
+                    a_items.append(pred * keep)
+                action_vel = torch.cat([v.reshape(-1) for v in a_items]).to(flat_latent.dtype)
+            parts.append(action_vel)
+
+        if sound_specs:
+            # Sound velocity: preds_sound are per-item [C, T], already zero on clean
+            # frames (unpack fills only noisy frames); zero on cond frames defensively.
+            sound_total = sum(spec.numel for spec in sound_specs)
+            sound_vel = torch.zeros(sound_total, device=flat_latent.device, dtype=flat_latent.dtype)
+            preds_s = out.get("preds_sound")
+            if preds_s is not None:
+                s_items: list[torch.Tensor] = []
+                for pred, cond_mask in zip(preds_s, packed.sound_condition_mask, strict=False):
+                    pred = pred.squeeze(0) if pred.dim() == 3 else pred  # [C, T]
+                    keep = (1.0 - cond_mask).reshape(1, -1).to(dtype=pred.dtype, device=pred.device)  # [1, T]
+                    s_items.append(pred * keep)
+                sound_vel = torch.cat([v.reshape(-1) for v in s_items]).to(flat_latent.dtype)
+            parts.append(sound_vel)
+
+        return vision_vel if len(parts) == 1 else torch.cat(parts)
+
+    cond_v = _run(cond_token_ids)
+    uncond_v = _run(uncond_token_ids)
+    v_pred = uncond_v + guidance * (cond_v - uncond_v)
+    if normalize_cfg:
+        scale = (torch.norm(cond_v) / (torch.norm(v_pred) + 1e-8)).clamp(min=0.0, max=1.0)
+        v_pred = v_pred * scale
+    return v_pred
+
+
+class Cosmos3DenoiseEngine:
+    """Stateless denoise driver tying CFG velocity to UniPC stepping.
+
+    Holds the transformer + scheduler + packing constants and runs the full
+    UniPC denoise loop. Kept separate from the pipeline so it can be exercised
+    in isolation (smoke + parity tests) with stub or real components.
+    """
+
+    def __init__(
+        self,
+        *,
+        transformer: Any,
+        scheduler: Any,
+        special_tokens: dict[str, int],
+        latent_patch_size: int,
+        temporal_modality_margin: int,
+        reset_spatial_ids: bool,
+        enable_fps_modulation: bool,
+        base_fps: float,
+        temporal_compression_factor: int,
+        include_end_of_generation_token: bool = False,
+    ) -> None:
+        self.transformer = transformer
+        self.scheduler = scheduler
+        self.special_tokens = special_tokens
+        self.latent_patch_size = latent_patch_size
+        self.temporal_modality_margin = temporal_modality_margin
+        self.reset_spatial_ids = reset_spatial_ids
+        self.enable_fps_modulation = enable_fps_modulation
+        self.base_fps = base_fps
+        self.temporal_compression_factor = temporal_compression_factor
+        self.include_end_of_generation_token = include_end_of_generation_token
+
+    def velocity(
+        self,
+        *,
+        flat_latent: torch.Tensor,
+        timestep: torch.Tensor,
+        guidance: float,
+        specs: list[Cosmos3VisionSpec],
+        cond_token_ids: list[int],
+        uncond_token_ids: list[int],
+        fps_per_item: list[float] | None = None,
+        sound_specs: list[Cosmos3SoundSpec] | None = None,
+        sound_fps_per_item: list[float] | None = None,
+        action_specs: list[Cosmos3ActionSpec] | None = None,
+        action_fps_per_item: list[float] | None = None,
+    ) -> torch.Tensor:
+        return cosmos3_get_cfg_velocity(
+            transformer=self.transformer,
+            flat_latent=flat_latent,
+            timestep=timestep,
+            guidance=guidance,
+            specs=specs,
+            cond_token_ids=cond_token_ids,
+            uncond_token_ids=uncond_token_ids,
+            special_tokens=self.special_tokens,
+            latent_patch_size=self.latent_patch_size,
+            temporal_modality_margin=self.temporal_modality_margin,
+            reset_spatial_ids=self.reset_spatial_ids,
+            enable_fps_modulation=self.enable_fps_modulation,
+            base_fps=self.base_fps,
+            temporal_compression_factor=self.temporal_compression_factor,
+            include_end_of_generation_token=self.include_end_of_generation_token,
+            fps_per_item=fps_per_item,
+            sound_specs=sound_specs,
+            sound_fps_per_item=sound_fps_per_item,
+            action_specs=action_specs,
+            action_fps_per_item=action_fps_per_item,
+        )
+
+    def denoise(
+        self,
+        *,
+        flat_latent: torch.Tensor,
+        timesteps: torch.Tensor,
+        guidance: float,
+        specs: list[Cosmos3VisionSpec],
+        cond_token_ids: list[int],
+        uncond_token_ids: list[int],
+        fps_per_item: list[float] | None = None,
+        progress_bar: Any | None = None,
+        sound_specs: list[Cosmos3SoundSpec] | None = None,
+        sound_fps_per_item: list[float] | None = None,
+        action_specs: list[Cosmos3ActionSpec] | None = None,
+        action_fps_per_item: list[float] | None = None,
+    ) -> torch.Tensor:
+        """Run the full UniPC denoise loop, returning the final flat latent.
+
+        For each timestep: compute the sequential-CFG velocity, then
+        ``scheduler.step(model_output=v, timestep, sample=latent.unsqueeze(0))``
+        (the framework steps with a leading batch axis), squeezing back to flat.
+        For t2vs the flat latent is ``[vision | sound]`` and the velocity covers
+        both; the scheduler steps the combined vector jointly.
+        """
+        latent = flat_latent
+        iterator = progress_bar(timesteps) if progress_bar is not None else timesteps
+        for t in iterator:
+            v_pred = self.velocity(
+                flat_latent=latent,
+                timestep=t.reshape(1),
+                guidance=guidance,
+                specs=specs,
+                cond_token_ids=cond_token_ids,
+                uncond_token_ids=uncond_token_ids,
+                fps_per_item=fps_per_item,
+                sound_specs=sound_specs,
+                sound_fps_per_item=sound_fps_per_item,
+                action_specs=action_specs,
+                action_fps_per_item=action_fps_per_item,
+            )
+            stepped = self.scheduler.step(
+                model_output=v_pred,
+                timestep=t,
+                sample=latent.unsqueeze(0),
+                return_dict=False,
+            )[0]
+            latent = stepped.squeeze(0)
+        return latent
+
+
+# ===========================================================================
+# Pipeline (ComposedPipelineBase)
+# ===========================================================================
+class Cosmos3OmniDiffusersPipeline(ComposedPipelineBase):
+    """Cosmos3 video generation pipeline (T2V / I2V / T2I).
+
+    Stage-based ``ComposedPipelineBase`` pipeline. The required modules
+    (``transformer`` / ``vae`` / ``scheduler`` / ``text_tokenizer``) are loaded
+    from the ``nvidia/Cosmos3-Nano`` checkpoint by the component loader. The
+    class name matches the checkpoint ``model_index.json`` ``_class_name`` so
+    the registry resolves it directly.
+
+    The denoise/CFG/VAE math is delegated to module-level helpers
+    (:func:`cosmos3_get_cfg_velocity`, :class:`Cosmos3DenoiseEngine`,
+    :func:`cosmos3_vae_encode` / :func:`cosmos3_vae_decode`) which are
+    framework-parity tested in ``tests/local_tests/cosmos3``.
+    """
+
+    is_video_pipeline = True
+    # ``vision_encoder`` / ``sound_tokenizer`` ship in the checkpoint but the
+    # video path does not need them; they are intentionally omitted here.
+    _required_config_modules = ["text_tokenizer", "vae", "transformer", "scheduler"]
+
+    # Engine-init flow_shift (T2V/I2V); T2I overrides to 3.0 per request.
+    _engine_init_flow_shift: float = 1.0
+    # Class-attribute defaults so ``__new__``-based unit tests can read these
+    # before ``initialize_pipeline`` runs.
+    scheduler: Any = None
+    _base_scheduler_config: Any = None
+    _current_flow_shift: float | None = None
+
+    @staticmethod
+    def _flow_scheduler_config(config: Any) -> dict[str, Any]:
+        """Coerce a loaded UniPC config to the framework's flow-matching setup.
+
+        The checkpoint ``scheduler_config.json`` carries diffusers-style fields
+        (``use_karras_sigmas=True``, ``sigma_min``/``sigma_max``, beta schedule)
+        that do not describe the framework sampler. The framework uses
+        ``FlowUniPCMultistepScheduler`` (pure flow matching: ``shift`` +
+        ``num_train_timesteps`` only). FastVideo's vendored UniPC checks
+        ``use_karras_sigmas`` *before* ``use_flow_sigmas``, so leaving karras on
+        builds diffusion-style sigmas and the denoise diverges to NaN. Force the
+        flow config here (parity-verified in ``test_cosmos3_scheduler_parity``).
+        """
+        cfg = dict(config)
+        cfg.update(
+            use_karras_sigmas=False,
+            use_exponential_sigmas=False,
+            use_beta_sigmas=False,
+            use_flow_sigmas=True,
+            prediction_type="flow_prediction",
+            predict_x0=True,
+            final_sigmas_type="zero",
+        )
+        return cfg
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None:
+        """Bind the loaded scheduler + snapshot its config so per-request
+        flow_shift rebuilds are cheap and the engine-init shift is applied."""
+        pipeline_config = fastvideo_args.pipeline_config
+        engine_shift = getattr(pipeline_config, "flow_shift", None)
+        if engine_shift is not None:
+            self._engine_init_flow_shift = float(engine_shift)
+        scheduler = self.get_module("scheduler")
+        if scheduler is not None:
+            # Rebuild from a flow-coerced config so the runtime scheduler matches
+            # the framework sampler (the loaded checkpoint config is diffusers-style).
+            flow_config = self._flow_scheduler_config(scheduler.config)
+            self.scheduler = UniPCMultistepScheduler.from_config(flow_config)
+            if isinstance(self.modules, dict):
+                self.modules["scheduler"] = self.scheduler
+            self._base_scheduler_config = self.scheduler.config
+            self._current_flow_shift = float(getattr(self.scheduler.config, "flow_shift", 1.0))
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        """Wire the Cosmos3 stages.
+
+        The whole text->latent->denoise->decode flow is custom (sequential CFG
+        with per-pass repacking), so a single :class:`Cosmos3DenoisingStage`
+        owns it. ``InputValidationStage`` runs first for the standard checks.
+        """
+        from fastvideo.pipelines.stages import InputValidationStage
+        from fastvideo.pipelines.stages.cosmos3_stages import Cosmos3DenoisingStage
+
+        self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
+        self.add_stage(
+            stage_name="denoising_stage",
+            stage=Cosmos3DenoisingStage(
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler"),
+                vae=self.get_module("vae"),
+                tokenizer=self.get_module("text_tokenizer"),
+                pipeline=self,
+            ),
+        )
+
+    # -- Scheduler control --------------------------------------------------
+
+    def _set_flow_shift(self, target_shift: float) -> None:
+        """Set UniPC ``flow_shift`` to ``target_shift``.
+
+        Lazily builds a default UniPC scheduler when called before
+        ``initialize_pipeline`` (e.g. the ``__new__``-based scheduler-parity
+        tests); otherwise rebuilds from the snapshotted base config only when
+        the target differs from the current shift.
+        """
+        target = float(target_shift)
+        base_config = self._base_scheduler_config
+        if base_config is None:
+            self.scheduler = UniPCMultistepScheduler(
+                num_train_timesteps=1000,
+                solver_order=2,
+                prediction_type="flow_prediction",
+                use_flow_sigmas=True,
+                flow_shift=target,
+            )
+            self._base_scheduler_config = self.scheduler.config
+            self._current_flow_shift = target
+            return
+        current = self._current_flow_shift
+        if current is not None and target == float(current):
+            return
+        self.scheduler = UniPCMultistepScheduler.from_config(base_config, flow_shift=target)
+        if isinstance(self.modules, dict):
+            self.modules["scheduler"] = self.scheduler
+        self._current_flow_shift = target
+
+    # -- Tokenization -------------------------------------------------------
+
+    def tokenize_caption(self, caption: str, *, is_video: bool = False, use_system_prompt: bool = False) -> list[int]:
+        return cosmos3_tokenize_caption(self.get_module("text_tokenizer"),
+                                        caption,
+                                        is_video=is_video,
+                                        use_system_prompt=use_system_prompt)
+
+
+# Entry point for the pipeline registry. The class name matches the checkpoint
+# ``model_index.json`` ``_class_name`` so ``resolve_pipeline_cls`` finds it.
+EntryClass = Cosmos3OmniDiffusersPipeline

fastvideo/pipelines/basic/cosmos3/presets.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 (Cosmos3-Nano) inference presets.
+
+Defaults track the official ``cosmos-framework`` ``sample_args`` for the video
+paths (``text2video`` / ``image2video``: guidance=6.0, num_steps=35, shift=10.0,
+fps=24, num_frames=189) and ``text2image`` (guidance=4.0, num_steps=50,
+shift=3.0). The default resolution is 16:9 at a VAE-aligned 704x1280 (spatial
+compression 16 -> 44x80 latent grid).
+"""
+from fastvideo.api.presets import InferencePreset, PresetStageSpec
+
+_DENOISE_STAGE = PresetStageSpec(
+    name="denoise",
+    kind="denoising",
+    description="Cosmos3 sequential-CFG UniPC denoising pass",
+    allowed_overrides=frozenset({
+        "num_inference_steps",
+        "guidance_scale",
+    }),
+)
+
+# Framework video negative prompt (Cosmos quality prompt).
+COSMOS3_VIDEO_NEGATIVE_PROMPT = (
+    "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, "
+    "over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, "
+    "underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, "
+    "jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, "
+    "fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. "
+    "Overall, the video is of poor quality.")
+
+COSMOS3_NANO = InferencePreset(
+    name="cosmos3_nano",
+    version=1,
+    model_family="cosmos3",
+    description="Cosmos3-Nano text-to-video",
+    workload_type="t2v",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults={
+        "height": 704,
+        "width": 1280,
+        "num_frames": 189,
+        "fps": 24,
+        "guidance_scale": 6.0,
+        "num_inference_steps": 35,
+        "negative_prompt": COSMOS3_VIDEO_NEGATIVE_PROMPT,
+    },
+)
+
+COSMOS3_NANO_I2V = InferencePreset(
+    name="cosmos3_nano_i2v",
+    version=1,
+    model_family="cosmos3",
+    description="Cosmos3-Nano image-to-video",
+    workload_type="i2v",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults={
+        "height": 704,
+        "width": 1280,
+        "num_frames": 189,
+        "fps": 24,
+        "guidance_scale": 6.0,
+        "num_inference_steps": 35,
+        "negative_prompt": COSMOS3_VIDEO_NEGATIVE_PROMPT,
+    },
+)
+
+COSMOS3_NANO_T2I = InferencePreset(
+    name="cosmos3_nano_t2i",
+    version=1,
+    model_family="cosmos3",
+    description="Cosmos3-Nano text-to-image",
+    workload_type="t2i",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults={
+        "height": 1024,
+        "width": 1024,
+        "num_frames": 1,
+        "fps": 24,
+        "guidance_scale": 4.0,
+        "num_inference_steps": 50,
+        "negative_prompt": "",
+    },
+)
+
+ALL_PRESETS = (COSMOS3_NANO, COSMOS3_NANO_I2V, COSMOS3_NANO_T2I)

fastvideo/pipelines/basic/cosmos3/sequence_packing.py

@@ -0,0 +1,549 @@
+# SPDX-License-Identifier: Apache-2.0
+"""FastVideo-native Cosmos3 sequence packing (video subset).
+
+Numerical-parity port of the official ``cosmos_framework`` data packer
+(``cosmos_framework.data.vfm.sequence_packing.pack_input_sequence``) restricted
+to the VIDEO generation path that the FastVideo Cosmos3 DiT consumes (T2V / I2V
+/ T2I). It builds, per sample, two splits:
+
+* a ``causal`` text split (prompt token ids, plus the trailing ``eos`` and
+  ``start_of_generation`` markers the framework appends when a generation
+  modality follows), and
+* a ``full`` vision split (VAE latent patch tokens).
+
+The 3D-MRoPE position ids ``[3, seq]`` are produced exactly like the framework:
+text tokens broadcast a single monotone id across the (t, h, w) axes, the
+temporal offset is bumped by ``temporal_modality_margin`` at the text->vision
+boundary, and vision tokens lay out a (T, H, W) grid with spatial ids reset per
+segment. Condition frames (I2V cond frame 0, T2I single conditioned frame, ...)
+are kept in the packed sequence and rope grid but excluded from the MSE-loss /
+timestep bookkeeping, mirroring the framework.
+
+The output ``Cosmos3PackedSequence`` maps 1:1 onto the
+``Cosmos3VFMTransformer.forward`` kwargs via :meth:`to_dit_kwargs`. This module
+is pure torch/python; it imports no diffusers/transformers model classes.
+
+Reference of record: ``cosmos_framework`` (NVIDIA), the parity oracle used by
+``tests/local_tests/cosmos3/test_cosmos3_packing_parity.py``.
+"""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+
+from fastvideo.models.dits.cosmos3 import (
+    compute_mrope_position_ids_text,
+    compute_mrope_position_ids_vision,
+)
+
+__all__ = [
+    "Cosmos3VisionItem",
+    "Cosmos3SampleInputs",
+    "Cosmos3PackedSequence",
+    "pack_cosmos3_video_sequence",
+]
+
+
+# ---------------------------------------------------------------------------
+# Inputs
+# ---------------------------------------------------------------------------
+@dataclass
+class Cosmos3VisionItem:
+    """One vision latent for a sample.
+
+    Args:
+        latent: VAE latent ``[C, T, H, W]`` (a leading batch axis of size 1 is
+            accepted and squeezed).
+        condition_frame_indexes: Latent-frame indices that are *conditioned*
+            (clean) rather than noisy. ``[]`` for T2V, ``[0]`` for I2V, and the
+            single conditioned frame for T2I.
+        fps: Frames-per-second for this clip; only used when
+            ``enable_fps_modulation`` is set.
+    """
+
+    latent: torch.Tensor
+    condition_frame_indexes: list[int] = field(default_factory=list)
+    fps: float | None = None
+
+
+@dataclass
+class Cosmos3SoundItem:
+    """One sound latent for a sample (t2vs).
+
+    Args:
+        latent: AVAE sound latent ``[C, T]`` (channels, temporal frames).
+        condition_frame_indexes: Latent-frame indices that are *conditioned*
+            (clean). ``[]`` for t2vs (all frames generated).
+        fps: Sound latent FPS (``sound_latent_fps``, e.g. 25); only used when
+            ``enable_fps_modulation`` is set.
+    """
+
+    latent: torch.Tensor
+    condition_frame_indexes: list[int] = field(default_factory=list)
+    fps: float | None = None
+
+
+@dataclass
+class Cosmos3ActionItem:
+    """One action latent for a sample (action-conditioned world model).
+
+    Args:
+        latent: Action latent ``[T, action_dim]`` (per-frame action vectors).
+        condition_frame_indexes: Frame indices kept clean (conditioning actions).
+        domain_id: Embodiment domain id (scalar / ``[1]``) for the
+            domain-aware action projection.
+        fps: Action FPS; only used when ``enable_fps_modulation`` is set.
+    """
+
+    latent: torch.Tensor
+    condition_frame_indexes: list[int] = field(default_factory=list)
+    domain_id: int = 0
+    fps: float | None = None
+
+
+@dataclass
+class Cosmos3SampleInputs:
+    """Per-sample packing inputs (text prompt + vision item, +sound, +action)."""
+
+    text_ids: list[int]
+    vision: Cosmos3VisionItem
+    timestep: float
+    sound: Cosmos3SoundItem | None = None
+    action: Cosmos3ActionItem | None = None
+
+
+# ---------------------------------------------------------------------------
+# Output
+# ---------------------------------------------------------------------------
+@dataclass
+class Cosmos3PackedSequence:
+    """Packed-sequence inputs consumed by ``Cosmos3VFMTransformer.forward``.
+
+    Field names mirror the framework ``PackedSequence`` (+ its ``vision``
+    ``ModalityData``) so the parity test can compare field-by-field.
+    """
+
+    # Sequence structure.
+    sample_lens: list[int]
+    split_lens: list[int]
+    attn_modes: list[str]
+    sequence_length: int
+    is_image_batch: bool
+
+    # Text modality.
+    text_ids: torch.Tensor
+    text_indexes: torch.Tensor
+    position_ids: torch.Tensor  # [3, sequence_length]
+
+    # Vision modality.
+    vision_tokens: list[torch.Tensor]
+    vision_token_shapes: list[tuple[int, int, int]]
+    vision_sequence_indexes: torch.Tensor
+    vision_timesteps: torch.Tensor
+    vision_mse_loss_indexes: torch.Tensor
+    vision_noisy_frame_indexes: list[torch.Tensor]
+    vision_condition_mask: list[torch.Tensor]
+    fps_vision: torch.Tensor | None = None
+
+    # Sound modality (t2vs); empty/None when no sound.
+    sound_tokens: list[torch.Tensor] = field(default_factory=list)
+    sound_token_shapes: list[tuple[int, int, int]] = field(default_factory=list)
+    sound_sequence_indexes: torch.Tensor | None = None
+    sound_timesteps: torch.Tensor | None = None
+    sound_mse_loss_indexes: torch.Tensor | None = None
+    sound_noisy_frame_indexes: list[torch.Tensor] = field(default_factory=list)
+    sound_condition_mask: list[torch.Tensor] = field(default_factory=list)
+    fps_sound: torch.Tensor | None = None
+
+    # Action modality (action-conditioned world model); empty/None when no action.
+    action_tokens: list[torch.Tensor] = field(default_factory=list)
+    action_token_shapes: list[tuple[int, ...]] = field(default_factory=list)
+    action_sequence_indexes: torch.Tensor | None = None
+    action_timesteps: torch.Tensor | None = None
+    action_mse_loss_indexes: torch.Tensor | None = None
+    action_noisy_frame_indexes: list[torch.Tensor] = field(default_factory=list)
+    action_condition_mask: list[torch.Tensor] = field(default_factory=list)
+    action_domain_id: list[torch.Tensor] = field(default_factory=list)
+
+    def to_dit_kwargs(self, device: torch.device | str | None = None) -> dict[str, Any]:
+        """Return the kwargs dict for ``Cosmos3VFMTransformer.forward``.
+
+        Packing is device-agnostic (ids/indexes/position-ids are built on CPU).
+        When ``device`` is given, every tensor input is moved to it so the DiT
+        forward runs on a single device (e.g. the model's GPU at inference).
+        """
+
+        def _mv(x: Any) -> Any:
+            return x.to(device) if (device is not None and torch.is_tensor(x)) else x
+
+        return dict(
+            text_ids=_mv(self.text_ids),
+            text_indexes=_mv(self.text_indexes),
+            position_ids=_mv(self.position_ids),
+            sequence_length=int(self.sequence_length),
+            split_lens=list(self.split_lens),
+            attn_modes=list(self.attn_modes),
+            vision_tokens=[_mv(t) for t in self.vision_tokens],
+            vision_token_shapes=list(self.vision_token_shapes),
+            vision_sequence_indexes=_mv(self.vision_sequence_indexes),
+            vision_timesteps=_mv(self.vision_timesteps),
+            vision_mse_loss_indexes=_mv(self.vision_mse_loss_indexes),
+            vision_noisy_frame_indexes=[_mv(t) for t in self.vision_noisy_frame_indexes],
+            fps_vision=self.fps_vision,
+            sound_tokens=[_mv(t) for t in self.sound_tokens],
+            sound_token_shapes=list(self.sound_token_shapes),
+            sound_sequence_indexes=_mv(self.sound_sequence_indexes),
+            sound_timesteps=_mv(self.sound_timesteps),
+            sound_mse_loss_indexes=_mv(self.sound_mse_loss_indexes),
+            sound_noisy_frame_indexes=[_mv(t) for t in self.sound_noisy_frame_indexes],
+            fps_sound=_mv(self.fps_sound),
+            action_tokens=[_mv(t) for t in self.action_tokens],
+            action_token_shapes=list(self.action_token_shapes),
+            action_sequence_indexes=_mv(self.action_sequence_indexes),
+            action_timesteps=_mv(self.action_timesteps),
+            action_mse_loss_indexes=_mv(self.action_mse_loss_indexes),
+            action_noisy_frame_indexes=[_mv(t) for t in self.action_noisy_frame_indexes],
+            action_domain_id=[_mv(t) for t in self.action_domain_id],
+        )
+
+
+# ---------------------------------------------------------------------------
+# Packing
+# ---------------------------------------------------------------------------
+def pack_cosmos3_video_sequence(
+    samples: list[Cosmos3SampleInputs],
+    special_tokens: dict[str, int],
+    *,
+    latent_patch_size: int = 2,
+    include_end_of_generation_token: bool = False,
+    temporal_modality_margin: int = 15_000,
+    reset_spatial_ids: bool = True,
+    enable_fps_modulation: bool = False,
+    base_fps: float = 24.0,
+    temporal_compression_factor: int = 4,
+    initial_mrope_temporal_offset: int | float = 0,
+) -> Cosmos3PackedSequence:
+    """Pack prompts + vision latents into the Cosmos3 DiT packed-sequence inputs.
+
+    Video subset of ``cosmos_framework`` ``pack_input_sequence`` under
+    ``unified_3d_mrope``: each sample is ``[causal text, full vision]``.
+
+    Args:
+        samples: Per-sample text prompt token ids + vision item + timestep.
+        special_tokens: Must contain ``eos_token_id`` and
+            ``start_of_generation`` (and ``end_of_generation`` if
+            ``include_end_of_generation_token``). ``bos_token_id`` is honored if
+            present (prepended) to match the framework.
+        latent_patch_size: Latent patch size used by the DiT.
+        include_end_of_generation_token: Append the framework's end-of-generation
+            marker after the vision split.
+        temporal_modality_margin: Temporal-offset bump applied at the
+            text->vision boundary (``unified_3d_mrope_temporal_modality_margin``).
+        reset_spatial_ids: Reset vision spatial ids to 0 per segment.
+        enable_fps_modulation: Use float, fps-scaled temporal positions.
+        base_fps: Base FPS used when ``enable_fps_modulation``.
+        temporal_compression_factor: VAE temporal compression factor.
+        initial_mrope_temporal_offset: Per-sample starting temporal offset.
+
+    Returns:
+        A :class:`Cosmos3PackedSequence`.
+    """
+    assert "eos_token_id" in special_tokens, "special_tokens must contain eos_token_id"
+    assert "start_of_generation" in special_tokens, "special_tokens must contain start_of_generation"
+    if latent_patch_size < 1:
+        raise ValueError(f"latent_patch_size must be >= 1, got {latent_patch_size}")
+
+    # Build-time accumulators (concatenated across samples).
+    sample_lens: list[int] = []
+    split_lens: list[int] = []
+    attn_modes: list[str] = []
+
+    text_ids: list[int] = []
+    text_indexes: list[int] = []
+    position_id_blocks: list[torch.Tensor] = []  # each [3, n]
+
+    vision_tokens: list[torch.Tensor] = []
+    vision_token_shapes: list[tuple[int, int, int]] = []
+    vision_sequence_indexes: list[int] = []
+    vision_timesteps: list[float] = []
+    vision_mse_loss_indexes: list[int] = []
+    vision_noisy_frame_indexes: list[torch.Tensor] = []
+    vision_condition_mask: list[torch.Tensor] = []
+    fps_values: list[float] = []
+
+    sound_tokens: list[torch.Tensor] = []
+    sound_token_shapes: list[tuple[int, int, int]] = []
+    sound_sequence_indexes: list[int] = []
+    sound_timesteps: list[float] = []
+    sound_mse_loss_indexes: list[int] = []
+    sound_noisy_frame_indexes: list[torch.Tensor] = []
+    sound_condition_mask: list[torch.Tensor] = []
+    sound_fps_values: list[float] = []
+
+    action_tokens: list[torch.Tensor] = []
+    action_token_shapes: list[tuple[int, ...]] = []
+    action_sequence_indexes: list[int] = []
+    action_timesteps: list[float] = []
+    action_mse_loss_indexes: list[int] = []
+    action_noisy_frame_indexes: list[torch.Tensor] = []
+    action_condition_mask: list[torch.Tensor] = []
+    action_domain_id: list[torch.Tensor] = []
+
+    curr = 0  # running position in the packed sequence
+    is_image_batch = True
+
+    for sample in samples:
+        temporal_offset: int | float = initial_mrope_temporal_offset
+        sample_len = 0
+
+        # ---- 1. Text split (causal) ----
+        if "bos_token_id" in special_tokens:
+            shifted_text_ids = [special_tokens["bos_token_id"], *sample.text_ids]
+        else:
+            shifted_text_ids = list(sample.text_ids)
+        # The video path always has a following generation modality, so the
+        # framework appends eos + start_of_generation.
+        shifted_text_ids = [*shifted_text_ids, special_tokens["eos_token_id"], special_tokens["start_of_generation"]]
+        text_split_len = len(shifted_text_ids)
+
+        text_ids.extend(shifted_text_ids)
+        text_indexes.extend(range(curr, curr + text_split_len))
+
+        text_mrope, temporal_offset = compute_mrope_position_ids_text(
+            num_tokens=text_split_len,
+            temporal_offset=int(temporal_offset),
+        )
+        position_id_blocks.append(text_mrope)
+
+        attn_modes.append("causal")
+        split_lens.append(text_split_len)
+        curr += text_split_len
+        sample_len += text_split_len
+
+        # End of text modality: bump temporal offset before vision.
+        temporal_offset += temporal_modality_margin
+        # Sound shares the vision temporal start (parallel temporal positions).
+        vision_start_temporal_offset = temporal_offset
+
+        # ---- 2. Vision split (full) ----
+        latent = sample.vision.latent
+        latent = latent.squeeze(0) if latent.dim() == 5 else latent  # [C, T, H, W]
+        _c, latent_t, latent_h, latent_w = latent.shape
+        patch_h = math.ceil(latent_h / latent_patch_size)
+        patch_w = math.ceil(latent_w / latent_patch_size)
+        num_vision_tokens = latent_t * patch_h * patch_w
+
+        vision_tokens.append(sample.vision.latent)
+        vision_token_shapes.append((latent_t, patch_h, patch_w))
+        vision_sequence_indexes.extend(range(curr, curr + num_vision_tokens))
+
+        condition_set = {idx for idx in sample.vision.condition_frame_indexes if 0 <= idx < latent_t}
+        cond_mask = torch.zeros((latent_t, 1, 1), device=latent.device, dtype=latent.dtype)
+        for frame_idx in condition_set:
+            cond_mask[frame_idx, 0, 0] = 1.0
+        vision_condition_mask.append(cond_mask)
+
+        noisy_frames = torch.tensor(
+            [idx for idx in range(latent_t) if idx not in condition_set],
+            device=latent.device,
+            dtype=torch.long,
+        )
+        vision_noisy_frame_indexes.append(noisy_frames)
+
+        # MSE-loss indices + per-token timesteps cover only the noisy frames.
+        frame_token_stride = patch_h * patch_w
+        for frame_idx in range(latent_t):
+            if frame_idx in condition_set:
+                continue
+            frame_start = curr + frame_idx * frame_token_stride
+            vision_mse_loss_indexes.extend(range(frame_start, frame_start + frame_token_stride))
+            vision_timesteps.extend([float(sample.timestep)] * frame_token_stride)
+
+        vision_fps = sample.vision.fps if enable_fps_modulation else None
+        if vision_fps is not None:
+            fps_values.append(float(vision_fps))
+        vision_mrope, temporal_offset = compute_mrope_position_ids_vision(
+            grid_t=latent_t,
+            grid_h=patch_h,
+            grid_w=patch_w,
+            temporal_offset=temporal_offset,
+            fps=vision_fps,
+            base_fps=base_fps,
+            temporal_compression_factor=temporal_compression_factor,
+            enable_fps_modulation=enable_fps_modulation,
+        )
+        position_id_blocks.append(vision_mrope)
+
+        curr += num_vision_tokens
+        sample_len += num_vision_tokens
+
+        # ---- 2a2. Action split: shares the vision "full" split ----
+        # Mirrors framework ``_pack_action_tokens``: action latent [T, D] -> T
+        # tokens (token shape (T,)), domain-aware, 3D-MRoPE at the vision temporal
+        # offset with ``start_frame_offset=1`` (parallel to vision; tcf=1; does
+        # not advance the offset).
+        action_split_len = 0
+        if sample.action is not None:
+            action_latent = sample.action.latent  # [T, D]
+            action_t = int(action_latent.shape[0])
+            action_split_len = action_t
+
+            action_tokens.append(action_latent)
+            action_token_shapes.append((action_t, ))
+            action_sequence_indexes.extend(range(curr, curr + action_t))
+            action_domain_id.append(torch.tensor([int(sample.action.domain_id)], dtype=torch.long))
+
+            a_cond_set = {idx for idx in sample.action.condition_frame_indexes if 0 <= idx < action_t}
+            a_cond_mask = torch.zeros((action_t, 1), device=action_latent.device, dtype=action_latent.dtype)
+            for fi in a_cond_set:
+                a_cond_mask[fi, 0] = 1.0
+            action_condition_mask.append(a_cond_mask)
+
+            a_noisy = torch.tensor([idx for idx in range(action_t) if idx not in a_cond_set],
+                                   device=action_latent.device,
+                                   dtype=torch.long)
+            action_noisy_frame_indexes.append(a_noisy)
+
+            for fi in range(action_t):
+                if fi in a_cond_set:
+                    continue
+                action_mse_loss_indexes.append(curr + fi)
+                action_timesteps.append(float(sample.timestep))
+
+            action_fps = sample.action.fps if enable_fps_modulation else None
+            action_mrope, _ = compute_mrope_position_ids_vision(
+                grid_t=action_t,
+                grid_h=1,
+                grid_w=1,
+                temporal_offset=vision_start_temporal_offset,
+                fps=action_fps,
+                base_fps=base_fps,
+                temporal_compression_factor=1,  # action is at frame rate
+                base_temporal_compression_factor=temporal_compression_factor,
+                enable_fps_modulation=enable_fps_modulation,
+                start_frame_offset=1,
+            )
+            position_id_blocks.append(action_mrope)
+            curr += action_t
+            sample_len += action_t
+
+        # ---- 2b. Sound split (t2vs): shares the vision "full" split ----
+        # Mirrors framework ``_pack_sound_tokens``: sound latent [C, T] -> T
+        # tokens (token shape (T,1,1)), packed right after vision, with 3D-MRoPE
+        # temporal positions starting at the vision temporal offset (parallel to
+        # vision, start_frame_offset=0, tcf=1) and NOT advancing it.
+        sound_split_len = 0
+        if sample.sound is not None:
+            sound_latent = sample.sound.latent
+            sound_latent = sound_latent.squeeze(0) if sound_latent.dim() == 3 else sound_latent  # [C, T]
+            _sc, sound_t = sound_latent.shape
+            sound_split_len = sound_t
+
+            sound_tokens.append(sound_latent)
+            sound_token_shapes.append((sound_t, 1, 1))
+            sound_sequence_indexes.extend(range(curr, curr + sound_t))
+
+            s_cond_set = {idx for idx in sample.sound.condition_frame_indexes if 0 <= idx < sound_t}
+            s_cond_mask = torch.zeros((sound_t, 1), device=sound_latent.device, dtype=sound_latent.dtype)
+            for fi in s_cond_set:
+                s_cond_mask[fi, 0] = 1.0
+            sound_condition_mask.append(s_cond_mask)
+
+            s_noisy = torch.tensor([idx for idx in range(sound_t) if idx not in s_cond_set],
+                                   device=sound_latent.device,
+                                   dtype=torch.long)
+            sound_noisy_frame_indexes.append(s_noisy)
+
+            for fi in range(sound_t):
+                if fi in s_cond_set:
+                    continue
+                sound_mse_loss_indexes.append(curr + fi)  # 1 token per sound frame
+                sound_timesteps.append(float(sample.timestep))
+
+            sound_fps = sample.sound.fps if enable_fps_modulation else None
+            if sound_fps is not None:
+                sound_fps_values.append(float(sound_fps))
+            sound_mrope, _ = compute_mrope_position_ids_vision(
+                grid_t=sound_t,
+                grid_h=1,
+                grid_w=1,
+                temporal_offset=vision_start_temporal_offset,
+                fps=sound_fps,
+                base_fps=base_fps,
+                temporal_compression_factor=1,  # sound latent already at sound_latent_fps
+                enable_fps_modulation=enable_fps_modulation,
+                start_frame_offset=0,
+            )
+            position_id_blocks.append(sound_mrope)
+            curr += sound_t
+            sample_len += sound_t
+
+        # ---- 3. Optional end-of-generation marker ----
+        eov_len = 0
+        if include_end_of_generation_token:
+            assert "end_of_generation" in special_tokens, ("special_tokens must contain end_of_generation when "
+                                                           "include_end_of_generation_token=True")
+            text_ids.append(special_tokens["end_of_generation"])
+            text_indexes.append(curr)
+            eov_dtype = torch.float32 if enable_fps_modulation else torch.long
+            eov_ids = torch.full((3, 1), temporal_offset, dtype=eov_dtype)
+            position_id_blocks.append(eov_ids)
+            temporal_offset += 1
+            curr += 1
+            eov_len = 1
+            sample_len += 1
+
+        # Vision + action + sound + any trailing eov marker share one "full" split.
+        attn_modes.append("full")
+        split_lens.append(num_vision_tokens + action_split_len + sound_split_len + eov_len)
+        sample_lens.append(sample_len)
+
+        if latent_t != 1:
+            is_image_batch = False
+
+    sequence_length = sum(sample_lens)
+
+    # position_ids: float iff any block is float (fps modulation path).
+    any_float = any(b.dtype.is_floating_point for b in position_id_blocks)
+    if any_float:
+        position_id_blocks = [b.to(torch.float32) for b in position_id_blocks]
+    position_ids = torch.cat(position_id_blocks, dim=1)  # [3, sequence_length]
+
+    timesteps_dtype = torch.float32
+    return Cosmos3PackedSequence(
+        sample_lens=sample_lens,
+        split_lens=split_lens,
+        attn_modes=attn_modes,
+        sequence_length=sequence_length,
+        is_image_batch=is_image_batch,
+        text_ids=torch.tensor(text_ids, dtype=torch.long),
+        text_indexes=torch.tensor(text_indexes, dtype=torch.long),
+        position_ids=position_ids,
+        vision_tokens=vision_tokens,
+        vision_token_shapes=vision_token_shapes,
+        vision_sequence_indexes=torch.tensor(vision_sequence_indexes, dtype=torch.long),
+        vision_timesteps=torch.tensor(vision_timesteps, dtype=timesteps_dtype),
+        vision_mse_loss_indexes=torch.tensor(vision_mse_loss_indexes, dtype=torch.long),
+        vision_noisy_frame_indexes=vision_noisy_frame_indexes,
+        vision_condition_mask=vision_condition_mask,
+        fps_vision=(torch.tensor(fps_values, dtype=torch.float32) if fps_values else None),
+        sound_tokens=sound_tokens,
+        sound_token_shapes=sound_token_shapes,
+        sound_sequence_indexes=(torch.tensor(sound_sequence_indexes, dtype=torch.long) if sound_tokens else None),
+        sound_timesteps=(torch.tensor(sound_timesteps, dtype=timesteps_dtype) if sound_tokens else None),
+        sound_mse_loss_indexes=(torch.tensor(sound_mse_loss_indexes, dtype=torch.long) if sound_tokens else None),
+        sound_noisy_frame_indexes=sound_noisy_frame_indexes,
+        sound_condition_mask=sound_condition_mask,
+        fps_sound=(torch.tensor(sound_fps_values, dtype=torch.float32) if sound_fps_values else None),
+        action_tokens=action_tokens,
+        action_token_shapes=action_token_shapes,
+        action_sequence_indexes=(torch.tensor(action_sequence_indexes, dtype=torch.long) if action_tokens else None),
+        action_timesteps=(torch.tensor(action_timesteps, dtype=timesteps_dtype) if action_tokens else None),
+        action_mse_loss_indexes=(torch.tensor(action_mse_loss_indexes, dtype=torch.long) if action_tokens else None),
+        action_noisy_frame_indexes=action_noisy_frame_indexes,
+        action_condition_mask=action_condition_mask,
+        action_domain_id=action_domain_id,
+    )

fastvideo/pipelines/basic/flux_2/__init__.py

@@ -1,7 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Flux2 pipeline module."""
-
-from fastvideo.pipelines.basic.flux_2.flux_2_pipeline import Flux2Pipeline
-from fastvideo.pipelines.basic.flux_2.flux_2_klein_pipeline import Flux2KleinPipeline
-
-__all__ = ["Flux2Pipeline", "Flux2KleinPipeline"]

fastvideo/pipelines/basic/flux_2/flux_2_klein_pipeline.py

@@ -1,17 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Copied and adapted from: https://github.com/sglang-ai/sglang
-"""
-Flux2 Klein image generation pipeline (distilled, 4-step, no guidance).
-"""
-
-from fastvideo.configs.pipelines.flux_2 import Flux2KleinPipelineConfig
-from fastvideo.pipelines.basic.flux_2.flux_2_pipeline import Flux2Pipeline
-
-
-class Flux2KleinPipeline(Flux2Pipeline):
-    """Flux2 Klein image diffusion pipeline (distilled, 4-step, no guidance)."""
-
-    pipeline_config_cls: type[Flux2KleinPipelineConfig] = Flux2KleinPipelineConfig
-
-
-EntryClass = Flux2KleinPipeline

fastvideo/pipelines/basic/flux_2/flux_2_latent_preparation.py

@@ -1,138 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-Flux2 latent preparation stage using packed 2x2 layout.
-
-Flux2 uses packed latents: transformer sees 128 channels (32*4) with half
-spatial resolution; after denoising we unpatchify to 32 channels and full
-spatial for VAE decode. This stage prepares (B, 128, T, H//2, W//2).
-"""
-
-import torch
-from diffusers.utils.torch_utils import randn_tensor
-
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.pipelines.stages.latent_preparation import LatentPreparationStage
-
-
-class Flux2LatentPreparationStage(LatentPreparationStage):
-    """
-    Latent preparation for Flux2: packed layout with half spatial dimensions.
-
-    Matches diffusers Flux2Pipeline.prepare_latents: shape is
-    (B, num_channels_latents, T, H_latent//2, W_latent//2) so the transformer
-    sees 128 channels and half spatial; after denoising we unpatchify to
-    (B, 32, H_latent, W_latent) before VAE.
-    """
-
-    def forward(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ) -> ForwardBatch:
-        """Prepare latents with Flux2 packed half-spatial shape."""
-        from fastvideo.distributed import get_local_torch_device
-
-        latent_num_frames = None
-        if hasattr(self, "adjust_video_length"):
-            latent_num_frames = self.adjust_video_length(batch, fastvideo_args)
-
-        if not batch.prompt_embeds:
-            if batch.keyboard_cond is not None:
-                batch_size = batch.keyboard_cond.shape[0]
-            elif batch.mouse_cond is not None:
-                batch_size = batch.mouse_cond.shape[0]
-            elif batch.image_embeds:
-                batch_size = batch.image_embeds[0].shape[0]
-            else:
-                batch_size = 1
-        elif isinstance(batch.prompt, list):
-            batch_size = len(batch.prompt)
-        elif batch.prompt is not None:
-            batch_size = 1
-        else:
-            batch_size = batch.prompt_embeds[0].shape[0]
-
-        batch_size *= batch.num_videos_per_prompt
-
-        if not batch.prompt_embeds:
-            transformer_dtype = next(self.transformer.parameters()).dtype
-            device = get_local_torch_device()
-            dummy_prompt = torch.zeros(
-                batch_size,
-                0,
-                self.transformer.hidden_size,
-                device=device,
-                dtype=transformer_dtype,
-            )
-            batch.prompt_embeds = [dummy_prompt]
-            batch.negative_prompt_embeds = []
-            batch.do_classifier_free_guidance = False
-
-        dtype = batch.prompt_embeds[0].dtype
-        device = get_local_torch_device()
-        generator = batch.generator
-        latents = batch.latents
-        num_frames = (latent_num_frames if latent_num_frames is not None else batch.num_frames)
-        height = batch.height
-        width = batch.width
-
-        if height is None or width is None:
-            raise ValueError("Height and width must be provided")
-
-        vae_arch = fastvideo_args.pipeline_config.vae_config.arch_config
-        scale = vae_arch.spatial_compression_ratio
-        # Flux2 packed: half spatial (2x2 patch packing)
-        latent_h = (height // scale) // 2
-        latent_w = (width // scale) // 2
-
-        if self.use_btchw_layout:
-            shape = (
-                batch_size,
-                num_frames,
-                self.transformer.num_channels_latents,
-                latent_h,
-                latent_w,
-            )
-            bcthw_shape = tuple(shape[i] for i in [0, 2, 1, 3, 4])
-        else:
-            shape = (
-                batch_size,
-                self.transformer.num_channels_latents,
-                num_frames,
-                latent_h,
-                latent_w,
-            )
-            bcthw_shape = shape
-
-        if isinstance(generator, list) and len(generator) != batch_size:
-            raise ValueError(f"You have passed a list of generators of length {len(generator)}, "
-                             f"but requested an effective batch size of {batch_size}.")
-
-        if latents is None:
-            latents = randn_tensor(
-                shape,
-                generator=generator,
-                device=device,
-                dtype=dtype,
-            )
-            if hasattr(self.scheduler, "init_noise_sigma"):
-                latents = latents * self.scheduler.init_noise_sigma
-        else:
-            latents = latents.to(device)
-            is_longcat_refine = (batch.refine_from is not None or batch.stage1_video is not None)
-            if (not is_longcat_refine) and hasattr(self.scheduler, "init_noise_sigma"):
-                latents = latents * self.scheduler.init_noise_sigma
-
-        batch.latents = latents
-        batch.raw_latent_shape = bcthw_shape
-        latent_ids = torch.cartesian_prod(
-            torch.arange(num_frames, device=device),
-            torch.arange(latent_h, device=device),
-            torch.arange(latent_w, device=device),
-            torch.arange(1, device=device),
-        )
-        batch.extra["flux2_img_ids"] = latent_ids.unsqueeze(0).expand(batch_size, -1, -1)
-        # Flux2 mu depends on image_seq_len; use packed spatial size
-        batch.n_tokens = latent_h * latent_w
-        return batch

fastvideo/pipelines/basic/flux_2/flux_2_pipeline.py

@@ -1,95 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Copied and adapted from: https://github.com/sglang-ai/sglang
-"""
-Flux2 image generation pipeline implementation.
-
-This module contains an implementation of the Flux2 image diffusion pipeline
-using the modular pipeline architecture.
-"""
-
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.logger import init_logger
-from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
-from fastvideo.pipelines.basic.flux_2.flux_2_latent_preparation import (
-    Flux2LatentPreparationStage, )
-from fastvideo.pipelines.basic.flux_2.flux_2_timestep_preparation import (
-    Flux2TimestepPreparationStage, )
-from fastvideo.pipelines.basic.flux_2.flux_2_text_encoding import (
-    Flux2TextEncodingStage, )
-from fastvideo.pipelines.stages import (
-    ConditioningStage,
-    DecodingStage,
-    DenoisingStage,
-    InputValidationStage,
-)
-
-logger = init_logger(__name__)
-
-
-class Flux2Pipeline(LoRAPipeline, ComposedPipelineBase):
-    """
-    Flux2 image diffusion pipeline with LoRA support.
-    """
-
-    _required_config_modules = [
-        "text_encoder",
-        "tokenizer",
-        "vae",
-        "transformer",
-        "scheduler",
-    ]
-
-    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
-        """Set up pipeline stages with proper dependency injection."""
-
-        self.add_stage(
-            stage_name="input_validation_stage",
-            stage=InputValidationStage(),
-        )
-
-        self.add_stage(
-            stage_name="prompt_encoding_stage",
-            stage=Flux2TextEncodingStage(
-                text_encoders=[self.get_module("text_encoder")],
-                tokenizers=[self.get_module("tokenizer")],
-            ),
-        )
-
-        self.add_stage(
-            stage_name="conditioning_stage",
-            stage=ConditioningStage(),
-        )
-
-        self.add_stage(
-            stage_name="latent_preparation_stage",
-            stage=Flux2LatentPreparationStage(
-                scheduler=self.get_module("scheduler"),
-                transformer=self.get_module("transformer", None),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="timestep_preparation_stage",
-            stage=Flux2TimestepPreparationStage(scheduler=self.get_module("scheduler"), ),
-        )
-
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=DenoisingStage(
-                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=DecodingStage(
-                vae=self.get_module("vae"),
-                pipeline=self,
-            ),
-        )
-
-
-EntryClass = Flux2Pipeline

fastvideo/pipelines/basic/flux_2/flux_2_text_encoding.py

@@ -1,161 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Flux2 text encoding stages."""
-from __future__ import annotations
-
-from typing import Any
-
-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.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.pipelines.stages.text_encoding import TextEncodingStage
-
-FLUX2_SYSTEM_MESSAGE = ("You are an AI that reasons about image descriptions. You give structured "
-                        "responses focusing on object relationships, object\nattribution and actions "
-                        "without speculation.")
-
-
-def _format_flux2_full_input(prompts: list[str], system_message: str) -> list[list[dict[str, Any]]]:
-    return [[
-        {
-            "role": "system",
-            "content": [{
-                "type": "text",
-                "text": system_message
-            }],
-        },
-        {
-            "role": "user",
-            "content": [{
-                "type": "text",
-                "text": prompt.replace("[IMG]", "")
-            }],
-        },
-    ] for prompt in prompts]
-
-
-def _prepare_flux2_text_ids(prompt_embeds: torch.Tensor) -> torch.Tensor:
-    batch_size, seq_len, _ = prompt_embeds.shape
-    text_ids = torch.cartesian_prod(
-        torch.arange(1, device=prompt_embeds.device),
-        torch.arange(1, device=prompt_embeds.device),
-        torch.arange(1, device=prompt_embeds.device),
-        torch.arange(seq_len, device=prompt_embeds.device),
-    )
-    return text_ids.unsqueeze(0).expand(batch_size, -1, -1)
-
-
-class Flux2TextEncodingStage(TextEncodingStage):
-    """Text encoding for Flux2 full and Klein variants."""
-
-    def _uses_embedded_guidance(self, fastvideo_args: FastVideoArgs) -> bool:
-        return getattr(fastvideo_args.pipeline_config, "embedded_cfg_scale", None) is not None
-
-    @torch.no_grad()
-    def forward(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ) -> ForwardBatch:
-        if self._uses_embedded_guidance(fastvideo_args):
-            batch.do_classifier_free_guidance = False
-            batch.negative_prompt_embeds = []
-
-        if batch.prompt_embeds is not None and len(batch.prompt_embeds) > 0:
-            if "flux2_txt_ids" not in batch.extra:
-                batch.extra["flux2_txt_ids"] = _prepare_flux2_text_ids(batch.prompt_embeds[0])
-            return batch
-
-        if getattr(fastvideo_args.pipeline_config, "flux2_text_encoder_type", "") != "mistral3":
-            return super().forward(batch, fastvideo_args)
-
-        assert batch.prompt is not None
-        prompt_embeds, attention_mask = self.encode_flux2_full_text(
-            batch.prompt,
-            fastvideo_args,
-            max_length=batch.max_sequence_length,
-        )
-        batch.prompt_embeds.append(prompt_embeds)
-        batch.extra["flux2_txt_ids"] = _prepare_flux2_text_ids(prompt_embeds)
-        if batch.prompt_attention_mask is not None:
-            batch.prompt_attention_mask.append(attention_mask)
-        return batch
-
-    @torch.no_grad()
-    def encode_flux2_full_text(
-        self,
-        text: str | list[str],
-        fastvideo_args: FastVideoArgs,
-        max_length: int | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        tokenizer = self.tokenizers[0]
-        text_encoder = self.text_encoders[0]
-        encoder_config = fastvideo_args.pipeline_config.text_encoder_configs[0]
-        arch_config = encoder_config.arch_config
-
-        prompts = [text] if isinstance(text, str) else text
-        max_sequence_length = max_length or getattr(arch_config, "text_len", 512) or 512
-        hidden_state_layers = getattr(
-            fastvideo_args.pipeline_config,
-            "text_encoder_out_layers",
-            (10, 20, 30),
-        )
-        system_message = getattr(
-            fastvideo_args.pipeline_config,
-            "flux2_system_message",
-            FLUX2_SYSTEM_MESSAGE,
-        )
-
-        messages = _format_flux2_full_input(prompts, system_message)
-        inputs = tokenizer.apply_chat_template(
-            messages,
-            add_generation_prompt=False,
-            tokenize=True,
-            return_dict=True,
-            return_tensors="pt",
-            padding="max_length",
-            truncation=True,
-            max_length=max_sequence_length,
-        )
-
-        try:
-            encoder_device = next(text_encoder.parameters()).device
-        except StopIteration:
-            encoder_device = get_local_torch_device()
-        encoder_dtype = getattr(text_encoder, "dtype", None)
-
-        input_ids = inputs["input_ids"].to(encoder_device)
-        attention_mask = inputs["attention_mask"].to(encoder_device)
-
-        forward_kwargs: dict[str, Any] = {
-            "input_ids": input_ids,
-            "attention_mask": attention_mask,
-            "output_hidden_states": True,
-            "use_cache": False,
-        }
-        if "pixel_values" in inputs:
-            forward_kwargs["pixel_values"] = inputs["pixel_values"].to(
-                device=encoder_device,
-                dtype=encoder_dtype or torch.bfloat16,
-            )
-        if "image_sizes" in inputs:
-            forward_kwargs["image_sizes"] = inputs["image_sizes"].to(encoder_device)
-
-        with set_forward_context(current_timestep=0, attn_metadata=None):
-            outputs = text_encoder(**forward_kwargs)
-
-        if outputs.hidden_states is None:
-            raise ValueError("Full Flux2 requires output_hidden_states=True from text encoder")
-
-        stacked = torch.stack([outputs.hidden_states[k] for k in hidden_state_layers], dim=1)
-        if encoder_dtype is not None:
-            stacked = stacked.to(dtype=encoder_dtype)
-        batch_size, num_layers, seq_len, hidden_dim = stacked.shape
-        prompt_embeds = stacked.permute(0, 2, 1, 3).reshape(
-            batch_size,
-            seq_len,
-            num_layers * hidden_dim,
-        )
-        return prompt_embeds, attention_mask

fastvideo/pipelines/basic/flux_2/flux_2_timestep_preparation.py

@@ -1,111 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Flux2-specific timestep preparation."""
-
-import inspect
-
-import numpy as np
-import torch
-
-from fastvideo.distributed import get_local_torch_device
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.pipelines.stages.timestep_preparation import TimestepPreparationStage
-
-
-def compute_empirical_mu(image_seq_len: int, num_steps: int) -> float:
-    """
-    Resolution-dependent mu for Flux2 flow-match scheduler.
-    From Black Forest Labs flux2 official repo: sampling.compute_empirical_mu.
-    """
-    a1, b1 = 8.73809524e-05, 1.89833333
-    a2, b2 = 0.00016927, 0.45666666
-
-    if image_seq_len > 4300:
-        return float(a2 * image_seq_len + b2)
-
-    m_200 = a2 * image_seq_len + b2
-    m_10 = a1 * image_seq_len + b1
-    a = (m_200 - m_10) / 190.0
-    b = m_200 - 200.0 * a
-    return float(a * num_steps + b)
-
-
-class Flux2TimestepPreparationStage(TimestepPreparationStage):
-    """Flux2 timestep preparation matching the Diffusers Flux2 schedule."""
-
-    def forward(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ) -> ForwardBatch:
-        scheduler = self.scheduler
-        device = get_local_torch_device()
-        num_inference_steps = batch.num_inference_steps
-        timesteps = batch.timesteps
-        sigmas = batch.sigmas
-        n_tokens = batch.n_tokens
-
-        extra_set_timesteps_kwargs = {}
-        if n_tokens is not None and "n_tokens" in inspect.signature(scheduler.set_timesteps).parameters:
-            extra_set_timesteps_kwargs["n_tokens"] = n_tokens
-
-        # Flux2/BFL: Diffusers' Flux2 pipeline passes a custom sigma grid and
-        # always supplies the resolution-dependent mu when the scheduler accepts
-        # it.
-        scheduler_config = getattr(scheduler, "config", None)
-        use_flow_sigmas = (getattr(scheduler_config, "use_flow_sigmas", False) if scheduler_config else False)
-        if timesteps is None and sigmas is None and not use_flow_sigmas:
-            sigmas = np.linspace(1.0, 1.0 / num_inference_steps, num_inference_steps)
-
-        if "mu" in inspect.signature(scheduler.set_timesteps).parameters:
-            if batch.n_tokens is not None:
-                image_seq_len = batch.n_tokens
-            else:
-                h = (batch.height if isinstance(batch.height, int) else (batch.height[0] if batch.height else None))
-                w = (batch.width if isinstance(batch.width, int) else (batch.width[0] if batch.width else None))
-                vae_config = getattr(fastvideo_args.pipeline_config, "vae_config", None)
-                if vae_config is not None:
-                    arch = getattr(vae_config, "arch_config", None)
-                    scale = (getattr(arch, "spatial_compression_ratio", 8) if arch else 8)
-                else:
-                    scale = 8
-                image_seq_len = ((h // scale) * (w // scale) if h is not None and w is not None else 256)
-            extra_set_timesteps_kwargs["mu"] = compute_empirical_mu(image_seq_len, num_inference_steps)
-
-        if timesteps is not None and sigmas is not None:
-            raise ValueError("Only one of `timesteps` or `sigmas` can be passed. "
-                             "Please choose one to set custom values")
-
-        if timesteps is not None:
-            accepts_timesteps = "timesteps" in inspect.signature(scheduler.set_timesteps).parameters
-            if not accepts_timesteps:
-                raise ValueError(f"The current scheduler class {scheduler.__class__}'s "
-                                 f"`set_timesteps` does not support custom timestep schedules.")
-            timesteps_for_scheduler = (timesteps.cpu() if isinstance(timesteps, torch.Tensor) else timesteps)
-            scheduler.set_timesteps(
-                timesteps=timesteps_for_scheduler,
-                device=device,
-                **extra_set_timesteps_kwargs,
-            )
-            timesteps = scheduler.timesteps
-        elif sigmas is not None:
-            accept_sigmas = "sigmas" in inspect.signature(scheduler.set_timesteps).parameters
-            if not accept_sigmas:
-                raise ValueError(f"The current scheduler class {scheduler.__class__}'s "
-                                 f"`set_timesteps` does not support custom sigmas schedules.")
-            scheduler.set_timesteps(
-                sigmas=sigmas,
-                device=device,
-                **extra_set_timesteps_kwargs,
-            )
-            timesteps = scheduler.timesteps
-        else:
-            scheduler.set_timesteps(
-                num_inference_steps,
-                device=device,
-                **extra_set_timesteps_kwargs,
-            )
-            timesteps = scheduler.timesteps
-
-        batch.timesteps = timesteps
-        return batch

fastvideo/pipelines/basic/flux_2/presets.py

@@ -1,75 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Flux2 model family pipeline presets.
-
-Each preset is a named inference preset that declares the user-facing
-stage topology, default sampling values, and which per-stage overrides
-are allowed.  Presets are registered explicitly from
-:func:`fastvideo.registry._register_presets`.
-"""
-from fastvideo.api.presets import InferencePreset, PresetStageSpec
-
-_DENOISE_STAGE = PresetStageSpec(
-    name="denoise",
-    kind="denoising",
-    description="Main denoising pass",
-    allowed_overrides=frozenset({
-        "num_inference_steps",
-        "guidance_scale",
-    }),
-)
-
-FLUX2_DEV = InferencePreset(
-    name="flux2_dev",
-    version=1,
-    model_family="flux2",
-    description="Flux2 full T2I with embedded guidance",
-    workload_type="t2i",
-    stage_schemas=(_DENOISE_STAGE, ),
-    defaults={
-        "height": 1024,
-        "width": 1024,
-        "num_frames": 1,
-        "fps": 1,
-        "seed": 0,
-        "guidance_scale": 4.0,
-        "num_inference_steps": 50,
-    },
-)
-
-FLUX2_KLEIN_4B = InferencePreset(
-    name="flux2_klein_4b",
-    version=1,
-    model_family="flux2",
-    description="Flux2 Klein 4B (distilled, 4-step, no guidance)",
-    workload_type="t2i",
-    stage_schemas=(_DENOISE_STAGE, ),
-    defaults={
-        "height": 1024,
-        "width": 1024,
-        "num_frames": 1,
-        "fps": 1,
-        "seed": 0,
-        "guidance_scale": 1.0,
-        "num_inference_steps": 4,
-    },
-)
-
-FLUX2_KLEIN_9B = InferencePreset(
-    name="flux2_klein_9b",
-    version=1,
-    model_family="flux2",
-    description="Flux2 Klein 9B (distilled, 4-step, no guidance)",
-    workload_type="t2i",
-    stage_schemas=(_DENOISE_STAGE, ),
-    defaults={
-        "height": 1024,
-        "width": 1024,
-        "num_frames": 1,
-        "fps": 1,
-        "seed": 0,
-        "guidance_scale": 1.0,
-        "num_inference_steps": 4,
-    },
-)
-
-ALL_PRESETS = (FLUX2_DEV, FLUX2_KLEIN_4B, FLUX2_KLEIN_9B)

fastvideo/pipelines/basic/wan/lucy_edit_pipeline.py

@@ -1,80 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Lucy Edit video editing pipeline.
-
-Lucy Edit uses a Wan2.2 5B transformer with an input video latent appended to
-the noisy latent channels. The stage topology is therefore closest to Wan V2V,
-but the model repo does not include CLIP image-encoder components.
-"""
-
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.logger import init_logger
-from fastvideo.pipelines.basic.wan.wan_v2v_pipeline import WanVideoToVideoPipeline
-from fastvideo.pipelines.stages import (
-    ConditioningStage,
-    DecodingStage,
-    DenoisingStage,
-    InputValidationStage,
-    LatentPreparationStage,
-    TextEncodingStage,
-    TimestepPreparationStage,
-    VideoVAEEncodingStage,
-)
-
-logger = init_logger(__name__)
-
-
-class LucyEditPipeline(WanVideoToVideoPipeline):
-    """FastVideo pipeline for decart-ai/Lucy-Edit-Dev."""
-
-    _required_config_modules = [
-        "text_encoder",
-        "tokenizer",
-        "vae",
-        "transformer",
-        "scheduler",
-    ]
-
-    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
-        self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
-
-        self.add_stage(
-            stage_name="prompt_encoding_stage",
-            stage=TextEncodingStage(
-                text_encoders=[self.get_module("text_encoder")],
-                tokenizers=[self.get_module("tokenizer")],
-            ),
-        )
-
-        self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
-
-        self.add_stage(
-            stage_name="timestep_preparation_stage",
-            stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
-        )
-
-        self.add_stage(
-            stage_name="latent_preparation_stage",
-            stage=LatentPreparationStage(
-                scheduler=self.get_module("scheduler"),
-                transformer=self.get_module("transformer"),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="video_latent_preparation_stage",
-            stage=VideoVAEEncodingStage(vae=self.get_module("vae")),
-        )
-
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=DenoisingStage(
-                transformer=self.get_module("transformer"),
-                transformer_2=self.get_module("transformer_2"),
-                scheduler=self.get_module("scheduler"),
-            ),
-        )
-
-        self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")))
-
-
-EntryClass = LucyEditPipeline

fastvideo/pipelines/basic/wan/presets.py

@@ -268,24 +268,6 @@ FAST_WAN_2_2_TI2V_5B = InferencePreset(
     },
 )
 
-LUCY_EDIT_DEV = InferencePreset(
-    name="lucy_edit_dev",
-    version=1,
-    model_family="wan",
-    description="Lucy Edit Dev 5B video editing",
-    workload_type="t2v",
-    stage_schemas=(_DENOISE_STAGE, ),
-    defaults={
-        "height": 480,
-        "width": 832,
-        "num_frames": 81,
-        "fps": 24,
-        "guidance_scale": 5.0,
-        "num_inference_steps": 50,
-        "negative_prompt": "",
-    },
-)
-
 # -------------------------------------------------------------------
 # Self-Forcing (causal) presets
 # -------------------------------------------------------------------
@@ -359,7 +341,6 @@ ALL_PRESETS = (
     FAST_WAN_T2V_480P,
     WAN_2_2_TI2V_5B,
     FAST_WAN_2_2_TI2V_5B,
-    LUCY_EDIT_DEV,
     SF_WAN_T2V_1_3B,
     SF_WAN_2_2_T2V_A14B,
     SF_WAN_2_2_I2V_A14B,

fastvideo/pipelines/composed_pipeline_base.py

@@ -380,8 +380,6 @@ class ComposedPipelineBase(ABC):
         model_index.pop("boundary_ratio", None)
         # used by Wan2.2 ti2v
         model_index.pop("expand_timesteps", None)
-        # HF metadata (e.g. Flux2 Klein is_distilled); not a loadable module
-        model_index.pop("is_distilled", None)
 
         # some sanity checks
         assert len(model_index) > 1, "model_index.json must contain at least one pipeline module"

fastvideo/pipelines/stages/cosmos3_stages.py

@@ -0,0 +1,331 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 video denoising stage.
+
+The Cosmos3 video path is monolithic by design: each CFG pass repacks the whole
+text+vision sequence (the conditional pass carries prompt tokens, the
+unconditional pass carries negative-prompt tokens), so the standard
+encode/condition/denoise/decode stage split does not apply. This single stage
+owns the full flow, delegating the framework-parity-tested math to
+``fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline``:
+
+  1. resolve mode (T2I / I2V / T2V) + per-mode defaults, set ``flow_shift``;
+  2. tokenize the prompt + negative prompt with the Qwen2 chat template;
+  3. VAE-encode the conditioning frame(s) for I2V / T2I (kept clean), build the
+     initial noise (clean condition frames + pure noise elsewhere);
+  4. run the UniPC denoise loop with sequential CFG
+     (``Cosmos3DenoiseEngine.denoise``);
+  5. VAE-decode + ``(1 + x) / 2`` clamp to [0, 1].
+
+This mirrors the framework ``Cosmos3OmniDiffusersPipeline.__call__``.
+"""
+from __future__ import annotations
+
+import os
+import weakref
+from typing import Any
+
+import torch
+from diffusers.utils.torch_utils import randn_tensor
+from tqdm import tqdm
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import (
+    Cosmos3DenoiseEngine,
+    Cosmos3SoundSpec,
+    Cosmos3VisionSpec,
+    _VaeNorm,
+    cosmos3_special_tokens,
+    cosmos3_tokenize_caption,
+    cosmos3_vae_decode,
+    cosmos3_vae_encode,
+)
+from fastvideo.pipelines.basic.cosmos3.presets import (
+    COSMOS3_VIDEO_NEGATIVE_PROMPT, )
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+
+logger = init_logger(__name__)
+
+
+class Cosmos3DenoisingStage(PipelineStage):
+    """Full Cosmos3 video denoise: tokenize + encode + denoise + decode."""
+
+    def __init__(self, *, transformer, scheduler, vae, tokenizer, pipeline=None) -> None:
+        self.transformer = transformer
+        self.scheduler = scheduler
+        self.vae = vae
+        self.tokenizer = tokenizer
+        self.pipeline = weakref.ref(pipeline) if pipeline is not None else None
+
+    # ------------------------------------------------------------------
+    # Geometry helpers
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _latent_frames(num_frames: int, temporal_factor: int) -> int:
+        return (int(num_frames) - 1) // int(temporal_factor) + 1
+
+    @staticmethod
+    def _flow_shift_for_resolution(height: int, width: int) -> float:
+        """UniPC ``flow_shift`` for a given pixel resolution.
+
+        Mirrors the framework's ``_RESOLUTION_SHIFT_DEFAULTS`` (8B VLM backbone,
+        which Cosmos3-Nano uses): the shift is keyed by the named resolution
+        bucket the (H, W) belongs to, regardless of task (T2V/I2V/T2I):
+
+            "256" -> 3.0,  "480" -> 5.0,  "704"/"720"/"768" -> 10.0
+
+        We invert the framework's ``{IMAGE,VIDEO}_RES_SIZE_INFO`` tables by the
+        longest side: <=320 is the 256 bucket, 640-832 the 480 bucket, and
+        960-1360 the 704/720/768 buckets.
+        """
+        long_side = max(int(height), int(width))
+        if long_side <= 480:
+            return 3.0
+        if long_side <= 896:
+            return 5.0
+        return 10.0
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        pipeline_config = fastvideo_args.pipeline_config
+        arch = pipeline_config.dit_config.arch_config
+        device = self.transformer.embed_tokens.weight.device
+        dtype = self.transformer.embed_tokens.weight.dtype
+
+        num_frames = int(batch.num_frames) if batch.num_frames is not None else 1
+        height = int(batch.height)
+        width = int(batch.width)
+        fps = float(batch.fps) if batch.fps is not None else float(arch.base_fps)
+        guidance = float(batch.guidance_scale)
+
+        is_t2i = num_frames == 1 and batch.preprocessed_image is None and batch.pil_image is None
+        is_i2v = (batch.preprocessed_image is not None or batch.pil_image is not None) and not is_t2i
+
+        # Resolution-based flow_shift, set on the owning pipeline (rebuilds the
+        # scheduler). The framework picks the UniPC shift purely from the named
+        # resolution bucket (``_RESOLUTION_SHIFT_DEFAULTS``), NOT from the task,
+        # so T2V/I2V/T2I at the same resolution share a shift.
+        pipe = self.pipeline() if self.pipeline is not None else None
+        flow_shift = self._flow_shift_for_resolution(height, width)
+        if pipe is not None and hasattr(pipe, "_set_flow_shift"):
+            pipe._set_flow_shift(flow_shift)
+            scheduler = pipe.scheduler
+        else:
+            scheduler = self.scheduler
+
+        # ---- Tokenize prompt + negative prompt ----
+        prompt = batch.prompt if isinstance(batch.prompt, str) else (batch.prompt[0] if batch.prompt else "")
+        negative_prompt = batch.negative_prompt
+        if negative_prompt is None:
+            negative_prompt = "" if is_t2i else COSMOS3_VIDEO_NEGATIVE_PROMPT
+        if isinstance(negative_prompt, list):
+            negative_prompt = negative_prompt[0] if negative_prompt else ""
+
+        special_tokens = cosmos3_special_tokens(self.tokenizer)
+        is_video = not is_t2i
+        cond_ids = cosmos3_tokenize_caption(self.tokenizer, prompt, is_video=is_video, use_system_prompt=False)
+        uncond_ids = cosmos3_tokenize_caption(self.tokenizer,
+                                              negative_prompt,
+                                              is_video=is_video,
+                                              use_system_prompt=False)
+
+        # ---- VAE normalization constants + geometry ----
+        norm = _VaeNorm.from_vae(self.vae, dtype)
+        temporal_factor = int(arch.temporal_compression_factor)
+        spatial_factor = int(self.vae.config.scale_factor_spatial)
+        latent_t = self._latent_frames(num_frames, temporal_factor)
+        latent_h = height // spatial_factor
+        latent_w = width // spatial_factor
+        latent_channel = int(arch.latent_channel)
+        latent_shape = (latent_channel, latent_t, latent_h, latent_w)
+
+        generator = batch.generator
+        if isinstance(generator, list):
+            generator = generator[0] if generator else None
+
+        # ---- Conditioning latent (I2V / T2I) + condition mask ----
+        condition_frame_indexes: list[int] = []
+        clean_latent: torch.Tensor | None = None
+        if is_i2v or (is_t2i and (batch.preprocessed_image is not None or batch.pil_image is not None)):
+            image = batch.preprocessed_image if batch.preprocessed_image is not None else batch.pil_image
+            cond_pixels = self._image_to_video_tensor(image, num_frames, height, width, device, dtype)
+            clean_latent = cosmos3_vae_encode(self.vae, cond_pixels, norm).squeeze(0).float()  # [C, T, H, W]
+            condition_frame_indexes = [0]
+
+        # ---- Initial noise (clean condition frames + pure noise elsewhere) ----
+        pure_noise = randn_tensor(latent_shape, generator=generator, device=device, dtype=dtype).float()
+        if clean_latent is not None:
+            cond_mask = torch.zeros((latent_t, 1, 1), device=device, dtype=pure_noise.dtype)
+            for idx in condition_frame_indexes:
+                if 0 <= idx < latent_t:
+                    cond_mask[idx, 0, 0] = 1.0
+            clean = clean_latent.to(device=device, dtype=pure_noise.dtype)
+            init_latent = cond_mask * clean + (1.0 - cond_mask) * pure_noise
+        else:
+            init_latent = pure_noise
+
+        spec = Cosmos3VisionSpec(
+            shape=latent_shape,
+            condition_frame_indexes=condition_frame_indexes,
+        )
+
+        # ---- Scheduler timesteps ----
+        scheduler.set_timesteps(int(batch.num_inference_steps), device=device)
+        timesteps = scheduler.timesteps
+
+        engine = Cosmos3DenoiseEngine(
+            transformer=self.transformer,
+            scheduler=scheduler,
+            special_tokens=special_tokens,
+            latent_patch_size=int(arch.latent_patch_size),
+            temporal_modality_margin=int(arch.temporal_modality_margin),
+            reset_spatial_ids=bool(arch.unified_3d_mrope_reset_spatial_ids),
+            enable_fps_modulation=bool(arch.enable_fps_modulation),
+            base_fps=float(arch.base_fps),
+            temporal_compression_factor=temporal_factor,
+            include_end_of_generation_token=False,
+        )
+
+        flat_latent = init_latent.reshape(-1)
+        fps_per_item = [fps] if bool(arch.enable_fps_modulation) else None
+
+        # ---- t2vs: jointly generate sound (combined [vision | sound] latent) ----
+        # Mirrors the framework: a placeholder audio sized to the video duration
+        # sets the sound latent length; sound shares the denoise/CFG with vision.
+        with_audio = is_video and os.environ.get("COSMOS3_T2VS", "") not in ("", "0")
+        sound_specs = None
+        sound_fps_per_item = None
+        sound_vae = None
+        sound_shape: tuple[int, int] | None = None
+        if with_audio:
+            sound_vae = self._get_sound_vae(pipe, device, dtype)
+            sound_dim = int(arch.sound_dim)
+            sound_latent_fps = float(arch.sound_latent_fps)
+            # Framework ``create_placeholder_audio`` + ``get_latent_num_samples``.
+            num_audio_samples = int(num_frames / fps * sound_vae.sample_rate)
+            sound_latent_t = max(1, sound_vae.get_latent_num_samples(num_audio_samples))
+            sound_shape = (sound_dim, sound_latent_t)
+            sound_noise = randn_tensor((sound_dim, sound_latent_t), generator=generator, device=device,
+                                       dtype=dtype).float()
+            flat_latent = torch.cat([flat_latent, sound_noise.reshape(-1)])
+            sound_specs = [Cosmos3SoundSpec(shape=sound_shape, condition_frame_indexes=[], fps=sound_latent_fps)]
+            sound_fps_per_item = [sound_latent_fps] if bool(arch.enable_fps_modulation) else None
+
+        final_flat = engine.denoise(
+            flat_latent=flat_latent,
+            timesteps=timesteps,
+            guidance=guidance,
+            specs=[spec],
+            cond_token_ids=cond_ids,
+            uncond_token_ids=uncond_ids,
+            fps_per_item=fps_per_item,
+            progress_bar=lambda it: tqdm(it, desc="Cosmos3 denoising"),
+            sound_specs=sound_specs,
+            sound_fps_per_item=sound_fps_per_item,
+        )
+
+        # ---- Decode vision: [C, T, H, W] -> pixels [B, 3, T, H, W] in [0, 1] ----
+        vision_flat = final_flat[:spec.numel]
+        result_latent = vision_flat.reshape(latent_shape).unsqueeze(0).to(device=device, dtype=dtype)
+        decoded = cosmos3_vae_decode(self.vae, result_latent, norm)  # [B, 3, T, H, W] in [-1, 1]
+        video = ((1.0 + decoded) / 2.0).clamp(0.0, 1.0)
+
+        batch.latents = result_latent
+        batch.output = video
+
+        # ---- Decode sound: AVAE latent [C, T] -> waveform [C, N] in [-1, 1] ----
+        if with_audio and sound_vae is not None and sound_shape is not None:
+            sound_latent = final_flat[spec.numel:].reshape(sound_shape).unsqueeze(0).to(device=device, dtype=dtype)
+            waveform = sound_vae.decode(sound_latent)  # [1, C_audio, N]
+            batch.extra["audio"] = waveform[0].detach().float().cpu()  # [C_audio, N]
+            batch.extra["audio_sample_rate"] = int(sound_vae.sample_rate)
+        return batch
+
+    # ------------------------------------------------------------------
+    # Image preprocessing
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _resize_and_center_crop(img: torch.Tensor, target_h: int, target_w: int) -> torch.Tensor:
+        """Aspect-ratio-preserving resize + center crop, matching the framework
+        (``cosmos_framework.inference.vision._resize_and_center_crop``)."""
+        import math
+
+        import torchvision.transforms.functional as TF
+        orig_h, orig_w = img.shape[-2], img.shape[-1]
+        scaling_ratio = max(target_w / orig_w, target_h / orig_h)
+        resize_h = int(math.ceil(scaling_ratio * orig_h))
+        resize_w = int(math.ceil(scaling_ratio * orig_w))
+        img = TF.resize(img, [resize_h, resize_w])
+        return TF.center_crop(img, [target_h, target_w])
+
+    @staticmethod
+    def _get_sound_vae(pipe: Any, device: torch.device, dtype: torch.dtype) -> Any:
+        """Lazily load + cache the Cosmos3 sound AVAE decoder from the checkpoint.
+
+        The video path does not load ``sound_tokenizer``; t2vs needs only its
+        decoder, so we load it on first use from ``<model_path>/sound_tokenizer``.
+        """
+        cached = getattr(pipe, "_sound_vae", None) if pipe is not None else None
+        if cached is not None:
+            return cached
+        from fastvideo.models.audio.cosmos3_avae import Cosmos3SoundVAE
+        model_path = pipe.model_path
+        sound_dir = os.path.join(model_path, "sound_tokenizer")
+        sound_vae = Cosmos3SoundVAE.from_pretrained(sound_dir, torch_dtype=dtype).to(device)
+        if pipe is not None:
+            pipe._sound_vae = sound_vae
+        return sound_vae
+
+    @classmethod
+    def _image_to_video_tensor(
+        cls,
+        image: Any,
+        num_frames: int,
+        height: int,
+        width: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Build the I2V conditioning pixel video ``[1, 3, T, H, W]`` in [-1, 1].
+
+        Faithful to the framework (``cosmos_framework.inference.vision``):
+        ``load_conditioning_image`` (aspect-preserving resize + center crop +
+        uint8 quantization, then ``/127.5 - 1``) followed by
+        ``build_conditioned_video_batch``, which fills frame 0 with the image and
+        **repeats the last conditioning frame** for the rest of the clip (a static
+        video), NOT zeros. The whole clip is VAE-encoded by the caller; only the
+        latent condition frame(s) are kept clean by the condition mask, but the
+        VAE is temporal, so the repeated (not zeroed) frames change the condition
+        latent — zero-filling here produces a wrong conditioning latent.
+        """
+        import numpy as np
+
+        if hasattr(image, "convert"):  # PIL.Image: framework-exact preprocessing.
+            arr = np.array(image.convert("RGB"))
+            img = torch.from_numpy(arr).permute(2, 0, 1).float()  # [3, H, W] in [0, 255]
+            # Resize + center crop + uint8 quantization, then -> [-1, 1]
+            # (load_conditioning_image / load_conditioning_image_pixels).
+            img = cls._resize_and_center_crop(img.unsqueeze(0), height, width).squeeze(0)
+            img = img.round().clamp(0, 255) / 127.5 - 1.0  # [3, H, W] in [-1, 1]
+        elif isinstance(image, torch.Tensor):  # already-preprocessed conditioning frame.
+            img = image.float()
+            if img.dim() == 5:  # [B,3,T,H,W]
+                img = img[0]
+            if img.dim() == 4:  # [3,T,H,W] or [B,3,H,W] -> first frame
+                img = img[:, 0]
+            if img.max() > 1.5:  # [0, 255] -> [-1, 1]; otherwise assume already [-1, 1].
+                img = img / 127.5 - 1.0
+            if img.shape[-2:] != (height, width):
+                img = cls._resize_and_center_crop(img.unsqueeze(0), height, width).squeeze(0)
+        else:
+            raise TypeError(f"Unsupported conditioning image type: {type(image)}")
+
+        # Static-repeat video (build_conditioned_video_batch: frame 0 = image,
+        # remaining frames repeat the last conditioning frame). The whole clip is
+        # VAE-encoded by the caller; only the latent condition frame(s) are kept
+        # clean by the condition mask, but the VAE is temporal, so the repeated
+        # (not zeroed) frames change the condition latent — zero-filling here
+        # produces a wrong conditioning latent.
+        img = img.to(device=device, dtype=dtype)
+        video = img.unsqueeze(0).unsqueeze(2).expand(1, 3, num_frames, height, width)
+        return video.contiguous()

fastvideo/pipelines/stages/decoding.py

@@ -47,23 +47,6 @@ class DecodingStage(PipelineStage):
         result.add_check("output", batch.output, [V.is_tensor, V.with_dims(5)])
         return result
 
-    def _is_flux2_packed(self, latents: torch.Tensor) -> bool:
-        """Detect Flux2 packed latents by checking channel count against VAE geometry.
-
-        Flux2 packs latent_channels into 2x2 spatial patches, so the DiT
-        operates on ``latent_channels * 4`` channels at half spatial resolution.
-        The VAE's ``post_quant_conv`` input dimension equals ``latent_channels``.
-        """
-        if not hasattr(self.vae, "bn"):
-            return False
-        pqc = getattr(self.vae, "post_quant_conv", None)
-        if pqc is None:
-            return False
-        vae_latent_ch = pqc.weight.shape[1]
-        packed_ch = vae_latent_ch * 4  # 2x2 patch packing
-        ch_dim = 1 if latents.ndim >= 4 else -1
-        return latents.shape[ch_dim] == packed_ch
-
     def _denormalize_latents(self, latents: torch.Tensor) -> torch.Tensor:
         """Convert normalized latents into the VAE's expected latent space."""
         # Some VAEs handle latent (de)normalization internally.
@@ -94,29 +77,6 @@ class DecodingStage(PipelineStage):
 
         return latents
 
-    @staticmethod
-    def _unpatchify_latents(latents: torch.Tensor) -> torch.Tensor:
-        """Inverse of 2x2 patch packing: ``(B, C*4, H', W') -> (B, C, 2*H', 2*W')``."""
-        batch_size, num_channels, height, width = latents.shape
-        latents = latents.reshape(batch_size, num_channels // (2 * 2), 2, 2, height, width)
-        latents = latents.permute(0, 1, 4, 2, 5, 3)
-        latents = latents.reshape(batch_size, num_channels // (2 * 2), height * 2, width * 2)
-        return latents
-
-    def _flux2_bn_denorm_and_unpatchify(self, latents: torch.Tensor) -> torch.Tensor:
-        """BN denormalize then unpatchify packed latents for VAE decode.
-
-        Handles any channel count (e.g. 64->16, 128->32) via 2x2 spatial unpack.
-        """
-        running_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype)
-        running_var = self.vae.bn.running_var.view(1, -1, 1, 1).to(latents.device, latents.dtype)
-        cfg = getattr(self.vae, "config", None)
-        arch = getattr(cfg, "arch_config", None) if cfg else None
-        eps = getattr(arch, "batch_norm_eps", None) or getattr(cfg, "batch_norm_eps", 1e-5)
-        bn_std = torch.sqrt(torch.clamp(running_var + eps, min=1e-6))
-        latents = latents * bn_std + running_mean
-        return self._unpatchify_latents(latents)
-
     @torch.no_grad()
     def decode(self, latents: torch.Tensor, fastvideo_args: FastVideoArgs) -> torch.Tensor:
         """
@@ -140,9 +100,7 @@ class DecodingStage(PipelineStage):
         vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
         vae_autocast_enabled = (vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
 
-        # Flux2: skip denormalize on packed latents; BN denorm runs below instead
-        if not (latents.ndim == 5 and self._is_flux2_packed(latents)):
-            latents = self._denormalize_latents(latents)
+        latents = self._denormalize_latents(latents)
 
         # Decode latents
         with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
@@ -152,27 +110,7 @@ class DecodingStage(PipelineStage):
             #     self.vae.enable_parallel()
             if not vae_autocast_enabled:
                 latents = latents.to(vae_dtype)
-            # Flux2's image VAE expects 4D (B, C, H, W); squeeze the singleton T
-            # only for Flux2 packed latents. Gated on `_is_flux2_packed` so video
-            # VAEs that legitimately decode 5D latents with T=1 are untouched.
-            squeezed_for_vae = False
-            if latents.ndim == 5 and latents.shape[2] == 1 and self._is_flux2_packed(latents):
-                latents = latents.squeeze(2)
-                squeezed_for_vae = True
-            # Flux2 packed: BN denorm + unpatchify for VAE decode.
-            # BN denorm is the complete inverse normalisation for Flux2 (no
-            # scaling_factor/shift_factor step), matching Diffusers.
-            if latents.ndim == 4 and self._is_flux2_packed(latents):
-                latents = self._flux2_bn_denorm_and_unpatchify(latents)
             image = self.vae.decode(latents)
-            # Unwrap diffusers-style DecoderOutput / tuple (Flux2 VAE returns a
-            # DecoderOutput). No-op for existing VAEs that return a plain tensor.
-            if hasattr(image, "sample"):
-                image = image.sample
-            elif isinstance(image, tuple | list):
-                image = image[0]
-            if squeezed_for_vae:
-                image = image.unsqueeze(2)
 
         # Normalize image to [0, 1] range
         image = (image / 2 + 0.5).clamp(0, 1)
@@ -263,13 +201,7 @@ class DecodingStage(PipelineStage):
                 pipeline.add_module("vae", self.vae)
             fastvideo_args.model_loaded["vae"] = True
 
-        if fastvideo_args.output_type == "latent":
-            frames = batch.latents
-            if frames.ndim == 5 and frames.shape[2] == 1 and self._is_flux2_packed(frames):
-                frames = self._flux2_bn_denorm_and_unpatchify(frames.squeeze(2))
-                frames = frames.unsqueeze(2)
-        else:
-            frames = self.decode(batch.latents, fastvideo_args)
+        frames = batch.latents if fastvideo_args.output_type == "latent" else self.decode(batch.latents, fastvideo_args)
 
         # decode trajectory latents if needed
         if batch.return_trajectory_decoded:

fastvideo/pipelines/stages/denoising.py

@@ -4,7 +4,6 @@ Denoising stage for diffusion pipelines.
 """
 
 import inspect
-import os
 import weakref
 from collections.abc import Iterable
 from typing import Any
@@ -104,37 +103,7 @@ class DenoisingStage(PipelineStage):
         # TODO(will): make the precision configurable for inference
         # target_dtype = PRECISION_TO_TYPE[fastvideo_args.precision]
         target_dtype = torch.bfloat16
-        # Flux2-only denoising compensations.
-        #
-        # `_is_flux` gates four behaviors that exist because Flux2's transformer
-        # forward() does things internally that the generic pipeline must undo or
-        # match. These are architectural facts about the Flux2 transformer, not
-        # tunable precision policies (the precision policies #5/#6 — prompt-embed
-        # casting and scheduler-step placement — were already moved to config:
-        # DiTArchConfig.cast_prompt_embeds_to_dit_dtype and
-        # PipelineConfig.scheduler_step_in_fp32).
-        #
-        # The four behaviors gated below:
-        #   1. env-var bf16-reduced-precision matmul disable (4-step Klein drift)
-        #   2. autocast disabled (Flux2 long-sequence attention breaks parity under autocast)
-        #   3. guidance: skip the external x1000 (Flux2 multiplies guidance by 1000 internally)
-        #   4. timestep: divide by 1000 with cast-before-divide (Flux2 multiplies timestep by 1000 internally)
-        #
-        # Contract: `prefix == "Flux"` is set ONLY by Flux2 (fastvideo/configs/
-        # models/dits/flux_2.py). No other model uses that prefix, so this exact
-        # match cannot false-positive. A future Flux variant that needs the same
-        # compensations must either set prefix == "Flux" too, OR (preferred) these
-        # gates should graduate to arch-config declarations like the precision
-        # policies above.
-        _is_flux = (getattr(fastvideo_args.pipeline_config.dit_config, "prefix", "") == "Flux")
-        if _is_flux and os.getenv("FASTVIDEO_FLUX2_DISABLE_BF16_REDUCED_PRECISION_REDUCTION",
-                                  "").lower() in {"1", "true", "yes"}:
-            # Gate 1: tighten bf16 matmul accumulation for the 4-step Klein model (opt-in via env var).
-            torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
-        # Gate 2: Flux2 runs its bf16 transformer WITHOUT autocast — autocast perturbs long-sequence attention enough to break 4-step latent parity.
-        autocast_enabled = ((target_dtype != torch.float32) and not fastvideo_args.disable_autocast and not _is_flux)
-        scheduler_fp32 = getattr(fastvideo_args.pipeline_config, "scheduler_step_in_fp32", False)
-        local_device = get_local_torch_device()
+        autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast
 
         # Get timesteps and calculate warmup steps
         timesteps = batch.timesteps
@@ -190,40 +159,13 @@ class DenoisingStage(PipelineStage):
             },
         )
 
-        for key in ("flux2_txt_ids", "flux2_img_ids"):
-            value = batch.extra.get(key)
-            if torch.is_tensor(value):
-                batch.extra[key] = value.to(device=local_device)
-
-        flux2_id_kwargs = self.prepare_extra_func_kwargs(
-            self.transformer.forward,
-            {
-                "txt_ids": batch.extra.get("flux2_txt_ids"),
-                "img_ids": batch.extra.get("flux2_img_ids"),
-            },
-        )
-
         # Get latents and embeddings
         latents = batch.latents
-        cast_embeds = getattr(fastvideo_args.pipeline_config.dit_config, "cast_prompt_embeds_to_dit_dtype", False)
-        if cast_embeds:
-            prompt_embeds = [
-                embed.to(device=local_device, dtype=target_dtype) if torch.is_tensor(embed) else embed
-                for embed in batch.prompt_embeds
-            ]
-        else:
-            prompt_embeds = batch.prompt_embeds
+        prompt_embeds = batch.prompt_embeds
         assert not torch.isnan(prompt_embeds[0]).any(), "prompt_embeds contains nan"
         if batch.do_classifier_free_guidance:
             neg_prompt_embeds = batch.negative_prompt_embeds
             assert neg_prompt_embeds is not None
-            if cast_embeds:
-                neg_prompt_embeds = [
-                    embed.to(device=local_device, dtype=target_dtype) if torch.is_tensor(embed) else embed
-                    for embed in neg_prompt_embeds
-                ]
-            else:
-                neg_prompt_embeds = batch.negative_prompt_embeds
             assert not torch.isnan(neg_prompt_embeds[0]).any(), "neg_prompt_embeds contains nan"
 
         # (Wan2.2) Calculate timestep to switch from high noise expert to low noise expert
@@ -270,34 +212,23 @@ class DenoisingStage(PipelineStage):
         # Initialize lists for ODE trajectory
         trajectory_timesteps: list[torch.Tensor] = []
         trajectory_latents: list[torch.Tensor] = []
-        is_lucy_edit = fastvideo_args.pipeline_config.lucy_edit_task
 
         # Hoisted out of the per-step loop: depends only on inputs that
         # are constant across denoising steps.
         use_meanflow = getattr(self.transformer.config, "use_meanflow", False)
-        # Gate 3: Flux2's transformer multiplies guidance by 1000 internally, so we
-        # skip the external *1000 pre-scaling for Flux models.
         embedded_cfg_scale = fastvideo_args.pipeline_config.embedded_cfg_scale
-        if _is_flux and embedded_cfg_scale is not None:
-            embedded_cfg_scale = batch.guidance_scale
         if embedded_cfg_scale is not None:
             guidance_expand = (torch.tensor(
                 [embedded_cfg_scale] * latents.shape[0],
                 dtype=torch.float32,
                 device=get_local_torch_device(),
-            ).to(target_dtype) * (1.0 if _is_flux else 1000.0))
+            ).to(target_dtype) * 1000.0)
         else:
             guidance_expand = None
         # V2V padding: zero-filled tensor concatenated with each step's
         # latent_model_input.  Shape is fixed by latents and is never
         # written to, so we allocate once.
         v2v_zero_pad = torch.zeros_like(latents) if batch.video_latent is not None else None
-        lucy_timestep_seq_len = None
-        if is_lucy_edit:
-            patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size
-            assert patch_size[0] == 1, "Lucy Edit timestep expansion assumes temporal patch size 1"
-            lucy_timestep_seq_len = (latents.shape[2] * (latents.shape[3] // patch_size[1]) *
-                                     (latents.shape[4] // patch_size[2]))
 
         # CFG gating / stale-uncond reuse setup (Adaptive Guidance LinearAG
         # variant, Castillo et al. 2023).  When envs.FASTVIDEO_CFG_GATE_STEP
@@ -378,25 +309,14 @@ class DenoisingStage(PipelineStage):
                 # Expand latents for V2V/I2V
                 latent_model_input = latents.to(target_dtype)
                 if batch.video_latent is not None:
-                    if is_lucy_edit:
-                        latent_model_input = torch.cat(
-                            [latent_model_input, batch.video_latent],
-                            dim=1,
-                        ).to(target_dtype)
-                    else:
-                        latent_model_input = torch.cat(
-                            [latent_model_input, batch.video_latent, v2v_zero_pad],
-                            dim=1,
-                        ).to(target_dtype)
+                    latent_model_input = torch.cat([latent_model_input, batch.video_latent, v2v_zero_pad],
+                                                   dim=1).to(target_dtype)
                 elif batch.image_latent is not None:
                     assert not fastvideo_args.pipeline_config.ti2v_task, "image latents should not be provided for TI2V task"
                     latent_model_input = torch.cat([latent_model_input, batch.image_latent], dim=1).to(target_dtype)
 
                 assert not torch.isnan(latent_model_input).any(), "latent_model_input contains nan"
-                if is_lucy_edit:
-                    assert lucy_timestep_seq_len is not None
-                    t_expand = t.repeat(latent_model_input.shape[0], lucy_timestep_seq_len)
-                elif fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
+                if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
                     timestep = torch.stack([t]).to(get_local_torch_device())
                     temp_ts = (mask2[0][0][:, ::2, ::2] * timestep).flatten()
                     temp_ts = torch.cat([temp_ts, temp_ts.new_ones(seq_len - temp_ts.size(0)) * timestep])
@@ -404,19 +324,7 @@ class DenoisingStage(PipelineStage):
                     t_expand = timestep.repeat(latent_model_input.shape[0], 1)
                 else:
                     t_expand = t.repeat(latent_model_input.shape[0])
-                # Gate 4: Flux2 transformer multiplies timestep by 1000 internally, so
-                # the pipeline must pass timestep/1000 (matching Diffusers).
-                # Diffusers casts to the latent dtype before the division; doing
-                # the division in fp32 first changes BF16 rounding for the final
-                # Klein timestep and breaks latent parity.
-                if _is_flux:
-                    t_expand = t_expand.to(
-                        device=get_local_torch_device(),
-                        dtype=latent_model_input.dtype,
-                    )
-                    t_expand = t_expand / 1000.0
-                else:
-                    t_expand = t_expand.to(get_local_torch_device())
+                t_expand = t_expand.to(get_local_torch_device())
 
                 if use_meanflow:
                     if i == len(timesteps) - 1:
@@ -495,7 +403,6 @@ class DenoisingStage(PipelineStage):
                             **action_kwargs,
                             **camera_kwargs,
                             **timesteps_r_kwarg,
-                            **flux2_id_kwargs,
                         )
 
                     if batch.do_classifier_free_guidance:
@@ -537,7 +444,6 @@ class DenoisingStage(PipelineStage):
                                     **action_kwargs,
                                     **camera_kwargs,
                                     **timesteps_r_kwarg,
-                                    **flux2_id_kwargs,
                                 )
                             _cfg_gate_fresh_uncond += 1
 
@@ -561,16 +467,12 @@ class DenoisingStage(PipelineStage):
                                 noise_pred_text,
                                 guidance_rescale=batch.guidance_rescale,
                             )
-                if scheduler_fp32:
-                    # Diffusers-style: fp32 Euler update outside autocast avoids BF16 drift.
+                    # Compute the previous noisy sample
                     latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
-                else:
-                    with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled):
-                        latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
-                if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
-                    latents = latents.squeeze(0)
-                    latents = (1. - mask2[0]) * z + mask2[0] * latents
-                    # latents = latents.unsqueeze(0)
+                    if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
+                        latents = latents.squeeze(0)
+                        latents = (1. - mask2[0]) * z + mask2[0] * latents
+                        # latents = latents.unsqueeze(0)
 
                 # save trajectory latents if needed
                 if batch.return_trajectory_latents:

fastvideo/pipelines/stages/image_encoding.py

@@ -629,10 +629,9 @@ class VideoVAEEncodingStage(ImageVAEEncodingStage):
             encoder_output = self.vae.encode(video_condition)
 
         generator = batch.generator
-        sample_mode = "argmax" if fastvideo_args.pipeline_config.lucy_edit_task else "sample"
-        if sample_mode == "sample" and generator is None:
-            raise ValueError("Generator must be provided for sampled video VAE encoding")
-        latent_condition = self.retrieve_latents(encoder_output, generator, sample_mode=sample_mode)
+        if generator is None:
+            raise ValueError("Generator must be provided")
+        latent_condition = self.retrieve_latents(encoder_output, generator)
 
         if (hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None):
             if isinstance(self.vae.shift_factor, torch.Tensor):

fastvideo/pipelines/stages/text_encoding.py

@@ -57,10 +57,6 @@ class TextEncodingStage(PipelineStage):
         assert len(self.tokenizers) == len(self.text_encoders)
         assert len(self.text_encoders) == len(fastvideo_args.pipeline_config.text_encoder_configs)
 
-        # Skip encoding if precomputed prompt_embeds were provided
-        if batch.prompt_embeds is not None and len(batch.prompt_embeds) > 0:
-            return batch
-
         # Encode positive prompt with all available encoders
         assert batch.prompt is not None
         prompt_text: str | list[str] = batch.prompt
@@ -222,8 +218,7 @@ class TextEncodingStage(PipelineStage):
                     # Qwen2-style tokenizers. Scoped via treat_empty_as_dot so
                     # models that legitimately use "" (e.g. negative_prompt="")
                     # are not affected.
-                    if isinstance(processed_text, str) and not processed_text.strip() and getattr(
-                            encoder_config, "treat_empty_as_dot", False):
+                    if not processed_text.strip() and getattr(encoder_config, "treat_empty_as_dot", False):
                         processed_text = "."
                     processed_texts.append(processed_text)
                 else:
@@ -240,27 +235,7 @@ class TextEncodingStage(PipelineStage):
             tok = getattr(tokenizer, "tokenizer", tokenizer)
 
             if encoder_config.is_chat_model:
-                already_chat_formatted = bool(processed_texts) and isinstance(processed_texts[0], list)
-                if already_chat_formatted:
-                    # Existing chat models (e.g. HunyuanVideo 1.5 / Qwen2.5-VL)
-                    # pre-format prompts into message lists upstream and rely on
-                    # the inner tokenizer + full tokenizer_kwargs (which include
-                    # add_generation_prompt). Preserve that original path exactly.
-                    text_inputs = tok.apply_chat_template(processed_texts, **tok_kwargs).to(target_device)
-                else:
-                    # Two-step approach matching Diffusers: format with chat
-                    # template first, then tokenize the resulting strings.
-                    formatted_texts = []
-                    for pt in processed_texts:
-                        messages = [{"role": "user", "content": pt}]
-                        formatted = tokenizer.apply_chat_template(
-                            messages,
-                            tokenize=False,
-                            add_generation_prompt=True,
-                            enable_thinking=False,
-                        )
-                        formatted_texts.append(formatted)
-                    text_inputs = tokenizer(formatted_texts, **tok_kwargs).to(target_device)
+                text_inputs = tok.apply_chat_template(processed_texts, **tok_kwargs).to(target_device)
             else:
                 text_inputs = tok(processed_texts, **tok_kwargs).to(target_device)
 

fastvideo/registry.py

@@ -20,6 +20,7 @@ from fastvideo.configs.pipelines.cosmos2_5 import (
     Cosmos25Config,
     Cosmos25_14BConfig,
 )
+from fastvideo.configs.pipelines.cosmos3 import Cosmos3Config
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
 from fastvideo.configs.pipelines.gen3c import Gen3CConfig
@@ -30,10 +31,6 @@ from fastvideo.configs.pipelines.hyworld import HYWorldConfig
 from fastvideo.configs.pipelines.lingbotworld import LingBotWorldI2V480PConfig
 from fastvideo.configs.pipelines.longcat import LongCatT2V480PConfig
 from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
-from fastvideo.configs.pipelines.flux_2 import (
-    Flux2KleinPipelineConfig,
-    Flux2PipelineConfig,
-)
 from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig
 from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig
 from fastvideo.configs.pipelines.turbodiffusion import (
@@ -44,7 +41,6 @@ from fastvideo.configs.pipelines.turbodiffusion import (
 from fastvideo.configs.pipelines.wan import (
     FastWan2_1_T2V_480P_Config,
     FastWan2_2_TI2V_5B_Config,
-    LucyEditDevConfig,
     SelfForcingWan2_2_T2V480PConfig,
     SelfForcingWanT2V480PConfig,
     WANV2VConfig,
@@ -329,45 +325,6 @@ def _register_configs() -> None:
         default_preset="stable_audio_open_small",
     )
 
-    def _is_flux2_klein(path: str) -> bool:
-        path_lower = path.lower()
-        return "flux.2-klein" in path_lower or "flux2-klein" in path_lower or "flux2klein" in path_lower
-
-    def _is_flux2_full(path: str) -> bool:
-        path_lower = path.lower()
-        is_flux2 = "flux.2" in path_lower or "flux2" in path_lower or "flux_2" in path_lower or "flux-2" in path_lower
-        return is_flux2 and "klein" not in path_lower
-
-    # Flux2 Klein (distilled, 4-step, no guidance)
-    register_configs(
-        sampling_param_cls=None,
-        pipeline_config_cls=Flux2KleinPipelineConfig,
-        workload_types=(WorkloadType.T2I, ),
-        hf_model_paths=[
-            "black-forest-labs/FLUX.2-klein-4B",
-            "black-forest-labs/FLUX.2-klein-9B",
-        ],
-        model_detectors=[
-            _is_flux2_klein,
-        ],
-        model_family="flux2",
-        default_preset="flux2_klein_4b",
-    )
-    # Flux2 (full, Mistral3 text encoder, embedded guidance)
-    register_configs(
-        sampling_param_cls=None,
-        pipeline_config_cls=Flux2PipelineConfig,
-        workload_types=(WorkloadType.T2I, ),
-        hf_model_paths=[
-            "black-forest-labs/FLUX.2-dev",
-        ],
-        model_detectors=[
-            _is_flux2_full,
-        ],
-        model_family="flux2",
-        default_preset="flux2_dev",
-    )
-
     # Hunyuan 1.5 (specific)
     register_configs(
         sampling_param_cls=None,
@@ -599,6 +556,22 @@ def _register_configs() -> None:
         default_preset="gen3c_cosmos_7b",
     )
 
+    # Cosmos 3 (must register before Cosmos 2.5 and generic Cosmos detectors
+    # so the cosmos3 path-detection takes precedence)
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=Cosmos3Config,
+        workload_types=(WorkloadType.T2V, ),
+        hf_model_paths=[
+            "nvidia/Cosmos3-Nano",
+        ],
+        model_detectors=[
+            lambda path: "cosmos3" in path.lower() or "cosmos-3" in path.lower(),
+        ],
+        model_family="cosmos3",
+        default_preset="cosmos3_nano",
+    )
+
     # Cosmos 2.5 (2B)
     register_configs(
         sampling_param_cls=None,
@@ -782,18 +755,6 @@ def _register_configs() -> None:
         model_family="wan",
         default_preset="fast_wan_2_2_ti2v_5b",
     )
-    register_configs(
-        sampling_param_cls=None,
-        pipeline_config_cls=LucyEditDevConfig,
-        workload_types=(),
-        hf_model_paths=[
-            "decart-ai/Lucy-Edit-Dev",
-            "decart-ai/Lucy-Edit-1.1-Dev",
-        ],
-        model_detectors=[lambda path: "lucy-edit" in path.lower()],
-        model_family="wan",
-        default_preset="lucy_edit_dev",
-    )
     register_configs(
         sampling_param_cls=None,
         pipeline_config_cls=Wan2_2_T2V_A14B_Config,
@@ -896,19 +857,15 @@ def get_model_info(
     if workload_type is None:
         workload_type = WorkloadType.T2V
 
-    config_info = _get_config_info(model_path, raise_on_missing=True)
-    assert config_info is not None, "config_info must be resolved"
-
-    if override_pipeline_cls_name:
-        pipeline_name = override_pipeline_cls_name
-        logger.info("Using override pipeline class name %s", pipeline_name)
+    if os.path.exists(model_path):
+        config = verify_model_config_and_directory(model_path)
     else:
-        if os.path.exists(model_path):
-            config = verify_model_config_and_directory(model_path)
-        else:
-            config = maybe_download_model_index(model_path)
+        config = maybe_download_model_index(model_path)
 
-        pipeline_name = config.get("_class_name")
+    pipeline_name = config.get("_class_name")
+    if override_pipeline_cls_name:
+        logger.info("Overriding pipeline class name from %s to %s", pipeline_name, override_pipeline_cls_name)
+        pipeline_name = override_pipeline_cls_name
 
     if pipeline_name is None:
         raise ValueError("Model config does not contain a _class_name attribute. "
@@ -917,6 +874,9 @@ def get_model_info(
     pipeline_registry = get_pipeline_registry(pipeline_type)
     pipeline_cls = pipeline_registry.resolve_pipeline_cls(pipeline_name, pipeline_type, workload_type)
 
+    config_info = _get_config_info(model_path, raise_on_missing=True)
+    assert config_info is not None, "config_info must be resolved"
+
     sampling_param_cls = config_info.sampling_param_cls or SamplingParam
 
     return ModelInfo(
@@ -949,6 +909,8 @@ def _register_presets() -> None:
     from fastvideo.api.presets import register_preset
     from fastvideo.pipelines.basic.cosmos.presets import (
         ALL_PRESETS as COSMOS_PRESETS, )
+    from fastvideo.pipelines.basic.cosmos3.presets import (
+        ALL_PRESETS as COSMOS3_PRESETS, )
     from fastvideo.pipelines.basic.gamecraft.presets import (
         ALL_PRESETS as GAMECRAFT_PRESETS, )
     from fastvideo.pipelines.basic.gen3c.presets import (
@@ -977,12 +939,10 @@ def _register_presets() -> None:
         ALL_PRESETS as TURBODIFFUSION_PRESETS, )
     from fastvideo.pipelines.basic.wan.presets import (
         ALL_PRESETS as WAN_PRESETS, )
-    from fastvideo.pipelines.basic.flux_2.presets import (
-        ALL_PRESETS as FLUX2_PRESETS, )
 
     all_preset_groups = (
         COSMOS_PRESETS,
-        FLUX2_PRESETS,
+        COSMOS3_PRESETS,
         GAMECRAFT_PRESETS,
         GEN3C_PRESETS,
         HUNYUAN_PRESETS,

fastvideo/tests/api/test_lucy_edit.py

@@ -1,51 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from fastvideo.api.presets import get_preset
-from fastvideo.api.sampling_param import SamplingParam
-from fastvideo.configs.pipelines.wan import LucyEditDevConfig
-from fastvideo.fastvideo_args import WorkloadType
-from fastvideo.pipelines.basic.wan.lucy_edit_pipeline import LucyEditPipeline
-from fastvideo.pipelines.pipeline_registry import PipelineType, get_pipeline_registry
-from fastvideo.registry import get_default_preset, get_pipeline_config_cls_from_name
-
-
-def test_lucy_edit_registry_and_preset() -> None:
-    import fastvideo.registry  # noqa: F401
-
-    preset = get_preset("lucy_edit_dev", "wan")
-    assert preset.model_family == "wan"
-    assert preset.defaults["height"] == 480
-    assert preset.defaults["width"] == 832
-    assert preset.defaults["num_frames"] == 81
-
-    config = LucyEditDevConfig()
-    assert config.lucy_edit_task is True
-    assert config.ti2v_task is False
-    assert config.dit_config.arch_config.out_channels == 48
-    assert config.dit_config.arch_config.in_channels == 96
-    assert config.vae_config.arch_config.z_dim == 48
-    assert config.dit_config.arch_config.in_channels == config.vae_config.arch_config.z_dim * 2
-    assert get_default_preset("decart-ai/Lucy-Edit-Dev") == "lucy_edit_dev"
-    assert get_default_preset("decart-ai/Lucy-Edit-1.1-Dev") == "lucy_edit_dev"
-    assert get_pipeline_config_cls_from_name("decart-ai/Lucy-Edit-Dev") is LucyEditDevConfig
-    assert get_pipeline_config_cls_from_name("decart-ai/Lucy-Edit-1.1-Dev") is LucyEditDevConfig
-
-    sampling_param = SamplingParam.from_pretrained("decart-ai/Lucy-Edit-Dev")
-    assert sampling_param.height == 480
-    assert sampling_param.width == 832
-    assert sampling_param.num_frames == 81
-    assert sampling_param.fps == 24
-    assert sampling_param.guidance_scale == 5.0
-    assert sampling_param.negative_prompt == ""
-
-    sampling_param_1_1 = SamplingParam.from_pretrained("decart-ai/Lucy-Edit-1.1-Dev")
-    assert sampling_param_1_1.height == 480
-    assert sampling_param_1_1.width == 832
-    assert sampling_param_1_1.num_frames == 81
-    assert sampling_param_1_1.fps == 24
-    assert sampling_param_1_1.guidance_scale == 5.0
-    assert sampling_param_1_1.negative_prompt == ""
-
-    # FastVideo has no V2V workload enum today; model_index dispatches Lucy by pipeline class name.
-    registry = get_pipeline_registry(PipelineType.BASIC)
-    assert registry.resolve_pipeline_cls("LucyEditPipeline", PipelineType.BASIC, WorkloadType.T2V) is LucyEditPipeline

fastvideo/tests/attention/test_video_sparse_attention_metadata.py

@@ -1,38 +0,0 @@
-import torch
-
-from fastvideo.attention.backends.video_sparse_attn import (
-    VideoSparseAttentionImpl,
-    VideoSparseAttentionMetadataBuilder,
-)
-
-
-def _build_metadata(cache_tile_buf: bool):
-    return VideoSparseAttentionMetadataBuilder().build(
-        current_timestep=0,
-        raw_latent_shape=(4, 4, 4),
-        patch_size=(1, 1, 1),
-        VSA_sparsity=0.5,
-        device=torch.device("cpu"),
-        cache_tile_buf=cache_tile_buf,
-    )
-
-
-def test_vsa_tile_does_not_cache_training_scratch_when_disabled():
-    metadata = _build_metadata(cache_tile_buf=False)
-    impl = object.__new__(VideoSparseAttentionImpl)
-    x = torch.ones(1, 64, 2, 2)
-
-    tiled = impl.tile(x, metadata)
-
-    assert tiled.shape == x.shape
-    assert metadata.tile_buf is None
-
-
-def test_vsa_tile_caches_scratch_by_default():
-    metadata = _build_metadata(cache_tile_buf=True)
-    impl = object.__new__(VideoSparseAttentionImpl)
-    x = torch.ones(1, 64, 2, 2)
-
-    tiled = impl.tile(x, metadata)
-
-    assert metadata.tile_buf is tiled

fastvideo/tests/layers/test_linear_mlp_quant_gating.py

@@ -1,49 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Regression coverage for PR #1390's S2-1 plumbing finding: the dormant FP4 shape-tracking path must stay
-gated off unless explicitly enabled, and the MLP quant_config=None default path must keep using ReplicatedLinear's
-unquantized fallback.
-"""
-from __future__ import annotations
-
-import torch
-
-from fastvideo.layers.linear import ReplicatedLinear, UnquantizedLinearMethod
-from fastvideo.layers.mlp import MLP
-
-
-def test_replicated_linear_shape_tracking_default_off() -> None:
-    ReplicatedLinear.reset_shape_tracking()
-    assert ReplicatedLinear.enable_shape_tracking is False
-
-    linear = ReplicatedLinear(input_size=8, output_size=4)
-    linear(torch.randn(2, 8))
-
-    assert len(ReplicatedLinear._shape_to_layer_types) == 0
-
-
-def test_replicated_linear_shape_tracking_enabled_records_unique_shapes() -> None:
-    ReplicatedLinear.reset_shape_tracking()
-    ReplicatedLinear.enable_shape_tracking = True
-    try:
-        linear = ReplicatedLinear(input_size=8, output_size=4)
-        linear(torch.randn(2, 8))
-        linear(torch.randn(3, 8))
-
-        assert len(ReplicatedLinear._shape_to_layer_types) == 2
-        for layer_types in ReplicatedLinear._shape_to_layer_types.values():
-            assert "ReplicatedLinear" in layer_types
-
-        ReplicatedLinear.reset_shape_tracking()
-        assert len(ReplicatedLinear._shape_to_layer_types) == 0
-    finally:
-        ReplicatedLinear.enable_shape_tracking = False
-
-
-def test_mlp_quant_config_none_uses_unquantized_path() -> None:
-    mlp = MLP(input_dim=8, mlp_hidden_dim=16)
-
-    assert isinstance(mlp.fc_in.quant_method, UnquantizedLinearMethod)
-    assert isinstance(mlp.fc_out.quant_method, UnquantizedLinearMethod)
-
-    output = mlp.forward(torch.randn(2, 8))
-    assert output.shape == (2, 8)

fastvideo/tests/modal/launch_l40s_job.py

@@ -1,531 +0,0 @@
-"""Launch an arbitrary FastVideo command on Modal GPUs.
-
-Examples:
-    python -m modal run fastvideo/tests/modal/launch_l40s_job.py --command "nvidia-smi" --install-extra none
-
-    python -m modal run fastvideo/tests/modal/launch_l40s_job.py \
-        --num-gpus 2 \
-        --install-extra test \
-        --command "pytest fastvideo/tests/vaes -vs"
-
-    python -m modal run fastvideo/tests/modal/launch_l40s_job.py \
-        --gpu-type H100 \
-        --num-gpus 1 \
-        --install-extra none \
-        --command "nvidia-smi"
-
-Use ``--no-wait`` with ``modal run --detach`` when the job should keep running
-after the local Modal client exits.
-"""
-
-import os
-import base64
-import shlex
-import shutil
-import subprocess
-import sys
-import time
-from collections.abc import Callable
-from typing import Any
-
-import modal
-
-app = modal.App("fastvideo-gpu-job")
-
-REPO_DIR = "/FastVideo"
-MODEL_VOLUME_NAME = os.environ.get("FASTVIDEO_MODAL_VOLUME", "hf-model-weights")
-IMAGE_VERSION = os.environ.get("IMAGE_VERSION", "latest")
-IMAGE_TAG = os.environ.get(
-    "FASTVIDEO_MODAL_IMAGE",
-    f"ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:{IMAGE_VERSION}",
-)
-SECRET_ENV_KEYS = (
-    "HF_API_KEY",
-    "HUGGINGFACE_HUB_TOKEN",
-    "HF_TOKEN",
-    "WANDB_API_KEY",
-    "WANDB_BASE_URL",
-    "WANDB_MODE",
-)
-
-print(f"Using image: {IMAGE_TAG}")
-print(f"Using Modal volume: {MODEL_VOLUME_NAME}")
-
-model_vol = modal.Volume.from_name(MODEL_VOLUME_NAME, create_if_missing=True)
-local_secrets = modal.Secret.from_dict({
-    key: os.environ[key]
-    for key in SECRET_ENV_KEYS
-    if os.environ.get(key)
-})
-
-image = (
-    modal.Image.from_registry(IMAGE_TAG, add_python="3.12")
-    .apt_install(
-        "cmake",
-        "pkg-config",
-        "build-essential",
-        "curl",
-        "git",
-        "libssl-dev",
-        "ffmpeg",
-    )
-    .run_commands("curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable")
-    .run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc")
-    .env({
-        "PATH": "/root/.cargo/bin:$PATH",
-        "HF_HOME": "/root/data/.cache",
-        "TOKENIZERS_PARALLELISM": "false",
-        "FASTVIDEO_ATTENTION_BACKEND": os.environ.get("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN"),
-    })
-)
-
-COMMON_FUNCTION_KWARGS = dict(
-    image=image,
-    timeout=86400,
-    secrets=[local_secrets],
-    volumes={"/root/data": model_vol},
-)
-
-
-def _run_local_git_command(args: list[str]) -> str:
-    result = subprocess.run(
-        ["git", *args],
-        check=True,
-        capture_output=True,
-        text=True,
-    )
-    return result.stdout.strip()
-
-
-def _run_local_git_command_allow_diff(args: list[str]) -> str:
-    result = subprocess.run(
-        ["git", *args],
-        check=False,
-        capture_output=True,
-        text=True,
-    )
-    if result.returncode not in (0, 1):
-        raise RuntimeError(result.stderr.strip() or f"git {' '.join(args)} failed")
-    return result.stdout
-
-
-def _split_patch_paths(patch_paths: str) -> list[str]:
-    return [path.strip() for path in patch_paths.split(",") if path.strip()]
-
-
-def _build_local_patch(patch_paths: str) -> str:
-    paths = _split_patch_paths(patch_paths)
-    diff_args = ["diff", "--binary"]
-    if paths:
-        diff_args.extend(["--", *paths])
-    patch_parts = [_run_local_git_command_allow_diff(diff_args)]
-
-    untracked_args = ["ls-files", "--others", "--exclude-standard"]
-    if paths:
-        untracked_args.extend(["--", *paths])
-    untracked = _run_local_git_command(untracked_args).splitlines()
-    for path in untracked:
-        if not os.path.isfile(path):
-            continue
-        patch_parts.append(_run_local_git_command_allow_diff(["diff", "--binary", "--no-index", "/dev/null", path]))
-
-    patch = "\n".join(part for part in patch_parts if part.strip())
-    if not patch.strip():
-        raise RuntimeError("Requested --apply-local-patch but no local diff was found.")
-    return patch
-
-
-def _apply_local_patch(patch_b64: str) -> None:
-    if not patch_b64:
-        return
-    patch = base64.b64decode(patch_b64.encode("ascii"))
-    print("Applying local workspace patch", flush=True)
-    result = subprocess.run(
-        ["git", "apply", "--binary", "--whitespace=nowarn", "-"],
-        cwd=REPO_DIR,
-        input=patch,
-        stdout=subprocess.PIPE,
-        stderr=sys.stderr,
-        check=False,
-    )
-    if result.stdout:
-        print(result.stdout.decode("utf-8", errors="replace"), end="", flush=True)
-    if result.returncode != 0:
-        raise RuntimeError(f"Failed to apply local patch with exit code {result.returncode}")
-
-
-def _normalize_git_repo_url(git_repo: str) -> str:
-    if git_repo.startswith("git@github.com:"):
-        return "https://github.com/" + git_repo[len("git@github.com:"):]
-    if git_repo.startswith("ssh://git@github.com/"):
-        return "https://github.com/" + git_repo[len("ssh://git@github.com/"):]
-    return git_repo
-
-
-def _resolve_git_repo(git_repo: str) -> str:
-    if git_repo.strip():
-        return _normalize_git_repo_url(git_repo.strip())
-
-    env_repo = os.environ.get("BUILDKITE_REPO", "").strip()
-    if env_repo:
-        return _normalize_git_repo_url(env_repo)
-
-    discovered_repo = _run_local_git_command(["config", "--get", "remote.origin.url"])
-    if discovered_repo:
-        return _normalize_git_repo_url(discovered_repo)
-
-    raise RuntimeError("Could not resolve git repo URL. Pass --git-repo or set BUILDKITE_REPO.")
-
-
-def _resolve_git_commit(git_commit: str) -> str:
-    if git_commit.strip():
-        return git_commit.strip()
-
-    env_commit = os.environ.get("BUILDKITE_COMMIT", "").strip()
-    if env_commit:
-        return env_commit
-
-    discovered_commit = _run_local_git_command(["rev-parse", "HEAD"])
-    if discovered_commit:
-        return discovered_commit
-
-    raise RuntimeError("Could not resolve git commit. Pass --git-commit or set BUILDKITE_COMMIT.")
-
-
-def _resolve_pull_request(pr_number: str) -> str:
-    if pr_number.strip():
-        return pr_number.strip()
-    env_pr = os.environ.get("BUILDKITE_PULL_REQUEST", "").strip()
-    if env_pr:
-        return env_pr
-    return "false"
-
-
-def _run(args: list[str], cwd: str | None = None, env: dict[str, str] | None = None) -> str:
-    print("$ " + " ".join(shlex.quote(arg) for arg in args), flush=True)
-    result = subprocess.run(
-        args,
-        cwd=cwd,
-        env=env,
-        check=True,
-        stdout=subprocess.PIPE,
-        stderr=sys.stderr,
-        text=True,
-    )
-    if result.stdout:
-        print(result.stdout, end="", flush=True)
-    return result.stdout.strip()
-
-
-def _run_shell(command: str, cwd: str, env: dict[str, str]) -> None:
-    print(f"$ {command}", flush=True)
-    result = subprocess.run(
-        ["/bin/bash", "-lc", command],
-        cwd=cwd,
-        env=env,
-        stdout=sys.stdout,
-        stderr=sys.stderr,
-        check=False,
-    )
-    if result.returncode != 0:
-        raise RuntimeError(f"Command failed with exit code {result.returncode}: {command}")
-
-
-def _parse_env_vars(env_vars: str) -> dict[str, str]:
-    parsed: dict[str, str] = {}
-    for item in env_vars.split(","):
-        item = item.strip()
-        if not item:
-            continue
-        if "=" not in item:
-            raise RuntimeError(f"Invalid env var override {item!r}; expected KEY=VALUE.")
-        key, value = item.split("=", 1)
-        parsed[key.strip()] = value.strip()
-    return parsed
-
-
-def _activate_remote_python_env(env: dict[str, str]) -> dict[str, str]:
-    venv_bin = "/opt/venv/bin"
-    if os.path.isdir(venv_bin):
-        env["VIRTUAL_ENV"] = "/opt/venv"
-        env["PATH"] = venv_bin + os.pathsep + env.get("PATH", "")
-    return env
-
-
-def _clone_checkout(git_repo: str, git_commit: str, pr_number: str) -> str:
-    last_clone_error: subprocess.CalledProcessError | None = None
-    for attempt in range(1, 4):
-        shutil.rmtree(REPO_DIR, ignore_errors=True)
-        try:
-            _run(
-                [
-                    "git",
-                    "-c",
-                    "http.version=HTTP/1.1",
-                    "clone",
-                    git_repo,
-                    REPO_DIR,
-                ],
-                cwd="/",
-            )
-            break
-        except subprocess.CalledProcessError as error:
-            last_clone_error = error
-            if attempt == 3:
-                raise
-            sleep_seconds = 5 * attempt
-            print(
-                f"git clone failed on attempt {attempt}; retrying in {sleep_seconds}s",
-                flush=True,
-            )
-            time.sleep(sleep_seconds)
-    if last_clone_error is not None and not os.path.isdir(REPO_DIR):
-        raise last_clone_error
-    if pr_number and pr_number != "false":
-        _run(["git", "fetch", "--prune", "origin", f"refs/pull/{pr_number}/head"], cwd=REPO_DIR)
-        _run(["git", "checkout", "FETCH_HEAD"], cwd=REPO_DIR)
-    else:
-        _run(["git", "checkout", git_commit], cwd=REPO_DIR)
-    _run(["git", "submodule", "update", "--init", "--recursive"], cwd=REPO_DIR)
-    return _run(["git", "rev-parse", "HEAD"], cwd=REPO_DIR)
-
-
-def _install_fastvideo(install_extra: str, env: dict[str, str]) -> None:
-    install_extra = install_extra.strip()
-    if install_extra.lower() in {"", "none", "skip", "false"}:
-        return
-    package = "." if install_extra == "." else f".[{install_extra}]"
-    _run_shell(
-        "source $HOME/.local/bin/env 2>/dev/null || true; "
-        "source /opt/venv/bin/activate 2>/dev/null || true; "
-        f"uv pip install -e {shlex.quote(package)}",
-        cwd=REPO_DIR,
-        env=env,
-    )
-
-
-def _build_kernel(env: dict[str, str]) -> None:
-    _run_shell(
-        "source $HOME/.local/bin/env 2>/dev/null || true; "
-        "source /opt/venv/bin/activate 2>/dev/null || true; "
-        "./build.sh",
-        cwd=os.path.join(REPO_DIR, "fastvideo-kernel"),
-        env=env,
-    )
-
-
-def _run_gpu_job(
-    command: str,
-    git_repo: str,
-    git_commit: str,
-    pr_number: str,
-    install_extra: str,
-    build_kernel: bool,
-    env_vars: str,
-    local_patch_b64: str,
-    commit_volume: bool,
-) -> dict[str, Any]:
-    remote_env = _activate_remote_python_env(os.environ.copy())
-    remote_env.update({
-        "HF_HOME": "/root/data/.cache",
-        "TOKENIZERS_PARALLELISM": "false",
-        "FASTVIDEO_ATTENTION_BACKEND": remote_env.get("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN"),
-    })
-    remote_env.update(_parse_env_vars(env_vars))
-
-    print(f"Cloning repository: {git_repo}")
-    print(f"Target commit: {git_commit}")
-    if pr_number and pr_number != "false":
-        print(f"Using PR ref: {pr_number}")
-    checked_out_commit = _clone_checkout(git_repo, git_commit, pr_number)
-    print(f"Checked out commit: {checked_out_commit}")
-    _apply_local_patch(local_patch_b64)
-
-    _install_fastvideo(install_extra, remote_env)
-    if build_kernel:
-        _build_kernel(remote_env)
-
-    try:
-        _run_shell(command, cwd=REPO_DIR, env=remote_env)
-    finally:
-        if commit_volume:
-            print("Committing Modal volume", flush=True)
-            model_vol.commit()
-    return {
-        "command": command,
-        "git_repo": git_repo,
-        "git_commit": checked_out_commit,
-        "install_extra": install_extra,
-        "build_kernel": build_kernel,
-        "local_patch_applied": bool(local_patch_b64),
-        "commit_volume": commit_volume,
-    }
-
-
-@app.function(gpu="L40S:1", **COMMON_FUNCTION_KWARGS)
-def run_l40s_1(
-    command: str,
-    git_repo: str,
-    git_commit: str,
-    pr_number: str,
-    install_extra: str,
-    build_kernel: bool,
-    env_vars: str,
-    local_patch_b64: str,
-    commit_volume: bool,
-):
-    return _run_gpu_job(command, git_repo, git_commit, pr_number, install_extra, build_kernel, env_vars,
-                        local_patch_b64, commit_volume)
-
-
-@app.function(gpu="L40S:2", **COMMON_FUNCTION_KWARGS)
-def run_l40s_2(
-    command: str,
-    git_repo: str,
-    git_commit: str,
-    pr_number: str,
-    install_extra: str,
-    build_kernel: bool,
-    env_vars: str,
-    local_patch_b64: str,
-    commit_volume: bool,
-):
-    return _run_gpu_job(command, git_repo, git_commit, pr_number, install_extra, build_kernel, env_vars,
-                        local_patch_b64, commit_volume)
-
-
-@app.function(gpu="L40S:4", **COMMON_FUNCTION_KWARGS)
-def run_l40s_4(
-    command: str,
-    git_repo: str,
-    git_commit: str,
-    pr_number: str,
-    install_extra: str,
-    build_kernel: bool,
-    env_vars: str,
-    local_patch_b64: str,
-    commit_volume: bool,
-):
-    return _run_gpu_job(command, git_repo, git_commit, pr_number, install_extra, build_kernel, env_vars,
-                        local_patch_b64, commit_volume)
-
-
-@app.function(gpu="L40S:8", **COMMON_FUNCTION_KWARGS)
-def run_l40s_8(
-    command: str,
-    git_repo: str,
-    git_commit: str,
-    pr_number: str,
-    install_extra: str,
-    build_kernel: bool,
-    env_vars: str,
-    local_patch_b64: str,
-    commit_volume: bool,
-):
-    return _run_gpu_job(command, git_repo, git_commit, pr_number, install_extra, build_kernel, env_vars,
-                        local_patch_b64, commit_volume)
-
-
-@app.function(gpu="H100:1", **COMMON_FUNCTION_KWARGS)
-def run_h100_1(
-    command: str,
-    git_repo: str,
-    git_commit: str,
-    pr_number: str,
-    install_extra: str,
-    build_kernel: bool,
-    env_vars: str,
-    local_patch_b64: str,
-    commit_volume: bool,
-):
-    return _run_gpu_job(command, git_repo, git_commit, pr_number, install_extra, build_kernel, env_vars,
-                        local_patch_b64, commit_volume)
-
-
-@app.function(gpu="H100:2", **COMMON_FUNCTION_KWARGS)
-def run_h100_2(
-    command: str,
-    git_repo: str,
-    git_commit: str,
-    pr_number: str,
-    install_extra: str,
-    build_kernel: bool,
-    env_vars: str,
-    local_patch_b64: str,
-    commit_volume: bool,
-):
-    return _run_gpu_job(command, git_repo, git_commit, pr_number, install_extra, build_kernel, env_vars,
-                        local_patch_b64, commit_volume)
-
-
-def _select_runner(gpu_type: str, num_gpus: int) -> Callable[..., Any]:
-    normalized_gpu_type = gpu_type.upper()
-    runners = {
-        ("L40S", 1): run_l40s_1,
-        ("L40S", 2): run_l40s_2,
-        ("L40S", 4): run_l40s_4,
-        ("L40S", 8): run_l40s_8,
-        ("H100", 1): run_h100_1,
-        ("H100", 2): run_h100_2,
-    }
-    try:
-        return runners[(normalized_gpu_type, num_gpus)]
-    except KeyError as error:
-        supported = ", ".join(f"{gpu}:{count}" for gpu, count in sorted(runners))
-        raise RuntimeError(f"Unsupported GPU request {gpu_type}:{num_gpus}. Supported requests: {supported}.") from error
-
-
-@app.local_entrypoint()
-def main(
-    command: str = "nvidia-smi",
-    gpu_type: str = "L40S",
-    num_gpus: int = 1,
-    git_repo: str = "",
-    git_commit: str = "",
-    pr_number: str = "",
-    install_extra: str = "dev",
-    build_kernel: bool = False,
-    env_vars: str = "",
-    apply_local_patch: bool = False,
-    patch_paths: str = "",
-    wait: bool = True,
-    commit_volume: bool = False,
-):
-    normalized_gpu_type = gpu_type.upper()
-    resolved_git_repo = _resolve_git_repo(git_repo)
-    resolved_git_commit = _resolve_git_commit(git_commit)
-    resolved_pr_number = _resolve_pull_request(pr_number)
-    runner = _select_runner(normalized_gpu_type, num_gpus)
-
-    print(f"Launching {normalized_gpu_type}:{num_gpus} job")
-    print(f"Command: {command}")
-    print(f"Repo: {resolved_git_repo}")
-    print(f"Commit: {resolved_git_commit}")
-    if resolved_pr_number and resolved_pr_number != "false":
-        print(f"PR ref: {resolved_pr_number}")
-    local_patch_b64 = ""
-    if apply_local_patch:
-        patch = _build_local_patch(patch_paths)
-        local_patch_b64 = base64.b64encode(patch.encode("utf-8")).decode("ascii")
-        print(f"Local patch payload: {len(patch)} bytes")
-
-    kwargs = dict(
-        command=command,
-        git_repo=resolved_git_repo,
-        git_commit=resolved_git_commit,
-        pr_number=resolved_pr_number,
-        install_extra=install_extra,
-        build_kernel=build_kernel,
-        env_vars=env_vars,
-        local_patch_b64=local_patch_b64,
-        commit_volume=commit_volume,
-    )
-    if wait:
-        result = runner.remote(**kwargs)
-        print(f"Completed {normalized_gpu_type} job: {result}")
-        return
-
-    function_call = runner.spawn(**kwargs)
-    print(f"Spawned Modal FunctionCall: {function_call.object_id}")
-    print("Poll later with:")
-    print(f"  python -c \"import modal; print(modal.FunctionCall.from_id('{function_call.object_id}').get())\"")

fastvideo/tests/ssim/test_flux2_similarity.py

@@ -1,114 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Latent-slice regression tests for Flux2 text-to-image variants.
-
-Flux2 currently has local parity coverage against the official/reference
-pipeline, but CI needs a small deterministic regression gate for seeded HF
-artefacts.  Pixel-space comparisons are unnecessarily brittle for this first
-slot, so the test follows the latent helper pattern used by LTX-2: generate a
-single-image latent with the production recipe, persist the generated latent,
-and compare a stable latent signature plus the full tensor against the device
-reference.
-
-The default and full-quality parameter maps intentionally carry the same
-recipe values for now.  The ``--ssim-full-quality`` flag still switches the
-reference tier through ``conftest.py``; separate full-quality recipes can be
-introduced after the initial Flux2 references have a stable CI window.
-"""
-
-from __future__ import annotations
-
-import os
-
-import pytest
-import torch
-
-from fastvideo.logger import init_logger
-from fastvideo.tests.ssim.inference_similarity_utils import (
-    resolve_inference_device_reference_folder,
-)
-from fastvideo.tests.ssim.latent_similarity_utils import (
-    run_text_to_latent_similarity_test,
-)
-
-logger = init_logger(__name__)
-
-REQUIRED_GPUS = 1
-
-device_reference_folder = resolve_inference_device_reference_folder(logger)
-
-FLUX2_MODEL_TO_PARAMS: dict[str, dict[str, object]] = {
-    "black-forest-labs/FLUX.2-klein-4B": {
-        "num_gpus": 1,
-        "model_path": "black-forest-labs/FLUX.2-klein-4B",
-        "height": 1024,
-        "width": 1024,
-        "num_frames": 1,
-        "num_inference_steps": 4,
-        "guidance_scale": 1.0,
-        "seed": 0,
-        "sp_size": 1,
-        "tp_size": 1,
-        "fps": 1,
-    },
-    "black-forest-labs/FLUX.2-klein-9B": {
-        "num_gpus": 1,
-        "model_path": "black-forest-labs/FLUX.2-klein-9B",
-        "height": 1024,
-        "width": 1024,
-        "num_frames": 1,
-        "num_inference_steps": 4,
-        "guidance_scale": 1.0,
-        "seed": 0,
-        "sp_size": 1,
-        "tp_size": 1,
-        "fps": 1,
-    },
-}
-
-FLUX2_FULL_QUALITY_MODEL_TO_PARAMS: dict[str, dict[str, object]] = {
-    model_id: dict(params)
-    for model_id, params in FLUX2_MODEL_TO_PARAMS.items()
-}
-
-TEST_PROMPTS: dict[str, str] = {
-    "black-forest-labs/FLUX.2-klein-4B": "a brushed steel espresso machine on a marble counter, morning window light",
-    "black-forest-labs/FLUX.2-klein-9B": "a brushed steel espresso machine on a marble counter, morning window light",
-}
-
-SLICE_COSINE_THRESHOLD = 0.96
-FULL_COSINE_THRESHOLD = 0.99
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="Flux2 SSIM test requires CUDA",
-)
-@pytest.mark.parametrize("attention_backend_name", ["TORCH_SDPA"])
-@pytest.mark.parametrize("model_id", list(FLUX2_MODEL_TO_PARAMS.keys()))
-def test_flux2_similarity(
-    attention_backend_name: str,
-    model_id: str,
-) -> None:
-    _ = run_text_to_latent_similarity_test(
-        logger=logger,
-        script_dir=os.path.dirname(os.path.abspath(__file__)),
-        device_reference_folder=device_reference_folder,
-        prompt=TEST_PROMPTS[model_id],
-        attention_backend_name=attention_backend_name,
-        model_id=model_id,
-        default_params_map=FLUX2_MODEL_TO_PARAMS,
-        full_quality_params_map=FLUX2_FULL_QUALITY_MODEL_TO_PARAMS,
-        slice_cosine_threshold=SLICE_COSINE_THRESHOLD,
-        full_cosine_threshold=FULL_COSINE_THRESHOLD,
-        init_kwargs_override={
-            "workload_type": "t2i",
-            "use_fsdp_inference": False,
-            "dit_cpu_offload": False,
-            "vae_cpu_offload": True,
-            "text_encoder_cpu_offload": True,
-            "pin_cpu_memory": False,
-            "override_pipeline_cls_name": (
-                "Flux2KleinPipeline" if "klein" in model_id.lower() else "Flux2Pipeline"
-            ),
-        },
-    )

fastvideo/train/callbacks/validation.py

@@ -362,6 +362,7 @@ class ValidationCallback(Callback):
 
         batch = ForwardBatch(
             **shallow_asdict(sampling_param),
+            latents=None,
             generator=self.validation_random_generator,
             n_tokens=n_tokens,
             eta=0.0,

fastvideo/train/models/wan/wan.py

@@ -401,7 +401,6 @@ class WanModel(ModelBase):
                 patch_size=patch_size,
                 VSA_sparsity=tc.vsa_sparsity,
                 device=self.device,
-                cache_tile_buf=False,
             )
         elif (envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN"):
             if (not is_vmoba_available() or VideoMobaAttentionMetadataBuilder is None):

fastvideo/training/distillation_pipeline.py

@@ -248,7 +248,13 @@ class DistillationPipeline(TrainingPipeline):
         ema_enabled = (self.training_args.ema_decay is not None) and (self.training_args.ema_decay > 0.0)
         if ema_enabled and (self.training_args.ema_start_step <= 0):
             # Only eager-construct from the cold init weights when averaging starts at step 0.
-            self._build_generator_emas(context="eager init, ema_start_step<=0")
+            self.generator_ema = EMA_FSDP(self.transformer, decay=self.training_args.ema_decay)
+            logger.info("Initialized generator EMA with decay=%s", self.training_args.ema_decay)
+
+            # Initialize EMA for transformer_2 if it exists
+            if self.transformer_2 is not None:
+                self.generator_ema_2 = EMA_FSDP(self.transformer_2, decay=self.training_args.ema_decay)
+                logger.info("Initialized generator EMA_2 with decay=%s", self.training_args.ema_decay)
         elif ema_enabled:
             # Defer construction to the lazy block in the train loop, which builds the EMA AT
             # ema_start_step from the already-trained weights. Eager-constructing here would anchor
@@ -898,45 +904,11 @@ class DistillationPipeline(TrainingPipeline):
         training_batch.total_loss = training_batch.generator_loss + training_batch.fake_score_loss
         return training_batch
 
-    def _build_generator_emas(self, context: str = "") -> None:
-        # Idempotently construct whichever generator EMA shadows are missing, per-expert and
-        # decoupled. Safe to call repeatedly; no-op once both exist or when EMA is disabled.
-        if (self.training_args.ema_decay is None) or (self.training_args.ema_decay <= 0.0):
-            return
-        suffix = f" [{context}]" if context else ""
-        if self.generator_ema is None:
-            self.generator_ema = EMA_FSDP(self.transformer, decay=self.training_args.ema_decay)
-            logger.info("Built generator EMA (decay=%s)%s", self.training_args.ema_decay, suffix)
-        if self.transformer_2 is not None and self.generator_ema_2 is None:
-            self.generator_ema_2 = EMA_FSDP(self.transformer_2, decay=self.training_args.ema_decay)
-            logger.info("Built generator EMA_2 (decay=%s)%s", self.training_args.ema_decay, suffix)
-
-    def _build_deferred_ema_for_resume(self) -> None:
-        # Build a deferred EMA before checkpoint load only when a saved shard exists, so the
-        # shadow reloads instead of being skipped and rebuilt fresh. Gating on shard existence
-        # matters: building when none exists would reintroduce cold-init contamination.
-        if (self.training_args.ema_decay is None) or (self.training_args.ema_decay <= 0.0):
-            return
-
-        ema_shard_dir = os.path.join(self.training_args.resume_from_checkpoint, "ema_local_shard")
-
-        if (self.generator_ema is None
-                and os.path.exists(os.path.join(ema_shard_dir, f"generator_ema_rank{self.global_rank}.pt"))):
-            self.generator_ema = EMA_FSDP(self.transformer, decay=self.training_args.ema_decay)
-            logger.info("Pre-built generator EMA for resume from existing shard")
-
-        if (self.transformer_2 is not None and self.generator_ema_2 is None
-                and os.path.exists(os.path.join(ema_shard_dir, f"generator_ema_2_rank{self.global_rank}.pt"))):
-            self.generator_ema_2 = EMA_FSDP(self.transformer_2, decay=self.training_args.ema_decay)
-            logger.info("Pre-built generator EMA_2 for resume from existing shard")
-
     def _resume_from_checkpoint(self) -> None:
         """Resume training from checkpoint with distillation models."""
 
         logger.info("Loading distillation checkpoint from %s", self.training_args.resume_from_checkpoint)
 
-        self._build_deferred_ema_for_resume()
-
         resumed_step = load_distillation_checkpoint(
             self.transformer,
             self.fake_score_transformer,
@@ -1345,8 +1317,15 @@ class DistillationPipeline(TrainingPipeline):
             self.current_trainstep = step
             training_batch.current_vsa_sparsity = current_vsa_sparsity
 
-            if step >= self.training_args.ema_start_step:
-                self._build_generator_emas(context=f"lazy @ step {step}")
+            if (step >= self.training_args.ema_start_step) and \
+                    (self.generator_ema is None) and (self.training_args.ema_decay > 0):
+                self.generator_ema = EMA_FSDP(self.transformer, decay=self.training_args.ema_decay)
+                logger.info("Created generator EMA at step %s with decay=%s", step, self.training_args.ema_decay)
+
+                # Create EMA for transformer_2 if it exists
+                if self.transformer_2 is not None and self.generator_ema_2 is None:
+                    self.generator_ema_2 = EMA_FSDP(self.transformer_2, decay=self.training_args.ema_decay)
+                    logger.info("Created generator EMA_2 at step %s with decay=%s", step, self.training_args.ema_decay)
 
             with torch.autocast("cuda", dtype=torch.bfloat16):
                 training_batch = self.train_one_step(training_batch)

fastvideo/training/matrixgame2_ar_diffusion_pipeline.py

@@ -405,6 +405,7 @@ class MatrixGame2ARDiffusionPipeline(TrainingPipeline):
         sampling_param.num_frames = num_frames
         batch = ForwardBatch(
             **shallow_asdict(sampling_param),
+            latents=None,
             generator=torch.Generator(device="cpu").manual_seed(self.seed),
             n_tokens=n_tokens,
             eta=0.0,

fastvideo/training/matrixgame2_ode_causal_pipeline.py

@@ -351,6 +351,7 @@ class MatrixGame2ODEInitTrainingPipeline(TrainingPipeline):
         sampling_param.num_frames = num_frames
         batch = ForwardBatch(
             **shallow_asdict(sampling_param),
+            latents=None,
             generator=torch.Generator(device="cpu").manual_seed(self.seed),
             n_tokens=n_tokens,
             eta=0.0,

fastvideo/training/matrixgame2_self_forcing_distillation_pipeline.py

@@ -853,6 +853,7 @@ class MatrixGame2SelfForcingDistillationPipeline(SelfForcingDistillationPipeline
         sampling_param.num_frames = num_frames
         batch = ForwardBatch(
             **shallow_asdict(sampling_param),
+            latents=None,
             generator=torch.Generator(device="cpu").manual_seed(self.seed),
             n_tokens=n_tokens,
             eta=0.0,

fastvideo/training/matrixgame2_training_pipeline.py

@@ -172,6 +172,7 @@ class MatrixGame2TrainingPipeline(TrainingPipeline):
         sampling_param.num_frames = num_frames
         batch = ForwardBatch(
             **shallow_asdict(sampling_param),
+            latents=None,
             generator=torch.Generator(device="cpu").manual_seed(self.seed),
             n_tokens=n_tokens,
             eta=0.0,

fastvideo/training/self_forcing_distillation_pipeline.py

@@ -852,8 +852,15 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
             self.current_trainstep = step
             training_batch.current_vsa_sparsity = current_vsa_sparsity
 
-            if step >= self.training_args.ema_start_step:
-                self._build_generator_emas(context=f"lazy @ step {step}")
+            if (step >= self.training_args.ema_start_step) and \
+                    (self.generator_ema is None) and (self.training_args.ema_decay > 0):
+                self.generator_ema = EMA_FSDP(self.transformer, decay=self.training_args.ema_decay)
+                logger.info("Created generator EMA at step %s with decay=%s", step, self.training_args.ema_decay)
+
+                # Create EMA for transformer_2 if it exists
+                if self.transformer_2 is not None and self.generator_ema_2 is None:
+                    self.generator_ema_2 = EMA_FSDP(self.transformer_2, decay=self.training_args.ema_decay)
+                    logger.info("Created generator EMA_2 at step %s with decay=%s", step, self.training_args.ema_decay)
 
             with torch.autocast("cuda", dtype=torch.bfloat16):
                 training_batch = self.train_one_step(training_batch)

fastvideo/training/training_pipeline.py

@@ -359,8 +359,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 current_timestep=training_batch.timesteps,
                 patch_size=patch_size,
                 VSA_sparsity=current_vsa_sparsity,
-                device=get_local_torch_device(),
-                cache_tile_buf=False)
+                device=get_local_torch_device())
         elif envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN":
             if not vmoba_available:
                 raise ImportError("FASTVIDEO_ATTENTION_BACKEND is set to VMOBA_ATTN, "
@@ -691,6 +690,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
         sampling_param.num_frames = num_frames
         batch = ForwardBatch(
             **shallow_asdict(sampling_param),
+            latents=None,
             generator=self.validation_random_generator,
             n_tokens=n_tokens,
             eta=0.0,

fastvideo/training/training_utils.py

@@ -613,10 +613,18 @@ def load_distillation_checkpoint(
                 end_time - begin_time,
                 local_main_process_only=False)
 
-    # Load EMA separately
+    # Load EMA separately if saved in rank0_full mode
     if generator_ema is not None:
         try:
-            if getattr(generator_ema, "mode", None) == "local_shard":
+            if getattr(generator_ema, "mode", None) == "rank0_full":
+                ema_path = os.path.join(checkpoint_path, "ema", "generator_ema.pt")
+                if rank == 0 and os.path.exists(ema_path):
+                    ema_state = torch.load(ema_path, map_location="cpu")
+                    generator_ema.load_state_dict(ema_state)
+                    logger.info("rank: %s, generator EMA (rank0_full) loaded from %s", rank, ema_path)
+                elif rank == 0:
+                    logger.info("rank: %s, generator EMA file not found at %s; skipping", rank, ema_path)
+            elif getattr(generator_ema, "mode", None) == "local_shard":
                 ema_path = os.path.join(checkpoint_path, "ema_local_shard", f"generator_ema_rank{rank}.pt")
                 if os.path.exists(ema_path):
                     ema_state = torch.load(ema_path, map_location="cpu")
@@ -624,28 +632,9 @@ def load_distillation_checkpoint(
                     logger.info("rank: %s, generator EMA shard (local_shard) loaded from %s", rank, ema_path)
                 else:
                     logger.info("rank: %s, generator EMA shard file not found at %s; skipping", rank, ema_path)
-            else:
-                logger.info("rank: %s, generator EMA mode %s not supported for resume; skipping", rank,
-                            getattr(generator_ema, "mode", None))
         except Exception as e:
             logger.warning("rank: %s, failed to load generator EMA: %s", rank, str(e))
 
-    # Load EMA_2 from its shard (symmetric with generator_ema above)
-    if generator_ema_2 is not None:
-        try:
-            if getattr(generator_ema_2, "mode", None) == "local_shard":
-                ema_2_path = os.path.join(checkpoint_path, "ema_local_shard", f"generator_ema_2_rank{rank}.pt")
-                if os.path.exists(ema_2_path):
-                    generator_ema_2.load_state_dict(torch.load(ema_2_path, map_location="cpu"))
-                    logger.info("rank: %s, generator EMA_2 shard (local_shard) loaded from %s", rank, ema_2_path)
-                else:
-                    logger.info("rank: %s, generator EMA_2 shard file not found at %s; skipping", rank, ema_2_path)
-            else:
-                logger.info("rank: %s, generator EMA_2 mode %s not supported for resume; skipping", rank,
-                            getattr(generator_ema_2, "mode", None))
-        except Exception as e:
-            logger.warning("rank: %s, failed to load generator EMA_2: %s", rank, str(e))
-
     # Load generator_2 distributed checkpoint (MoE support)
     if generator_transformer_2 is not None:
         generator_2_dcp_dir = os.path.join(checkpoint_path, "distributed_checkpoint", "generator_2")
@@ -676,6 +665,18 @@ def load_distillation_checkpoint(
                         rank,
                         end_time - begin_time,
                         local_main_process_only=False)
+
+            # Load EMA_2 state if available and generator_ema_2 is provided
+            if generator_ema_2 is not None:
+                try:
+                    ema_2_state = generator_2_states.get("ema")
+                    if ema_2_state is not None:
+                        generator_ema_2.load_state_dict(ema_2_state)
+                        logger.info("rank: %s, generator_2 EMA state loaded successfully", rank)
+                    else:
+                        logger.info("rank: %s, no EMA_2 state found in checkpoint", rank)
+                except Exception as e:
+                    logger.warning("rank: %s, failed to load EMA_2 state: %s", rank, str(e))
         else:
             logger.info("rank: %s, generator_2 checkpoint not found, skipping", rank)
 

fastvideo/training/wan_i2v_distillation_pipeline.py

@@ -114,6 +114,7 @@ class WanI2VDistillationPipeline(DistillationPipeline):
         sampling_param.num_frames = num_frames
         batch = ForwardBatch(
             **shallow_asdict(sampling_param),
+            latents=None,
             generator=torch.Generator(device="cpu").manual_seed(self.seed),
             n_tokens=n_tokens,
             eta=0.0,

fastvideo/training/wan_i2v_training_pipeline.py

@@ -157,6 +157,7 @@ class WanI2VTrainingPipeline(TrainingPipeline):
         sampling_param.num_frames = num_frames
         batch = ForwardBatch(
             **shallow_asdict(sampling_param),
+            latents=None,
             generator=torch.Generator(device="cpu").manual_seed(self.seed),
             n_tokens=n_tokens,
             eta=0.0,

fastvideo/worker/multiproc_executor.py

@@ -138,9 +138,7 @@ class MultiprocExecutor(Executor):
         result_batch = ForwardBatch(data_type=forward_batch.data_type,
                                     output=output,
                                     logging_info=logging_info,
-                                    extra=extra,
-                                    trajectory_latents=responses[0].get("trajectory_latents"),
-                                    trajectory_timesteps=responses[0].get("trajectory_timesteps"))
+                                    extra=extra)
 
         return result_batch
 
@@ -701,8 +699,6 @@ class WorkerMultiprocProc:
                             "output_batch": result,
                             "logging_info": logging_info,
                             "extra": extra,
-                            "trajectory_latents": output_batch.trajectory_latents,
-                            "trajectory_timesteps": output_batch.trajectory_timesteps,
                         })
                     else:
                         result = self.worker.execute_method(method, *args, **kwargs)

fastvideo/worker/ray_distributed_executor.py

@@ -307,8 +307,6 @@ class RayDistributedExecutor(Executor):
             data_type=forward_batch.data_type,
             output=output,
             logging_info=logging_info,
-            trajectory_latents=responses[0].trajectory_latents,
-            trajectory_timesteps=responses[0].trajectory_timesteps,
         )
         return result_batch
 

pyproject.toml

@@ -26,7 +26,7 @@ dependencies = [
     "sentencepiece>=0.2.0",
     "timm>=1.0.11",
     "peft>=0.15.0",
-    "diffusers>=0.38.0",
+    "diffusers>=0.33.1",
     "torch==2.11.0",
     "torchvision",
     "torchaudio",

scripts/checkpoint_conversion/convert_flux2_klein.py

@@ -1,527 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-License-Identifier: Apache-2.0
-# pyright: reportAny=false, reportExplicitAny=false, reportUnknownMemberType=false
-# pyright: reportUnknownVariableType=false, reportUnusedCallResult=false
-"""Convert FLUX.2 Klein weights into a FastVideo-loadable layout.
-
-The published ``black-forest-labs/FLUX.2-klein-4B`` repo contains two useful
-transformer surfaces:
-
-* ``flux-2-klein-4b.safetensors``: BFL's compact raw transformer checkpoint.
-  Its double-stream attention projections are fused as ``img_attn.qkv`` and
-  ``txt_attn.qkv``.
-* ``transformer/diffusion_pytorch_model.safetensors``: Diffusers/FastVideo
-  names. These keys already match ``Flux2Transformer2DModel.state_dict()``.
-
-By default this script prefers the raw root checkpoint when present, converts it
-to the FastVideo native transformer key surface, reloads/resaves the VAE, copies
-the HF-backed Qwen3 text encoder/tokenizer and scheduler, and emits a standard
-Diffusers-style FastVideo repo:
-
-    <dst>/
-      model_index.json
-      transformer/{config.json,diffusion_pytorch_model*.safetensors}
-      vae/{config.json,diffusion_pytorch_model*.safetensors}
-      text_encoder/...
-      tokenizer/...
-      scheduler/scheduler_config.json
-
-Qwen3 and Mistral3 are intentionally copied as Transformers passthrough
-components in the standard layout: ``qwen3.py`` and ``mistral3.py`` both route
-Flux2 text encoding through ``from_pretrained_local()`` for exact HF parity.
-
-Example:
-    python scripts/checkpoint_conversion/convert_flux2_klein.py \
-        --src black-forest-labs/FLUX.2-klein-4B \
-        --dst converted_weights/flux2-klein-4b-fastvideo \
-        --overwrite
-"""
-
-from __future__ import annotations
-
-import argparse
-import json
-import os
-import re
-import shutil
-from collections import OrderedDict
-from pathlib import Path
-from typing import Any
-
-import torch
-from safetensors import safe_open
-
-try:
-    from huggingface_hub import save_torch_state_dict, snapshot_download
-except ImportError:
-    save_torch_state_dict = None
-    snapshot_download = None
-
-
-DEFAULT_REPO_ID = "black-forest-labs/FLUX.2-klein-4B"
-RAW_TRANSFORMER_FILENAME = "flux-2-klein-4b.safetensors"
-DIFFUSION_WEIGHTS_BASENAME = "diffusion_pytorch_model"
-
-BASE_SNAPSHOT_ALLOW_PATTERNS = (
-    "model_index.json",
-    "transformer/config.json",
-    "vae/config.json",
-    "vae/*.safetensors",
-    "vae/*.safetensors.index.json",
-    "text_encoder/*",
-    "tokenizer/*",
-    "scheduler/*",
-)
-
-PASSTHROUGH_SUBFOLDERS = ("text_encoder", "tokenizer", "scheduler")
-
-DEFAULT_MODEL_INDEX: dict[str, Any] = {
-    "_class_name": "Flux2KleinPipeline",
-    "_diffusers_version": "0.37.0.dev0",
-    "is_distilled": True,
-    "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"],
-    "text_encoder": ["transformers", "Qwen3ForCausalLM"],
-    "tokenizer": ["transformers", "Qwen2TokenizerFast"],
-    "transformer": ["diffusers", "Flux2Transformer2DModel"],
-    "vae": ["diffusers", "AutoencoderKLFlux2"],
-}
-
-DEFAULT_SCHEDULER_CONFIG: dict[str, Any] = {
-    "_class_name": "FlowMatchEulerDiscreteScheduler",
-    "_diffusers_version": "0.37.0.dev0",
-    "base_image_seq_len": 256,
-    "base_shift": 0.5,
-    "invert_sigmas": False,
-    "max_image_seq_len": 4096,
-    "max_shift": 1.15,
-    "num_train_timesteps": 1000,
-    "shift": 3.0,
-    "shift_terminal": None,
-    "stochastic_sampling": False,
-    "time_shift_type": "exponential",
-    "use_beta_sigmas": False,
-    "use_dynamic_shifting": True,
-    "use_exponential_sigmas": False,
-    "use_karras_sigmas": False,
-}
-
-TRANSFORMER_REQUIRED_KEYS = (
-    "x_embedder.weight",
-    "context_embedder.weight",
-    "time_guidance_embed.timestep_embedder.linear_1.weight",
-    "time_guidance_embed.timestep_embedder.linear_2.weight",
-    "double_stream_modulation_img.linear.weight",
-    "double_stream_modulation_txt.linear.weight",
-    "single_stream_modulation.linear.weight",
-    "norm_out.linear.weight",
-    "proj_out.weight",
-)
-
-VAE_REQUIRED_KEYS = (
-    "encoder.conv_in.weight",
-    "decoder.conv_out.weight",
-    "quant_conv.weight",
-    "post_quant_conv.weight",
-    "bn.running_mean",
-    "bn.running_var",
-)
-
-
-class ConversionError(RuntimeError):
-    pass
-
-
-def _resolve_hf_token() -> str | bool | None:
-    try:
-        from fastvideo.utils import resolve_hf_token
-
-        token = resolve_hf_token()
-        return token if token else None
-    except Exception:
-        return os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN") or None
-
-
-def _snapshot_allow_patterns(transformer_source: str) -> list[str]:
-    patterns: list[str] = list(BASE_SNAPSHOT_ALLOW_PATTERNS)
-    if transformer_source == "diffusers":
-        patterns.extend(("transformer/*.safetensors", "transformer/*.safetensors.index.json"))
-    else:
-        patterns.append(RAW_TRANSFORMER_FILENAME)
-    return patterns
-
-
-def _resolve_src(src: str, revision: str | None, cache_dir: str | None, transformer_source: str) -> Path:
-    local = Path(src).expanduser()
-    if local.exists():
-        return local
-
-    if snapshot_download is None:
-        raise ConversionError("huggingface_hub is required when --src is a repo id")
-
-    return Path(
-        snapshot_download(
-            repo_id=src,
-            revision=revision,
-            cache_dir=cache_dir,
-            token=_resolve_hf_token(),
-            allow_patterns=_snapshot_allow_patterns(transformer_source),
-        )
-    )
-
-
-def _read_json(path: Path) -> dict[str, Any]:
-    with path.open("r", encoding="utf-8") as f:
-        return json.load(f)
-
-
-def _write_json(path: Path, payload: dict[str, Any]) -> None:
-    path.parent.mkdir(parents=True, exist_ok=True)
-    with path.open("w", encoding="utf-8") as f:
-        json.dump(payload, f, indent=2)
-        f.write("\n")
-
-
-def _prepare_output_dir(dst: Path, overwrite: bool) -> None:
-    if dst.exists() and any(dst.iterdir()):
-        if not overwrite:
-            raise FileExistsError(f"Output directory is not empty: {dst}. Pass --overwrite to replace it.")
-        shutil.rmtree(dst)
-    dst.mkdir(parents=True, exist_ok=True)
-
-
-def _find_safetensors_files(component_dir: Path, basename: str = DIFFUSION_WEIGHTS_BASENAME) -> list[Path]:
-    if component_dir.is_file():
-        return [component_dir]
-
-    index_path = component_dir / f"{basename}.safetensors.index.json"
-    if index_path.exists():
-        index = _read_json(index_path)
-        return sorted({component_dir / shard for shard in index["weight_map"].values()})
-
-    single = component_dir / f"{basename}.safetensors"
-    if single.exists():
-        return [single]
-
-    return sorted(component_dir.glob("*.safetensors"))
-
-
-def _load_safetensors(files: list[Path]) -> OrderedDict[str, torch.Tensor]:
-    if not files:
-        raise FileNotFoundError("No safetensors files found")
-
-    state: OrderedDict[str, torch.Tensor] = OrderedDict()
-    for file in files:
-        with safe_open(str(file), framework="pt", device="cpu") as handle:
-            for key in handle.keys():
-                if key in state:
-                    raise ConversionError(f"Duplicate tensor key {key!r} while reading {file}")
-                state[key] = handle.get_tensor(key)
-    return state
-
-
-def _write_state_dict(
-    state: OrderedDict[str, torch.Tensor],
-    output_dir: Path,
-    max_shard_size: str,
-) -> None:
-    if save_torch_state_dict is None:
-        raise ConversionError("huggingface_hub.save_torch_state_dict is required to write sharded safetensors")
-    output_dir.mkdir(parents=True, exist_ok=True)
-    save_torch_state_dict(
-        state,
-        output_dir,
-        filename_pattern=f"{DIFFUSION_WEIGHTS_BASENAME}" + "{suffix}.safetensors",
-        max_shard_size=max_shard_size,
-        metadata={"format": "pt"},
-        safe_serialization=True,
-    )
-
-
-def _split_qkv(weight: torch.Tensor, source_key: str) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    if weight.shape[0] % 3 != 0:
-        raise ConversionError(f"Expected first dim divisible by 3 for {source_key}, got {tuple(weight.shape)}")
-    q_weight, k_weight, v_weight = torch.chunk(weight, 3, dim=0)
-    return q_weight, k_weight, v_weight
-
-
-def _convert_raw_transformer_key(
-    key: str,
-    value: torch.Tensor,
-    output: OrderedDict[str, torch.Tensor],
-) -> bool:
-    literal_map = {
-        "img_in.weight": "x_embedder.weight",
-        "txt_in.weight": "context_embedder.weight",
-        "time_in.in_layer.weight": "time_guidance_embed.timestep_embedder.linear_1.weight",
-        "time_in.out_layer.weight": "time_guidance_embed.timestep_embedder.linear_2.weight",
-        "double_stream_modulation_img.lin.weight": "double_stream_modulation_img.linear.weight",
-        "double_stream_modulation_txt.lin.weight": "double_stream_modulation_txt.linear.weight",
-        "single_stream_modulation.lin.weight": "single_stream_modulation.linear.weight",
-        "final_layer.adaLN_modulation.1.weight": "norm_out.linear.weight",
-        "final_layer.linear.weight": "proj_out.weight",
-    }
-    if key in literal_map:
-        output[literal_map[key]] = value
-        return True
-
-    match = re.fullmatch(r"double_blocks\.(\d+)\.(img|txt)_attn\.qkv\.weight", key)
-    if match:
-        block, stream = match.groups()
-        q_weight, k_weight, v_weight = _split_qkv(value, key)
-        if stream == "img":
-            prefix = f"transformer_blocks.{block}.attn"
-            output[f"{prefix}.to_q.weight"] = q_weight
-            output[f"{prefix}.to_k.weight"] = k_weight
-            output[f"{prefix}.to_v.weight"] = v_weight
-        else:
-            prefix = f"transformer_blocks.{block}.attn"
-            output[f"{prefix}.add_q_proj.weight"] = q_weight
-            output[f"{prefix}.add_k_proj.weight"] = k_weight
-            output[f"{prefix}.add_v_proj.weight"] = v_weight
-        return True
-
-    double_rewrites = (
-        (r"double_blocks\.(\d+)\.img_attn\.proj\.weight", r"transformer_blocks.\1.attn.to_out.0.weight"),
-        (r"double_blocks\.(\d+)\.txt_attn\.proj\.weight", r"transformer_blocks.\1.attn.to_add_out.weight"),
-        (r"double_blocks\.(\d+)\.img_attn\.norm\.query_norm\.scale", r"transformer_blocks.\1.attn.norm_q.weight"),
-        (r"double_blocks\.(\d+)\.img_attn\.norm\.key_norm\.scale", r"transformer_blocks.\1.attn.norm_k.weight"),
-        (r"double_blocks\.(\d+)\.txt_attn\.norm\.query_norm\.scale", r"transformer_blocks.\1.attn.norm_added_q.weight"),
-        (r"double_blocks\.(\d+)\.txt_attn\.norm\.key_norm\.scale", r"transformer_blocks.\1.attn.norm_added_k.weight"),
-        (r"double_blocks\.(\d+)\.img_mlp\.0\.weight", r"transformer_blocks.\1.ff.linear_in.weight"),
-        (r"double_blocks\.(\d+)\.img_mlp\.2\.weight", r"transformer_blocks.\1.ff.linear_out.weight"),
-        (r"double_blocks\.(\d+)\.txt_mlp\.0\.weight", r"transformer_blocks.\1.ff_context.linear_in.weight"),
-        (r"double_blocks\.(\d+)\.txt_mlp\.2\.weight", r"transformer_blocks.\1.ff_context.linear_out.weight"),
-        (r"single_blocks\.(\d+)\.linear1\.weight", r"single_transformer_blocks.\1.attn.to_qkv_mlp_proj.weight"),
-        (r"single_blocks\.(\d+)\.linear2\.weight", r"single_transformer_blocks.\1.attn.to_out.weight"),
-        (r"single_blocks\.(\d+)\.norm\.query_norm\.scale", r"single_transformer_blocks.\1.attn.norm_q.weight"),
-        (r"single_blocks\.(\d+)\.norm\.key_norm\.scale", r"single_transformer_blocks.\1.attn.norm_k.weight"),
-    )
-    for pattern, replacement in double_rewrites:
-        if re.fullmatch(pattern, key):
-            output[re.sub(pattern, replacement, key)] = value
-            return True
-
-    return False
-
-
-def convert_raw_transformer(raw_state: OrderedDict[str, torch.Tensor]) -> OrderedDict[str, torch.Tensor]:
-    converted: OrderedDict[str, torch.Tensor] = OrderedDict()
-    unexpected: list[str] = []
-    for key, value in raw_state.items():
-        if not _convert_raw_transformer_key(key, value, converted):
-            unexpected.append(key)
-
-    if unexpected:
-        raise ConversionError(
-            "Unmapped raw Flux2 Klein transformer keys:\n" + "\n".join(f"  - {key}" for key in unexpected)
-        )
-    return converted
-
-
-def convert_diffusers_transformer(state: OrderedDict[str, torch.Tensor]) -> OrderedDict[str, torch.Tensor]:
-    converted: OrderedDict[str, torch.Tensor] = OrderedDict()
-    for key, value in state.items():
-        if key.startswith("transformer."):
-            key = key[len("transformer."):]
-        converted[key] = value
-    return converted
-
-
-def convert_vae(state: OrderedDict[str, torch.Tensor]) -> OrderedDict[str, torch.Tensor]:
-    return OrderedDict(state.items())
-
-
-def _infer_block_count(keys: set[str], prefix: str) -> int:
-    count = 0
-    for key in keys:
-        if key.startswith(prefix):
-            suffix = key[len(prefix):]
-            first = suffix.split(".", 1)[0]
-            if first.isdigit():
-                count = max(count, int(first) + 1)
-    return count
-
-
-def _validate_required_keys(component: str, state: OrderedDict[str, torch.Tensor], required: tuple[str, ...]) -> None:
-    missing = [key for key in required if key not in state]
-    if missing:
-        raise ConversionError(f"{component} conversion missing required keys: {missing}")
-
-
-def _validate_transformer(state: OrderedDict[str, torch.Tensor]) -> None:
-    _validate_required_keys("transformer", state, TRANSFORMER_REQUIRED_KEYS)
-    raw_leftovers = [key for key in state if key.startswith(("double_blocks.", "single_blocks."))]
-    if raw_leftovers:
-        raise ConversionError(f"Raw transformer keys leaked into output: {raw_leftovers[:8]}")
-    keys = set(state)
-    double_blocks = _infer_block_count(keys, "transformer_blocks.")
-    single_blocks = _infer_block_count(keys, "single_transformer_blocks.")
-    print(
-        "  transformer keys: " +
-        f"{len(state)} total, {double_blocks} double blocks, {single_blocks} single blocks"
-    )
-
-
-def _validate_vae(state: OrderedDict[str, torch.Tensor]) -> None:
-    _validate_required_keys("vae", state, VAE_REQUIRED_KEYS)
-    print(f"  vae keys: {len(state)} total")
-
-
-def _component_config(src_dir: Path, component: str, class_name: str) -> dict[str, Any]:
-    config_path = src_dir / component / "config.json"
-    if not config_path.exists():
-        raise FileNotFoundError(f"Missing {component} config: {config_path}")
-    config = _read_json(config_path)
-    config["_class_name"] = class_name
-    config.pop("_name_or_path", None)
-    return config
-
-
-def _copy_passthrough_subfolder(src_dir: Path, dst_dir: Path, subfolder: str) -> None:
-    src = src_dir / subfolder
-    if not src.is_dir():
-        raise FileNotFoundError(f"Missing passthrough subfolder: {src}")
-    dst = dst_dir / subfolder
-    shutil.copytree(src, dst)
-    print(f"  copied {subfolder}/")
-
-
-def _copy_or_write_scheduler(src_dir: Path, dst_dir: Path) -> None:
-    src = src_dir / "scheduler"
-    dst = dst_dir / "scheduler"
-    if src.is_dir():
-        shutil.copytree(src, dst)
-        print("  copied scheduler/")
-        return
-    _write_json(dst / "scheduler_config.json", DEFAULT_SCHEDULER_CONFIG)
-    print("  wrote default scheduler/scheduler_config.json")
-
-
-def _build_model_index(src_dir: Path, source_label: str) -> dict[str, Any]:
-    index_path = src_dir / "model_index.json"
-    if index_path.exists():
-        index = _read_json(index_path)
-    else:
-        index = dict(DEFAULT_MODEL_INDEX)
-    index.update(
-        {
-            "_class_name": "Flux2KleinPipeline",
-            "is_distilled": True,
-            "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"],
-            "text_encoder": ["transformers", "Qwen3ForCausalLM"],
-            "tokenizer": ["transformers", "Qwen2TokenizerFast"],
-            "transformer": ["diffusers", "Flux2Transformer2DModel"],
-            "vae": ["diffusers", "AutoencoderKLFlux2"],
-            "_fastvideo_converted_from": source_label,
-        }
-    )
-    return index
-
-
-def _select_transformer_source(src_dir: Path, mode: str) -> tuple[str, Path]:
-    if src_dir.is_file():
-        if mode == "diffusers":
-            raise FileNotFoundError(f"Cannot use --transformer-source diffusers with a file source: {src_dir}")
-        return "raw", src_dir
-
-    raw_path = src_dir / RAW_TRANSFORMER_FILENAME
-    diffusers_dir = src_dir / "transformer"
-
-    if mode == "raw":
-        if not raw_path.exists():
-            raise FileNotFoundError(f"Requested raw transformer but missing {raw_path}")
-        return "raw", raw_path
-    if mode == "diffusers":
-        if not diffusers_dir.is_dir():
-            raise FileNotFoundError(f"Requested diffusers transformer but missing {diffusers_dir}")
-        return "diffusers", diffusers_dir
-    if raw_path.exists():
-        return "raw", raw_path
-    if diffusers_dir.is_dir():
-        return "diffusers", diffusers_dir
-    raise FileNotFoundError(f"Missing transformer weights under {src_dir}")
-
-
-def convert(src: str, dst: str, revision: str | None, cache_dir: str | None, max_shard_size: str,
-            transformer_source: str, overwrite: bool) -> None:
-    resolved_src = _resolve_src(src, revision=revision, cache_dir=cache_dir, transformer_source=transformer_source)
-    src_dir = resolved_src.parent if resolved_src.is_file() else resolved_src
-    dst_dir = Path(dst).expanduser()
-    _prepare_output_dir(dst_dir, overwrite=overwrite)
-
-    print(f"src: {src_dir}")
-    print(f"dst: {dst_dir}")
-
-    print("\n[1/5] Converting transformer:")
-    source_kind, source_path = _select_transformer_source(resolved_src, transformer_source)
-    if source_kind == "raw":
-        print(f"  using raw BFL transformer: {source_path.name}")
-        transformer_state = convert_raw_transformer(_load_safetensors([source_path]))
-    else:
-        print("  using Diffusers transformer subfolder")
-        transformer_state = convert_diffusers_transformer(_load_safetensors(_find_safetensors_files(source_path)))
-    _validate_transformer(transformer_state)
-    transformer_dir = dst_dir / "transformer"
-    _write_state_dict(transformer_state, transformer_dir, max_shard_size=max_shard_size)
-    _write_json(transformer_dir / "config.json", _component_config(src_dir, "transformer", "Flux2Transformer2DModel"))
-
-    print("\n[2/5] Converting VAE:")
-    vae_dir = src_dir / "vae"
-    vae_state = convert_vae(_load_safetensors(_find_safetensors_files(vae_dir)))
-    _validate_vae(vae_state)
-    output_vae_dir = dst_dir / "vae"
-    _write_state_dict(vae_state, output_vae_dir, max_shard_size=max_shard_size)
-    _write_json(output_vae_dir / "config.json", _component_config(src_dir, "vae", "AutoencoderKLFlux2"))
-
-    print("\n[3/5] Copying HF-backed encoders/tokenizer/scheduler:")
-    # Current FastVideo Flux2 text encoders intentionally use HF passthrough:
-    # Qwen3ForCausalLM.from_pretrained_local() and Mistral3ForConditionalGeneration.from_pretrained_local().
-    # Rewriting their weights to native fused QKV names would make the standard loader call the wrong HF class.
-    _copy_passthrough_subfolder(src_dir, dst_dir, "text_encoder")
-    _copy_passthrough_subfolder(src_dir, dst_dir, "tokenizer")
-    _copy_or_write_scheduler(src_dir, dst_dir)
-
-    print("\n[4/5] Writing model_index.json:")
-    _write_json(dst_dir / "model_index.json", _build_model_index(src_dir, src))
-    print("  wrote model_index.json")
-
-    print("\n[5/5] Summary:")
-    if source_kind == "raw":
-        print("  transformer mapping: raw BFL -> FastVideo native")
-    else:
-        print("  transformer mapping: Diffusers/FastVideo identity")
-    print("  vae mapping: Diffusers/FastVideo identity")
-    print("  text encoder mapping: HF Qwen3 passthrough (FastVideo loader uses from_pretrained_local)")
-    print(f"  max shard size: {max_shard_size}")
-    print("Done.")
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description=(__doc__ or "").split("\n\n", 1)[0])
-    parser.add_argument("--src", default=DEFAULT_REPO_ID, help="HF repo id or local Flux2 Klein snapshot directory")
-    parser.add_argument("--dst", required=True, help="Output directory for the FastVideo Diffusers-style repo")
-    parser.add_argument("--revision", default=None, help="Optional HF revision for --src repo id")
-    parser.add_argument("--cache-dir", default=None, help="Optional Hugging Face cache directory")
-    parser.add_argument("--max-shard-size", default="5GB", help="Maximum output shard size")
-    parser.add_argument(
-        "--transformer-source",
-        choices=("auto", "raw", "diffusers"),
-        default="auto",
-        help="auto prefers flux-2-klein-4b.safetensors when present, else transformer/",
-    )
-    parser.add_argument("--overwrite", action="store_true", help="Replace a non-empty output directory")
-    args = parser.parse_args()
-
-    convert(
-        src=args.src,
-        dst=args.dst,
-        revision=args.revision,
-        cache_dir=args.cache_dir,
-        max_shard_size=args.max_shard_size,
-        transformer_source=args.transformer_source,
-        overwrite=args.overwrite,
-    )
-
-
-if __name__ == "__main__":
-    main()

scripts/checkpoint_conversion/cosmos3_convert.py

@@ -0,0 +1,87 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 checkpoint strict-load verifier (no weight conversion required).
+
+The published ``nvidia/Cosmos3-Nano`` checkpoint is diffusers-format and its
+transformer weight keys map 1:1 (identity) onto FastVideo's native
+``Cosmos3VFMTransformer`` parameters -- ``needs_conversion=no``. There is no
+remap to apply; the checkpoint loads directly.
+
+This utility verifies strict-load completeness (every checkpoint key has a
+matching DiT parameter of the right shape, and every DiT parameter is provided
+by the checkpoint) without allocating the full ~30 GB model, by reading
+safetensors headers and instantiating the DiT on the ``meta`` device.
+
+Usage:
+    python scripts/checkpoint_conversion/cosmos3_convert.py \
+        --transformer official_weights/cosmos3/transformer
+"""
+from __future__ import annotations
+
+import argparse
+import glob
+import os
+import re
+
+import torch
+from safetensors import safe_open
+
+from fastvideo.configs.models.dits.cosmos3 import Cosmos3VideoConfig
+from fastvideo.models.dits.cosmos3 import Cosmos3VFMTransformer
+
+
+def checkpoint_key_shapes(transformer_dir: str) -> dict[str, tuple[int, ...]]:
+    """Read ``{key: shape}`` from a sharded safetensors transformer dir."""
+    shards = sorted(glob.glob(os.path.join(transformer_dir, "*.safetensors")))
+    if not shards:
+        raise FileNotFoundError(f"no .safetensors found in {transformer_dir}")
+    shapes: dict[str, tuple[int, ...]] = {}
+    for shard in shards:
+        with safe_open(shard, framework="pt") as handle:
+            for key in handle.keys():
+                shapes[key] = tuple(handle.get_slice(key).get_shape())
+    return shapes
+
+
+def verify_strict_load(transformer_dir: str) -> None:
+    """Raise SystemExit if the checkpoint does not strict-load into the DiT."""
+    ckpt = checkpoint_key_shapes(transformer_dir)
+    cfg = Cosmos3VideoConfig()
+    with torch.device("meta"):
+        dit = Cosmos3VFMTransformer(cfg, hf_config={})
+    params = {name: tuple(p.shape) for name, p in dit.named_parameters()}
+    buffers = {name for name, _ in dit.named_buffers()}
+    name_map: dict[str, str] = cfg.arch_config.param_names_mapping
+
+    def remap(key: str) -> str:
+        for pattern, replacement in name_map.items():
+            if re.match(pattern, key):
+                return re.sub(pattern, replacement, key)
+        return key
+
+    mapped = {remap(key): shape for key, shape in ckpt.items()}
+    unexpected = sorted(set(mapped) - set(params) - buffers)
+    missing = sorted(set(params) - set(mapped))
+    mismatched = [(k, mapped[k], params[k]) for k in (set(mapped) & set(params)) if mapped[k] != params[k]]
+
+    if unexpected or missing or mismatched:
+        raise SystemExit("strict-load FAILED: "
+                         f"unexpected={unexpected[:10]} missing={missing[:10]} "
+                         f"shape_mismatch={mismatched[:10]}")
+    print(f"strict-load OK: {len(ckpt)} checkpoint keys map 1:1 onto "
+          f"{len(params)} DiT params (identity; needs_conversion=no)")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--transformer",
+        default=os.path.join("official_weights", "cosmos3", "transformer"),
+        help="path to the checkpoint transformer/ directory",
+    )
+    args = parser.parse_args()
+    verify_strict_load(args.transformer)
+
+
+if __name__ == "__main__":
+    main()

scripts/inference/inference_flux2_klein.yaml

@@ -1,23 +0,0 @@
-# Run with:
-#   FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA fastvideo generate --config scripts/inference/inference_flux2_klein.yaml
-generator:
-  model_path: black-forest-labs/FLUX.2-klein-4B
-  engine:
-    num_gpus: 1
-    offload:
-      dit: false
-      vae: false
-      pin_cpu_memory: false
-  pipeline:
-    workload_type: t2i
-request:
-  prompt: "a photo of a banana on a wooden table, studio lighting"
-  sampling:
-    seed: 0
-    num_frames: 1
-    height: 1024
-    width: 1024
-    num_inference_steps: 4
-    guidance_scale: 1.0
-  output:
-    output_path: outputs/flux2-klein/

tests/local_tests/README.md

@@ -23,7 +23,6 @@ tests/local_tests/<family>/
 
 | Family | Workload | Dir |
 |---|---|---|
-| Flux2 | T2I | [`flux2/`](./flux2/) |
 | Hunyuan GameCraft | T2V / I2V | [`gamecraft/`](./gamecraft/) |
 | GEN3C | T2V | [`gen3c/`](./gen3c/) |
 | Kandinsky-5 | T2V | [`kandinsky5/`](./kandinsky5/) |

tests/local_tests/cosmos3/AUDIO_PORT_PLAN.md

@@ -0,0 +1,77 @@
+# Cosmos3 Audio (PR2) — Port Plan
+
+Branch: `feat/cosmos3-audio` (stacked on `feat/cosmos3-i2v`, which has T2V/I2V/T2I).
+Goal: text-to-video+sound (**t2vs**) — generate synchronized audio alongside video.
+
+## How the framework does audio (studied 2026-06-07)
+
+- **Sound tokenizer = AVAE** (`cosmos_framework/model/vfm/tokenizers/audio/avae.py`
+  + `avae_utils/`, ~2268 lines): a 48 kHz **stereo** neural audio codec.
+  - checkpoint: `official_weights/cosmos3/sound_tokenizer/` (`model_type:
+    autoencoder_v2`, ~1.9 GB). enc=`spec_convnext` (enc_dim 192, latent_dim 128,
+    n_fft 64), dec=`oobleck` (dec_dim 320, strides [2,4,5,6,8]), VAE bottleneck,
+    `snakebeta` activations, hop_size 1920.
+  - interface: `encode(audio[1,C,N]) -> latent`, `decode(latent) -> audio`,
+    `get_latent_num_samples(N)`, `sample_rate=48000`, `audio_channels=2`,
+    `sound_latent_fps=25`.
+- **DiT sound pathway** (`cosmos3_vfm_network.py`, 136 sound/audio refs): the MoT
+  has `sound2llm` / `llm2sound` / `sound_modality_embed` + `pack_sound_latents`
+  and joint vision+sound denoising (`preds_sound`, sound `condition_mask`, sound
+  noise init `cond_mask*x0 + (1-cond_mask)*noise`, velocity `pred*(1-cond_mask)`).
+  - FastVideo's native DiT ALREADY constructs the dormant heads
+    (`audio_proj_in`/`audio_proj_out`/`audio_modality_embed`, gated on
+    `arch.sound_gen`) for strict-load — the forward just doesn't use them yet.
+- **Inference flow** (`cosmos_framework/inference/sound.py`): t2vs builds a
+  zero **placeholder audio** sized to the video duration (sets sound latent
+  length), `inject_sound_into_batch` upgrades the SequencePlan to has_sound,
+  the omni model denoises vision+sound jointly, then AVAE-decodes the sound
+  latent and `mux_audio_into_video` (PyAV, AAC) muxes it into the mp4
+  (`save_sound` writes a WAV).
+
+## Components (each: native port + framework parity test, per methodology)
+
+1. **AVAE codec** — `fastvideo/models/.../cosmos3_avae.py` + config. Port
+   encoder/decoder/bottleneck/snake. Parity: tiny AVAE, framework weights copied
+   in, bit-exact `decode` (and `encode`) on CPU/fp32. **(largest piece)**
+2. **DiT sound pathway** — activate the dormant heads in `forward`; port
+   `pack_sound_latents` + sound token scatter/proj/modality-embed/velocity.
+   Parity: extend the DiT harness with sound tokens.
+3. **Sound sequence packing** — extend `sequence_packing.py` with the sound
+   modality (positions, attn mode, condition mask). Parity vs framework
+   `pack_input_sequence` with sound.
+4. **Pipeline (t2vs)** — placeholder audio -> joint denoise -> split ->
+   AVAE-decode sound -> mux into mp4 / save wav. Extend `Cosmos3DenoisingStage`
+   + a sound-decode/mux stage.
+5. **FastVideo AV infra** — audio in `OutputConfig` / a mux stage (check what
+   exists; `cosmos_framework.inference.sound.mux_audio_into_video` is the ref).
+
+## Open decisions
+- **D1 (AVAE approach)** — full native port (methodology-consistent; ~2.3k lines)
+  vs a documented lazy-wrapper around the framework AVAE (faster; but pulls heavy
+  deps and bends the "native + no-framework-at-runtime" rule). Default per
+  methodology: native port.
+- **D2 (scope)** — t2vs (T+video+sound) first; defer audio-conditioned / v2vs.
+- **D3** — confirm FastVideo can mux/emit audio (output format).
+
+## Status
+- [x] Branch forked, framework audio path studied, plan written.
+- [x] D1: native port (user-chosen). D2: t2vs first.
+- [x] **AVAE sound decoder (component 1) — DONE** (commit `5f81fb3d5`). Key
+  finding: the checkpoint is decoder-only in AutoencoderOobleck naming with
+  SnakeBeta + weight_g/v == FastVideo's native `OobleckVAE` decoder. Reused it
+  (+ `output_padding=stride%2` for the odd stride 5); `Cosmos3SoundVAE`
+  decoder-only wrapper; bit-exact parity vs the framework OobleckDecoder
+  (`test_cosmos3_avae_parity`); real 1.9 GB checkpoint strict-loads, decodes
+  [1,64,25] -> [1,2,48000] (1 s @ 48 kHz stereo).
+- [x] **DiT sound pathway (component 2) — DONE** (commit `005d6684a`). Activated
+  the dormant audio heads in the forward (`_encode_sound`/`_decode_sound` mirror);
+  `preds_vision` + `preds_sound` bit-exact (max=mean=0.0).
+- [x] **Sound sequence packing (component 3) — DONE** (commit `005d6684a`).
+  `Cosmos3SoundItem` + sound fields; sound shares the vision "full" split with
+  parallel MRoPE. Field-by-field + position_ids exact vs framework.
+- [x] **t2vs pipeline + AV mux (components 4-5) — DONE** (commit `3d8355129`).
+  Joint [vision|sound] denoise, AVAE-decode, stereo 48 kHz AAC mux. t2vs CFG
+  velocity parity max=mean=0.0; real-weights run produces coherent video + real
+  audio (mean -10.2 dB). Example `basic_cosmos3_t2vs_new_api.py`.
+
+**PR2 (audio/t2vs) COMPLETE** — every component bit-exact vs the framework.

tests/local_tests/cosmos3/PORT_STATUS.md

@@ -0,0 +1,124 @@
+# Cosmos3 Port Status
+
+## Summary
+
+- model_family: `cosmos3`
+- workload_types: `T2V, I2V, T2I` supported by `WorkloadType` today; full-omni target also needs audio (AV), VLM reasoning, and action-conditioning, which require framework extensions (Q002, Q003).
+- official_ref: `https://github.com/NVIDIA/cosmos-framework` — diffusers backend `diffusers_cosmos3.pipeline.Cosmos3OmniDiffusersPipeline`; HF `nvidia/Cosmos3-Nano`.
+- official_ref_dir: `cosmos-framework` (symlink -> `/home/william5lin/FastVideo/cosmos-framework`, commit `003d66d4`)
+- hf_weights_path: `nvidia/Cosmos3-Nano`
+- local_weights_dir: `official_weights/cosmos3` (symlink -> `/home/william5lin/FastVideo/official_weights/cosmos3`, 33 GiB / 67 files)
+- source_layout: `diffusers`
+- local_tests_readme: `tests/local_tests/cosmos3/README.md`
+
+## Current Phase
+
+- phase: `PR1 (video core) + PR2 (audio/t2vs) COMPLETE, real-weights verified. Video core (T2V/I2V/T2I) + audio (AVAE decode / DiT sound pathway / sound packing / t2vs CFG velocity) all framework-parity verified bit-exact (suite 130 passed, 0 skipped). t2vs real-weights on B200 produces coherent video + real stereo 48kHz audio. Branch chain: feat/cosmos3-tier-a-port (T2V) -> feat/cosmos3-i2v (I2V+T2I) -> feat/cosmos3-audio (t2vs). Next: PR3 action / PR4 reasoning.`
+- status: `in_progress`
+- owner: `orchestrator`
+- last_updated: `2026-06-07`
+- env: `fv-cosmos3` (conda clone of fv-main; `fastvideo` editable repointed to this worktree). Run tests from the worktree cwd with this env's python.
+- branch: rebased onto `origin/main` @ `1c627a3f9` (was 33 behind, merge-base 2026-05-22); now 6 commits ahead; `fastvideo` imports clean; Tier-A `13 passed, 2 skipped`.
+
+## Component Matrix
+
+| Component | Type | Reuse/Port | Official Definition | Official Instantiation | FastVideo Target | Prototype | Conversion | Parity | Open Issues |
+|---|---|---|---|---|---|---|---|---|---|
+| transformer | dit | port | `diffusers_cosmos3/transformer.py:Cosmos3OmniTransformer` (model_type `qwen3_vl_text`, MoT + MRoPE) | `model_index.json: transformer`; `cosmos_framework/model/vfm/mot/cosmos3_vfm_network.py`, `omni_mot_model.py` | `fastvideo/models/dits/cosmos3.py` (branch: `Cosmos3VFMTransformer`+`Cosmos3LanguageModel` — reconcile to `Cosmos3OmniTransformer`) | skeleton | not_started | scaffold_skip | I001 |
+| vae | vae | reuse | diffusers `AutoencoderKLWan` | `model_index.json: vae` | reuse Wan VAE (`fastvideo/models/vaes/`, cf. `cosmos25wanvae.py`) | not_started | passthrough? | not_started | Q001 |
+| scheduler | generic | reuse (flow-coerced) | framework `FlowUniPCMultistepScheduler` (`cosmos_framework/.../fm_solvers_unipc.py`; checkpoint ships diffusers-style config) | `model_index.json: scheduler`; `cosmos_framework/.../samplers/unipc.py:UniPCSampler` | FastVideo-native `UniPCMultistepScheduler` (flow config), coerced in `initialize_pipeline` | done | n/a | framework-parity DONE (`test_cosmos3_scheduler_parity`: timesteps bit-exact, sigmas ~1e-8, trajectory <~1e-6) | I003 (resolved) |
+| text_tokenizer | tokenizer | reuse | transformers `Qwen2TokenizerFast` | `model_index.json: text_tokenizer` | reuse (tokenizer = allowed third-party) | not_started | passthrough | scaffold_skip (`test_cosmos3_tokenizer_chat_template`) | - |
+| vision_encoder | encoder | port | transformers `Qwen3VLVisionModel` | `model_index.json: vision_encoder` | new encoder bucket OR documented lazy-wrapper | not_started | not_started | not_started | Q002 |
+| sound_tokenizer | generic/vae | port (decode) | framework AVAE `LatentAutoEncoderV2` (`avae_utils`); checkpoint is decoder-only AutoencoderOobleck-named w/ SnakeBeta | `model_index.json: sound_tokenizer` | reuse FastVideo native `OobleckVAE` decoder + `Cosmos3SoundVAE` wrapper (`models/audio/cosmos3_avae.py`) | done (decode) | n/a | DECODE bit-exact vs framework (`test_cosmos3_avae_parity`); real ckpt strict-loads | PR2 (branch feat/cosmos3-audio) |
+
+## Conversion State
+
+- conversion_script: `scripts/checkpoint_conversion/cosmos3_convert.py` (branch has it, 246 lines, built vs vllm-omni — repoint/verify vs diffusers checkpoint)
+- converted_weights_dir: `converted_weights/cosmos3` (n/a while needs_conversion=no)
+- source_layout: `diffusers`
+- needs_conversion: `no` (HF already diffusers-format; verify FastVideo loaders consume directly)
+- strict_load_status: `not_run`
+- passthrough_components: `vae (AutoencoderKLWan), scheduler (UniPC), text_tokenizer (Qwen2)` likely passthrough
+- retry_history: `none`
+
+## Parity Commands
+
+| Scope | Command | Last Result | Notes |
+|---|---|---|---|
+| Tier-A scaffold | `cd <worktree> && <fv-cosmos3 python> -m pytest tests/local_tests/cosmos3/ -q` | `13 passed, 2 skipped` (2026-06-06, post-rebase) | 2 skips: Cosmos3 tokenizer/_tokenize_prompt not yet wired on pipeline |
+| component | `pytest tests/local_tests/<bucket>/test_cosmos3_<component>_parity.py -v -s` | `not_run` | after env activation + native prototypes |
+| pipeline | `pytest tests/local_tests/pipelines/test_cosmos3_pipeline_parity.py -v -s` | `not_run` | |
+
+## Open Questions
+
+| ID | Question | Owner | Needed By Phase | Status | Resolution |
+|---|---|---|---|---|---|
+| Q001 | Does Cosmos3 VAE (`AutoencoderKLWan`) match FastVideo's existing Wan VAE config/instantiation exactly (z_dim, scale factors, latents_mean/std)? | orchestrator | 3 (reuse gate) | open | |
+| Q002 | `vision_encoder` (`Qwen3VLVisionModel`): native port vs documented lazy-wrapper exception? Needed for I2V/reasoning. | user/orchestrator | 3 | open | |
+| Q003 | `sound_tokenizer` (`Cosmos3AVAEAudioTokenizer`) + audio output requires `WorkloadType` AV + audio regression metric. | user | 0/10 | open | full-omni scope chosen 2026-06-06; infra extensions pending |
+
+## Issues And Blockers
+
+| ID | Phase | Component | Severity | Issue | Evidence | Owner | Status | Resolution |
+|---|---|---|---|---|---|---|---|---|
+| I001 | port | transformer | high | Branch DiT (`Cosmos3VFMTransformer`+`Cosmos3LanguageModel`) built vs vllm-omni #3454; official checkpoint loads `Cosmos3OmniTransformer` (diffusers shim). Class/structure reconciliation required. | `model_index.json`; `diffusers_cosmos3/transformer.py`; branch commit `52bb65f49` | orchestrator | resolved | DiT rewritten to checkpoint layout (single `layers` dual-pathway, BaseDiT-conformant); bit-identical framework parity (3d_rope + unified_3d_mrope), commits 59a4a571c/7c4633295 |
+| I002 | all | tests | medium | Tier-A conftest+tests mirror vllm-omni line-by-line (stubs, `vllm_omni...guardrails`). Must be repointed to `diffusers_cosmos3` / official structures. | `tests/local_tests/cosmos3/conftest.py` | orchestrator | open | |
+| I003 | inference | scheduler | high | First real-weights T2V was all-black: checkpoint `scheduler_config.json` sets `use_karras_sigmas=true`; vendored UniPC checks karras before `use_flow_sigmas` -> diffusion (beta) sigmas -> `scheduler.step` -> NaN latents. DiT/CFG velocity was clean. The scheduler had never been parity-tested vs the framework (`test_cosmos3_denoise_cfg_parity` used diffusers UniPC on both sides). | `result_latent` NaN at denoise step 0 (v_pred clean); ffprobe 3 KB black mp4 | orchestrator | resolved | Coerce loaded config to flow setup in `initialize_pipeline`; switch pipeline+tests to native UniPC (no diffusers at runtime); add `test_cosmos3_scheduler_parity` vs framework `FlowUniPCMultistepScheduler`; repoint denoise_cfg oracle to the framework scheduler. Commit 255311cf2 |
+
+## Escape Hatches
+
+| ID | Phase | Decision Type | Question | Recommended Option | Status | Resolution |
+|---|---|---|---|---|---|---|
+| E001 | prep | dependency/env | Shared `fv-main` env has `fastvideo` editable-installed from the MAIN worktree; the cosmos3 worktree's `fastvideo` is not importable (PEP660 finder overrides PYTHONPATH), so Tier-A tests skip. How to activate the worktree's `fastvideo` for verification without disrupting ~24 other worktrees sharing the env? | Dedicated conda env for the cosmos3 worktree | resolved | Created fv-cosmos3 (clone of fv-main); repointed fastvideo editable to worktree; run from worktree cwd. Branch also rebased onto origin/main to fix stale import. |
+
+## Decisions
+
+| Date | Decision | Rationale | Impact |
+|---|---|---|---|
+| 2026-06-06 | Reference source of truth = official diffusers (`Cosmos3OmniDiffusersPipeline` + `cosmos-framework`/`diffusers-cosmos3`), not vllm-omni #3454 | Official weights now public & diffusers-format; the artifact users actually load | Repoint DiT/pipeline/conversion/tests off vllm-omni (I001, I002) |
+| 2026-06-06 | Resume in worktree `/home/william5lin/FastVideo_cosmos3_port`; weights+reference symlinked (no copy) | Preserve 2,492 lines of Tier-A work; avoid 33 GB duplication | Verification needs worktree `fastvideo` active (E001) |
+| 2026-06-06 | Scope = full omni (video + audio + reasoning + action) | User choice (revised from branch's original video-only scope) | Adds `vision_encoder`, `sound_tokenizer` ports + `WorkloadType` AV + audio metric |
+| 2026-06-06 | Downloaded full 34.9 GB (33 GiB) `nvidia/Cosmos3-Nano` | Unblocks May-22 `PENDING` weight status (HF was 401, now public) | Real parity now possible |
+| 2026-06-06 | Rebased branch onto origin/main (33 commits); resolved registry.py conflict by reconstructing from main + cosmos3 import/entry | Branch was stale; fastvideo failed to import (main removed MatrixGameI2V480PConfig) | Branch imports clean; Tier-A 13 passed/2 skipped |
+| 2026-06-06 | Reference = cosmos_framework ONLY (full omni); diffusers shim dropped even for video | User directive (Phase 1 found diffusers __call__ is video-only; sound/action/reasoning live only in the framework) | Larger port; ref DiT = `Cosmos3VFMNetwork`/`Cosmos3VFMNetworkConfig` (not diffusers `Cosmos3OmniTransformer`); core model imports in fv-cosmos3 with light deps; TE only in optional dot_product_attention |
+
+## Handoff Notes
+
+- Prep (weights/reference/env editable installs) done in MAIN worktree; symlinked into this worktree. Env installs (`diffusers-cosmos3`, `cosmos-framework`) are in shared `fv-main`.
+- Next: resolve E001 (env), then Phase 1 reference study of `diffusers_cosmos3` pipeline/transformer, then Phase 3 reuse gate (VAE/scheduler/tokenizer) + component dispatch (transformer, vision_encoder, sound_tokenizer).
+- diffusers 0.36.0 imports the shim OK; checkpoint saved with 0.37.1 — watch `from_pretrained` needs (bump within FastVideo's `diffusers>=0.33.1` pin if required).
+
+### PR1 (video core) progress — 2026-06-06
+- Arch config 1:1 with checkpoint, committed `9567efdf0`.
+- Framework parity-reference harness committed `dd97efda3`: `tests/local_tests/cosmos3/test_cosmos3_reference_forward.py` builds a tiny `Cosmos3VFMNetwork` on CPU/float32 (SDPA monkeypatch; flash2/3/natten are CUDA-only) and forwards `packed_seq -> {last_hidden_state, preds_vision}`. 23 tests pass in fv-cosmos3. This is the ground-truth side for DiT parity. Run: `cd <worktree> && <fv-cosmos3 py> -m pytest tests/local_tests/cosmos3/test_cosmos3_reference_forward.py -q`.
+- THREE naming conventions to bridge:
+  1. framework-native (`Cosmos3VFMNetwork`): `language_model.model.layers.{i}.self_attn.{q,k,v,o}_proj(+ _moe_gen)`, `{q,k}_norm(+_moe_gen)`, `mlp(+_moe_gen)`, `vae2llm`/`llm2vae`, `time_embedder.mlp.{0,2}`.
+  2. diffusers checkpoint (on disk, what we load): `layers.{i}.self_attn.{to_q,to_k,to_v,to_out}` + `{add_q,add_k,add_v}_proj`/`to_add_out`, `{norm_q,norm_k,norm_added_q,norm_added_k}`, `mlp`/`mlp_moe_gen`, `proj_in`/`proj_out`, `time_embedder.linear_{1,2}`.
+  3. FastVideo DiT (our choice). Conversion maps (2)->(3); the DiT parity test copies (1)->(3).
+- BaseDiT signature is `__init__(self, config: DiTConfig, hf_config: dict)`; the branch `Cosmos3VFMTransformer` uses `fastvideo_args`/SimpleNamespace and does NOT conform — rewrite to conform + match the checkpoint key surface (single `layers` dual-pathway, not split language_model/gen_layers).
+- Native layers (per cosmos2_5): `ReplicatedLinear`/`MLP`/`RMSNorm` (fastvideo.layers.*), `LocalAttention`/`DistributedAttention` (fastvideo.attention), `apply_rotary_emb` (use_real_unbind_dim=-2 for Cosmos). EntryClass at module bottom; class attrs bound from config; 3D-MRoPE has no reusable util — adapt Cosmos25RotaryPosEmbed.
+- NEXT: write native `fastvideo/models/dits/cosmos3.py` + fastvideo-vs-framework forward parity test (copy framework weights into the FastVideo DiT, compare outputs), then conversion script (diffusers checkpoint -> FastVideo) + strict-load, then video pipeline/packing.
+
+### PR1 (video core) acceptance — real-weights E2E — 2026-06-07
+- First real-weights T2V (`examples/inference/basic/basic_cosmos3_new_api.py`, `COSMOS3_MODEL_PATH=official_weights/cosmos3`) ran mechanically but produced an all-black 3 KB mp4. Instrumenting the denoise loop showed `v_pred` clean at step 0 but `scheduler.step` -> NaN. Root cause I003: checkpoint `scheduler_config.json` is diffusers-style (`use_karras_sigmas=true`), and the vendored UniPC checks karras before `use_flow_sigmas` -> diffusion (beta) sigmas instead of flow sigmas -> NaN. The framework actually samples with `FlowUniPCMultistepScheduler` (pure flow: `shift` + `num_train_timesteps`).
+- Fix (commit `255311cf2`): coerce the loaded scheduler to the flow setup in `Cosmos3OmniDiffusersPipeline.initialize_pipeline`; use FastVideo's native UniPC (not diffusers) in pipeline + tests. Added `test_cosmos3_scheduler_parity.py` (native UniPC flow-config vs framework `FlowUniPCMultistepScheduler`: timesteps bit-exact, sigmas ~1e-8, full trajectory <~1e-6 over shift in {10,3}, steps in {4,10,35}). Repointed `test_cosmos3_denoise_cfg_parity` oracle to the framework scheduler (it previously compared diffusers-vs-diffusers, so the scheduler was never checked against the framework).
+- Also wired the remaining integration glue (registry alias `Cosmos3OmniTransformer`->`Cosmos3VFMTransformer`; `text_tokenizer`->TokenizerLoader; scheduler config param-filtering; DiT `materialize_non_persistent_buffers` + compute-dtype casts; packing device-move in `to_dit_kwargs`; empty text-preprocess).
+- Verified: 1280x704, 29 frames, 35 steps on a single B200 -> coherent golden-retriever-in-meadow video matching the prompt (no NaNs; per-frame pixel std ~58; visible temporal motion). Full cosmos3 suite: 95 passed, 0 skipped.
+- NEXT: PR2 audio (`sound_tokenizer` AVAE) / PR3 action / PR4 reasoning. Optional: I2V/T2I real-weights spot-checks; force-push branch (needs explicit OK).
+
+### PR1 (video core) — I2V real-weights — 2026-06-07 (branch feat/cosmos3-i2v)
+- Forked `feat/cosmos3-i2v` off `feat/cosmos3-tier-a-port` (stacked, includes the T2V + scheduler fix).
+- Studied the framework I2V path: `cosmos_framework.inference.vision.load_conditioning_image` (aspect-preserving resize + center crop + uint8 quantize -> `/127.5-1`) + `build_conditioned_video_batch` (frame 0 = image, remaining frames REPEAT the last conditioning frame -> static video), then VAE-encode; `condition_frame_indexes=[0]` (latent). Condition frames kept clean during sampling exactly as FastVideo already does: init noise `cond_mask*x0 + (1-cond_mask)*noise` (`omni_mot_model._prepare_inference_data`) + velocity zeroed `pred*(1-cond_mask)` each step (`_get_velocity`), no re-injection.
+- Bug found + fixed (commit `bd8d604fb`): FastVideo's `_image_to_video_tensor` ZERO-filled the non-condition frames; the temporal Wan VAE (4x) makes latent frame 0 depend on several pixel frames, so zero-fill -> wrong conditioning latent. Rewrote it to repeat-fill + framework resize/crop/quantize.
+- Parity: `test_cosmos3_i2v_conditioning_parity.py` vs framework `load_conditioning_image` + repeat-fill — bit-exact (max abs diff 0.0) across aspect/size/frame cases. Existing `test_cosmos3_denoise_cfg_parity` already covers the I2V cond-mask + velocity math (i2v case).
+- Example: `examples/inference/basic/basic_cosmos3_i2v_new_api.py` (`InputConfig(image_path=...)`, default `assets/images/cyclist.jpg`).
+- Verified on B200 (1280x704, 29f, 35 steps, real weights): output frame 0 reproduces the conditioning cyclist image; later frames show coherent forward motion down the trail following the prompt. Full suite 98 passed, 0 skipped.
+- NEXT: optional T2I real-weights spot-check; then PR2 audio / PR3 action / PR4 reasoning.
+
+### PR1 (video core) — T2I real-weights + resolution-based flow_shift — 2026-06-07 (branch feat/cosmos3-i2v)
+- Studied framework T2I: tokenization uses `vlm_config.use_system_prompt` which is `false` in the checkpoint (config.json:199) — matches FastVideo's hardcoded `use_system_prompt=False` for all modes (no divergence). Canonical T2I is 960x960 (inputs/omni/t2i.json), single-frame (num_frames=1).
+- Bug found + fixed (commit `604dc2637`): the stage chose `flow_shift` by task (`3.0 if is_t2i else 10.0`), but the framework picks it purely by the named resolution bucket (`OmniSampleArgs._RESOLUTION_SHIFT_DEFAULTS`, 8B backbone: 256->3.0, 480->5.0, 720/768->10.0; model default resolution "720"). Task-based only matched T2V@720 / T2I@256 by luck; canonical T2I@960x960 is the "720" bucket -> 10.0, so `is_t2i->3.0` was wrong. Replaced with `_flow_shift_for_resolution(h,w)` (longest-side bucketing), applied to all tasks.
+- Parity: `test_cosmos3_flow_shift_parity.py` checks the mapping vs framework `{VIDEO,IMAGE}_RES_SIZE_INFO` x `_RESOLUTION_SHIFT_DEFAULTS` (8B rows, 20 cases). Also hardened `_image_to_video_tensor` tensor branch to respect the [-1,1] convention (PIL path stays framework-exact).
+- Example: `examples/inference/basic/basic_cosmos3_t2i_new_api.py` (num_frames=1, 960x960).
+- Verified on B200 (real weights, 35 steps): coherent red-panda image matching the prompt, flow_shift=10.0. Full suite 118 passed, 0 skipped.
+- Video core (T2V/I2V/T2I) is now complete and real-weights verified. NEXT: PR2 audio (sound_tokenizer AVAE + audio output) on a new stacked branch.

tests/local_tests/cosmos3/README.md

@@ -0,0 +1,71 @@
+# Cosmos3 local parity workspace
+
+## Overview
+
+This workspace tracks the FastVideo Cosmos3 port. Live port state, component matrix,
+decisions, and blockers live in `PORT_STATUS.md`.
+
+- **Reference (2026-06-06): official NVIDIA `cosmos-framework` diffusers backend** —
+  `Cosmos3OmniDiffusersPipeline` from the `diffusers-cosmos3` shim — loading the
+  now-public `nvidia/Cosmos3-Nano` checkpoint.
+- **Scope: full omni** — T2V / I2V / T2I, audio (sound generation), VLM reasoning,
+  and action-conditioning.
+- The original Tier-A scaffold was written against vllm-omni PR #3454 before official
+  weights were public; it is being repointed to the diffusers reference (see I001/I002
+  in `PORT_STATUS.md`).
+
+## Reference code
+
+Primary (official):
+
+- Local: `cosmos-framework/` (symlink -> `/home/william5lin/FastVideo/cosmos-framework`,
+  commit `003d66d4`); GitHub <https://github.com/NVIDIA/cosmos-framework>
+- diffusers shim `cosmos-framework/packages/diffusers-cosmos3/diffusers_cosmos3/`:
+  - `pipeline.py` — `Cosmos3OmniDiffusersPipeline`
+  - `transformer.py` — `Cosmos3OmniTransformer`
+  - `sequence_packing.py`
+- framework model code: `cosmos_framework/model/vfm/mot/cosmos3_vfm_network.py`,
+  `cosmos_framework/model/vfm/omni_mot_model.py`
+- Installed editable in shared `fv-main`: `diffusers-cosmos3`, `cosmos-framework`
+  (both `--no-deps`).
+
+Original Tier-A reference (superseded, kept for diffing during repoint):
+
+- vllm-omni PR #3454 <https://github.com/vllm-project/vllm-omni/pull/3454>, pinned
+  `8536f5b1`, checkout `/home/william5lin/cosmos3-reference`.
+- The current `conftest.py` + tests still mirror this suite line-by-line.
+
+## Weight status
+
+DOWNLOADED (2026-06-06). `nvidia/Cosmos3-Nano` is now public and diffusers-format
+(the 2026-05-22 `401` is resolved).
+
+- Local: `official_weights/cosmos3/` (symlink -> main worktree; 33 GiB, 67 files,
+  `model_index.json` present)
+- Source: `nvidia/Cosmos3-Nano`, default revision; `source_layout=diffusers`,
+  `needs_conversion=no`
+- `model_index` class: `Cosmos3OmniDiffusersPipeline` (diffusers 0.37.1)
+- Token: not required (public repo)
+
+Components (from `model_index.json`): `transformer` (`Cosmos3OmniTransformer`),
+`vae` (`AutoencoderKLWan`), `scheduler` (`UniPCMultistepScheduler`),
+`text_tokenizer` (`Qwen2TokenizerFast`), `vision_encoder` (`Qwen3VLVisionModel`),
+`sound_tokenizer` (`Cosmos3AVAEAudioTokenizer`).
+
+## Running the Tier-A scaffold
+
+```bash
+PYTHONPATH=/home/william5lin/FastVideo_cosmos3_port \
+  python -m pytest tests/local_tests/cosmos3/ -q
+```
+
+NOTE: as of 2026-06-06 these report `15 skipped` because the shared `fv-main` env's
+editable `fastvideo` resolves to the MAIN worktree (a PEP660 finder overrides
+`PYTHONPATH`), so the worktree's cosmos3 modules are not importable. Tracked as E001
+in `PORT_STATUS.md`.
+
+## SSIM placeholder
+
+No SSIM references seeded yet. Add SSIM coverage only after a FastVideo inference path
+can load the Cosmos3 weights and generate stable T2V/I2V/T2I outputs. Audio quality
+uses a separate metric (not SSIM); see `PORT_STATUS.md` Q003.

tests/local_tests/cosmos3/conftest.py

@@ -0,0 +1,214 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Shared fixtures for the Cosmos3 native-pipeline local tests.
+
+These fixtures build the FastVideo-native Cosmos3 pipeline
+(``fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline.Cosmos3OmniDiffusersPipeline``)
+via ``__new__`` and wire it with tiny stub components so the runtime call graph
+(sequential CFG, condition-frame masking, mode dispatch) can be exercised on CPU
+without real weights or ``cosmos_framework``.
+
+The stub transformer implements the native DiT's packed-input contract
+(``{"preds_vision": [[1, C, T, H, W], ...]}``) and records, per call, the first
+``text_ids`` token so tests can assert the cond/uncond pass order.
+"""
+from __future__ import annotations
+
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import torch
+from fastvideo.models.schedulers.scheduling_unipc_multistep import (
+    UniPCMultistepScheduler,
+)
+from torch import nn
+
+_LATENT_CHANNEL = 16
+_LATENT_PATCH_SIZE = 2
+_SPATIAL_FACTOR = 8
+_TEMPORAL_FACTOR = 4
+
+
+def pytest_configure(config: pytest.Config) -> None:
+    """Register the ``local`` marker used by sibling test files."""
+    config.addinivalue_line(
+        "markers",
+        "local: marker for local-only parity/scaffold tests (skipped in CI)",
+    )
+
+
+# ---------------------------------------------------------------------------
+# Stub transformer: records cond/uncond call order; bounded preds_vision.
+# ---------------------------------------------------------------------------
+class StubCosmos3Transformer(nn.Module):
+    """Records each forward's first ``text_ids`` token + returns preds_vision.
+
+    ``preds_vision`` is keyed by the first text token (so the conditional and
+    unconditional passes return different velocities) and is zero on
+    conditioning frames, matching the real DiT's unpatchify output.
+    """
+
+    def __init__(self, latent_channel: int = _LATENT_CHANNEL) -> None:
+        super().__init__()
+        self.latent_channel = latent_channel
+        self.embed_tokens = nn.Embedding(64, 8)
+        self.calls: list[dict[str, Any]] = []
+
+    def forward(self, **kwargs: Any) -> dict[str, Any]:
+        token_ids = kwargs["text_ids"]
+        token = int(token_ids.reshape(-1)[0].item()) if token_ids.numel() else 0
+        self.calls.append({"token": token, "kwargs": dict(kwargs)})
+        scale = 0.01 * (1.0 + (token % 7))
+        preds: list[torch.Tensor] = []
+        for latent, _shape, nfi in zip(kwargs["vision_tokens"], kwargs["vision_token_shapes"],
+                                       kwargs["vision_noisy_frame_indexes"]):
+            lat = latent.squeeze(0) if latent.dim() == 5 else latent  # [C, T, H, W]
+            out = torch.zeros_like(lat)
+            if nfi.numel() > 0:
+                out[:, nfi] = scale * torch.tanh(lat[:, nfi])
+            preds.append(out.unsqueeze(0))
+        return {"preds_vision": preds}
+
+
+class _StubLatentDist:
+
+    def __init__(self, latents: torch.Tensor) -> None:
+        self._latents = latents
+
+    def mode(self) -> torch.Tensor:
+        return self._latents
+
+
+class StubCosmos3VAE:
+    """Deterministic VAE shaped by the Wan scale factors."""
+
+    def __init__(self, z_dim: int = _LATENT_CHANNEL) -> None:
+        self.config = SimpleNamespace(
+            z_dim=z_dim,
+            scale_factor_temporal=_TEMPORAL_FACTOR,
+            scale_factor_spatial=_SPATIAL_FACTOR,
+            latents_mean=[0.0] * z_dim,
+            latents_std=[1.0] * z_dim,
+        )
+
+    def encode(self, video: torch.Tensor):
+        b, _c, t, h, w = video.shape
+        lt = (t - 1) // self.config.scale_factor_temporal + 1
+        lh = h // self.config.scale_factor_spatial
+        lw = w // self.config.scale_factor_spatial
+        return _StubLatentDist(torch.ones(b, self.config.z_dim, lt, lh, lw, dtype=video.dtype, device=video.device))
+
+    def decode(self, z: torch.Tensor):
+        b, _c, lt, lh, lw = z.shape
+        t = (lt - 1) * self.config.scale_factor_temporal + 1
+        h = lh * self.config.scale_factor_spatial
+        w = lw * self.config.scale_factor_spatial
+        sig = torch.nan_to_num(torch.tanh(z[:, :1, :1, :1, :1])).reshape(b, 1, 1, 1, 1)
+        return torch.clamp(torch.zeros(b, 3, t, h, w, dtype=z.dtype, device=z.device) + sig, -1.0, 1.0)
+
+
+class StubQwen2Tokenizer:
+    """Qwen2-shaped chat tokenizer stub (special tokens + chat template)."""
+
+    eos_token_id = 62
+    _SPECIAL = {"<|vision_start|>": 60, "<|vision_end|>": 61}
+
+    def convert_tokens_to_ids(self, token: str) -> int:
+        return self._SPECIAL[token]
+
+    def apply_chat_template(self, conversations, *, tokenize=True, add_generation_prompt=True, add_vision_id=False):
+        user = next((c["content"] for c in conversations if c["role"] == "user"), "")
+        n = max(1, min(8, len(user) % 8 + 1))
+        return [10 + (i % 40) for i in range(n)]
+
+
+def make_scheduler(flow_shift: float = 10.0) -> UniPCMultistepScheduler:
+    return UniPCMultistepScheduler(
+        num_train_timesteps=1000,
+        solver_order=2,
+        prediction_type="flow_prediction",
+        use_flow_sigmas=True,
+        flow_shift=flow_shift,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Pipeline factory — builds the native pipeline via __new__ + stub modules.
+# ---------------------------------------------------------------------------
+@pytest.fixture
+def make_cosmos3_pipeline():
+    """Return a factory building the native Cosmos3 pipeline wired with stubs."""
+
+    def _make(**overrides: Any):
+        from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import (  # noqa: F401
+            Cosmos3OmniDiffusersPipeline, )
+
+        pipe = Cosmos3OmniDiffusersPipeline.__new__(Cosmos3OmniDiffusersPipeline)
+        scheduler = make_scheduler()
+        pipe.modules = {
+            "transformer": StubCosmos3Transformer(),
+            "vae": StubCosmos3VAE(),
+            "scheduler": scheduler,
+            "text_tokenizer": StubQwen2Tokenizer(),
+        }
+        pipe.scheduler = scheduler
+        pipe._base_scheduler_config = scheduler.config
+        pipe._current_flow_shift = float(scheduler.config.flow_shift)
+        pipe._engine_init_flow_shift = 10.0
+        for key, value in overrides.items():
+            setattr(pipe, key, value)
+        return pipe
+
+    return _make
+
+
+@pytest.fixture
+def make_cosmos3_stage():
+    """Return a factory building a ``Cosmos3DenoisingStage`` bound to a pipeline."""
+
+    def _make(pipeline):
+        from fastvideo.pipelines.stages.cosmos3_stages import Cosmos3DenoisingStage
+
+        return Cosmos3DenoisingStage(
+            transformer=pipeline.modules["transformer"],
+            scheduler=pipeline.modules["scheduler"],
+            vae=pipeline.modules["vae"],
+            tokenizer=pipeline.modules["text_tokenizer"],
+            pipeline=pipeline,
+        )
+
+    return _make
+
+
+def make_forward_batch(*, num_frames: int, height: int, width: int, image: Any = None, **overrides: Any):
+    """Build a tiny ``ForwardBatch`` for the Cosmos3 stage."""
+    from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+    values: dict[str, Any] = dict(
+        data_type="video",
+        prompt="a calm ocean at sunrise",
+        negative_prompt="",
+        height=height,
+        width=width,
+        num_frames=num_frames,
+        fps=24,
+        num_inference_steps=2,
+        guidance_scale=6.0,
+        generator=torch.Generator("cpu").manual_seed(0),
+        preprocessed_image=image,
+    )
+    values.update(overrides)
+    return ForwardBatch(**values)
+
+
+def make_fastvideo_args():
+    """Build minimal ``fastvideo_args`` (only ``pipeline_config`` is read)."""
+    from fastvideo.configs.pipelines.cosmos3 import Cosmos3Config
+
+    cfg = Cosmos3Config()
+    arch = cfg.dit_config.arch_config
+    arch.latent_channel = _LATENT_CHANNEL
+    arch.latent_patch_size = _LATENT_PATCH_SIZE
+    arch.temporal_compression_factor = _TEMPORAL_FACTOR
+    arch.enable_fps_modulation = False
+    return SimpleNamespace(pipeline_config=cfg)

tests/local_tests/cosmos3/test_cosmos3_action_parity.py

@@ -0,0 +1,271 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Numerical-parity test: FastVideo Cosmos3 action pathway vs the framework.
+
+Covers the action (multi-embodiment world-model) modality at the DiT level:
+
+  * **action packing** — native ``pack_cosmos3_video_sequence`` with a
+    ``Cosmos3ActionItem`` vs framework ``pack_input_sequence`` with
+    ``has_action``: action tokens share the vision "full" split, with ``(T,)``
+    shapes, a ``(T,1)`` condition mask, and 3D-MRoPE temporal positions at the
+    vision offset with ``start_frame_offset=1`` (parallel to vision); and
+  * **DiT action forward** — the dormant domain-aware ``action_proj_in`` /
+    ``action_proj_out`` (``DomainAwareLinear``) + ``action_modality_embed`` heads,
+    now activated, with a per-token embodiment ``domain_id``.
+
+Framework model + pack is the parity ORACLE (CPU/float32 via SDPA monkey-patch).
+We assert the native packer matches the framework field-by-field, then that
+``preds_vision`` AND ``preds_action`` match the framework forward.
+
+Run:
+    cd <worktree> && <fv-cosmos3 python> -m pytest \
+        tests/local_tests/cosmos3/test_cosmos3_action_parity.py -q -s
+"""
+from __future__ import annotations
+
+import pytest
+import torch
+
+cosmos_framework = pytest.importorskip(
+    "cosmos_framework",
+    reason="cosmos_framework not installed; run in fv-cosmos3 env.",
+)
+
+from .test_cosmos3_dit_parity import (  # noqa: E402
+    _fastvideo_inputs_from_packed_seq,
+    _framework_to_fastvideo_state_dict,
+)
+from .test_cosmos3_dit_parity_mrope import (  # noqa: E402
+    _ACTION_DIM,
+    _LATENT_CHANNEL,
+    _LATENT_PATCH_SIZE,
+    _RESET_SPATIAL_IDS,
+    _TCF,
+    _TEMPORAL_MODALITY_MARGIN,
+    _build_tiny_cosmos3_mrope,
+    _build_tiny_fastvideo_dit_mrope,
+)
+from .test_cosmos3_reference_forward import _apply_sdpa_patches  # noqa: E402
+
+pytestmark = [pytest.mark.local]
+_apply_sdpa_patches()
+
+_SPECIAL_TOKENS = {"start_of_generation": 60, "end_of_generation": 61, "eos_token_id": 62}
+
+
+def _copy_weights_with_action(vfm, dit) -> None:
+    """Copy backbone + vision weights AND the domain-aware action heads."""
+    mapped = _framework_to_fastvideo_state_dict(vfm, num_layers=dit.num_hidden_layers)
+    src = dict(vfm.named_parameters())
+    mapped["action_proj_in.fc.weight"] = src["action2llm.fc.weight"].detach().clone()
+    mapped["action_proj_in.bias.weight"] = src["action2llm.bias.weight"].detach().clone()
+    mapped["action_proj_out.fc.weight"] = src["llm2action.fc.weight"].detach().clone()
+    mapped["action_proj_out.bias.weight"] = src["llm2action.bias.weight"].detach().clone()
+    mapped["action_modality_embed"] = src["action_modality_embed"].detach().clone()
+    dst = dict(dit.named_parameters())
+    with torch.no_grad():
+        for name, tensor in mapped.items():
+            assert name in dst, f"DiT missing param {name!r}"
+            assert dst[name].shape == tensor.shape, f"shape mismatch {name}"
+            dst[name].copy_(tensor.to(dst[name].dtype))
+
+
+def _framework_pack_action(*, text_ids, vision, action, cond_vision, cond_action, domain_id, timestep,
+                           is_image_batch):
+    from cosmos_framework.data.vfm.sequence_packing import (
+        GenerationDataClean,
+        SequencePlan,
+        pack_input_sequence,
+    )
+
+    gen = GenerationDataClean(
+        batch_size=1,
+        is_image_batch=is_image_batch,
+        x0_tokens_vision=[vision],
+        fps_vision=None,
+        num_vision_items_per_sample=[1],
+        x0_tokens_action=[action],
+        fps_action=None,
+        action_domain_id=[torch.tensor([domain_id], dtype=torch.long)],
+    )
+    plans = [SequencePlan(
+        has_text=True, has_vision=True, has_action=True,
+        condition_frame_indexes_vision=list(cond_vision),
+        condition_frame_indexes_action=list(cond_action),
+    )]
+    ps = pack_input_sequence(
+        sequence_plans=plans,
+        input_text_indexes=[list(text_ids)],
+        gen_data_clean=gen,
+        input_timesteps=torch.tensor([timestep], dtype=torch.float32),
+        special_tokens=_SPECIAL_TOKENS,
+        latent_patch_size=_LATENT_PATCH_SIZE,
+        include_end_of_generation_token=False,
+        position_embedding_type="unified_3d_mrope",
+        unified_3d_mrope_reset_spatial_ids=_RESET_SPATIAL_IDS,
+        unified_3d_mrope_temporal_modality_margin=_TEMPORAL_MODALITY_MARGIN,
+        enable_fps_modulation=False,
+        base_fps=24.0,
+        temporal_compression_factor=_TCF,
+    )
+    # The framework sets action.domain_id on the packed sequence from
+    # gen_data_clean inside the model (_get_velocity); mirror that for the oracle.
+    if ps.action is not None:
+        ps.action.domain_id = [torch.tensor([domain_id], dtype=torch.long)]
+    return ps
+
+
+def _fastvideo_pack_action(*, text_ids, vision, action, cond_vision, cond_action, domain_id, timestep):
+    from fastvideo.pipelines.basic.cosmos3.sequence_packing import (
+        Cosmos3ActionItem,
+        Cosmos3SampleInputs,
+        Cosmos3VisionItem,
+        pack_cosmos3_video_sequence,
+    )
+
+    samples = [Cosmos3SampleInputs(
+        text_ids=list(text_ids),
+        vision=Cosmos3VisionItem(latent=vision, condition_frame_indexes=list(cond_vision)),
+        action=Cosmos3ActionItem(latent=action, condition_frame_indexes=list(cond_action), domain_id=domain_id),
+        timestep=float(timestep),
+    )]
+    return pack_cosmos3_video_sequence(
+        samples, _SPECIAL_TOKENS,
+        latent_patch_size=_LATENT_PATCH_SIZE, include_end_of_generation_token=False,
+        temporal_modality_margin=_TEMPORAL_MODALITY_MARGIN, reset_spatial_ids=_RESET_SPATIAL_IDS,
+        enable_fps_modulation=False, base_fps=24.0, temporal_compression_factor=_TCF,
+    )
+
+
+def _fv_inputs_with_action(ps) -> dict:
+    kw = _fastvideo_inputs_from_packed_seq(ps)
+    a = ps.action
+    kw.update(
+        action_tokens=list(a.tokens),
+        action_token_shapes=[tuple(x) for x in a.token_shapes],
+        action_sequence_indexes=a.sequence_indexes,
+        action_timesteps=a.timesteps,
+        action_mse_loss_indexes=a.mse_loss_indexes,
+        action_noisy_frame_indexes=list(a.noisy_frame_indexes),
+        action_domain_id=list(a.domain_id),
+    )
+    return kw
+
+
+def _diffs(a, b):
+    d = (a - b).abs()
+    return d.max().item(), d.mean().item()
+
+
+# (grid_t, lh, lw, action_t, n_text, cond_vision, cond_action, domain_id)
+_CASES = [
+    pytest.param(2, 4, 4, 6, 4, [], [], 0, id="a2v_2x2x2_act6_dom0"),
+    pytest.param(3, 8, 4, 9, 5, [], [], 7, id="a2v_3x4x2_act9_dom7"),
+    pytest.param(2, 4, 4, 5, 5, [0], [0], 3, id="ai2v_cond_act5_dom3"),
+]
+
+
+class TestCosmos3ActionParity:
+
+    def _build(self, num_layers=2, seed_model=42):
+        vfm = _build_tiny_cosmos3_mrope(seed=seed_model, num_layers=num_layers, action_gen=True)
+        dit = _build_tiny_fastvideo_dit_mrope(num_layers=num_layers)
+        _copy_weights_with_action(vfm, dit)
+        return vfm, dit
+
+    def _make_inputs(self, grid_t, lh, lw, act_t, n_text, cond_v, cond_a, dom, seed=7):
+        torch.manual_seed(seed)
+        return dict(
+            text_ids=torch.randint(0, 60, (n_text,)).tolist(),
+            vision=torch.randn(1, _LATENT_CHANNEL, grid_t, lh, lw),
+            action=torch.randn(act_t, _ACTION_DIM),  # [T, D]
+            cond_vision=cond_v, cond_action=cond_a, domain_id=dom, timestep=500.0,
+        )
+
+    @pytest.mark.parametrize(("grid_t", "lh", "lw", "act_t", "n_text", "cond_v", "cond_a", "dom"), _CASES)
+    def test_action_packing_matches_framework(self, grid_t, lh, lw, act_t, n_text, cond_v, cond_a, dom):
+        ins = self._make_inputs(grid_t, lh, lw, act_t, n_text, cond_v, cond_a, dom)
+        fw = _framework_pack_action(is_image_batch=(grid_t == 1), **ins)
+        fv = _fastvideo_pack_action(**ins)
+        assert fv.split_lens == list(fw.split_lens), f"split_lens fv={fv.split_lens} fw={list(fw.split_lens)}"
+        assert fv.attn_modes == list(fw.attn_modes)
+        assert int(fv.sequence_length) == int(fw.sequence_length)
+        torch.testing.assert_close(fv.position_ids, fw.position_ids, rtol=0, atol=0)
+        a = fw.action
+        torch.testing.assert_close(fv.action_sequence_indexes, a.sequence_indexes.to(torch.long), rtol=0, atol=0)
+        assert fv.action_token_shapes == [tuple(x) for x in a.token_shapes]
+        torch.testing.assert_close(fv.action_timesteps.to(torch.float32), a.timesteps.to(torch.float32))
+        torch.testing.assert_close(fv.action_mse_loss_indexes, a.mse_loss_indexes.to(torch.long), rtol=0, atol=0)
+        for x, y in zip(fv.action_noisy_frame_indexes, a.noisy_frame_indexes):
+            torch.testing.assert_close(x.to(torch.long), y.to(torch.long), rtol=0, atol=0)
+        print(f"\n[action_packing {grid_t}x{lh}x{lw} act={act_t} dom={dom}] position_ids + action fields exact")
+
+    @pytest.mark.parametrize(("grid_t", "lh", "lw", "act_t", "n_text", "cond_v", "cond_a", "dom"), _CASES)
+    def test_action_dit_forward_matches_framework(self, grid_t, lh, lw, act_t, n_text, cond_v, cond_a, dom):
+        vfm, dit = self._build()
+        ins = self._make_inputs(grid_t, lh, lw, act_t, n_text, cond_v, cond_a, dom)
+        fw_pack = _framework_pack_action(is_image_batch=(grid_t == 1), **ins)
+        fv_pack = _fastvideo_pack_action(**ins)
+        with torch.no_grad():
+            fw_out = vfm(packed_seq=fw_pack)
+            fv_out = dit(**fv_pack.to_dit_kwargs())
+            fv_on_fw = dit(**_fv_inputs_with_action(fw_pack))
+        pv_mx, pv_mn = _diffs(fv_out["preds_vision"][0], fw_out["preds_vision"][0])
+        pa_mx, pa_mn = _diffs(fv_out["preds_action"][0], fw_out["preds_action"][0])
+        paf_mx, paf_mn = _diffs(fv_on_fw["preds_action"][0], fw_out["preds_action"][0])
+        print(f"\n[action_dit {grid_t}x{lh}x{lw} act={act_t} dom={dom}] "
+              f"preds_vision max={pv_mx:.3e} mean={pv_mn:.3e} | "
+              f"preds_action max={pa_mx:.3e} mean={pa_mn:.3e} | "
+              f"preds_action(fwpack) max={paf_mx:.3e} mean={paf_mn:.3e}")
+        assert fv_out["preds_action"][0].shape == fw_out["preds_action"][0].shape
+        torch.testing.assert_close(fv_out["preds_vision"][0], fw_out["preds_vision"][0], atol=1e-4, rtol=1e-3)
+        torch.testing.assert_close(fv_out["preds_action"][0], fw_out["preds_action"][0], atol=1e-4, rtol=1e-3)
+        torch.testing.assert_close(fv_on_fw["preds_action"][0], fw_out["preds_action"][0], atol=1e-4, rtol=1e-3)
+
+    @pytest.mark.parametrize(("grid_t", "lh", "lw", "act_t", "n_text", "cond_v", "cond_a", "dom"), _CASES)
+    def test_action_cfg_velocity_matches_framework(self, grid_t, lh, lw, act_t, n_text, cond_v, cond_a, dom):
+        """Combined [vision|action] sequential-CFG velocity (action pipeline glue)
+        matches a framework-DiT oracle."""
+        from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import (
+            Cosmos3ActionSpec,
+            Cosmos3VisionSpec,
+            cosmos3_get_cfg_velocity,
+        )
+
+        vfm, dit = self._build()
+        vlat_shape = (_LATENT_CHANNEL, grid_t, lh, lw)
+        action_shape = (act_t, _ACTION_DIM)
+        torch.manual_seed(3)
+        cond_ids = torch.randint(0, 60, (n_text,)).tolist()
+        uncond_ids = torch.randint(0, 60, (max(1, n_text - 1),)).tolist()
+        vis_numel = int(torch.tensor(vlat_shape).prod())
+        act_numel = int(torch.tensor(action_shape).prod())
+        flat = torch.randn(vis_numel + act_numel)
+        guidance, ts = 6.0, 500.0
+
+        def _fw_velocity(ids):
+            vision = flat[:vis_numel].reshape(vlat_shape).unsqueeze(0)
+            action = flat[vis_numel:].reshape(action_shape)
+            ps = _framework_pack_action(text_ids=ids, vision=vision, action=action, cond_vision=cond_v,
+                                        cond_action=cond_a, domain_id=dom, timestep=ts,
+                                        is_image_batch=(grid_t == 1))
+            with torch.no_grad():
+                out = vfm(packed_seq=ps)
+            pv = out["preds_vision"][0].squeeze(0)  # [C,T,H,W] (zero on clean)
+            pa = out["preds_action"][0]  # [T,D] (zero on clean)
+            return torch.cat([pv.reshape(-1), pa.reshape(-1)])
+
+        fw_cond, fw_uncond = _fw_velocity(cond_ids), _fw_velocity(uncond_ids)
+        fw_v = fw_uncond + guidance * (fw_cond - fw_uncond)
+        fv_v = cosmos3_get_cfg_velocity(
+            transformer=dit, flat_latent=flat, timestep=torch.tensor([ts]), guidance=guidance,
+            specs=[Cosmos3VisionSpec(shape=vlat_shape, condition_frame_indexes=list(cond_v))],
+            action_specs=[Cosmos3ActionSpec(shape=action_shape, condition_frame_indexes=list(cond_a), domain_id=dom)],
+            cond_token_ids=cond_ids, uncond_token_ids=uncond_ids,
+            special_tokens=_SPECIAL_TOKENS, latent_patch_size=_LATENT_PATCH_SIZE,
+            temporal_modality_margin=_TEMPORAL_MODALITY_MARGIN, reset_spatial_ids=_RESET_SPATIAL_IDS,
+            enable_fps_modulation=False, base_fps=24.0, temporal_compression_factor=_TCF,
+        )
+        assert fv_v.shape == fw_v.shape, f"shape fv={fv_v.shape} fw={fw_v.shape}"
+        mx, mn = _diffs(fv_v, fw_v)
+        print(f"\n[action_cfg_velocity {grid_t}x{lh}x{lw} act={act_t} dom={dom}] max={mx:.3e} mean={mn:.3e}")
+        torch.testing.assert_close(fv_v, fw_v, atol=1e-4, rtol=1e-3)

tests/local_tests/cosmos3/test_cosmos3_avae_parity.py

@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Numerical-parity test: FastVideo Cosmos3 sound decoder vs the framework AVAE.
+
+The Cosmos3 ``sound_tokenizer`` is an AVAE (audio VAE). Its shipped diffusers
+checkpoint is **decoder-only** (``decoder.*``; the SpectrogramConvNeXt encoder is
+not exported) in diffusers ``AutoencoderOobleck`` naming, but with **SnakeBeta**
+activations (alpha+beta, logscale) and ``weight_g``/``weight_v`` weight-norm —
+i.e. exactly FastVideo's existing native ``OobleckVAE`` decoder
+(``fastvideo/models/vaes/oobleck.py``). t2vs only needs DECODE (generate sound
+latents -> waveform), so this pins the decoder.
+
+The framework decoder
+(``cosmos_framework.model.vfm.tokenizers.audio.avae_utils.models.OobleckDecoder``,
+``nn.Sequential`` naming, ``output_padding=stride%2`` on the transpose convs) is
+the parity ORACLE. We build a tiny framework decoder, map its weights into the
+FastVideo decoder (Sequential -> conv1/block.N/res_unitM/snake1/conv2), and
+assert bit-exact decode. Strides include an ODD value (5, as in the real config
+``[2,4,5,6,8]``) to exercise the ``output_padding`` path that diverged before.
+
+CPU / float32. Run:
+    cd <worktree> && <fv-cosmos3 python> -m pytest \
+        tests/local_tests/cosmos3/test_cosmos3_avae_parity.py -q
+"""
+from __future__ import annotations
+
+import pytest
+import torch
+
+# The official framework provides the parity oracle.
+_fw_models = pytest.importorskip(
+    "cosmos_framework.model.vfm.tokenizers.audio.avae_utils.models",
+    reason="cosmos_framework not installed; run in fv-cosmos3 env.",
+)
+from cosmos_framework.model.vfm.tokenizers.audio.avae_utils.env import (  # noqa: E402
+    AttrDict,
+)
+
+from fastvideo.models.vaes.oobleck import OobleckDecoder as FvOobleckDecoder  # noqa: E402
+
+pytestmark = [pytest.mark.local]
+
+FwOobleckDecoder = _fw_models.OobleckDecoder
+
+
+def _framework_decoder(dec_dim, vocoder_input_dim, dec_c_mults, dec_strides):
+    """Framework OobleckDecoder (the parity oracle), non-causal / no-antialias."""
+    h = AttrDict({
+        "vocoder_input_dim": vocoder_input_dim,
+        "input_channels": 1,
+        "stereo": True,  # 2 audio channels
+        "dec_dim": dec_dim,
+        "dec_c_mults": dec_c_mults,
+        "dec_strides": dec_strides,
+        "dec_use_snake": True,
+        "dec_use_nearest_upsample": False,
+        "dec_anti_aliasing": False,
+        "causal": False,
+        "dec_use_tanh_at_final": False,
+        "padding_mode": "zeros",
+    })
+    return FwOobleckDecoder(h).eval()
+
+
+def _framework_to_fastvideo_decoder_state(fw_decoder, num_blocks):
+    """Map framework Sequential decoder weights -> FastVideo decoder names.
+
+    framework: layers.0=first conv; layers.{1..K}=OobleckDecoderBlock
+    (.layers.0 snake, .1 conv_t, .{2,3,4} ResidualUnit{.layers.0 snake,
+    .1 conv, .2 snake, .3 conv}); layers.{1+K}=final snake; layers.{2+K}=final conv.
+    FastVideo: conv1; block.{b}.{snake1,conv_t1,res_unit{1,2,3}.{snake1,conv1,snake2,conv2}};
+    snake1; conv2. Snake alpha/beta: framework [C] -> FastVideo [1,C,1].
+    """
+    out = {}
+    for k, v in fw_decoder.state_dict().items():
+        p = k.split(".")
+        li = int(p[1])
+        if li == 0:
+            nk = "conv1." + ".".join(p[2:])
+        elif li == 1 + num_blocks:
+            nk = "snake1." + ".".join(p[2:])
+        elif li == 2 + num_blocks:
+            nk = "conv2." + ".".join(p[2:])
+        else:
+            b = li - 1
+            sub = int(p[3])
+            if sub == 0:
+                nk = f"block.{b}.snake1." + ".".join(p[4:])
+            elif sub == 1:
+                nk = f"block.{b}.conv_t1." + ".".join(p[4:])
+            else:
+                r = sub - 2  # ResidualUnit index 0..2
+                m = {0: "snake1", 1: "conv1", 2: "snake2", 3: "conv2"}[int(p[5])]
+                nk = f"block.{b}.res_unit{r + 1}.{m}." + ".".join(p[6:])
+        if nk.endswith(".alpha") or nk.endswith(".beta"):
+            v = v.reshape(1, -1, 1)
+        out[nk] = v
+    return out
+
+
+# (dec_dim, vocoder_input_dim, dec_c_mults, dec_strides) — tiny; strides incl odd.
+_CASES = [
+    pytest.param(4, 8, [1, 2], [5, 2], id="odd_stride5"),
+    pytest.param(6, 8, [1, 2, 4], [2, 5, 6], id="real_stride_pattern_tiny"),
+    pytest.param(4, 4, [1, 2], [4, 8], id="even_strides"),
+]
+
+
+class TestCosmos3AVAEParity:
+
+    @pytest.mark.parametrize(("dec_dim", "vin", "cmults", "strides"), _CASES)
+    def test_decode_matches_framework(self, dec_dim, vin, cmults, strides):
+        torch.manual_seed(0)
+        fw = _framework_decoder(dec_dim, vin, cmults, strides)
+        fv = FvOobleckDecoder(
+            channels=dec_dim,
+            input_channels=vin,
+            audio_channels=2,
+            upsampling_ratios=list(reversed(strides)),  # framework reverses dec_strides
+            channel_multiples=cmults,
+        ).eval()
+        state = _framework_to_fastvideo_decoder_state(fw, num_blocks=len(strides))
+        fv.load_state_dict(state, strict=True)  # exact name + shape match
+
+        z = torch.randn(1, vin, 5)
+        with torch.no_grad():
+            a = fw(z)
+            b = fv(z)
+        assert a.shape == b.shape, f"shape: fw={a.shape} fv={b.shape}"
+        max_abs = (a - b).abs().max().item()
+        print(f"\n[avae_decode dim={dec_dim} strides={strides}] max abs diff = {max_abs:.3e}")
+        torch.testing.assert_close(b, a, atol=1e-6, rtol=1e-5)

tests/local_tests/cosmos3/test_cosmos3_denoise_cfg_parity.py

@@ -0,0 +1,342 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Numerical-parity test: FastVideo Cosmos3 denoise/CFG glue vs the framework.
+
+The DiT forward and the sequence-packing are already framework-parity-verified
+(``test_cosmos3_dit_parity*`` / ``test_cosmos3_packing_parity``). This test pins
+the remaining glue that the native pipeline adds — the SEQUENTIAL classifier-free
+guidance velocity and one UniPC scheduler step — against the framework math
+(``diffusers_cosmos3.pipeline.Cosmos3OmniDiffusersPipeline.get_cfg_velocity`` /
+``__call__``):
+
+  * for one denoise step, replicate the framework's ``get_cfg_velocity`` exactly
+    on top of the OFFICIAL ``Cosmos3VFMNetwork`` forward (oracle): a conditional
+    pass (prompt tokens) and an unconditional pass (negative-prompt tokens),
+    each masking the prediction on conditioning frames
+    (``pred * (1 - condition_mask)``), then ``v = uncond + g*(cond - uncond)``;
+  * run FastVideo's :func:`cosmos3_get_cfg_velocity` with the native DiT (the
+    framework weights copied in) + the native packer, and assert the velocity
+    matches the oracle;
+  * take one ``UniPCMultistepScheduler.step`` on each (the actual checkpoint
+    scheduler) and assert the stepped latent matches;
+  * drive :meth:`Cosmos3DenoiseEngine.denoise` for >= 2 steps and assert it
+    equals the manual framework step-by-step loop.
+
+CPU / float32, via the reference SDPA monkey-patch. The official model is the
+parity ORACLE.
+
+Run:
+    cd <worktree> && <fv-cosmos3 python> -m pytest \
+        tests/local_tests/cosmos3/test_cosmos3_denoise_cfg_parity.py -q
+"""
+from __future__ import annotations
+
+import pytest
+import torch
+
+# The official framework provides the parity oracle.
+cosmos_framework = pytest.importorskip(
+    "cosmos_framework",
+    reason="cosmos_framework not installed; run in fv-cosmos3 env.",
+)
+
+from .test_cosmos3_dit_parity import _copy_weights  # noqa: E402
+from .test_cosmos3_dit_parity_mrope import (  # noqa: E402
+    _LATENT_CHANNEL,
+    _LATENT_PATCH_SIZE,
+    _RESET_SPATIAL_IDS,
+    _TCF,
+    _TEMPORAL_MODALITY_MARGIN,
+    _build_tiny_cosmos3_mrope,
+    _build_tiny_fastvideo_dit_mrope,
+)
+from .test_cosmos3_reference_forward import _apply_sdpa_patches  # noqa: E402
+from .test_cosmos3_scheduler_parity import (  # noqa: E402
+    _fastvideo_scheduler,
+    _framework_scheduler,
+)
+
+pytestmark = [pytest.mark.local]
+
+_apply_sdpa_patches()
+
+# Tiny special tokens (< tiny vocab_size=64), video path appends eos + sog.
+_SPECIAL_TOKENS = {
+    "start_of_generation": 60,
+    "end_of_generation": 61,
+    "eos_token_id": 62,
+}
+
+# Cosmos3 video flow_shift; framework scheduler is the parity oracle, FastVideo's
+# vendored UniPC (flow config) is the unit under test.
+_FLOW_SHIFT = 10.0
+
+
+# ---------------------------------------------------------------------------
+# Framework-oracle CFG velocity (replicates pipeline.get_cfg_velocity math).
+# ---------------------------------------------------------------------------
+def _framework_pack(*, text_ids, vision_latent, cond_frames, timestep):
+    from cosmos_framework.data.vfm.sequence_packing import (
+        GenerationDataClean,
+        SequencePlan,
+        pack_input_sequence,
+    )
+
+    # vision_latent is [1, C, T, H, W]; temporal dim is axis 2.
+    gen_data_clean = GenerationDataClean(
+        batch_size=1,
+        is_image_batch=(vision_latent.shape[2] == 1),
+        x0_tokens_vision=[vision_latent],
+        fps_vision=None,
+        num_vision_items_per_sample=[1],
+    )
+    plans = [SequencePlan(has_text=True, has_vision=True, condition_frame_indexes_vision=list(cond_frames))]
+    return pack_input_sequence(
+        sequence_plans=plans,
+        input_text_indexes=[list(text_ids)],
+        gen_data_clean=gen_data_clean,
+        input_timesteps=torch.tensor([timestep], dtype=torch.float32),
+        special_tokens=_SPECIAL_TOKENS,
+        latent_patch_size=_LATENT_PATCH_SIZE,
+        include_end_of_generation_token=False,
+        position_embedding_type="unified_3d_mrope",
+        unified_3d_mrope_reset_spatial_ids=_RESET_SPATIAL_IDS,
+        unified_3d_mrope_temporal_modality_margin=_TEMPORAL_MODALITY_MARGIN,
+        enable_fps_modulation=False,
+        base_fps=24.0,
+        temporal_compression_factor=_TCF,
+    )
+
+
+def _framework_inputs(ps):
+    """Framework PackedSequence -> framework Cosmos3VFMNetwork forward kwargs."""
+    return dict(packed_seq=ps)
+
+
+def _framework_cfg_velocity(
+    *,
+    vfm,
+    flat_latent: torch.Tensor,
+    timestep: torch.Tensor,
+    guidance: float,
+    vision_shape: tuple[int, int, int, int],
+    cond_frames: list[int],
+    cond_ids: list[int],
+    uncond_ids: list[int],
+) -> torch.Tensor:
+    """Replicate the framework ``get_cfg_velocity`` on the oracle model.
+
+    Single vision item; sequential cond then uncond pass; mask condition
+    frames; ``v = uncond + g*(cond - uncond)``.
+    """
+    timestep_value = float(timestep.reshape(()).item())
+    vision_latent = flat_latent.reshape(vision_shape)  # [C, T, H, W]
+
+    def _run(text_ids: list[int]) -> torch.Tensor:
+        ps = _framework_pack(
+            text_ids=text_ids,
+            # The framework packer expects a 5D [1, C, T, H, W] latent.
+            vision_latent=vision_latent.unsqueeze(0),
+            cond_frames=cond_frames,
+            timestep=timestep_value,
+        )
+        out = vfm(**_framework_inputs(ps))
+        preds = out.get("preds_vision")
+        cond_mask = ps.vision.condition_mask[0]  # [T] or [T,1,1]
+        if preds is None:
+            return torch.zeros_like(flat_latent)
+        pred = preds[0].squeeze(0)  # [C, T, H, W]
+        keep = (1.0 - cond_mask.reshape(-1, 1, 1)).to(dtype=pred.dtype, device=pred.device)
+        velocity = pred * keep if keep.sum() > 0 else torch.zeros_like(pred)
+        return velocity.reshape(-1)
+
+    cond_v = _run(cond_ids)
+    uncond_v = _run(uncond_ids)
+    return uncond_v + guidance * (cond_v - uncond_v)
+
+
+# ---------------------------------------------------------------------------
+# Cases: T2V (no cond), I2V (cond frame 0), single-frame T2I.
+# ---------------------------------------------------------------------------
+_CASES = [
+    pytest.param(2, 4, 4, 6, [], id="t2v_2x2x2"),
+    pytest.param(3, 8, 4, 5, [0], id="i2v_3x4x2_cond0"),
+    pytest.param(1, 8, 8, 4, [], id="t2i_1x4x4"),
+]
+
+
+def _build_models(num_layers: int = 2, seed_model: int = 42):
+    vfm = _build_tiny_cosmos3_mrope(seed=seed_model, num_layers=num_layers)
+    dit = _build_tiny_fastvideo_dit_mrope(num_layers=num_layers)
+    _copy_weights(vfm, dit)
+    return vfm, dit
+
+
+def _fastvideo_velocity(dit, *, flat_latent, timestep, guidance, vision_shape, cond_frames, cond_ids, uncond_ids):
+    from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import (
+        Cosmos3VisionSpec,
+        cosmos3_get_cfg_velocity,
+    )
+
+    spec = Cosmos3VisionSpec(shape=vision_shape, condition_frame_indexes=list(cond_frames))
+    return cosmos3_get_cfg_velocity(
+        transformer=dit,
+        flat_latent=flat_latent,
+        timestep=timestep,
+        guidance=guidance,
+        specs=[spec],
+        cond_token_ids=cond_ids,
+        uncond_token_ids=uncond_ids,
+        special_tokens=_SPECIAL_TOKENS,
+        latent_patch_size=_LATENT_PATCH_SIZE,
+        temporal_modality_margin=_TEMPORAL_MODALITY_MARGIN,
+        reset_spatial_ids=_RESET_SPATIAL_IDS,
+        enable_fps_modulation=False,
+        base_fps=24.0,
+        temporal_compression_factor=_TCF,
+    )
+
+
+class TestCosmos3DenoiseCFGParity:
+
+    @pytest.mark.parametrize(("grid_t", "latent_h", "latent_w", "n_text", "cond"), _CASES)
+    def test_cfg_velocity_matches_framework(self, grid_t, latent_h, latent_w, n_text, cond):
+        vfm, dit = _build_models()
+        torch.manual_seed(0)
+        cond_ids = torch.randint(0, 60, (n_text,)).tolist()
+        uncond_ids = torch.randint(0, 60, (max(1, n_text - 1),)).tolist()
+        vision_shape = (_LATENT_CHANNEL, grid_t, latent_h, latent_w)
+        flat_latent = torch.randn(int(torch.tensor(vision_shape).prod()))
+        timestep = torch.tensor([[500.0]])  # framework expects [1,1]; we reshape to scalar
+        guidance = 6.0
+
+        fw_v = _framework_cfg_velocity(
+            vfm=vfm,
+            flat_latent=flat_latent,
+            timestep=timestep,
+            guidance=guidance,
+            vision_shape=vision_shape,
+            cond_frames=cond,
+            cond_ids=cond_ids,
+            uncond_ids=uncond_ids,
+        )
+        fv_v = _fastvideo_velocity(
+            dit,
+            flat_latent=flat_latent,
+            timestep=timestep,
+            guidance=guidance,
+            vision_shape=vision_shape,
+            cond_frames=cond,
+            cond_ids=cond_ids,
+            uncond_ids=uncond_ids,
+        )
+        assert fw_v.shape == fv_v.shape, f"shape: fw={fw_v.shape} fv={fv_v.shape}"
+        max_abs = (fw_v - fv_v).abs().max().item()
+        print(f"\n[cfg_velocity {grid_t}x{latent_h}x{latent_w} cond={cond}] max abs diff = {max_abs:.3e}")
+        torch.testing.assert_close(fv_v, fw_v, atol=1e-4, rtol=1e-3)
+
+    def test_one_unipc_step_matches_framework(self):
+        """CFG velocity + one UniPC step: FastVideo == framework math."""
+        vfm, dit = _build_models()
+        grid_t, latent_h, latent_w = 2, 4, 4
+        vision_shape = (_LATENT_CHANNEL, grid_t, latent_h, latent_w)
+        torch.manual_seed(3)
+        cond_ids = torch.randint(0, 60, (5,)).tolist()
+        uncond_ids = torch.randint(0, 60, (4,)).tolist()
+        flat_latent = torch.randn(int(torch.tensor(vision_shape).prod()))
+        guidance = 6.0
+
+        fw_sched = _framework_scheduler(4, _FLOW_SHIFT)
+        fv_sched = _fastvideo_scheduler(4, _FLOW_SHIFT)
+        t = fw_sched.timesteps[0]
+
+        fw_v = _framework_cfg_velocity(
+            vfm=vfm,
+            flat_latent=flat_latent,
+            timestep=t.reshape(1, 1),
+            guidance=guidance,
+            vision_shape=vision_shape,
+            cond_frames=[],
+            cond_ids=cond_ids,
+            uncond_ids=uncond_ids,
+        )
+        fw_stepped = fw_sched.step(model_output=fw_v, timestep=t, sample=flat_latent.unsqueeze(0),
+                                   return_dict=False)[0].squeeze(0)
+
+        fv_v = _fastvideo_velocity(
+            dit,
+            flat_latent=flat_latent,
+            timestep=t.reshape(1),
+            guidance=guidance,
+            vision_shape=vision_shape,
+            cond_frames=[],
+            cond_ids=cond_ids,
+            uncond_ids=uncond_ids,
+        )
+        fv_stepped = fv_sched.step(model_output=fv_v, timestep=t, sample=flat_latent.unsqueeze(0),
+                                   return_dict=False)[0].squeeze(0)
+
+        max_abs = (fw_stepped - fv_stepped).abs().max().item()
+        print(f"\n[unipc_step] max abs diff = {max_abs:.3e}")
+        torch.testing.assert_close(fv_stepped, fw_stepped, atol=1e-4, rtol=1e-3)
+
+    def test_full_denoise_loop_matches_framework(self):
+        """Cosmos3DenoiseEngine.denoise (>= 2 steps) == framework step-by-step."""
+        from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import (
+            Cosmos3DenoiseEngine,
+            Cosmos3VisionSpec,
+        )
+
+        vfm, dit = _build_models()
+        grid_t, latent_h, latent_w = 2, 4, 4
+        vision_shape = (_LATENT_CHANNEL, grid_t, latent_h, latent_w)
+        torch.manual_seed(5)
+        cond_ids = torch.randint(0, 60, (5,)).tolist()
+        uncond_ids = torch.randint(0, 60, (4,)).tolist()
+        flat_latent = torch.randn(int(torch.tensor(vision_shape).prod()))
+        guidance = 6.0
+        num_steps = 3
+
+        # Manual framework loop (oracle).
+        fw_sched = _framework_scheduler(num_steps, _FLOW_SHIFT)
+        fw_latent = flat_latent.clone()
+        for t in fw_sched.timesteps:
+            v = _framework_cfg_velocity(
+                vfm=vfm,
+                flat_latent=fw_latent,
+                timestep=t.reshape(1, 1),
+                guidance=guidance,
+                vision_shape=vision_shape,
+                cond_frames=[],
+                cond_ids=cond_ids,
+                uncond_ids=uncond_ids,
+            )
+            fw_latent = fw_sched.step(model_output=v, timestep=t, sample=fw_latent.unsqueeze(0),
+                                      return_dict=False)[0].squeeze(0)
+
+        # FastVideo engine loop.
+        fv_sched = _fastvideo_scheduler(num_steps, _FLOW_SHIFT)
+        engine = Cosmos3DenoiseEngine(
+            transformer=dit,
+            scheduler=fv_sched,
+            special_tokens=_SPECIAL_TOKENS,
+            latent_patch_size=_LATENT_PATCH_SIZE,
+            temporal_modality_margin=_TEMPORAL_MODALITY_MARGIN,
+            reset_spatial_ids=_RESET_SPATIAL_IDS,
+            enable_fps_modulation=False,
+            base_fps=24.0,
+            temporal_compression_factor=_TCF,
+        )
+        spec = Cosmos3VisionSpec(shape=vision_shape, condition_frame_indexes=[])
+        fv_latent = engine.denoise(
+            flat_latent=flat_latent.clone(),
+            timesteps=fv_sched.timesteps,
+            guidance=guidance,
+            specs=[spec],
+            cond_token_ids=cond_ids,
+            uncond_token_ids=uncond_ids,
+        )
+
+        assert fv_latent.shape == fw_latent.shape
+        max_abs = (fw_latent - fv_latent).abs().max().item()
+        print(f"\n[full_denoise {num_steps} steps] final latent max abs diff = {max_abs:.3e}")
+        torch.testing.assert_close(fv_latent, fw_latent, atol=1e-4, rtol=1e-3)

tests/local_tests/cosmos3/test_cosmos3_dit_parity.py

@@ -0,0 +1,256 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Numerical-parity test: FastVideo Cosmos3 DiT vs official ``Cosmos3VFMNetwork``.
+
+Builds a tiny official-framework ``Cosmos3VFMNetwork`` AND a tiny FastVideo
+``Cosmos3VFMTransformer`` from the SAME tiny config, copies the framework
+weights into the FastVideo DiT via an explicit framework->fastvideo name map,
+runs BOTH forwards on identical deterministic inputs (CPU / float32), and
+asserts ``torch.allclose`` on the vision prediction output (``preds_vision``)
+and the per-token ``last_hidden_state``.
+
+The official model is the parity ORACLE. It runs on CPU / float32 via the SDPA
+monkey-patch in ``test_cosmos3_reference_forward`` (flash2/flash3/natten are
+CUDA-only). The FastVideo DiT runs natively on CPU with plain SDPA.
+
+Run:
+    cd <worktree> && <fv-cosmos3 python> -m pytest \
+        tests/local_tests/cosmos3/test_cosmos3_dit_parity.py -q
+"""
+from __future__ import annotations
+
+import pytest
+import torch
+
+# The official framework provides the parity oracle.
+cosmos_framework = pytest.importorskip(
+    "cosmos_framework",
+    reason="cosmos_framework not installed; run in fv-cosmos3 env.",
+)
+
+# Reuse the reference harness's tiny-model builder + SDPA monkey-patch.
+from .test_cosmos3_reference_forward import (  # noqa: E402
+    _apply_sdpa_patches,
+    _build_tiny_cosmos3,
+    _build_tiny_packed_seq,
+)
+
+pytestmark = [pytest.mark.local]
+
+# Ensure the CPU/float32 SDPA patches are installed (idempotent).
+_apply_sdpa_patches()
+
+
+# ---------------------------------------------------------------------------
+# Tiny config shared by both models (must match _build_tiny_cosmos3).
+# ---------------------------------------------------------------------------
+def _build_tiny_fastvideo_dit() -> "Cosmos3VFMTransformer":  # noqa: F821
+    from fastvideo.configs.models.dits.cosmos3 import (
+        Cosmos3ArchConfig,
+        Cosmos3VideoConfig,
+    )
+    from fastvideo.models.dits.cosmos3 import Cosmos3VFMTransformer
+
+    arch = Cosmos3ArchConfig(
+        hidden_size=16,
+        num_hidden_layers=1,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        head_dim=8,
+        intermediate_size=32,
+        vocab_size=64,
+        rms_norm_eps=1e-6,
+        attention_bias=False,
+        latent_patch_size=2,
+        latent_channel=16,
+        rope_theta=5_000_000.0,
+        mrope_section=[24, 20, 20],
+        position_embedding_type="3d_rope",
+        base_fps=24.0,
+        temporal_compression_factor=4,
+        enable_fps_modulation=False,
+        # Dormant heads present in the checkpoint surface (constructed for
+        # strict-load parity; not exercised by this video-path forward).
+        action_gen=True,
+        action_dim=64,
+        max_action_dim=64,
+        num_embodiment_domains=32,
+        sound_gen=True,
+        sound_dim=64,
+    )
+    cfg = Cosmos3VideoConfig(arch_config=arch)
+    model = Cosmos3VFMTransformer(cfg, hf_config={})
+    return model.to(torch.float32).eval()
+
+
+# ---------------------------------------------------------------------------
+# Framework -> FastVideo weight name map.
+# ---------------------------------------------------------------------------
+def _framework_to_fastvideo_state_dict(vfm, num_layers: int) -> dict[str, torch.Tensor]:
+    """Translate framework param names into the FastVideo DiT param names.
+
+    Framework (Cosmos3VFMNetwork):
+      language_model.model.{embed_tokens,norm,norm_moe_gen}
+      language_model.lm_head
+      language_model.model.layers.{i}.self_attn.{q,k,v,o}_proj(+ _moe_gen)
+      language_model.model.layers.{i}.self_attn.{q,k}_norm(+ _moe_gen)
+      language_model.model.layers.{i}.{mlp,mlp_moe_gen}.{gate,up,down}_proj
+      language_model.model.layers.{i}.{input,post_attention}_layernorm(+ _moe_gen)
+      vae2llm / llm2vae / time_embedder.mlp.{0,2}
+
+    FastVideo (Cosmos3VFMTransformer):
+      embed_tokens / norm / norm_moe_gen / lm_head
+      layers.{i}.self_attn.{to_q,to_k,to_v,to_out} (und)
+      layers.{i}.self_attn.{add_q,add_k,add_v}_proj / to_add_out (gen)
+      layers.{i}.self_attn.{norm_q,norm_k,norm_added_q,norm_added_k}
+      layers.{i}.{mlp,mlp_moe_gen}.{gate,up,down}_proj
+      layers.{i}.{input,post_attention}_layernorm(+ _moe_gen)
+      proj_in / proj_out / time_embedder.linear_{1,2}
+    """
+    src = dict(vfm.named_parameters())
+    out: dict[str, torch.Tensor] = {}
+
+    def take(name: str) -> torch.Tensor:
+        return src[name].detach().clone()
+
+    # ---- Top-level backbone ----
+    out["embed_tokens.weight"] = take("language_model.model.embed_tokens.weight")
+    out["norm.weight"] = take("language_model.model.norm.weight")
+    out["norm_moe_gen.weight"] = take("language_model.model.norm_moe_gen.weight")
+    out["lm_head.weight"] = take("language_model.lm_head.weight")
+
+    # ---- Vision adapters ----
+    out["proj_in.weight"] = take("vae2llm.weight")
+    out["proj_in.bias"] = take("vae2llm.bias")
+    out["proj_out.weight"] = take("llm2vae.weight")
+    out["proj_out.bias"] = take("llm2vae.bias")
+
+    # ---- Timestep embedder (mlp.0/mlp.2 -> linear_1/linear_2) ----
+    out["time_embedder.linear_1.weight"] = take("time_embedder.mlp.0.weight")
+    out["time_embedder.linear_1.bias"] = take("time_embedder.mlp.0.bias")
+    out["time_embedder.linear_2.weight"] = take("time_embedder.mlp.2.weight")
+    out["time_embedder.linear_2.bias"] = take("time_embedder.mlp.2.bias")
+
+    # ---- Per layer ----
+    und_attn = {"q_proj": "to_q", "k_proj": "to_k", "v_proj": "to_v", "o_proj": "to_out"}
+    gen_attn = {
+        "q_proj_moe_gen": "add_q_proj",
+        "k_proj_moe_gen": "add_k_proj",
+        "v_proj_moe_gen": "add_v_proj",
+        "o_proj_moe_gen": "to_add_out",
+    }
+    und_norm = {"q_norm": "norm_q", "k_norm": "norm_k"}
+    gen_norm = {"q_norm_moe_gen": "norm_added_q", "k_norm_moe_gen": "norm_added_k"}
+
+    for i in range(num_layers):
+        fw = f"language_model.model.layers.{i}"
+        fv = f"layers.{i}"
+        for s, d in und_attn.items():
+            out[f"{fv}.self_attn.{d}.weight"] = take(f"{fw}.self_attn.{s}.weight")
+        for s, d in gen_attn.items():
+            out[f"{fv}.self_attn.{d}.weight"] = take(f"{fw}.self_attn.{s}.weight")
+        for s, d in und_norm.items():
+            out[f"{fv}.self_attn.{d}.weight"] = take(f"{fw}.self_attn.{s}.weight")
+        for s, d in gen_norm.items():
+            out[f"{fv}.self_attn.{d}.weight"] = take(f"{fw}.self_attn.{s}.weight")
+        for mlp in ("mlp", "mlp_moe_gen"):
+            for proj in ("gate_proj", "up_proj", "down_proj"):
+                out[f"{fv}.{mlp}.{proj}.weight"] = take(f"{fw}.{mlp}.{proj}.weight")
+        for ln in ("input_layernorm", "input_layernorm_moe_gen", "post_attention_layernorm",
+                   "post_attention_layernorm_moe_gen"):
+            out[f"{fv}.{ln}.weight"] = take(f"{fw}.{ln}.weight")
+
+    return out
+
+
+def _copy_weights(vfm, dit) -> None:
+    """Copy framework weights into the FastVideo DiT (shape-checked)."""
+    mapped = _framework_to_fastvideo_state_dict(vfm, num_layers=dit.num_hidden_layers)
+    dst = dict(dit.named_parameters())
+    # Every mapped tensor must land on an existing FastVideo param with a matching shape.
+    for name, tensor in mapped.items():
+        assert name in dst, f"FastVideo DiT missing param for mapped key {name!r}"
+        assert dst[name].shape == tensor.shape, (f"shape mismatch for {name}: "
+                                                 f"dit={tuple(dst[name].shape)} fw={tuple(tensor.shape)}")
+    with torch.no_grad():
+        for name, tensor in mapped.items():
+            dst[name].copy_(tensor.to(dst[name].dtype))
+
+
+def _fastvideo_inputs_from_packed_seq(ps) -> dict:
+    """Build the FastVideo DiT forward kwargs from a framework PackedSequence."""
+    v = ps.vision
+    return dict(
+        text_ids=ps.text_ids,
+        text_indexes=ps.text_indexes,
+        position_ids=ps.position_ids,
+        sequence_length=int(ps.sequence_length),
+        split_lens=list(ps.split_lens),
+        attn_modes=list(ps.attn_modes),
+        vision_tokens=list(v.tokens),
+        vision_token_shapes=list(v.token_shapes),
+        vision_sequence_indexes=v.sequence_indexes,
+        vision_timesteps=v.timesteps,
+        vision_mse_loss_indexes=v.mse_loss_indexes,
+        vision_noisy_frame_indexes=list(v.noisy_frame_indexes),
+        fps_vision=None,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+class TestCosmos3DiTParity:
+
+    def _run_both(self, seed_model: int = 42, seed_data: int = 7):
+        vfm = _build_tiny_cosmos3(seed=seed_model)
+        dit = _build_tiny_fastvideo_dit()
+        _copy_weights(vfm, dit)
+        ps = _build_tiny_packed_seq(n_text=4, seed=seed_data)
+
+        with torch.no_grad():
+            fw_out = vfm(packed_seq=ps)
+            fv_out = dit(**_fastvideo_inputs_from_packed_seq(ps))
+        return fw_out, fv_out
+
+    def test_weight_map_is_complete(self):
+        """The framework->fastvideo map must cover EVERY FastVideo DiT parameter
+        that is exercised by the video path (i.e. all non-dormant params).
+
+        Dormant action/audio heads have no framework counterpart in this tiny
+        vision-only setup, so they are excluded from the copy; everything else
+        must be covered.
+        """
+        vfm = _build_tiny_cosmos3(seed=42)
+        dit = _build_tiny_fastvideo_dit()
+        mapped = set(_framework_to_fastvideo_state_dict(vfm, num_layers=dit.num_hidden_layers))
+        dit_params = set(n for n, _ in dit.named_parameters())
+        dormant = {
+            n
+            for n in dit_params
+            if n.startswith(("action_", "audio_"))
+        }
+        uncovered = dit_params - mapped - dormant
+        assert not uncovered, f"FastVideo DiT params not covered by weight map: {sorted(uncovered)}"
+
+    def test_preds_vision_parity(self):
+        fw_out, fv_out = self._run_both()
+        fw_pv = fw_out["preds_vision"][0]  # [1, C, T, H, W]
+        fv_pv = fv_out["preds_vision"][0]
+        assert fw_pv.shape == fv_pv.shape, f"shape mismatch: fw={fw_pv.shape} fv={fv_pv.shape}"
+        max_abs = (fw_pv - fv_pv).abs().max().item()
+        print(f"\n[preds_vision] max abs diff = {max_abs:.3e}")
+        torch.testing.assert_close(fv_pv, fw_pv, atol=1e-4, rtol=1e-3)
+
+    def test_last_hidden_state_parity(self):
+        fw_out, fv_out = self._run_both()
+        fw_lhs = fw_out["last_hidden_state"]  # [N, hidden]
+        fv_lhs = fv_out["last_hidden_state"]
+        assert fw_lhs.shape == fv_lhs.shape, f"shape mismatch: fw={fw_lhs.shape} fv={fv_lhs.shape}"
+        max_abs = (fw_lhs - fv_lhs).abs().max().item()
+        print(f"\n[last_hidden_state] max abs diff = {max_abs:.3e}")
+        torch.testing.assert_close(fv_lhs, fw_lhs, atol=1e-4, rtol=1e-3)
+
+    def test_parity_holds_across_seeds(self):
+        """Re-running with a different random init still matches (not a fluke)."""
+        fw_out, fv_out = self._run_both(seed_model=99, seed_data=13)
+        torch.testing.assert_close(fv_out["preds_vision"][0], fw_out["preds_vision"][0], atol=1e-4, rtol=1e-3)

tests/local_tests/cosmos3/test_cosmos3_dit_parity_mrope.py

@@ -0,0 +1,371 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Numerical-parity test: FastVideo Cosmos3 DiT vs ``Cosmos3VFMNetwork`` (mRoPE).
+
+Companion to ``test_cosmos3_dit_parity.py`` (which covers ``3d_rope``). This
+module exercises the rotary mode the REAL ``nvidia/Cosmos3-Nano`` checkpoint
+uses: ``position_embedding_type="unified_3d_mrope"`` with the real-checkpoint
+settings (``mrope_section=[24,20,20]``, ``mrope_interleaved=True``,
+``rope_theta=5e6``, ``unified_3d_mrope_reset_spatial_ids=True``,
+``temporal_modality_margin=15000``).
+
+Under unified 3D mRoPE there is NO additive latent position embedding
+(``latent_pos_embed is None``); all positional information rides on the
+per-token 3D (T, H, W) rotary embedding. The packed-sequence ``position_ids``
+are therefore shape ``[3, seq_len]``, built exactly like the framework data
+packer (``cosmos_framework.data.vfm.sequence_packing``):
+
+  * text tokens broadcast one monotone id across all three axes
+    (``get_3d_mrope_ids_text_tokens``),
+  * the temporal offset is bumped by ``temporal_modality_margin`` at the
+    text->vision boundary,
+  * vision tokens lay out a (T, H, W) grid with spatial ids reset per segment
+    (``get_3d_mrope_ids_vae_tokens`` with ``reset_spatial_indices=True``).
+
+Both models are built tiny from the SAME config, framework weights are copied
+into the FastVideo DiT (reusing the ``3d_rope`` test's weight map — the
+transformer key surface is identical across rotary modes), and BOTH forwards
+run on identical deterministic CPU / float32 inputs. The official model is the
+parity ORACLE (run on CPU via the SDPA monkey-patch in
+``test_cosmos3_reference_forward``).
+
+Run:
+    cd <worktree> && <fv-cosmos3 python> -m pytest \
+        tests/local_tests/cosmos3/test_cosmos3_dit_parity_mrope.py -q
+"""
+from __future__ import annotations
+
+import pytest
+import torch
+
+# The official framework provides the parity oracle.
+cosmos_framework = pytest.importorskip(
+    "cosmos_framework",
+    reason="cosmos_framework not installed; run in fv-cosmos3 env.",
+)
+
+# Reuse the reference harness's SDPA monkey-patch and the 3d_rope parity
+# test's weight-copy + input-builder helpers (key surface is rotary-agnostic).
+from .test_cosmos3_dit_parity import (  # noqa: E402
+    _copy_weights,
+    _fastvideo_inputs_from_packed_seq,
+    _framework_to_fastvideo_state_dict,
+)
+from .test_cosmos3_reference_forward import _apply_sdpa_patches  # noqa: E402
+
+pytestmark = [pytest.mark.local]
+
+# Ensure the CPU/float32 SDPA patches are installed (idempotent).
+_apply_sdpa_patches()
+
+# Real-checkpoint unified_3d_mrope settings (tiny model, real rope constants).
+_ROPE_THETA = 5_000_000.0
+_MROPE_SECTION = [24, 20, 20]
+_MROPE_INTERLEAVED = True
+_RESET_SPATIAL_IDS = True
+_TEMPORAL_MODALITY_MARGIN = 15_000
+_LATENT_PATCH_SIZE = 2
+_LATENT_CHANNEL = 16
+_TCF = 4  # temporal compression factor
+
+
+# ---------------------------------------------------------------------------
+# Tiny model builders (framework + FastVideo) with unified_3d_mrope.
+# ---------------------------------------------------------------------------
+_SOUND_DIM = 64
+_SOUND_LATENT_FPS = 25
+_ACTION_DIM = 64
+_NUM_EMBODIMENT_DOMAINS = 32
+
+
+def _build_tiny_cosmos3_mrope(seed: int = 42, num_layers: int = 2, sound_gen: bool = False,
+                              action_gen: bool = False):
+    """Tiny framework ``Cosmos3VFMNetwork`` with ``unified_3d_mrope``.
+
+    ``rope_theta`` / ``rope_scaling`` (carrying ``mrope_section`` +
+    ``mrope_interleaved``) are threaded through the materialized text config;
+    ``position_embedding_type="unified_3d_mrope"`` leaves ``latent_pos_embed``
+    as ``None`` so positions ride solely on the 3D rotary embedding.
+
+    ``sound_gen=True`` additionally builds the sound MoT heads (``sound2llm`` /
+    ``llm2sound`` / ``sound_modality_embed``) for the t2vs parity test.
+    """
+    from cosmos_framework.model.vfm.mot.cosmos3_vfm_network import (
+        Cosmos3VFMNetwork,
+        Cosmos3VFMNetworkConfig,
+    )
+    from cosmos_framework.model.vfm.mot.unified_mot import (
+        Qwen3MoTConfig,
+        Qwen3VLTextForCausalLM,
+    )
+    from cosmos_framework.model.vfm.vlm.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig
+
+    tiny_text_dict = dict(
+        model_type="qwen3_vl_text",
+        vocab_size=64,
+        hidden_size=16,
+        intermediate_size=32,
+        num_hidden_layers=num_layers,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        head_dim=8,
+        rms_norm_eps=1e-6,
+        attention_bias=False,
+        attention_dropout=0.0,
+        rope_theta=_ROPE_THETA,
+        rope_scaling={
+            "rope_type": "default",
+            "mrope_section": _MROPE_SECTION,
+            "mrope_interleaved": _MROPE_INTERLEAVED,
+        },
+        max_position_embeddings=262144,
+    )
+    mot_cfg = Qwen3MoTConfig(
+        config_dict=tiny_text_dict,
+        qk_norm_for_text=True,
+        qk_norm_for_diffusion=True,
+        include_visual=False,
+    )
+    tiny_vlm_cfg = Qwen3VLConfig(text_config=tiny_text_dict)
+    sound_kwargs = dict(
+        sound_gen=True,
+        sound_dim=_SOUND_DIM,
+        temporal_compression_factor_sound=1,
+        sound_latent_fps=_SOUND_LATENT_FPS,
+    ) if sound_gen else {}
+    action_kwargs = dict(
+        action_gen=True,
+        action_dim=_ACTION_DIM,
+        num_embodiment_domains=_NUM_EMBODIMENT_DOMAINS,
+    ) if action_gen else {}
+    vfm_cfg = Cosmos3VFMNetworkConfig(
+        vision_gen=True,
+        vlm_config=tiny_vlm_cfg,
+        latent_patch_size=_LATENT_PATCH_SIZE,
+        latent_downsample_factor=8,
+        latent_channel_size=_LATENT_CHANNEL,
+        position_embedding_type="unified_3d_mrope",
+        max_latent_h=16,
+        max_latent_w=16,
+        max_latent_t=8,
+        temporal_compression_factor_vision=_TCF,
+        **sound_kwargs,
+        **action_kwargs,
+    )
+    torch.manual_seed(seed)
+    lm = Qwen3VLTextForCausalLM(config=mot_cfg)
+    vfm = Cosmos3VFMNetwork(language_model=lm, config=vfm_cfg)
+    # inv_freq is a non-persistent buffer; init it on CPU (mirrors from_pretrained).
+    vfm.language_model.model.rotary_emb.init_weights(buffer_device=None)
+    vfm.eval()
+    return vfm
+
+
+def _build_tiny_fastvideo_dit_mrope(num_layers: int = 2):
+    from fastvideo.configs.models.dits.cosmos3 import (
+        Cosmos3ArchConfig,
+        Cosmos3VideoConfig,
+    )
+    from fastvideo.models.dits.cosmos3 import Cosmos3VFMTransformer
+
+    arch = Cosmos3ArchConfig(
+        hidden_size=16,
+        num_hidden_layers=num_layers,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        head_dim=8,
+        intermediate_size=32,
+        vocab_size=64,
+        rms_norm_eps=1e-6,
+        attention_bias=False,
+        latent_patch_size=_LATENT_PATCH_SIZE,
+        latent_channel=_LATENT_CHANNEL,
+        rope_theta=_ROPE_THETA,
+        mrope_section=_MROPE_SECTION,
+        mrope_interleaved=_MROPE_INTERLEAVED,
+        unified_3d_mrope_reset_spatial_ids=_RESET_SPATIAL_IDS,
+        temporal_modality_margin=_TEMPORAL_MODALITY_MARGIN,
+        position_embedding_type="unified_3d_mrope",
+        base_fps=24.0,
+        temporal_compression_factor=_TCF,
+        enable_fps_modulation=False,
+        # Dormant heads present in the checkpoint surface (constructed for
+        # strict-load parity; not exercised by this video-path forward).
+        action_gen=True,
+        action_dim=64,
+        max_action_dim=64,
+        num_embodiment_domains=32,
+        sound_gen=True,
+        sound_dim=64,
+    )
+    cfg = Cosmos3VideoConfig(arch_config=arch)
+    model = Cosmos3VFMTransformer(cfg, hf_config={})
+    return model.to(torch.float32).eval()
+
+
+# ---------------------------------------------------------------------------
+# [3, seq_len] mRoPE position-id builder (mirrors the framework data packer).
+# ---------------------------------------------------------------------------
+def _build_mrope_position_ids(n_text: int, grid_t: int, patch_h: int, patch_w: int) -> torch.Tensor:
+    """Build ``[3, seq_len]`` (T, H, W) mRoPE ids for one text+vision sample.
+
+    Reproduces ``pack_input_sequence`` for a single causal-text + full-vision
+    sample: monotone text ids on all axes, ``+temporal_modality_margin`` at the
+    text->vision boundary, then a reset-spatial (T, H, W) vision grid.
+    """
+    from cosmos_framework.data.vfm.sequence_packing import (
+        get_3d_mrope_ids_text_tokens,
+        get_3d_mrope_ids_vae_tokens,
+    )
+
+    offset: int | float = 0
+    text_ids, offset = get_3d_mrope_ids_text_tokens(num_tokens=n_text, temporal_offset=offset)
+    # End of text modality: add the boundary margin before vision.
+    offset += _TEMPORAL_MODALITY_MARGIN
+    vision_ids, offset = get_3d_mrope_ids_vae_tokens(
+        grid_t=grid_t,
+        grid_h=patch_h,
+        grid_w=patch_w,
+        temporal_offset=offset,
+        reset_spatial_indices=_RESET_SPATIAL_IDS,
+        fps=None,  # integer positions (enable_fps_modulation=False)
+        temporal_compression_factor=_TCF,
+    )
+    return torch.cat([text_ids, vision_ids], dim=1)  # [3, seq_len]
+
+
+def _build_tiny_packed_seq_mrope(
+    *,
+    n_text: int = 6,
+    grid_t: int = 2,
+    latent_h: int = 4,
+    latent_w: int = 4,
+    seed: int = 7,
+):
+    """Minimal PackedSequence with ``[3, seq]`` mRoPE position ids.
+
+    Vision latent ``[C, grid_t, latent_h, latent_w]`` patchifies (patch=2) to a
+    ``(grid_t, latent_h/2, latent_w/2)`` token grid; all frames are noisy.
+    """
+    from cosmos_framework.data.vfm.sequence_packing import ModalityData, PackedSequence
+
+    patch_h = latent_h // _LATENT_PATCH_SIZE
+    patch_w = latent_w // _LATENT_PATCH_SIZE
+    n_vision = grid_t * patch_h * patch_w
+    total_len = n_text + n_vision
+
+    torch.manual_seed(seed)
+    vision_tensor = torch.randn(_LATENT_CHANNEL, grid_t, latent_h, latent_w)
+    text_ids = torch.randint(0, 64, (n_text,))
+    position_ids = _build_mrope_position_ids(n_text, grid_t, patch_h, patch_w)  # [3, total_len]
+
+    noisy_frame_indexes = torch.arange(grid_t, dtype=torch.long)  # all frames noisy
+    vision_mod = ModalityData(
+        sequence_indexes=torch.arange(n_text, total_len, dtype=torch.long),
+        timesteps=torch.full((n_vision,), 500.0),
+        mse_loss_indexes=torch.arange(n_text, total_len, dtype=torch.long),
+        token_shapes=[(grid_t, patch_h, patch_w)],
+        tokens=[vision_tensor],
+        condition_mask=[torch.zeros(grid_t, dtype=torch.long)],  # 0 = noisy
+        noisy_frame_indexes=[noisy_frame_indexes],
+    )
+    packed_seq = PackedSequence(
+        sample_lens=[total_len],
+        split_lens=[n_text, n_vision],
+        attn_modes=["causal", "full"],
+        is_image_batch=(grid_t == 1),
+        sequence_length=total_len,
+        text_ids=text_ids,
+        text_indexes=torch.arange(n_text, dtype=torch.long),
+        position_ids=position_ids,
+        vision=vision_mod,
+    )
+    return packed_seq
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+# (grid_t, latent_h, latent_w): a single image, a small video, and a taller
+# video, to exercise the spatial mRoPE overwrite + gen<->gen full attention.
+_GRIDS = [
+    pytest.param(1, 8, 8, id="image_1x4x4"),
+    pytest.param(2, 4, 4, id="video_2x2x2"),
+    pytest.param(3, 8, 4, id="video_3x4x2"),
+]
+
+
+class TestCosmos3DiTParityMRoPE:
+
+    def _run_both(
+        self,
+        *,
+        grid_t: int,
+        latent_h: int,
+        latent_w: int,
+        seed_model: int = 42,
+        seed_data: int = 7,
+        num_layers: int = 2,
+    ):
+        vfm = _build_tiny_cosmos3_mrope(seed=seed_model, num_layers=num_layers)
+        dit = _build_tiny_fastvideo_dit_mrope(num_layers=num_layers)
+        _copy_weights(vfm, dit)
+        ps = _build_tiny_packed_seq_mrope(
+            n_text=6, grid_t=grid_t, latent_h=latent_h, latent_w=latent_w, seed=seed_data
+        )
+        with torch.no_grad():
+            fw_out = vfm(packed_seq=ps)
+            fv_out = dit(**_fastvideo_inputs_from_packed_seq(ps))
+        return fw_out, fv_out
+
+    def test_position_ids_are_3xN_mrope(self):
+        """The packed mRoPE ids are ``[3, seq_len]`` with the text->vision margin."""
+        ps = _build_tiny_packed_seq_mrope(n_text=6, grid_t=2, latent_h=4, latent_w=4)
+        pos = ps.position_ids
+        assert pos.ndim == 2 and pos.shape[0] == 3, f"expected [3, N], got {tuple(pos.shape)}"
+        assert pos.shape[1] == int(ps.sequence_length)
+        # Text axis is monotone 0..5 on all 3 rows; vision temporal jumps by the margin.
+        assert pos[0, :6].tolist() == [0, 1, 2, 3, 4, 5]
+        assert pos[1, :6].tolist() == [0, 1, 2, 3, 4, 5]
+        assert pos[2, :6].tolist() == [0, 1, 2, 3, 4, 5]
+        # First vision token temporal id == last_text_id (5) + margin + 1.
+        assert pos[0, 6].item() == 5 + _TEMPORAL_MODALITY_MARGIN + 1
+        # Reset spatial: first vision token H/W ids are 0.
+        assert pos[1, 6].item() == 0 and pos[2, 6].item() == 0
+
+    def test_no_additive_latent_pos_embed(self):
+        """unified_3d_mrope must NOT build an additive latent position embedding."""
+        dit = _build_tiny_fastvideo_dit_mrope()
+        assert dit.position_embedding_type == "unified_3d_mrope"
+        assert dit.latent_pos_embed is None
+
+    @pytest.mark.parametrize(("grid_t", "latent_h", "latent_w"), _GRIDS)
+    def test_preds_vision_parity(self, grid_t, latent_h, latent_w):
+        fw_out, fv_out = self._run_both(grid_t=grid_t, latent_h=latent_h, latent_w=latent_w)
+        fw_pv = fw_out["preds_vision"][0]  # [1, C, T, H, W]
+        fv_pv = fv_out["preds_vision"][0]
+        assert fw_pv.shape == fv_pv.shape, f"shape mismatch: fw={fw_pv.shape} fv={fv_pv.shape}"
+        max_abs = (fw_pv - fv_pv).abs().max().item()
+        print(f"\n[preds_vision mrope {grid_t}x{latent_h}x{latent_w}] max abs diff = {max_abs:.3e}")
+        torch.testing.assert_close(fv_pv, fw_pv, atol=1e-4, rtol=1e-3)
+
+    @pytest.mark.parametrize(("grid_t", "latent_h", "latent_w"), _GRIDS)
+    def test_last_hidden_state_parity(self, grid_t, latent_h, latent_w):
+        fw_out, fv_out = self._run_both(grid_t=grid_t, latent_h=latent_h, latent_w=latent_w)
+        fw_lhs = fw_out["last_hidden_state"]  # [N, hidden]
+        fv_lhs = fv_out["last_hidden_state"]
+        assert fw_lhs.shape == fv_lhs.shape, f"shape mismatch: fw={fw_lhs.shape} fv={fv_lhs.shape}"
+        max_abs = (fw_lhs - fv_lhs).abs().max().item()
+        print(f"\n[last_hidden_state mrope {grid_t}x{latent_h}x{latent_w}] max abs diff = {max_abs:.3e}")
+        torch.testing.assert_close(fv_lhs, fw_lhs, atol=1e-4, rtol=1e-3)
+
+    def test_parity_holds_across_seeds(self):
+        """A different random init still matches bit-for-bit (not a fluke)."""
+        fw_out, fv_out = self._run_both(
+            grid_t=2, latent_h=4, latent_w=4, seed_model=99, seed_data=13
+        )
+        torch.testing.assert_close(
+            fv_out["preds_vision"][0], fw_out["preds_vision"][0], atol=1e-4, rtol=1e-3
+        )
+        torch.testing.assert_close(
+            fv_out["last_hidden_state"], fw_out["last_hidden_state"], atol=1e-4, rtol=1e-3
+        )