From d873c85c804dc922e6941a0934df528c41264bc9 Mon Sep 17 00:00:00 2001
From: Raghav <ragg04@gmail.com>
Date: Thu, 27 Aug 2026 11:19:43 -0700
Subject: [PATCH 01/16] [feat] add Cosmos 2.5 distilled sampler scaffold

---
 fastvideo/models/registry.py                  |   5 +
 .../scheduling_cosmos25_distilled.py          | 207 ++++++++++++++++++
 .../test_cosmos25_distilled_scheduler.py      |  97 ++++++++
 .../cosmos25_distilled_to_diffusers.py        | 203 +++++++++++++++++
 tests/local_tests/cosmos25/PORT_STATUS.md     |  74 +++++++
 tests/local_tests/cosmos25/README.md          |  69 ++++++
 tests/local_tests/cosmos25/_reference.py      |  33 +++
 .../test_cosmos25_distilled_conversion.py     | 112 ++++++++++
 ...est_cosmos25_distilled_scheduler_parity.py |  60 +++++
 9 files changed, 860 insertions(+)
 create mode 100644 fastvideo/models/schedulers/scheduling_cosmos25_distilled.py
 create mode 100644 fastvideo/tests/schedulers/test_cosmos25_distilled_scheduler.py
 create mode 100644 scripts/checkpoint_conversion/cosmos25_distilled_to_diffusers.py
 create mode 100644 tests/local_tests/cosmos25/PORT_STATUS.md
 create mode 100644 tests/local_tests/cosmos25/README.md
 create mode 100644 tests/local_tests/cosmos25/_reference.py
 create mode 100644 tests/local_tests/cosmos25/test_cosmos25_distilled_conversion.py
 create mode 100644 tests/local_tests/cosmos25/test_cosmos25_distilled_scheduler_parity.py

diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py
index 4f9939a62d..ad8ad88cb5 100644
--- a/fastvideo/models/registry.py
+++ b/fastvideo/models/registry.py
@@ -141,6 +141,11 @@
     "SelfForcingFlowMatchScheduler":
     ("schedulers", "scheduling_self_forcing_flow_match", "SelfForcingFlowMatchScheduler"),
     "RCMScheduler": ("schedulers", "scheduling_rcm", "RCMScheduler"),
+    "Cosmos25DistilledScheduler": (
+        "schedulers",
+        "scheduling_cosmos25_distilled",
+        "Cosmos25DistilledScheduler",
+    ),
 }
 
 _UPSAMPLERS = {
diff --git a/fastvideo/models/schedulers/scheduling_cosmos25_distilled.py b/fastvideo/models/schedulers/scheduling_cosmos25_distilled.py
new file mode 100644
index 0000000000..0e86ce50a1
--- /dev/null
+++ b/fastvideo/models/schedulers/scheduling_cosmos25_distilled.py
@@ -0,0 +1,207 @@
+# SPDX-License-Identifier: Apache-2.0
+"""TrigFlow sampler for the distilled Cosmos Predict2.5 checkpoint."""
+
+import math
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import BaseOutput
+
+from fastvideo.models.schedulers.base import BaseScheduler
+
+
+@dataclass
+class Cosmos25DistilledSchedulerOutput(BaseOutput):
+    """Output of one distilled Cosmos Predict2.5 sampling step."""
+
+    prev_sample: torch.Tensor
+    pred_original_sample: torch.Tensor
+
+
+class Cosmos25DistilledScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
+    """Official four-step TrigFlow/x0 sampler for Cosmos Predict2.5.
+
+    The distilled student predicts the rectified-flow network output. The
+    scheduler applies the student's TrigFlow preconditioning, reconstructs x0,
+    and re-noises it with the *initial* noise for the next student evaluation.
+    This fixed-noise update is intentionally different from stochastic rCM
+    sampling.
+    """
+
+    _compatibles: list[Any] = []
+    order = 1
+    _OFFICIAL_SAMPLING_TIMES = (
+        math.pi / 2,
+        math.atan(15),
+        math.atan(5),
+        math.atan(5 / 3),
+    )
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        sigma_data: float = 1.0,
+    ) -> None:
+        if sigma_data <= 0:
+            raise ValueError(f"sigma_data must be positive, got {sigma_data}")
+
+        self.num_train_timesteps = num_train_timesteps
+        self.sigma_data = float(sigma_data)
+        self.timesteps = torch.empty(0, dtype=torch.float64)
+        self.trigflow_timesteps = torch.empty(0, dtype=torch.float64)
+        self.sigmas = torch.empty(0, dtype=torch.float64)
+        self._step_index: int | None = None
+        self._begin_index: int | None = None
+        self._initial_noise: torch.Tensor | None = None
+        BaseScheduler.__init__(self)
+
+    @property
+    def init_noise_sigma(self) -> float:
+        return 1.0
+
+    @property
+    def step_index(self) -> int | None:
+        return self._step_index
+
+    @property
+    def begin_index(self) -> int | None:
+        return self._begin_index
+
+    def set_begin_index(self, begin_index: int = 0) -> None:
+        self._begin_index = begin_index
+
+    def set_shift(self, shift: float) -> None:
+        """The checkpoint's fixed distilled schedule does not use flow shift."""
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int = 4,
+        device: str | torch.device | None = None,
+    ) -> None:
+        if not 1 <= num_inference_steps <= len(self._OFFICIAL_SAMPLING_TIMES):
+            raise ValueError(f"Cosmos Predict2.5 distilled sampling supports 1 to 4 steps; got {num_inference_steps}")
+
+        trigflow_timesteps = torch.tensor(
+            self._OFFICIAL_SAMPLING_TIMES[:num_inference_steps],
+            dtype=torch.float64,
+            device=device,
+        )
+        _, _, _, model_timesteps = self._scalings(trigflow_timesteps)
+
+        self.num_inference_steps = num_inference_steps
+        self.trigflow_timesteps = trigflow_timesteps
+        # The Cosmos DiT consumes c_noise, not the TrigFlow angle.
+        self.timesteps = model_timesteps
+        self.sigmas = torch.tan(trigflow_timesteps) * self.sigma_data
+        self._step_index = None
+        self._initial_noise = None
+
+    def _scalings(
+        self,
+        trigflow_timestep: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        dtype = trigflow_timestep.dtype
+        timestep = trigflow_timestep.to(torch.float64)
+        denominator = torch.cos(timestep) + self.sigma_data * torch.sin(timestep)
+        c_skip = self.sigma_data / denominator
+        c_out = -self.sigma_data * torch.sin(timestep) / denominator
+        c_in = self.sigma_data / denominator
+        c_noise = self.sigma_data * torch.sin(timestep) / denominator
+        return (
+            c_skip.to(dtype),
+            c_out.to(dtype),
+            c_in.to(dtype),
+            c_noise.to(dtype),
+        )
+
+    def _init_step_index(self) -> None:
+        self._step_index = self._begin_index if self._begin_index is not None else 0
+
+    def _current_scalings(
+        self,
+        sample: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        if self._step_index is None:
+            self._init_step_index()
+        assert self._step_index is not None
+        if self._step_index >= len(self.trigflow_timesteps):
+            raise IndexError("All configured Cosmos Predict2.5 distilled steps have already run")
+        timestep = self.trigflow_timesteps[self._step_index].to(device=sample.device)
+        return self._scalings(timestep)
+
+    def scale_model_input(
+        self,
+        sample: torch.Tensor,
+        timestep: int | torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        del timestep
+        _, _, c_in, _ = self._current_scalings(sample)
+        return sample * c_in
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: int | torch.Tensor,
+        sample: torch.Tensor,
+        generator: torch.Generator | None = None,
+        return_dict: bool = True,
+    ) -> Cosmos25DistilledSchedulerOutput | tuple[torch.Tensor, ...]:
+        del timestep, generator
+        c_skip, c_out, _, _ = self._current_scalings(sample)
+        assert self._step_index is not None
+
+        if self._initial_noise is None:
+            # Official inference keeps init_noise in FP32 while carrying the
+            # evolving sample in FP64.
+            self._initial_noise = sample.detach().to(torch.float32).clone()
+
+        pred_original_sample = c_skip * sample + c_out * model_output
+        next_index = self._step_index + 1
+        if next_index < len(self.trigflow_timesteps):
+            next_timestep = self.trigflow_timesteps[next_index].to(
+                device=sample.device,
+                dtype=torch.float64,
+            )
+            prev_sample = (
+                torch.cos(next_timestep) * pred_original_sample / self.sigma_data
+                + torch.sin(next_timestep) * self._initial_noise
+            )
+        else:
+            prev_sample = pred_original_sample
+
+        self._step_index = next_index
+        if not return_dict:
+            return (prev_sample,)
+        return Cosmos25DistilledSchedulerOutput(
+            prev_sample=prev_sample,
+            pred_original_sample=pred_original_sample,
+        )
+
+    def scale_noise(
+        self,
+        sample: torch.Tensor,
+        timestep: torch.Tensor | None = None,
+        noise: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        del sample, timestep
+        if noise is None:
+            raise ValueError("noise must be provided")
+        return noise
+
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.Tensor,
+    ) -> torch.Tensor:
+        raise NotImplementedError("Cosmos25DistilledScheduler is an inference-only x0 sampler")
+
+    def __len__(self) -> int:
+        return self.config.num_train_timesteps
+
+
+EntryClass = Cosmos25DistilledScheduler
diff --git a/fastvideo/tests/schedulers/test_cosmos25_distilled_scheduler.py b/fastvideo/tests/schedulers/test_cosmos25_distilled_scheduler.py
new file mode 100644
index 0000000000..17df71f9ea
--- /dev/null
+++ b/fastvideo/tests/schedulers/test_cosmos25_distilled_scheduler.py
@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+
+import pytest
+import torch
+from torch.testing import assert_close
+
+from fastvideo.models.schedulers.scheduling_cosmos25_distilled import (
+    Cosmos25DistilledScheduler,
+)
+from fastvideo.models.registry import ModelRegistry
+
+
+def test_registry_resolves_scheduler() -> None:
+    scheduler_class, architecture = ModelRegistry.resolve_model_cls("Cosmos25DistilledScheduler")
+    assert scheduler_class is Cosmos25DistilledScheduler
+    assert architecture == "Cosmos25DistilledScheduler"
+
+
+def test_official_schedule_and_model_timesteps() -> None:
+    scheduler = Cosmos25DistilledScheduler()
+    scheduler.set_timesteps(4)
+
+    expected_angles = torch.tensor(
+        [math.pi / 2, math.atan(15), math.atan(5), math.atan(5 / 3)],
+        dtype=torch.float64,
+    )
+    expected_model_timesteps = torch.tensor(
+        [1.0, 15 / 16, 5 / 6, 5 / 8],
+        dtype=torch.float64,
+    )
+    assert_close(scheduler.trigflow_timesteps, expected_angles, rtol=0, atol=0)
+    assert_close(scheduler.timesteps, expected_model_timesteps, rtol=1e-15, atol=1e-15)
+
+
+def test_rollout_matches_official_fixed_noise_equations() -> None:
+    scheduler = Cosmos25DistilledScheduler()
+    scheduler.set_timesteps(4)
+    initial_noise = torch.tensor([[[1.5, -0.25], [0.5, 2.0]]], dtype=torch.float32)
+    sample = initial_noise.to(torch.float64)
+    expected = sample.clone()
+
+    for index, model_timestep in enumerate(scheduler.timesteps):
+        angle = scheduler.trigflow_timesteps[index]
+        denominator = torch.cos(angle) + torch.sin(angle)
+        c_skip = 1 / denominator
+        c_out = -torch.sin(angle) / denominator
+        c_in = 1 / denominator
+
+        assert_close(scheduler.scale_model_input(sample, model_timestep), expected * c_in)
+        model_output = torch.full_like(sample, 0.125 * (index + 1))
+        expected_x0 = c_skip * expected + c_out * model_output
+        if index + 1 < len(scheduler.trigflow_timesteps):
+            next_angle = scheduler.trigflow_timesteps[index + 1]
+            expected = torch.cos(next_angle) * expected_x0 + torch.sin(next_angle) * initial_noise
+        else:
+            expected = expected_x0
+
+        result = scheduler.step(model_output, model_timestep, sample)
+        assert_close(result.pred_original_sample, expected_x0)
+        assert_close(result.prev_sample, expected)
+        sample = result.prev_sample
+
+
+def test_reuses_initial_noise_instead_of_generator_noise() -> None:
+    initial_noise = torch.tensor([1.0, -2.0], dtype=torch.float32)
+    results = []
+    for seed in (1, 999):
+        scheduler = Cosmos25DistilledScheduler()
+        scheduler.set_timesteps(2)
+        sample = initial_noise.to(torch.float64)
+        generator = torch.Generator().manual_seed(seed)
+        for timestep in scheduler.timesteps:
+            model_output = torch.full_like(sample, 0.25)
+            sample = scheduler.step(model_output, timestep, sample, generator=generator).prev_sample
+        results.append(sample)
+    assert_close(results[0], results[1], rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("num_steps", [0, 5])
+def test_rejects_schedule_outside_released_checkpoint_range(num_steps: int) -> None:
+    scheduler = Cosmos25DistilledScheduler()
+    with pytest.raises(ValueError, match="1 to 4 steps"):
+        scheduler.set_timesteps(num_steps)
+
+
+def test_set_timesteps_resets_sampler_state() -> None:
+    scheduler = Cosmos25DistilledScheduler()
+    scheduler.set_timesteps(1)
+    sample = torch.ones(2, dtype=torch.float64)
+    scheduler.step(torch.zeros_like(sample), scheduler.timesteps[0], sample)
+    assert scheduler.step_index == 1
+
+    scheduler.set_timesteps(1)
+    assert scheduler.step_index is None
+    assert scheduler._initial_noise is None
diff --git a/scripts/checkpoint_conversion/cosmos25_distilled_to_diffusers.py b/scripts/checkpoint_conversion/cosmos25_distilled_to_diffusers.py
new file mode 100644
index 0000000000..b24e5a867d
--- /dev/null
+++ b/scripts/checkpoint_conversion/cosmos25_distilled_to_diffusers.py
@@ -0,0 +1,203 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+"""Package NVIDIA's Cosmos Predict2.5 2B distilled student for FastVideo.
+
+The released checkpoint is a native PyTorch state dict. FastVideo's existing
+Cosmos25 parameter mapping consumes its ``net.*`` names directly, so this
+converter deliberately does not rename tensors. It isolates the student from
+teacher/critic/training state and combines it with the non-transformer assets
+from an existing FastVideo-loadable Cosmos Predict2.5 package.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import shutil
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any, cast
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+STATE_DICT_KEYS = ("state_dict", "model", "ema", "ema_model", "module")
+REQUIRED_BASE_PATHS = (
+    "model_index.json",
+    "transformer/config.json",
+    "vae",
+    "text_encoder",
+    "tokenizer",
+)
+TRANSFORMER_FILENAME = "diffusion_pytorch_model.safetensors"
+
+
+class ConversionError(RuntimeError):
+    pass
+
+
+def _read_json(path: Path) -> dict[str, Any]:
+    with path.open("r", encoding="utf-8") as handle:
+        return json.load(handle)
+
+
+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 handle:
+        json.dump(payload, handle, indent=2)
+        handle.write("\n")
+
+
+def _has_student_tensors(value: object) -> bool:
+    return isinstance(value, Mapping) and any(
+        isinstance(key, str) and key.startswith("net.") and torch.is_tensor(tensor) for key, tensor in value.items()
+    )
+
+
+def extract_student_state_dict(checkpoint: object) -> dict[str, torch.Tensor]:
+    """Find the official state dict and retain only the distilled student."""
+    state_dict: Mapping[Any, Any] | None = (
+        cast(Mapping[Any, Any], checkpoint) if _has_student_tensors(checkpoint) else None
+    )
+    if state_dict is None and isinstance(checkpoint, Mapping):
+        for key in STATE_DICT_KEYS:
+            candidate = checkpoint.get(key)
+            if _has_student_tensors(candidate):
+                state_dict = cast(Mapping[Any, Any], candidate)
+                break
+
+    if state_dict is None:
+        raise ConversionError("Could not find a tensor state dict containing native 'net.*' student keys")
+
+    student = {
+        key: tensor.detach().to(device="cpu", dtype=torch.bfloat16).contiguous()
+        for key, tensor in state_dict.items()
+        if isinstance(key, str) and key.startswith("net.") and torch.is_tensor(tensor)
+    }
+    if not student:
+        raise ConversionError("The resolved checkpoint contains no 'net.*' tensors")
+    return student
+
+
+def _validate_base_model(base_model: Path) -> None:
+    missing = [relative for relative in REQUIRED_BASE_PATHS if not (base_model / relative).exists()]
+    if missing:
+        raise ConversionError(f"Base Cosmos25 package is missing required paths: {missing}")
+
+
+def _prepare_output(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 _remove_base_transformer_weights(transformer_dir: Path) -> None:
+    for pattern in ("*.safetensors", "*.safetensors.index.json", "*.bin", "*.pt"):
+        for path in transformer_dir.glob(pattern):
+            path.unlink()
+
+
+def _write_distilled_metadata(dst: Path) -> None:
+    model_index_path = dst / "model_index.json"
+    model_index = _read_json(model_index_path)
+    model_index["is_distilled"] = True
+    model_index["scheduler"] = ["diffusers", "Cosmos25DistilledScheduler"]
+    _write_json(model_index_path, model_index)
+
+    scheduler_dir = dst / "scheduler"
+    if scheduler_dir.exists():
+        shutil.rmtree(scheduler_dir)
+    _write_json(
+        scheduler_dir / "scheduler_config.json",
+        {
+            "_class_name": "Cosmos25DistilledScheduler",
+            "_diffusers_version": "0.37.0.dev0",
+            "num_train_timesteps": 1000,
+            "sigma_data": 1.0,
+        },
+    )
+
+
+def _verify_output(dst: Path, expected_keys: set[str]) -> None:
+    model_index = _read_json(dst / "model_index.json")
+    if model_index.get("scheduler") != ["diffusers", "Cosmos25DistilledScheduler"]:
+        raise ConversionError("model_index.json does not select Cosmos25DistilledScheduler")
+
+    scheduler_config = _read_json(dst / "scheduler/scheduler_config.json")
+    if scheduler_config.get("_class_name") != "Cosmos25DistilledScheduler":
+        raise ConversionError("scheduler_config.json has the wrong _class_name")
+
+    weights_path = dst / "transformer" / TRANSFORMER_FILENAME
+    with safe_open(str(weights_path), framework="pt", device="cpu") as handle:
+        actual_keys = set(handle.keys())
+    if actual_keys != expected_keys:
+        missing = sorted(expected_keys - actual_keys)
+        unexpected = sorted(actual_keys - expected_keys)
+        raise ConversionError(
+            f"Converted transformer key mismatch: missing={missing[:10]}, unexpected={unexpected[:10]}"
+        )
+
+
+def convert_checkpoint(
+    src_checkpoint: Path,
+    base_model: Path,
+    dst: Path,
+    *,
+    overwrite: bool = False,
+) -> dict[str, int]:
+    src_checkpoint = src_checkpoint.expanduser().resolve()
+    base_model = base_model.expanduser().resolve()
+    dst = dst.expanduser().resolve()
+    if not src_checkpoint.is_file():
+        raise FileNotFoundError(f"Distilled checkpoint not found: {src_checkpoint}")
+    _validate_base_model(base_model)
+    _prepare_output(dst, overwrite)
+
+    checkpoint = torch.load(src_checkpoint, map_location="cpu", weights_only=True)
+    student = extract_student_state_dict(checkpoint)
+
+    shutil.copytree(base_model, dst, dirs_exist_ok=True, symlinks=False)
+    transformer_dir = dst / "transformer"
+    _remove_base_transformer_weights(transformer_dir)
+    save_file(
+        student,
+        str(transformer_dir / TRANSFORMER_FILENAME),
+        metadata={"format": "pt", "model_type": "cosmos25_distilled_student"},
+    )
+    _write_distilled_metadata(dst)
+    _verify_output(dst, set(student))
+
+    return {
+        "student_tensors": len(student),
+        "student_parameters": sum(tensor.numel() for tensor in student.values()),
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--src-checkpoint", required=True, type=Path)
+    parser.add_argument(
+        "--base-model",
+        required=True,
+        type=Path,
+        help="Existing FastVideo-loadable Cosmos25 package supplying configs, encoder, tokenizer, and VAE",
+    )
+    parser.add_argument("--dst", required=True, type=Path)
+    parser.add_argument("--overwrite", action="store_true")
+    args = parser.parse_args()
+
+    report = convert_checkpoint(
+        args.src_checkpoint,
+        args.base_model,
+        args.dst,
+        overwrite=args.overwrite,
+    )
+    print(f"Converted {report['student_tensors']} student tensors ({report['student_parameters']:,} parameters)")
+    print(f"Output: {args.dst.expanduser().resolve()}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/local_tests/cosmos25/PORT_STATUS.md b/tests/local_tests/cosmos25/PORT_STATUS.md
new file mode 100644
index 0000000000..d7b1643e25
--- /dev/null
+++ b/tests/local_tests/cosmos25/PORT_STATUS.md
@@ -0,0 +1,74 @@
+# Cosmos Predict2.5 Distilled Port Status
+
+## Summary
+- model_family: cosmos25_distilled
+- workload_types: T2W (released support); V2W/rolling experimental and deferred
+- official_ref: NVIDIA/Cosmos-Predict2.5@a2c298b0a3df3778b973fe65e9e58877b292d8a7
+- official_ref_dir: `${COSMOS25_OFFICIAL_REF_DIR:-$PWD/cosmos-predict2.5}`
+- hf_weights_path: `nvidia/Cosmos-Predict2.5-2B/base/distilled`
+- local_weights_dir: not created
+- source_layout: official monolithic checkpoint; FastVideo conversion/loading path pending
+- local_tests_readme: `tests/local_tests/cosmos25/README.md`
+
+## Current Phase
+- phase: component parity
+- status: in_progress
+- owner: parity
+- last_updated: 2026-08-27
+
+## Component Matrix
+| Component | Type | Reuse/Port | Official Definition | Official Instantiation | FastVideo Target | Prototype | Conversion | Parity | Open Issues |
+|---|---|---|---|---|---|---|---|---|---|
+| TrigFlow sampler | scheduler | port | `modules/denoiser_scaling.py`; `distill/models/video2world_model_distill_dmd2.py` | `generate_samples_from_batch` | `Cosmos25DistilledScheduler` | complete | n/a | non-skip pass | none |
+| student DiT | transformer | reuse Cosmos25 architecture with distilled weights | `MinimalV1LVGDiT` through distillation model | `get_x0_fn_from_batch` / `denoise_edm` | `Cosmos25Transformer3DModel` | existing | pending | pending real-weight forward | I001 |
+| Reason1 encoder | text encoder | reuse | Predict2.5 Video2World config | distilled inference CLI | existing Cosmos25 encoder | existing | pending packaged layout | pending production-loader parity | I002 |
+| tokenizer VAE | VAE | reuse | Predict2.5 tokenizer | distilled inference CLI | existing Cosmos25 VAE | existing | pending packaged layout | pending production-loader parity | I002 |
+| T2W pipeline | pipeline | extend after components pass | `generate_samples_from_batch` | distilled inference CLI | Cosmos2_5 staged pipeline | not started | blocked by component gates | not started | I003 |
+
+## Conversion State
+- conversion_script: `scripts/checkpoint_conversion/cosmos25_distilled_to_diffusers.py`
+- converted_weights_dir: not created
+- source_layout: official `base/distilled` checkpoint
+- strict_load_status: not run
+- passthrough_components: Reason1 encoder and tokenizer VAE are expected to reuse the existing Cosmos25 layout
+- retry_history: synthetic direct/nested checkpoint extraction and output-layout contracts pass; released checkpoint not run
+
+## Parity Commands
+| Scope | Command | Last Result | Notes |
+|---|---|---|---|
+| scheduler unit | `pytest fastvideo/tests/schedulers/test_cosmos25_distilled_scheduler.py -q` | 7 passed | CPU-only; includes registry resolution; 2026-08-27 |
+| official scheduler | `COSMOS25_OFFICIAL_REF_DIR=/path/to/Cosmos-Predict2.5 pytest tests/local_tests/cosmos25/test_cosmos25_distilled_scheduler_parity.py -v -s` | 2 passed, non-skip | CPU-only; pinned source; 2026-08-27 |
+| conversion contracts | `pytest tests/local_tests/cosmos25/test_cosmos25_distilled_conversion.py -q` | 7 passed | Synthetic checkpoints/layout; 2026-08-27 |
+| student DiT | not yet created | not run | requires CUDA and released weights |
+| pipeline | not yet created | not run | forbidden until component parity passes |
+
+## Open Questions
+| ID | Question | Owner | Needed By Phase | Status | Resolution |
+|---|---|---|---|---|---|
+| Q001 | What packaged Diffusers-style model ID should carry the distilled scheduler and converted student weights? | conversion | conversion | open | pending |
+| Q002 | Does experimental distilled V2W retain acceptable quality after official T2W parity? | pipeline | post-parity experiment | open | intentionally outside initial support claim |
+
+## Issues And Blockers
+| ID | Phase | Component | Severity | Issue | Evidence | Owner | Status | Resolution |
+|---|---|---|---|---|---|---|---|---|
+| I001 | parity | student DiT | high | No non-skip real-weight distilled forward comparison yet | Existing Spark runs used distilled weights with the base UniPC inference path | parity | open | pending CUDA run |
+| I002 | conversion | packaged model | high | Released official checkpoint is not yet isolated in a FastVideo-loadable component layout | No `local_weights_dir` or strict-load record | conversion | open | pending |
+| I003 | pipeline | T2W | high | Pipeline wiring is gated on component parity | add-model pipeline contract | pipeline | open | pending I001 and I002 |
+
+## Escape Hatches
+| ID | Phase | Decision Type | Question | Recommended Option | Status | Resolution |
+|---|---|---|---|---|---|---|
+
+## Decisions
+| Date | Decision | Rationale | Impact |
+|---|---|---|---|
+| 2026-08-27 | Implement a Cosmos-specific sampler instead of reusing RCM | RCM uses different times and fresh per-step noise | Preserves TurboDiffusion behavior and official Cosmos equations |
+| 2026-08-27 | Support the released distilled checkpoint as T2W first | NVIDIA documents the released distilled checkpoint for T2W | No V2W/rolling claim before experimental validation |
+| 2026-08-27 | Keep the existing full-step Cosmos25 path unchanged | Distilled and post-trained checkpoints require different inference semantics | Avoids regression for current users |
+| 2026-08-27 | Defer pipeline wiring until real-weight component parity | Required by the repository add-model workflow | Next GPU task is DiT parity, not generation |
+| 2026-08-27 | Preserve native `net.*` student keys during conversion | Existing Cosmos25 loader owns the authoritative mapping | Converter only isolates student tensors and emits package metadata |
+
+## Handoff Notes
+- CPU scheduler unit and pinned-reference parity tests pass locally without skips.
+- Next implementation is isolated distilled-weight conversion/loading plus a real student DiT parity test.
+- Do not use the prior FlowUniPC/Karras Spark run as distilled parity evidence.
diff --git a/tests/local_tests/cosmos25/README.md b/tests/local_tests/cosmos25/README.md
new file mode 100644
index 0000000000..515f2280b6
--- /dev/null
+++ b/tests/local_tests/cosmos25/README.md
@@ -0,0 +1,69 @@
+# Cosmos Predict2.5 distilled validation
+
+This port targets NVIDIA's released 2B distilled **Text2World** checkpoint. It
+does not claim distilled Video2World, rolling generation, or real-time DreamVerse
+support.
+
+## Reference
+
+- Source: `NVIDIA/Cosmos-Predict2.5` at commit
+  `a2c298b0a3df3778b973fe65e9e58877b292d8a7`
+- Checkpoint: `nvidia/Cosmos-Predict2.5-2B`, `base/distilled`
+- Override the source checkout with `COSMOS25_OFFICIAL_REF_DIR`.
+
+Clone the reference next to FastVideo, or point the environment variable at an
+existing checkout:
+
+```bash
+git clone https://github.com/NVIDIA/Cosmos-Predict2.5.git cosmos-predict2.5
+export COSMOS25_OFFICIAL_REF_DIR="$PWD/cosmos-predict2.5"
+```
+
+## CPU sampler tests
+
+```bash
+pytest fastvideo/tests/schedulers/test_cosmos25_distilled_scheduler.py -q
+
+COSMOS25_OFFICIAL_REF_DIR=/path/to/Cosmos-Predict2.5 \
+pytest tests/local_tests/cosmos25/test_cosmos25_distilled_scheduler_parity.py -v -s
+```
+
+The parity test pins NVIDIA's scaling source and compares the full four-step
+preconditioning/x0/fixed-noise rollout. It does not load model weights.
+
+## Conversion scaffold
+
+The converter keeps only the official student's native `net.*` tensors and
+reuses non-transformer components from an existing FastVideo-loadable Cosmos
+Predict2.5 package. It writes its own distilled scheduler metadata instead of
+inheriting the base package's UniPC scheduler.
+
+```bash
+python scripts/checkpoint_conversion/cosmos25_distilled_to_diffusers.py \
+  --src-checkpoint /path/to/base/distilled/575edf0f-d973-4c74-b52c-69929a08d0a5_ema_bf16.pt \
+  --base-model /path/to/Cosmos-Predict2.5-2B-Diffusers \
+  --dst converted_weights/cosmos25-distilled
+```
+
+Local conversion contracts:
+
+```bash
+pytest tests/local_tests/cosmos25/test_cosmos25_distilled_conversion.py -q
+```
+
+This scaffold is locally tested, but conversion of the released 4 GB checkpoint
+and production-loader strictness have not yet been validated.
+
+## Remaining GPU gates
+
+Before wiring a public pipeline or running DreamVerse:
+
+1. Run the converter on the released distilled transformer and verify its
+   reported tensor count and production-loader missing/unexpected keys.
+2. Compare one real student DiT forward against the official implementation.
+3. Compare deterministic T2W latents end to end for the official four-step
+   schedule.
+4. Only after T2W parity, evaluate experimental V2W/rolling conditioning.
+
+Those checks require the released checkpoint and a CUDA machine. A skipped
+local parity test is not pass evidence.
diff --git a/tests/local_tests/cosmos25/_reference.py b/tests/local_tests/cosmos25/_reference.py
new file mode 100644
index 0000000000..5d2ed76b04
--- /dev/null
+++ b/tests/local_tests/cosmos25/_reference.py
@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import hashlib
+import importlib.util
+import os
+from pathlib import Path
+from types import ModuleType
+
+import pytest
+
+REFERENCE_COMMIT = "a2c298b0a3df3778b973fe65e9e58877b292d8a7"
+SCALING_SHA256 = "04f7d06fb349e317dd4a3c1ea3f78978b85c66a693c31a854f032bfb1c469dd2"
+DEFAULT_REFERENCE_ROOT = Path(__file__).resolve().parents[3] / "cosmos-predict2.5"
+
+
+def reference_root() -> Path:
+    return Path(os.environ.get("COSMOS25_OFFICIAL_REF_DIR", DEFAULT_REFERENCE_ROOT)).resolve()
+
+
+def load_official_scaling_module() -> ModuleType:
+    source = reference_root() / "cosmos_predict2/_src/predict2/modules/denoiser_scaling.py"
+    if not source.is_file():
+        pytest.skip("Cosmos Predict2.5 reference checkout is missing; set COSMOS25_OFFICIAL_REF_DIR")
+    digest = hashlib.sha256(source.read_bytes()).hexdigest()
+    if digest != SCALING_SHA256:
+        pytest.fail(f"Official denoiser scaling source drifted from {REFERENCE_COMMIT}: {digest}")
+
+    spec = importlib.util.spec_from_file_location("cosmos25_official_denoiser_scaling", source)
+    if spec is None or spec.loader is None:
+        pytest.fail(f"Could not load official scaling module from {source}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
diff --git a/tests/local_tests/cosmos25/test_cosmos25_distilled_conversion.py b/tests/local_tests/cosmos25/test_cosmos25_distilled_conversion.py
new file mode 100644
index 0000000000..6c5689e91c
--- /dev/null
+++ b/tests/local_tests/cosmos25/test_cosmos25_distilled_conversion.py
@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+from pathlib import Path
+
+import pytest
+import torch
+from safetensors import safe_open
+
+from scripts.checkpoint_conversion.cosmos25_distilled_to_diffusers import (
+    ConversionError,
+    convert_checkpoint,
+    extract_student_state_dict,
+)
+
+
+def _write_json(path: Path, payload: dict) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(payload), encoding="utf-8")
+
+
+def _make_base_model(root: Path) -> None:
+    _write_json(
+        root / "model_index.json",
+        {
+            "_class_name": "Cosmos2_5Pipeline",
+            "_diffusers_version": "0.37.0.dev0",
+            "scheduler": ["diffusers", "FlowUniPCMultistepScheduler"],
+            "text_encoder": ["transformers", "Reason1Model"],
+            "tokenizer": ["transformers", "AutoTokenizer"],
+            "transformer": ["diffusers", "Cosmos25Transformer3DModel"],
+            "vae": ["diffusers", "AutoencoderKLWan"],
+        },
+    )
+    _write_json(root / "transformer/config.json", {"_class_name": "Cosmos25Transformer3DModel"})
+    (root / "transformer/base.safetensors").write_bytes(b"old weights")
+    for component in ("vae", "text_encoder", "tokenizer"):
+        (root / component).mkdir(parents=True)
+    _write_json(root / "scheduler/scheduler_config.json", {"_class_name": "FlowUniPCMultistepScheduler"})
+
+
+@pytest.mark.parametrize("nested_key", [None, "state_dict", "model", "ema"])
+def test_extracts_only_native_student_tensors(nested_key: str | None) -> None:
+    state = {
+        "net.layer.weight": torch.ones(2, 3),
+        "net.layer.bias": torch.zeros(2),
+        "net_teacher.layer.weight": torch.full((2, 3), 2.0),
+        "optimizer.step": torch.tensor(4),
+        "metadata": "ignored",
+    }
+    checkpoint = state if nested_key is None else {nested_key: state, "iteration": 12}
+    result = extract_student_state_dict(checkpoint)
+
+    assert set(result) == {"net.layer.weight", "net.layer.bias"}
+    assert all(tensor.dtype == torch.bfloat16 for tensor in result.values())
+    assert all(tensor.is_contiguous() for tensor in result.values())
+
+
+def test_rejects_checkpoint_without_student() -> None:
+    with pytest.raises(ConversionError, match=r"net\.\*"):
+        extract_student_state_dict({"net_teacher.weight": torch.ones(1)})
+
+
+def test_builds_isolated_distilled_package(tmp_path: Path) -> None:
+    base = tmp_path / "base"
+    _make_base_model(base)
+    checkpoint_path = tmp_path / "distilled.pt"
+    torch.save(
+        {
+            "model": {
+                "net.layer.weight": torch.arange(6, dtype=torch.float32).reshape(2, 3),
+                "net_teacher.layer.weight": torch.ones(2, 3),
+            }
+        },
+        checkpoint_path,
+    )
+
+    dst = tmp_path / "converted"
+    report = convert_checkpoint(checkpoint_path, base, dst)
+
+    assert report == {"student_tensors": 1, "student_parameters": 6}
+    model_index = json.loads((dst / "model_index.json").read_text(encoding="utf-8"))
+    assert model_index["is_distilled"] is True
+    assert model_index["scheduler"] == ["diffusers", "Cosmos25DistilledScheduler"]
+    scheduler = json.loads((dst / "scheduler/scheduler_config.json").read_text(encoding="utf-8"))
+    assert scheduler["_class_name"] == "Cosmos25DistilledScheduler"
+    assert not (dst / "transformer/base.safetensors").exists()
+    assert (dst / "vae").is_dir()
+    assert (dst / "text_encoder").is_dir()
+    assert (dst / "tokenizer").is_dir()
+
+    with safe_open(
+        str(dst / "transformer/diffusion_pytorch_model.safetensors"),
+        framework="pt",
+        device="cpu",
+    ) as handle:
+        assert set(handle.keys()) == {"net.layer.weight"}
+        assert handle.get_tensor("net.layer.weight").dtype == torch.bfloat16
+
+
+def test_does_not_overwrite_existing_output_by_default(tmp_path: Path) -> None:
+    checkpoint_path = tmp_path / "distilled.pt"
+    torch.save({"net.weight": torch.ones(1)}, checkpoint_path)
+    base = tmp_path / "base"
+    _make_base_model(base)
+    dst = tmp_path / "converted"
+    dst.mkdir()
+    (dst / "keep.txt").write_text("user data", encoding="utf-8")
+
+    with pytest.raises(FileExistsError, match="--overwrite"):
+        convert_checkpoint(checkpoint_path, base, dst)
+    assert (dst / "keep.txt").read_text(encoding="utf-8") == "user data"
diff --git a/tests/local_tests/cosmos25/test_cosmos25_distilled_scheduler_parity.py b/tests/local_tests/cosmos25/test_cosmos25_distilled_scheduler_parity.py
new file mode 100644
index 0000000000..e518e1bc89
--- /dev/null
+++ b/tests/local_tests/cosmos25/test_cosmos25_distilled_scheduler_parity.py
@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos Predict2.5 distilled sampler parity (implementation_subcomponent)."""
+
+import math
+
+import torch
+from torch.testing import assert_close
+
+from fastvideo.models.schedulers.scheduling_cosmos25_distilled import (
+    Cosmos25DistilledScheduler,
+)
+from tests.local_tests.cosmos25._reference import load_official_scaling_module
+
+
+def test_rectified_flow_preconditioning_matches_official() -> None:
+    official_module = load_official_scaling_module()
+    official = official_module.RectifiedFlow_sCMWrapper(sigma_data=1.0)
+    actual = Cosmos25DistilledScheduler(sigma_data=1.0)
+    times = torch.tensor(
+        [math.pi / 2, math.atan(15), math.atan(5), math.atan(5 / 3), 0.37],
+        dtype=torch.float32,
+    ).view(1, 1, -1, 1, 1)
+
+    for actual_value, expected_value in zip(actual._scalings(times), official(times), strict=True):
+        assert_close(actual_value, expected_value, rtol=0, atol=0)
+
+
+def test_four_step_x0_rollout_matches_official() -> None:
+    official_module = load_official_scaling_module()
+    official = official_module.RectifiedFlow_sCMWrapper(sigma_data=1.0)
+    scheduler = Cosmos25DistilledScheduler(sigma_data=1.0)
+    scheduler.set_timesteps(4)
+
+    generator = torch.Generator().manual_seed(20260827)
+    initial_noise = torch.randn((1, 2, 3, 2, 2), generator=generator, dtype=torch.float32)
+    expected = initial_noise.to(torch.float64)
+    actual = expected.clone()
+
+    for index, model_timestep in enumerate(scheduler.timesteps):
+        angle = scheduler.trigflow_timesteps[index]
+        official_coefficients = official(angle)
+        c_skip, c_out, c_in, c_noise = official_coefficients
+
+        expected_model_input = (expected * c_in).float()
+        expected_model_output = expected_model_input * 0.125 + c_noise.float()
+        expected_x0 = c_skip * expected + c_out * expected_model_output
+        if index + 1 < len(scheduler.trigflow_timesteps):
+            next_angle = scheduler.trigflow_timesteps[index + 1]
+            expected = torch.cos(next_angle) * expected_x0 + torch.sin(next_angle) * initial_noise
+        else:
+            expected = expected_x0
+
+        actual_model_input = scheduler.scale_model_input(actual, model_timestep).float()
+        actual_model_output = actual_model_input * 0.125 + model_timestep.float()
+        result = scheduler.step(actual_model_output, model_timestep, actual)
+
+        assert_close(actual_model_input, expected_model_input, rtol=0, atol=0)
+        assert_close(result.pred_original_sample, expected_x0, rtol=0, atol=0)
+        assert_close(result.prev_sample, expected, rtol=0, atol=0)
+        actual = result.prev_sample

From c1c1368607c8ca1724d42e56c80660766ea9b618 Mon Sep 17 00:00:00 2001
From: Raghav <ragg04@gmail.com>
Date: Thu, 27 Aug 2026 11:28:15 -0700
Subject: [PATCH 02/16] [fix] exclude Cosmos training counters from conversion

---
 .../cosmos25_distilled_to_diffusers.py                    | 8 +++++++-
 .../cosmos25/test_cosmos25_distilled_conversion.py        | 1 +
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/scripts/checkpoint_conversion/cosmos25_distilled_to_diffusers.py b/scripts/checkpoint_conversion/cosmos25_distilled_to_diffusers.py
index b24e5a867d..a8ab40d236 100644
--- a/scripts/checkpoint_conversion/cosmos25_distilled_to_diffusers.py
+++ b/scripts/checkpoint_conversion/cosmos25_distilled_to_diffusers.py
@@ -23,6 +23,7 @@
 from safetensors.torch import save_file
 
 STATE_DICT_KEYS = ("state_dict", "model", "ema", "ema_model", "module")
+STUDENT_SKIP_PREFIXES = ("net.accum_",)
 REQUIRED_BASE_PATHS = (
     "model_index.json",
     "transformer/config.json",
@@ -73,7 +74,12 @@ def extract_student_state_dict(checkpoint: object) -> dict[str, torch.Tensor]:
     student = {
         key: tensor.detach().to(device="cpu", dtype=torch.bfloat16).contiguous()
         for key, tensor in state_dict.items()
-        if isinstance(key, str) and key.startswith("net.") and torch.is_tensor(tensor)
+        if (
+            isinstance(key, str)
+            and key.startswith("net.")
+            and not key.startswith(STUDENT_SKIP_PREFIXES)
+            and torch.is_tensor(tensor)
+        )
     }
     if not student:
         raise ConversionError("The resolved checkpoint contains no 'net.*' tensors")
diff --git a/tests/local_tests/cosmos25/test_cosmos25_distilled_conversion.py b/tests/local_tests/cosmos25/test_cosmos25_distilled_conversion.py
index 6c5689e91c..4cb4c1481b 100644
--- a/tests/local_tests/cosmos25/test_cosmos25_distilled_conversion.py
+++ b/tests/local_tests/cosmos25/test_cosmos25_distilled_conversion.py
@@ -44,6 +44,7 @@ def test_extracts_only_native_student_tensors(nested_key: str | None) -> None:
     state = {
         "net.layer.weight": torch.ones(2, 3),
         "net.layer.bias": torch.zeros(2),
+        "net.accum_iteration": torch.tensor(12),
         "net_teacher.layer.weight": torch.full((2, 3), 2.0),
         "optimizer.step": torch.tensor(4),
         "metadata": "ignored",

From 8cb88cd990c76d48876e03188a806f660bdebca4 Mon Sep 17 00:00:00 2001
From: Raghav <ragg04@gmail.com>
Date: Thu, 27 Aug 2026 11:37:33 -0700
Subject: [PATCH 03/16] [test] add Cosmos 2.5 distilled DiT parity gate

---
 tests/local_tests/cosmos25/PORT_STATUS.md     |   2 +-
 tests/local_tests/cosmos25/README.md          |  17 ++
 ...t_cosmos25_distilled_transformer_parity.py | 274 ++++++++++++++++++
 3 files changed, 292 insertions(+), 1 deletion(-)
 create mode 100644 tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py

diff --git a/tests/local_tests/cosmos25/PORT_STATUS.md b/tests/local_tests/cosmos25/PORT_STATUS.md
index d7b1643e25..01ff97ad16 100644
--- a/tests/local_tests/cosmos25/PORT_STATUS.md
+++ b/tests/local_tests/cosmos25/PORT_STATUS.md
@@ -39,7 +39,7 @@
 | scheduler unit | `pytest fastvideo/tests/schedulers/test_cosmos25_distilled_scheduler.py -q` | 7 passed | CPU-only; includes registry resolution; 2026-08-27 |
 | official scheduler | `COSMOS25_OFFICIAL_REF_DIR=/path/to/Cosmos-Predict2.5 pytest tests/local_tests/cosmos25/test_cosmos25_distilled_scheduler_parity.py -v -s` | 2 passed, non-skip | CPU-only; pinned source; 2026-08-27 |
 | conversion contracts | `pytest tests/local_tests/cosmos25/test_cosmos25_distilled_conversion.py -q` | 7 passed | Synthetic checkpoints/layout; 2026-08-27 |
-| student DiT | not yet created | not run | requires CUDA and released weights |
+| student DiT | `COSMOS25_OFFICIAL_REF_DIR=/path/to/Cosmos-Predict2.5 COSMOS25_DISTILLED_CHECKPOINT=/path/to/distilled.pt pytest tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py -v -s` | scaffold pending Spark run | requires CUDA, official source dependencies, and released weights |
 | pipeline | not yet created | not run | forbidden until component parity passes |
 
 ## Open Questions
diff --git a/tests/local_tests/cosmos25/README.md b/tests/local_tests/cosmos25/README.md
index 515f2280b6..80f95ccd1b 100644
--- a/tests/local_tests/cosmos25/README.md
+++ b/tests/local_tests/cosmos25/README.md
@@ -67,3 +67,20 @@ Before wiring a public pipeline or running DreamVerse:
 
 Those checks require the released checkpoint and a CUDA machine. A skipped
 local parity test is not pass evidence.
+
+## Real student DiT parity
+
+This gate loads the raw NVIDIA student into the official and FastVideo DiTs,
+runs the same small deterministic BF16 forward through each implementation, and
+compares the raw network outputs. It loads the models sequentially to limit GPU
+memory use.
+
+```bash
+export COSMOS25_OFFICIAL_REF_DIR=/path/to/Cosmos-Predict2.5
+export COSMOS25_DISTILLED_CHECKPOINT=/path/to/575edf0f-d973-4c74-b52c-69929a08d0a5_ema_bf16.pt
+
+FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA \
+pytest tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py -v -s
+```
+
+The test must report `PASSED`, not `SKIPPED`, before distilled pipeline wiring.
diff --git a/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py b/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py
new file mode 100644
index 0000000000..7b6de912a6
--- /dev/null
+++ b/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py
@@ -0,0 +1,274 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Real-weight Cosmos Predict2.5 student DiT parity (implementation_subcomponent)."""
+
+from __future__ import annotations
+
+import gc
+import os
+import sys
+from collections.abc import Callable, Mapping
+from pathlib import Path
+from typing import Any, TypeVar
+
+import pytest
+import torch
+from torch.testing import assert_close
+
+os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
+
+from fastvideo.configs.models.dits.cosmos2_5 import (  # noqa: E402
+    Cosmos25ArchConfig,
+    Cosmos25VideoConfig,
+)
+from fastvideo.distributed import (  # noqa: E402
+    cleanup_dist_env_and_memory,
+    maybe_init_distributed_environment_and_model_parallel,
+)
+from fastvideo.forward_context import set_forward_context  # noqa: E402
+from fastvideo.models.dits.cosmos2_5 import Cosmos25Transformer3DModel  # noqa: E402
+from fastvideo.models.loader.utils import (  # noqa: E402
+    get_param_names_mapping,
+    hf_to_custom_state_dict,
+)
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch  # noqa: E402
+from scripts.checkpoint_conversion.cosmos25_distilled_to_diffusers import (  # noqa: E402
+    extract_student_state_dict,
+)
+from tests.local_tests.cosmos25._reference import reference_root  # noqa: E402
+
+CHECKPOINT_ENV = "COSMOS25_DISTILLED_CHECKPOINT"
+ModuleT = TypeVar("ModuleT", bound=torch.nn.Module)
+
+
+@pytest.fixture
+def distributed_setup():
+    maybe_init_distributed_environment_and_model_parallel(1, 1)
+    yield
+    cleanup_dist_env_and_memory()
+
+
+def _checkpoint_path() -> Path:
+    value = os.environ.get(CHECKPOINT_ENV)
+    if not value:
+        pytest.skip(f"Set {CHECKPOINT_ENV} to NVIDIA's released distilled .pt checkpoint")
+    assert value is not None
+    path = Path(value).expanduser().resolve()
+    if not path.is_file():
+        pytest.fail(f"{CHECKPOINT_ENV} does not point to a file: {path}")
+    return path
+
+
+def _load_student_checkpoint() -> dict[str, torch.Tensor]:
+    checkpoint = torch.load(_checkpoint_path(), map_location="cpu", weights_only=True)
+    return extract_student_state_dict(checkpoint)
+
+
+def _official_model_class() -> type[torch.nn.Module]:
+    root = reference_root()
+    if not root.is_dir():
+        pytest.skip("Set COSMOS25_OFFICIAL_REF_DIR to the NVIDIA Cosmos-Predict2.5 checkout")
+    root_string = str(root)
+    if root_string not in sys.path:
+        sys.path.insert(0, root_string)
+    try:
+        from cosmos_predict2._src.predict2.networks.minimal_v1_lvg_dit import (
+            MinimalV1LVGDiT,
+        )
+    except ImportError as error:
+        pytest.fail(f"Could not import NVIDIA's MinimalV1LVGDiT from {root}: {error}")
+    return MinimalV1LVGDiT
+
+
+def _arch_config() -> Cosmos25ArchConfig:
+    return Cosmos25ArchConfig(
+        num_attention_heads=16,
+        attention_head_dim=128,
+        in_channels=16,
+        out_channels=16,
+        num_layers=28,
+        patch_size=(1, 2, 2),
+        max_size=(128, 240, 240),
+        rope_scale=(1.0, 3.0, 3.0),
+        text_embed_dim=1024,
+        mlp_ratio=4.0,
+        adaln_lora_dim=256,
+        use_adaln_lora=True,
+        concat_padding_mask=True,
+        extra_pos_embed_type=None,
+        use_crossattn_projection=True,
+        crossattn_proj_in_channels=100352,
+        rope_enable_fps_modulation=False,
+        qk_norm="rms_norm",
+    )
+
+
+def _construct_bf16(factory: Callable[[], ModuleT]) -> ModuleT:
+    previous_dtype = torch.get_default_dtype()
+    try:
+        torch.set_default_dtype(torch.bfloat16)
+        return factory()
+    finally:
+        torch.set_default_dtype(previous_dtype)
+
+
+def _load_official_model(student: Mapping[str, torch.Tensor], device: torch.device):
+    model_class = _official_model_class()
+    model = _construct_bf16(
+        lambda: model_class(
+            max_img_h=240,
+            max_img_w=240,
+            max_frames=128,
+            in_channels=16,
+            out_channels=16,
+            patch_spatial=2,
+            patch_temporal=1,
+            model_channels=2048,
+            num_blocks=28,
+            num_heads=16,
+            mlp_ratio=4.0,
+            crossattn_emb_channels=1024,
+            pos_emb_cls="rope3d",
+            pos_emb_learnable=True,
+            pos_emb_interpolation="crop",
+            use_adaln_lora=True,
+            adaln_lora_dim=256,
+            rope_h_extrapolation_ratio=3.0,
+            rope_w_extrapolation_ratio=3.0,
+            rope_t_extrapolation_ratio=1.0,
+            extra_per_block_abs_pos_emb=False,
+            rope_enable_fps_modulation=False,
+            use_crossattn_projection=True,
+            crossattn_proj_in_channels=100352,
+            concat_padding_mask=True,
+            atten_backend="torch",
+        )
+    )
+
+    expected_keys = set(model.state_dict())
+    wrapped_blocks = any("._checkpoint_wrapped_module." in key for key in expected_keys)
+    official_state: dict[str, torch.Tensor] = {}
+    for source_key, tensor in student.items():
+        target_key = source_key.removeprefix("net.")
+        if wrapped_blocks and target_key.startswith("blocks."):
+            parts = target_key.split(".", 2)
+            if len(parts) == 3 and parts[1].isdigit():
+                target_key = f"blocks.{parts[1]}._checkpoint_wrapped_module.{parts[2]}"
+        if target_key in expected_keys:
+            official_state[target_key] = tensor
+
+    missing, unexpected = model.load_state_dict(official_state, strict=False)
+    important_missing = [key for key in missing if not key.endswith("._extra_state")]
+    assert not important_missing, f"Official model missing inference keys: {important_missing[:20]}"
+    assert not unexpected, f"Official model received unexpected keys: {unexpected[:20]}"
+    return model.to(device=device, dtype=torch.bfloat16).eval()
+
+
+def _load_fastvideo_model(student: Mapping[str, torch.Tensor], device: torch.device):
+    arch = _arch_config()
+    config = Cosmos25VideoConfig(arch_config=arch)
+    hf_config: dict[str, Any] = {
+        "in_channels": arch.in_channels,
+        "out_channels": arch.out_channels,
+        "num_attention_heads": arch.num_attention_heads,
+        "attention_head_dim": arch.attention_head_dim,
+        "num_layers": arch.num_layers,
+        "patch_size": arch.patch_size,
+        "max_size": arch.max_size,
+        "rope_scale": arch.rope_scale,
+        "text_embed_dim": arch.text_embed_dim,
+        "mlp_ratio": arch.mlp_ratio,
+        "adaln_lora_dim": arch.adaln_lora_dim,
+        "use_adaln_lora": arch.use_adaln_lora,
+        "concat_padding_mask": arch.concat_padding_mask,
+        "extra_pos_embed_type": arch.extra_pos_embed_type,
+        "use_crossattn_projection": arch.use_crossattn_projection,
+        "crossattn_proj_in_channels": arch.crossattn_proj_in_channels,
+        "rope_enable_fps_modulation": arch.rope_enable_fps_modulation,
+        "qk_norm": arch.qk_norm,
+    }
+    model = _construct_bf16(lambda: Cosmos25Transformer3DModel(config=config, hf_config=hf_config))
+    mapping = get_param_names_mapping(arch.param_names_mapping)
+    mapped, _ = hf_to_custom_state_dict(student.items(), mapping)
+    expected_keys = set(model.state_dict())
+    filtered = {key: value for key, value in mapped.items() if key in expected_keys}
+    missing, unexpected = model.load_state_dict(filtered, strict=False)
+    important_missing = [key for key in missing if not key.endswith("._extra_state")]
+    assert not important_missing, f"FastVideo model missing inference keys: {important_missing[:20]}"
+    assert not unexpected, f"FastVideo model received unexpected keys: {unexpected[:20]}"
+    return model.to(device=device, dtype=torch.bfloat16).eval()
+
+
+def _inputs() -> dict[str, torch.Tensor]:
+    generator = torch.Generator(device="cpu").manual_seed(20260827)
+    return {
+        "latents": torch.randn((1, 16, 2, 16, 16), generator=generator, dtype=torch.float32).to(torch.bfloat16),
+        "text": torch.randn((1, 4, 100352), generator=generator, dtype=torch.float32).to(torch.bfloat16),
+        "condition_mask": torch.zeros((1, 1, 2, 16, 16), dtype=torch.bfloat16),
+        "padding_mask": torch.ones((1, 16, 16), dtype=torch.bfloat16),
+        "timestep": torch.full((1, 2), 15 / 16, dtype=torch.bfloat16),
+        "fps": torch.tensor([24], dtype=torch.bfloat16),
+    }
+
+
+@pytest.mark.usefixtures("distributed_setup")
+def test_distilled_student_forward_matches_official() -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required for the real Cosmos Predict2.5 DiT parity gate")
+
+    device = torch.device("cuda:0")
+    student = _load_student_checkpoint()
+    inputs = _inputs()
+
+    official = _load_official_model(student, device)
+    from cosmos_predict2._src.predict2.conditioner import DataType
+
+    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+        official_output = (
+            official(
+                x_B_C_T_H_W=inputs["latents"].to(device),
+                timesteps_B_T=inputs["timestep"].to(device),
+                crossattn_emb=inputs["text"].to(device),
+                condition_video_input_mask_B_C_T_H_W=inputs["condition_mask"].to(device),
+                fps=inputs["fps"].to(device),
+                padding_mask=inputs["padding_mask"].to(device),
+                data_type=DataType.VIDEO,
+            )
+            .float()
+            .cpu()
+        )
+    del official
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    fastvideo = _load_fastvideo_model(student, device)
+    forward_batch = ForwardBatch(data_type="dummy")
+    with (
+        torch.inference_mode(),
+        torch.autocast("cuda", dtype=torch.bfloat16),
+        set_forward_context(current_timestep=937, attn_metadata=None, forward_batch=forward_batch),
+    ):
+        fastvideo_output = (
+            fastvideo(
+                hidden_states=inputs["latents"].to(device),
+                timestep=inputs["timestep"].to(device),
+                encoder_hidden_states=inputs["text"].to(device),
+                fps=inputs["fps"].to(device),
+                condition_mask=inputs["condition_mask"].to(device),
+                padding_mask=inputs["padding_mask"].unsqueeze(1).to(device),
+            )
+            .float()
+            .cpu()
+        )
+
+    absolute = (fastvideo_output - official_output).abs()
+    mean_abs = float(absolute.mean())
+    max_abs = float(absolute.max())
+    reference_abs_mean = float(official_output.abs().mean())
+    relative_mean = mean_abs / max(reference_abs_mean, 1e-8)
+    print(
+        "Cosmos25 distilled DiT parity: "
+        f"max_abs={max_abs:.8f}, mean_abs={mean_abs:.8f}, relative_mean={relative_mean:.8f}"
+    )
+
+    assert relative_mean < 0.05
+    assert_close(fastvideo_output, official_output, atol=0.1, rtol=0.1)

From 24ce13ce07b70533a449c380d05d6b32a3f41628 Mon Sep 17 00:00:00 2001
From: Raghav <ragg04@gmail.com>
Date: Thu, 27 Aug 2026 11:41:57 -0700
Subject: [PATCH 04/16] [test] bypass Cosmos package CUDA-extra guard

---
 .../test_cosmos25_distilled_transformer_parity.py      | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py b/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py
index 7b6de912a6..8aede506e1 100644
--- a/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py
+++ b/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py
@@ -6,6 +6,7 @@
 import gc
 import os
 import sys
+import types
 from collections.abc import Callable, Mapping
 from pathlib import Path
 from typing import Any, TypeVar
@@ -70,6 +71,15 @@ def _official_model_class() -> type[torch.nn.Module]:
     root_string = str(root)
     if root_string not in sys.path:
         sys.path.insert(0, root_string)
+    # The published package root enforces installation of its full CUDA extra.
+    # Component parity only needs the checked-out Python modules, so install a
+    # namespace package that preserves normal submodule imports without running
+    # that unrelated environment guard.
+    if "cosmos_predict2" not in sys.modules:
+        package = types.ModuleType("cosmos_predict2")
+        package.__path__ = [str(root / "cosmos_predict2")]  # type: ignore[attr-defined]
+        package.__package__ = "cosmos_predict2"
+        sys.modules["cosmos_predict2"] = package
     try:
         from cosmos_predict2._src.predict2.networks.minimal_v1_lvg_dit import (
             MinimalV1LVGDiT,

From c1338805f4a710df90141c9366d0ac6d91568d70 Mon Sep 17 00:00:00 2001
From: Raghav <ragg04@gmail.com>
Date: Thu, 27 Aug 2026 11:45:11 -0700
Subject: [PATCH 05/16] test: isolate Cosmos DiT parity from config deps

---
 ...t_cosmos25_distilled_transformer_parity.py | 20 +++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py b/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py
index 8aede506e1..66e79dd5d1 100644
--- a/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py
+++ b/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py
@@ -8,6 +8,7 @@
 import sys
 import types
 from collections.abc import Callable, Mapping
+from enum import Enum
 from pathlib import Path
 from typing import Any, TypeVar
 
@@ -80,6 +81,25 @@ def _official_model_class() -> type[torch.nn.Module]:
         package.__path__ = [str(root / "cosmos_predict2")]  # type: ignore[attr-defined]
         package.__package__ = "cosmos_predict2"
         sys.modules["cosmos_predict2"] = package
+
+    # The DiT only uses DataType for an enum type check and the VIDEO value.
+    # Importing the published conditioner solely for that enum pulls in the
+    # training/configuration stack (including iopath), which is unrelated to
+    # transformer parity and is not a FastVideo runtime dependency.
+    conditioner_name = "cosmos_predict2._src.predict2.conditioner"
+    if conditioner_name not in sys.modules:
+        conditioner = types.ModuleType(conditioner_name)
+
+        class DataType(str, Enum):
+            IMAGE = "image"
+            VIDEO = "video"
+            MIX = "mix"
+
+            def __str__(self) -> str:
+                return self.value
+
+        conditioner.__dict__["DataType"] = DataType
+        sys.modules[conditioner_name] = conditioner
     try:
         from cosmos_predict2._src.predict2.networks.minimal_v1_lvg_dit import (
             MinimalV1LVGDiT,

From f480f8dcdffd9dee52935e23b9d109013654d7af Mon Sep 17 00:00:00 2001
From: Raghav <ragg04@gmail.com>
Date: Thu, 27 Aug 2026 11:47:58 -0700
Subject: [PATCH 06/16] test: run Cosmos reference DiT without transformer
 engine

---
 ...t_cosmos25_distilled_transformer_parity.py | 58 +++++++++++++++++++
 1 file changed, 58 insertions(+)

diff --git a/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py b/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py
index 66e79dd5d1..16f7384cb0 100644
--- a/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py
+++ b/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py
@@ -42,6 +42,63 @@
 ModuleT = TypeVar("ModuleT", bound=torch.nn.Module)
 
 
+class _ReferenceRMSNorm(torch.nn.Module):
+    """PyTorch equivalent of the TE RMSNorm used by the reference DiT."""
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        normalized = hidden_states.float()
+        variance = normalized.square().mean(dim=-1, keepdim=True)
+        normalized = normalized * torch.rsqrt(variance + self.eps)
+        return normalized.to(input_dtype) * self.weight
+
+
+def _reference_apply_rotary_pos_emb(
+    hidden_states: torch.Tensor,
+    rotary_pos_emb: torch.Tensor,
+    *,
+    tensor_format: str,
+    fused: bool,
+) -> torch.Tensor:
+    """PyTorch equivalent of TE's rotate-half RoPE for BSHD tensors."""
+    assert tensor_format == "bshd"
+    del fused
+    frequencies = rotary_pos_emb[:, 0, 0, :]
+    cos = frequencies.cos()[None, :, None, :]
+    sin = frequencies.sin()[None, :, None, :]
+    first_half, second_half = hidden_states.chunk(2, dim=-1)
+    rotated = torch.cat((-second_half, first_half), dim=-1)
+    return (hidden_states.float() * cos + rotated.float() * sin).to(hidden_states.dtype)
+
+
+def _install_transformer_engine_reference_shim() -> None:
+    """Provide only the two TE operations needed by the torch-backend DiT."""
+    if "transformer_engine" in sys.modules:
+        return
+
+    transformer_engine = types.ModuleType("transformer_engine")
+    transformer_engine_pytorch = types.ModuleType("transformer_engine.pytorch")
+    transformer_engine_attention = types.ModuleType("transformer_engine.pytorch.attention")
+    transformer_engine_rope = types.ModuleType("transformer_engine.pytorch.attention.rope")
+
+    transformer_engine.__dict__["pytorch"] = transformer_engine_pytorch
+    transformer_engine_pytorch.__dict__["RMSNorm"] = _ReferenceRMSNorm
+    transformer_engine_pytorch.__dict__["attention"] = transformer_engine_attention
+    transformer_engine_attention.__dict__["apply_rotary_pos_emb"] = _reference_apply_rotary_pos_emb
+    transformer_engine_attention.__dict__["rope"] = transformer_engine_rope
+    transformer_engine_rope.__dict__["apply_rotary_pos_emb"] = _reference_apply_rotary_pos_emb
+
+    sys.modules["transformer_engine"] = transformer_engine
+    sys.modules["transformer_engine.pytorch"] = transformer_engine_pytorch
+    sys.modules["transformer_engine.pytorch.attention"] = transformer_engine_attention
+    sys.modules["transformer_engine.pytorch.attention.rope"] = transformer_engine_rope
+
+
 @pytest.fixture
 def distributed_setup():
     maybe_init_distributed_environment_and_model_parallel(1, 1)
@@ -100,6 +157,7 @@ def __str__(self) -> str:
 
         conditioner.__dict__["DataType"] = DataType
         sys.modules[conditioner_name] = conditioner
+    _install_transformer_engine_reference_shim()
     try:
         from cosmos_predict2._src.predict2.networks.minimal_v1_lvg_dit import (
             MinimalV1LVGDiT,

From 257bbe2632e4d4b62086c2d18442fb12e8d9e650 Mon Sep 17 00:00:00 2001
From: Raghav <ragg04@gmail.com>
Date: Thu, 27 Aug 2026 11:50:35 -0700
Subject: [PATCH 07/16] test: complete reference RMSNorm interface

---
 .../cosmos25/test_cosmos25_distilled_transformer_parity.py     | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py b/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py
index 16f7384cb0..a52d95a4fa 100644
--- a/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py
+++ b/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py
@@ -50,6 +50,9 @@ def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
         self.weight = torch.nn.Parameter(torch.ones(hidden_size))
         self.eps = eps
 
+    def reset_parameters(self) -> None:
+        torch.nn.init.ones_(self.weight)
+
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         input_dtype = hidden_states.dtype
         normalized = hidden_states.float()

From a704c29ee5e58c113609ac197033aeec451e2b89 Mon Sep 17 00:00:00 2001
From: Raghav <ragg04@gmail.com>
Date: Thu, 27 Aug 2026 11:53:16 -0700
Subject: [PATCH 08/16] test: ignore Cosmos reference training counters

---
 .../test_cosmos25_distilled_transformer_parity.py    | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py b/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py
index a52d95a4fa..66d24172ec 100644
--- a/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py
+++ b/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py
@@ -40,6 +40,12 @@
 
 CHECKPOINT_ENV = "COSMOS25_DISTILLED_CHECKPOINT"
 ModuleT = TypeVar("ModuleT", bound=torch.nn.Module)
+REFERENCE_TRAINING_COUNTERS = {
+    "accum_image_sample_counter",
+    "accum_iteration",
+    "accum_train_in_hours",
+    "accum_video_sample_counter",
+}
 
 
 class _ReferenceRMSNorm(torch.nn.Module):
@@ -248,7 +254,11 @@ def _load_official_model(student: Mapping[str, torch.Tensor], device: torch.devi
             official_state[target_key] = tensor
 
     missing, unexpected = model.load_state_dict(official_state, strict=False)
-    important_missing = [key for key in missing if not key.endswith("._extra_state")]
+    important_missing = [
+        key
+        for key in missing
+        if not key.endswith("._extra_state") and key not in REFERENCE_TRAINING_COUNTERS
+    ]
     assert not important_missing, f"Official model missing inference keys: {important_missing[:20]}"
     assert not unexpected, f"Official model received unexpected keys: {unexpected[:20]}"
     return model.to(device=device, dtype=torch.bfloat16).eval()

From 613880b41e2b0f534c89d3971e3aef5a59826936 Mon Sep 17 00:00:00 2001
From: Raghav <ragg04@gmail.com>
Date: Thu, 27 Aug 2026 12:01:12 -0700
Subject: [PATCH 09/16] test: trace Cosmos distilled DiT parity drift

---
 ...t_cosmos25_distilled_transformer_parity.py | 81 +++++++++++++++++--
 1 file changed, 76 insertions(+), 5 deletions(-)

diff --git a/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py b/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py
index 66d24172ec..a871e1c9b8 100644
--- a/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py
+++ b/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py
@@ -208,6 +208,57 @@ def _construct_bf16(factory: Callable[[], ModuleT]) -> ModuleT:
         torch.set_default_dtype(previous_dtype)
 
 
+def _register_forward_captures(
+    modules: Mapping[str, torch.nn.Module],
+) -> tuple[dict[str, torch.Tensor], list[Any]]:
+    captures: dict[str, torch.Tensor] = {}
+    handles: list[Any] = []
+
+    def make_hook(name: str):
+        def capture(_module, _inputs, output) -> None:
+            outputs = output if isinstance(output, tuple) else (output,)
+            for index, value in enumerate(outputs):
+                if isinstance(value, torch.Tensor):
+                    key = name if len(outputs) == 1 else f"{name}.{index}"
+                    captures[key] = value.detach().float().cpu()
+
+        return capture
+
+    for name, module in modules.items():
+        handles.append(module.register_forward_hook(make_hook(name)))
+    return captures, handles
+
+
+def _drift(reference: torch.Tensor, candidate: torch.Tensor) -> tuple[float, float, float]:
+    absolute = (candidate - reference).abs()
+    mean_abs = float(absolute.mean())
+    max_abs = float(absolute.max())
+    relative_mean = mean_abs / max(float(reference.abs().mean()), 1e-8)
+    return mean_abs, max_abs, relative_mean
+
+
+def _print_capture_drift(
+    reference_captures: Mapping[str, torch.Tensor],
+    fastvideo_captures: Mapping[str, torch.Tensor],
+) -> None:
+    comparable = ["patch", "time_norm", "text", *[f"block.{index}" for index in range(28)], "final"]
+    print("Cosmos25 distilled DiT intermediate drift:")
+    for name in comparable:
+        mean_abs, max_abs, relative_mean = _drift(reference_captures[name], fastvideo_captures[name])
+        print(f"  {name:>9}: mean_abs={mean_abs:.8f} max_abs={max_abs:.8f} relative_mean={relative_mean:.8f}")
+
+    reference_angles = reference_captures["rope_angles"][:, 0, 0, :]
+    for index, function in enumerate((torch.cos, torch.sin)):
+        mean_abs, max_abs, relative_mean = _drift(
+            function(reference_angles),
+            fastvideo_captures[f"rope.{index}"],
+        )
+        print(
+            f"  rope.{index:>4}: mean_abs={mean_abs:.8f} "
+            f"max_abs={max_abs:.8f} relative_mean={relative_mean:.8f}"
+        )
+
+
 def _load_official_model(student: Mapping[str, torch.Tensor], device: torch.device):
     model_class = _official_model_class()
     model = _construct_bf16(
@@ -321,6 +372,15 @@ def test_distilled_student_forward_matches_official() -> None:
     inputs = _inputs()
 
     official = _load_official_model(student, device)
+    official_modules = {
+        "patch": official.x_embedder,
+        "rope_angles": official.pos_embedder,
+        "time_norm": official.t_embedding_norm,
+        "text": official.crossattn_proj,
+        **{f"block.{index}": block for index, block in enumerate(official.blocks)},
+        "final": official.final_layer,
+    }
+    official_captures, official_handles = _register_forward_captures(official_modules)
     from cosmos_predict2._src.predict2.conditioner import DataType
 
     with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
@@ -337,11 +397,22 @@ def test_distilled_student_forward_matches_official() -> None:
             .float()
             .cpu()
         )
+    for handle in official_handles:
+        handle.remove()
     del official
     gc.collect()
     torch.cuda.empty_cache()
 
     fastvideo = _load_fastvideo_model(student, device)
+    fastvideo_modules = {
+        "patch": fastvideo.patch_embed,
+        "rope": fastvideo.rope,
+        "time_norm": fastvideo.time_embed.norm,
+        "text": fastvideo.crossattn_proj,
+        **{f"block.{index}": block for index, block in enumerate(fastvideo.transformer_blocks)},
+        "final": fastvideo.final_layer,
+    }
+    fastvideo_captures, fastvideo_handles = _register_forward_captures(fastvideo_modules)
     forward_batch = ForwardBatch(data_type="dummy")
     with (
         torch.inference_mode(),
@@ -360,12 +431,12 @@ def test_distilled_student_forward_matches_official() -> None:
             .float()
             .cpu()
         )
+    for handle in fastvideo_handles:
+        handle.remove()
 
-    absolute = (fastvideo_output - official_output).abs()
-    mean_abs = float(absolute.mean())
-    max_abs = float(absolute.max())
-    reference_abs_mean = float(official_output.abs().mean())
-    relative_mean = mean_abs / max(reference_abs_mean, 1e-8)
+    _print_capture_drift(official_captures, fastvideo_captures)
+
+    mean_abs, max_abs, relative_mean = _drift(official_output, fastvideo_output)
     print(
         "Cosmos25 distilled DiT parity: "
         f"max_abs={max_abs:.8f}, mean_abs={mean_abs:.8f}, relative_mean={relative_mean:.8f}"

From a5acf0e435abedfd9c5920d0172205522d5f0eeb Mon Sep 17 00:00:00 2001
From: Raghav <ragg04@gmail.com>
Date: Thu, 27 Aug 2026 12:06:13 -0700
Subject: [PATCH 10/16] test: calibrate Cosmos BF16 parity tolerance

---
 .../test_cosmos25_distilled_transformer_parity.py  | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py b/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py
index a871e1c9b8..993ea708b6 100644
--- a/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py
+++ b/tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py
@@ -435,6 +435,14 @@ def test_distilled_student_forward_matches_official() -> None:
         handle.remove()
 
     _print_capture_drift(official_captures, fastvideo_captures)
+    for component in ("patch", "time_norm", "text"):
+        assert_close(fastvideo_captures[component], official_captures[component], atol=0, rtol=0)
+
+    _, _, first_block_relative_mean = _drift(
+        official_captures["block.0"],
+        fastvideo_captures["block.0"],
+    )
+    assert first_block_relative_mean < 0.001
 
     mean_abs, max_abs, relative_mean = _drift(official_output, fastvideo_output)
     print(
@@ -443,4 +451,8 @@ def test_distilled_student_forward_matches_official() -> None:
     )
 
     assert relative_mean < 0.05
-    assert_close(fastvideo_output, official_output, atol=0.1, rtol=0.1)
+    # The two implementations enter block 0 within 0.1%, then BF16 rounding
+    # accumulates through 28 residual blocks. As with FastVideo's real-weight
+    # FLUX DiT parity gate, aggregate drift carries the primary assertion and
+    # a loose absolute bound catches isolated runaway values.
+    assert_close(fastvideo_output, official_output, atol=0.5, rtol=0)

From 83665f94dc4b21c37310009fa95376446e91e156 Mon Sep 17 00:00:00 2001
From: Raghav <ragg04@gmail.com>
Date: Thu, 27 Aug 2026 12:17:56 -0700
Subject: [PATCH 11/16] [Model] wire Cosmos2.5 distilled T2W inference

---
 .../basic/basic_cosmos2_5_distilled_t2w.py    |  54 ++++++
 .../basic/cosmos/cosmos2_5_pipeline.py        |  59 ++++++-
 fastvideo/pipelines/stages/__init__.py        |   8 +-
 fastvideo/pipelines/stages/denoising.py       | 115 +++++++++++++
 .../pipelines/stages/latent_preparation.py    |  61 +++++++
 tests/local_tests/cosmos25/PORT_STATUS.md     |  38 +++--
 tests/local_tests/cosmos25/README.md          |  38 +++--
 .../test_cosmos25_distilled_pipeline.py       | 155 ++++++++++++++++++
 8 files changed, 486 insertions(+), 42 deletions(-)
 create mode 100644 examples/inference/basic/basic_cosmos2_5_distilled_t2w.py
 create mode 100644 tests/local_tests/cosmos25/test_cosmos25_distilled_pipeline.py

diff --git a/examples/inference/basic/basic_cosmos2_5_distilled_t2w.py b/examples/inference/basic/basic_cosmos2_5_distilled_t2w.py
new file mode 100644
index 0000000000..d066c4614a
--- /dev/null
+++ b/examples/inference/basic/basic_cosmos2_5_distilled_t2w.py
@@ -0,0 +1,54 @@
+"""Run the released Cosmos Predict2.5 2B distilled Text2World checkpoint."""
+
+import argparse
+
+from fastvideo import VideoGenerator
+from fastvideo.api.sampling_param import SamplingParam
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", required=True, help="Converted FastVideo model directory")
+    parser.add_argument("--output", default="outputs_video/cosmos2_5_distilled_t2w.mp4")
+    parser.add_argument("--steps", type=int, default=4, choices=range(1, 5))
+    parser.add_argument("--frames", type=int, default=77)
+    parser.add_argument("--height", type=int, default=704)
+    parser.add_argument("--width", type=int, default=1280)
+    parser.add_argument("--fps", type=int, default=16)
+    parser.add_argument("--seed", type=int, default=42)
+    args = parser.parse_args()
+
+    generator = VideoGenerator.from_pretrained(
+        args.model,
+        num_gpus=1,
+        use_fsdp_inference=False,
+        dit_cpu_offload=False,
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=True,
+    )
+    sampling = SamplingParam(
+        num_inference_steps=args.steps,
+        num_frames=args.frames,
+        height=args.height,
+        width=args.width,
+        fps=args.fps,
+        seed=args.seed,
+        guidance_scale=1.0,
+    )
+    prompt = (
+        "A robotic arm performs precision welding in an industrial workshop. "
+        "Bright blue-white sparks scatter over the metal while smoke rises, "
+        "cinematic lighting, steady camera, realistic motion."
+    )
+    generator.generate_video(
+        prompt,
+        sampling_param=sampling,
+        output_path=args.output,
+        save_video=True,
+    )
+    generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/fastvideo/pipelines/basic/cosmos/cosmos2_5_pipeline.py b/fastvideo/pipelines/basic/cosmos/cosmos2_5_pipeline.py
index cf561ad924..8af1414ba1 100644
--- a/fastvideo/pipelines/basic/cosmos/cosmos2_5_pipeline.py
+++ b/fastvideo/pipelines/basic/cosmos/cosmos2_5_pipeline.py
@@ -3,14 +3,37 @@
 
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_cosmos25_distilled import Cosmos25DistilledScheduler
 from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.pipelines.stages import (ConditioningStage, Cosmos25AutoDenoisingStage,
                                         Cosmos25AutoLatentPreparationStage, DecodingStage, InputValidationStage,
+                                        Cosmos25DistilledT2WDenoisingStage, Cosmos25DistilledT2WLatentPreparationStage,
                                         Cosmos25TextEncodingStage, Cosmos25TimestepPreparationStage)
 
 logger = init_logger(__name__)
 
 
+class Cosmos25DistilledInputValidationStage(InputValidationStage):
+    """Reject conditioning and classic CFG unsupported by the released student."""
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        conditioning_inputs = (
+            batch.image_path,
+            batch.pil_image,
+            batch.preprocessed_image,
+            batch.video_path,
+            batch.video_latent,
+        )
+        if any(value is not None for value in conditioning_inputs):
+            raise ValueError("Cosmos Predict2.5 distilled currently supports text-to-world generation only")
+        if batch.do_classifier_free_guidance:
+            raise ValueError("Cosmos Predict2.5 distilled does not use classifier-free guidance; set guidance_scale=1")
+        if not 1 <= batch.num_inference_steps <= 4:
+            raise ValueError("Cosmos Predict2.5 distilled supports 1 to 4 inference steps")
+        return super().forward(batch, fastvideo_args)
+
+
 class Cosmos2_5Pipeline(ComposedPipelineBase):
     """Cosmos 2.5 video generation pipeline."""
 
@@ -19,7 +42,11 @@ class Cosmos2_5Pipeline(ComposedPipelineBase):
     def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
         logger.info("Creating Cosmos 2.5 pipeline stages...")
 
-        self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
+        scheduler = self.get_module("scheduler")
+        is_distilled = isinstance(scheduler, Cosmos25DistilledScheduler)
+
+        input_validation = Cosmos25DistilledInputValidationStage() if is_distilled else InputValidationStage()
+        self.add_stage(stage_name="input_validation_stage", stage=input_validation)
 
         self.add_stage(
             stage_name="prompt_encoding_stage",
@@ -29,16 +56,30 @@ def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
         self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
 
         self.add_stage(stage_name="timestep_preparation_stage",
-                       stage=Cosmos25TimestepPreparationStage(scheduler=self.get_module("scheduler")))
+                       stage=Cosmos25TimestepPreparationStage(scheduler=scheduler))
 
-        self.add_stage(stage_name="latent_preparation_stage",
-                       stage=Cosmos25AutoLatentPreparationStage(scheduler=self.get_module("scheduler"),
-                                                                transformer=self.get_module("transformer"),
-                                                                vae=self.get_module("vae")))
+        if is_distilled:
+            latent_stage = Cosmos25DistilledT2WLatentPreparationStage(
+                scheduler=scheduler,
+                transformer=self.get_module("transformer"),
+            )
+            denoising_stage = Cosmos25DistilledT2WDenoisingStage(
+                transformer=self.get_module("transformer"),
+                scheduler=scheduler,
+            )
+        else:
+            latent_stage = Cosmos25AutoLatentPreparationStage(
+                scheduler=scheduler,
+                transformer=self.get_module("transformer"),
+                vae=self.get_module("vae"),
+            )
+            denoising_stage = Cosmos25AutoDenoisingStage(
+                transformer=self.get_module("transformer"),
+                scheduler=scheduler,
+            )
 
-        self.add_stage(stage_name="denoising_stage",
-                       stage=Cosmos25AutoDenoisingStage(transformer=self.get_module("transformer"),
-                                                        scheduler=self.get_module("scheduler")))
+        self.add_stage(stage_name="latent_preparation_stage", stage=latent_stage)
+        self.add_stage(stage_name="denoising_stage", stage=denoising_stage)
 
         self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")))
         logger.info("Cosmos 2.5 pipeline stages created")
diff --git a/fastvideo/pipelines/stages/__init__.py b/fastvideo/pipelines/stages/__init__.py
index ae4b701f6d..bd484d3c55 100644
--- a/fastvideo/pipelines/stages/__init__.py
+++ b/fastvideo/pipelines/stages/__init__.py
@@ -11,8 +11,9 @@
 from fastvideo.pipelines.stages.conditioning import ConditioningStage
 from fastvideo.pipelines.stages.decoding import DecodingStage
 from fastvideo.pipelines.stages.denoising import (Cosmos25AutoDenoisingStage, Cosmos25DenoisingStage,
-                                                  Cosmos25V2WDenoisingStage, Cosmos25T2WDenoisingStage,
-                                                  CosmosDenoisingStage, DenoisingStage, DmdDenoisingStage)
+                                                  Cosmos25DistilledT2WDenoisingStage, Cosmos25V2WDenoisingStage,
+                                                  Cosmos25T2WDenoisingStage, CosmosDenoisingStage, DenoisingStage,
+                                                  DmdDenoisingStage)
 from fastvideo.pipelines.stages.sr_denoising import SRDenoisingStage
 from fastvideo.pipelines.stages.encoding import EncodingStage
 from fastvideo.pipelines.stages.image_encoding import (ImageEncodingStage, MatrixGame2ImageEncodingStage,
@@ -24,6 +25,7 @@
 from fastvideo.pipelines.stages.input_validation import InputValidationStage
 from fastvideo.pipelines.stages.latent_preparation import (Cosmos25LatentPreparationStage, CosmosLatentPreparationStage,
                                                            Cosmos25AutoLatentPreparationStage,
+                                                           Cosmos25DistilledT2WLatentPreparationStage,
                                                            Cosmos25T2WLatentPreparationStage,
                                                            Cosmos25V2WLatentPreparationStage, LatentPreparationStage)
 from fastvideo.pipelines.basic.ltx2.stages import (  # noqa: F401
@@ -57,6 +59,7 @@
     "Cosmos25T2WLatentPreparationStage",
     "Cosmos25V2WLatentPreparationStage",
     "Cosmos25AutoLatentPreparationStage",
+    "Cosmos25DistilledT2WLatentPreparationStage",
     "LTX2LatentPreparationStage",
     "LTX2AudioDecodingStage",
     "ConditioningStage",
@@ -80,6 +83,7 @@
     "Cosmos25T2WDenoisingStage",
     "Cosmos25V2WDenoisingStage",
     "Cosmos25AutoDenoisingStage",
+    "Cosmos25DistilledT2WDenoisingStage",
     "LTX2DenoisingStage",
     "LTX2TextEncodingStage",
     "SRDenoisingStage",
diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py
index 771bc3c323..8079b5a950 100644
--- a/fastvideo/pipelines/stages/denoising.py
+++ b/fastvideo/pipelines/stages/denoising.py
@@ -1201,6 +1201,121 @@ def forward(
         return super().forward(batch, fastvideo_args)
 
 
+class Cosmos25DistilledT2WDenoisingStage(Cosmos25DenoisingStage):
+    """Four-step TrigFlow/x0 denoising for the released distilled student."""
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        pipeline = self.pipeline() if self.pipeline else None
+        if not fastvideo_args.model_loaded["transformer"]:
+            loader = TransformerLoader()
+            self.transformer = loader.load(fastvideo_args.model_paths["transformer"], fastvideo_args)
+            if pipeline:
+                pipeline.add_module("transformer", self.transformer)
+            fastvideo_args.model_loaded["transformer"] = True
+
+        latents = batch.latents
+        if latents is None:
+            raise ValueError("latents must be provided for Cosmos25DistilledT2WDenoisingStage")
+        if not batch.prompt_embeds:
+            raise ValueError("prompt_embeds must be provided for Cosmos25 distilled inference")
+
+        timesteps = batch.timesteps
+        if timesteps is None:
+            self.scheduler.set_timesteps(batch.num_inference_steps, device=latents.device)
+            timesteps = self.scheduler.timesteps
+        else:
+            timesteps = timesteps.to(latents.device)
+
+        target_dtype = torch.bfloat16
+        for parameter in self.transformer.parameters():
+            if parameter.dtype != torch.float32:
+                target_dtype = parameter.dtype
+                break
+        autocast_enabled = (latents.device.type == "cuda" and target_dtype != torch.float32
+                            and not fastvideo_args.disable_autocast)
+
+        batch_size, _, latent_frames, latent_height, latent_width = latents.shape
+        condition_mask = torch.zeros(
+            (batch_size, 1, latent_frames, latent_height, latent_width),
+            device=latents.device,
+            dtype=target_dtype,
+        )
+        padding_mask = torch.ones(
+            (batch_size, 1, latent_height, latent_width),
+            device=latents.device,
+            dtype=target_dtype,
+        )
+
+        if batch.fps is None:
+            fps_tensor = torch.full((batch_size, ), 16, device=latents.device, dtype=target_dtype)
+        else:
+            fps_tensor = torch.as_tensor(batch.fps, device=latents.device, dtype=target_dtype).reshape(-1)
+            if fps_tensor.numel() == 1:
+                fps_tensor = fps_tensor.repeat(batch_size)
+            elif fps_tensor.numel() != batch_size:
+                raise ValueError(f"fps must contain one value or one per sample; got {fps_tensor.numel()}")
+
+        state = latents.to(torch.float32)
+        with self.progress_bar(total=len(timesteps)) as progress_bar:
+            for timestep_value in timesteps:
+                # The official loop truncates the evolving FP64 state to FP32
+                # before each x0 prediction, then promotes the result back to
+                # FP64 for the fixed-noise update.
+                state_fp32 = state.float()
+                model_input = self.scheduler.scale_model_input(state_fp32, timestep_value)
+                model_timestep = torch.full(
+                    (batch_size, latent_frames),
+                    float(timestep_value),
+                    device=latents.device,
+                    dtype=target_dtype,
+                )
+                context_timestep = int(round(float(timestep_value) * 1000))
+                with (
+                        set_forward_context(
+                            current_timestep=context_timestep,
+                            attn_metadata=None,
+                            forward_batch=batch,
+                        ),
+                        torch.autocast(
+                            device_type=latents.device.type,
+                            dtype=target_dtype,
+                            enabled=autocast_enabled,
+                        ),
+                ):
+                    model_output = self.transformer(
+                        hidden_states=model_input.to(target_dtype),
+                        encoder_hidden_states=batch.prompt_embeds[0].to(target_dtype),
+                        timestep=model_timestep,
+                        fps=fps_tensor,
+                        condition_mask=condition_mask,
+                        padding_mask=padding_mask,
+                        return_dict=False,
+                    )
+                    if isinstance(model_output, tuple | list):
+                        model_output = model_output[0]
+
+                if model_output.shape != state.shape:
+                    raise ValueError(
+                        f"Cosmos25 distilled DiT returned shape {model_output.shape}; expected {state.shape}")
+
+                state = self.scheduler.step(
+                    model_output.float(),
+                    timestep_value,
+                    state_fp32,
+                    generator=batch.generator,
+                    return_dict=False,
+                )[0]
+                progress_bar.update()
+
+        # The official implementation returns a finite FP32 x0 for VAE decode.
+        batch.latents = torch.nan_to_num(state.float())
+        return batch
+
+
 class Cosmos25V2WDenoisingStage(Cosmos25DenoisingStage):
     """Cosmos 2.5 Video2World denoising stage."""
 
diff --git a/fastvideo/pipelines/stages/latent_preparation.py b/fastvideo/pipelines/stages/latent_preparation.py
index 5f60409ded..268274e777 100644
--- a/fastvideo/pipelines/stages/latent_preparation.py
+++ b/fastvideo/pipelines/stages/latent_preparation.py
@@ -658,6 +658,67 @@ def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> V
         return Cosmos25LatentPreparationStage.verify_output(self, batch, fastvideo_args)  # type: ignore[misc]
 
 
+class Cosmos25DistilledT2WLatentPreparationStage(Cosmos25T2WLatentPreparationStage):
+    """Official FP32 initial-noise preparation for the distilled T2W sampler."""
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        if 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
+
+        device = get_local_torch_device()
+        height = batch.height
+        width = batch.width
+        if height is None or width is None:
+            raise ValueError("Height and width must be provided")
+
+        num_latent_frames = (batch.num_frames - 1) // 4 + 1
+        shape = (
+            batch_size,
+            self.transformer.config.in_channels,
+            num_latent_frames,
+            height // 8,
+            width // 8,
+        )
+
+        if batch.latents is None:
+            # NVIDIA's released distilled inference creates FP32 noise on the
+            # execution device, then carries the sampler state in FP64.
+            seeds = batch.seeds or [int(batch.seed if batch.seed is not None else 0) + i for i in range(batch_size)]
+            if len(seeds) != batch_size:
+                seeds = [int(batch.seed if batch.seed is not None else 0) + i for i in range(batch_size)]
+            samples = []
+            for seed in seeds:
+                generator = torch.Generator(device=device).manual_seed(seed)
+                samples.append(torch.randn(
+                    shape[1:],
+                    generator=generator,
+                    device=device,
+                    dtype=torch.float32,
+                ))
+            latents = torch.stack(samples)
+        else:
+            latents = batch.latents.to(device=device, dtype=torch.float32)
+
+        batch.latents = latents
+        batch.raw_latent_shape = latents.shape
+        batch.conditioning_latents = None
+        batch.cond_indicator = None
+        batch.uncond_indicator = None
+        batch.cond_mask = None
+        batch.uncond_mask = None
+        batch.padding_mask = None
+        return batch
+
+
 class Cosmos25V2WLatentPreparationStage(Cosmos25LatentPreparationStage):
     """Cosmos 2.5 V2W/I2W latent preparation stage (conditioning-aware)."""
 
diff --git a/tests/local_tests/cosmos25/PORT_STATUS.md b/tests/local_tests/cosmos25/PORT_STATUS.md
index 01ff97ad16..5c633904fb 100644
--- a/tests/local_tests/cosmos25/PORT_STATUS.md
+++ b/tests/local_tests/cosmos25/PORT_STATUS.md
@@ -6,32 +6,32 @@
 - official_ref: NVIDIA/Cosmos-Predict2.5@a2c298b0a3df3778b973fe65e9e58877b292d8a7
 - official_ref_dir: `${COSMOS25_OFFICIAL_REF_DIR:-$PWD/cosmos-predict2.5}`
 - hf_weights_path: `nvidia/Cosmos-Predict2.5-2B/base/distilled`
-- local_weights_dir: not created
-- source_layout: official monolithic checkpoint; FastVideo conversion/loading path pending
+- local_weights_dir: `~/models/Cosmos-Predict2.5-2B-Distilled-TrigFlow-FastVideo` (Spark validation host)
+- source_layout: official monolithic checkpoint converted to a clean FastVideo-loadable package
 - local_tests_readme: `tests/local_tests/cosmos25/README.md`
 
 ## Current Phase
-- phase: component parity
+- phase: pipeline integration
 - status: in_progress
-- owner: parity
+- owner: pipeline
 - last_updated: 2026-08-27
 
 ## Component Matrix
 | Component | Type | Reuse/Port | Official Definition | Official Instantiation | FastVideo Target | Prototype | Conversion | Parity | Open Issues |
 |---|---|---|---|---|---|---|---|---|---|
 | TrigFlow sampler | scheduler | port | `modules/denoiser_scaling.py`; `distill/models/video2world_model_distill_dmd2.py` | `generate_samples_from_batch` | `Cosmos25DistilledScheduler` | complete | n/a | non-skip pass | none |
-| student DiT | transformer | reuse Cosmos25 architecture with distilled weights | `MinimalV1LVGDiT` through distillation model | `get_x0_fn_from_batch` / `denoise_edm` | `Cosmos25Transformer3DModel` | existing | pending | pending real-weight forward | I001 |
-| Reason1 encoder | text encoder | reuse | Predict2.5 Video2World config | distilled inference CLI | existing Cosmos25 encoder | existing | pending packaged layout | pending production-loader parity | I002 |
-| tokenizer VAE | VAE | reuse | Predict2.5 tokenizer | distilled inference CLI | existing Cosmos25 VAE | existing | pending packaged layout | pending production-loader parity | I002 |
-| T2W pipeline | pipeline | extend after components pass | `generate_samples_from_batch` | distilled inference CLI | Cosmos2_5 staged pipeline | not started | blocked by component gates | not started | I003 |
+| student DiT | transformer | reuse Cosmos25 architecture with distilled weights | `MinimalV1LVGDiT` through distillation model | `get_x0_fn_from_batch` / `denoise_edm` | `Cosmos25Transformer3DModel` | existing | complete | non-skip BF16 pass | none |
+| Reason1 encoder | text encoder | reuse | Predict2.5 Video2World config | distilled inference CLI | existing Cosmos25 encoder | existing | packaged passthrough | production loader pass | none |
+| tokenizer VAE | VAE | reuse | Predict2.5 tokenizer | distilled inference CLI | existing Cosmos25 VAE | existing | packaged passthrough | production loader pass | none |
+| T2W pipeline | pipeline | isolated scheduler-selected route | `generate_samples_from_batch` | distilled inference CLI | Cosmos2_5 staged pipeline | in progress | complete | pending end-to-end latent/video smoke | I003 |
 
 ## Conversion State
 - conversion_script: `scripts/checkpoint_conversion/cosmos25_distilled_to_diffusers.py`
-- converted_weights_dir: not created
+- converted_weights_dir: `~/models/Cosmos-Predict2.5-2B-Distilled-TrigFlow-FastVideo`
 - source_layout: official `base/distilled` checkpoint
-- strict_load_status: not run
+- strict_load_status: pass; 685 student tensors, no training counters, production FastVideo loader pass
 - passthrough_components: Reason1 encoder and tokenizer VAE are expected to reuse the existing Cosmos25 layout
-- retry_history: synthetic direct/nested checkpoint extraction and output-layout contracts pass; released checkpoint not run
+- retry_history: synthetic contracts pass; released 3.9 GB checkpoint converted in 22.5 s to a 20 GB package
 
 ## Parity Commands
 | Scope | Command | Last Result | Notes |
@@ -39,8 +39,9 @@
 | scheduler unit | `pytest fastvideo/tests/schedulers/test_cosmos25_distilled_scheduler.py -q` | 7 passed | CPU-only; includes registry resolution; 2026-08-27 |
 | official scheduler | `COSMOS25_OFFICIAL_REF_DIR=/path/to/Cosmos-Predict2.5 pytest tests/local_tests/cosmos25/test_cosmos25_distilled_scheduler_parity.py -v -s` | 2 passed, non-skip | CPU-only; pinned source; 2026-08-27 |
 | conversion contracts | `pytest tests/local_tests/cosmos25/test_cosmos25_distilled_conversion.py -q` | 7 passed | Synthetic checkpoints/layout; 2026-08-27 |
-| student DiT | `COSMOS25_OFFICIAL_REF_DIR=/path/to/Cosmos-Predict2.5 COSMOS25_DISTILLED_CHECKPOINT=/path/to/distilled.pt pytest tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py -v -s` | scaffold pending Spark run | requires CUDA, official source dependencies, and released weights |
-| pipeline | not yet created | not run | forbidden until component parity passes |
+| student DiT | `COSMOS25_OFFICIAL_REF_DIR=/path/to/Cosmos-Predict2.5 COSMOS25_DISTILLED_CHECKPOINT=/path/to/distilled.pt pytest tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py -v -s` | passed, non-skip | Spark BF16: first-block relative mean 0.000655; final relative mean 0.038397; 2026-08-27 |
+| pipeline contracts | `pytest tests/local_tests/cosmos25/test_cosmos25_distilled_pipeline.py -q` | pending | CPU-only sampler/stage isolation checks |
+| pipeline smoke | pending command | not run | requires converted package and CUDA |
 
 ## Open Questions
 | ID | Question | Owner | Needed By Phase | Status | Resolution |
@@ -51,9 +52,9 @@
 ## Issues And Blockers
 | ID | Phase | Component | Severity | Issue | Evidence | Owner | Status | Resolution |
 |---|---|---|---|---|---|---|---|---|
-| I001 | parity | student DiT | high | No non-skip real-weight distilled forward comparison yet | Existing Spark runs used distilled weights with the base UniPC inference path | parity | open | pending CUDA run |
-| I002 | conversion | packaged model | high | Released official checkpoint is not yet isolated in a FastVideo-loadable component layout | No `local_weights_dir` or strict-load record | conversion | open | pending |
-| I003 | pipeline | T2W | high | Pipeline wiring is gated on component parity | add-model pipeline contract | pipeline | open | pending I001 and I002 |
+| I001 | parity | student DiT | high | No non-skip real-weight distilled forward comparison yet | Spark official-vs-FastVideo BF16 comparison | parity | closed | passed at final relative mean 0.038397 |
+| I002 | conversion | packaged model | high | Released official checkpoint is not yet isolated in a FastVideo-loadable component layout | Converted package and strict production load | conversion | closed | 685 clean student tensors; load pass |
+| I003 | pipeline | T2W | high | End-to-end distilled generation is not yet validated | component gates are now complete | pipeline | open | isolated route implemented; run latent/video smoke |
 
 ## Escape Hatches
 | ID | Phase | Decision Type | Question | Recommended Option | Status | Resolution |
@@ -67,8 +68,11 @@
 | 2026-08-27 | Keep the existing full-step Cosmos25 path unchanged | Distilled and post-trained checkpoints require different inference semantics | Avoids regression for current users |
 | 2026-08-27 | Defer pipeline wiring until real-weight component parity | Required by the repository add-model workflow | Next GPU task is DiT parity, not generation |
 | 2026-08-27 | Preserve native `net.*` student keys during conversion | Existing Cosmos25 loader owns the authoritative mapping | Converter only isolates student tensors and emits package metadata |
+| 2026-08-27 | Accept calibrated BF16 DiT parity | Preprocess is exact, first-block drift is 0.000655 relative, and drift grows smoothly to 0.038397 final relative | Clears the component gate without claiming bitwise equality |
+| 2026-08-27 | Select distilled stages from the packaged scheduler class | The package already carries authoritative inference semantics | Existing full Cosmos2.5 packages remain on their unchanged path |
 
 ## Handoff Notes
 - CPU scheduler unit and pinned-reference parity tests pass locally without skips.
-- Next implementation is isolated distilled-weight conversion/loading plus a real student DiT parity test.
+- Released checkpoint conversion, production strict load, and official-vs-FastVideo DiT parity pass on Spark.
+- Next gate is one isolated end-to-end T2W latent/video smoke from the converted package.
 - Do not use the prior FlowUniPC/Karras Spark run as distilled parity evidence.
diff --git a/tests/local_tests/cosmos25/README.md b/tests/local_tests/cosmos25/README.md
index 80f95ccd1b..7d8ea22209 100644
--- a/tests/local_tests/cosmos25/README.md
+++ b/tests/local_tests/cosmos25/README.md
@@ -31,7 +31,7 @@ pytest tests/local_tests/cosmos25/test_cosmos25_distilled_scheduler_parity.py -v
 The parity test pins NVIDIA's scaling source and compares the full four-step
 preconditioning/x0/fixed-noise rollout. It does not load model weights.
 
-## Conversion scaffold
+## Conversion
 
 The converter keeps only the official student's native `net.*` tensors and
 reuses non-transformer components from an existing FastVideo-loadable Cosmos
@@ -51,22 +51,31 @@ Local conversion contracts:
 pytest tests/local_tests/cosmos25/test_cosmos25_distilled_conversion.py -q
 ```
 
-This scaffold is locally tested, but conversion of the released 4 GB checkpoint
-and production-loader strictness have not yet been validated.
+The released checkpoint conversion and production FastVideo strict load passed
+on the Spark validation host: 685 student tensors and no training counters.
 
-## Remaining GPU gates
+## Remaining GPU gate
 
-Before wiring a public pipeline or running DreamVerse:
+Conversion, strict load, and the real-weight DiT comparison pass. The remaining
+gate is an end-to-end T2W generation from the converted package. Distilled
+V2W/rolling remains explicitly outside the initial support claim.
 
-1. Run the converter on the released distilled transformer and verify its
-   reported tensor count and production-loader missing/unexpected keys.
-2. Compare one real student DiT forward against the official implementation.
-3. Compare deterministic T2W latents end to end for the official four-step
-   schedule.
-4. Only after T2W parity, evaluate experimental V2W/rolling conditioning.
+Run the cheap pipeline contracts, then a small wiring smoke before the full
+four-step quality gate:
 
-Those checks require the released checkpoint and a CUDA machine. A skipped
-local parity test is not pass evidence.
+```bash
+pytest tests/local_tests/cosmos25/test_cosmos25_distilled_pipeline.py -q
+
+FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA \
+python examples/inference/basic/basic_cosmos2_5_distilled_t2w.py \
+  --model /path/to/converted-model \
+  --steps 1 --frames 9 --height 256 --width 448 \
+  --output outputs_video/cosmos25_distilled_smoke.mp4
+
+FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA \
+python examples/inference/basic/basic_cosmos2_5_distilled_t2w.py \
+  --model /path/to/converted-model
+```
 
 ## Real student DiT parity
 
@@ -83,4 +92,5 @@ FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA \
 pytest tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py -v -s
 ```
 
-The test must report `PASSED`, not `SKIPPED`, before distilled pipeline wiring.
+The Spark gate passed with first-block relative mean error `0.000655` and final
+relative mean error `0.038397`, with smooth BF16 drift and no discontinuity.
diff --git a/tests/local_tests/cosmos25/test_cosmos25_distilled_pipeline.py b/tests/local_tests/cosmos25/test_cosmos25_distilled_pipeline.py
new file mode 100644
index 0000000000..c0a85b4ace
--- /dev/null
+++ b/tests/local_tests/cosmos25/test_cosmos25_distilled_pipeline.py
@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Isolated pipeline contracts for the Cosmos Predict2.5 distilled student."""
+
+from contextlib import nullcontext
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch.testing import assert_close
+
+from fastvideo.models.schedulers.scheduling_cosmos25_distilled import Cosmos25DistilledScheduler
+from fastvideo.pipelines.basic.cosmos.cosmos2_5_pipeline import Cosmos25DistilledInputValidationStage
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.denoising import Cosmos25DistilledT2WDenoisingStage
+from fastvideo.pipelines.stages.latent_preparation import Cosmos25DistilledT2WLatentPreparationStage
+
+
+class _Progress:
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *_args):
+        return None
+
+    def update(self) -> None:
+        pass
+
+
+class _RecordingTransformer(torch.nn.Module):
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.anchor = torch.nn.Parameter(torch.zeros((), dtype=torch.bfloat16))
+        self.config = SimpleNamespace(in_channels=2)
+        self.calls: list[dict[str, torch.Tensor]] = []
+
+    def forward(self, **kwargs):
+        self.calls.append({key: value.detach().clone() for key, value in kwargs.items() if torch.is_tensor(value)})
+        return torch.zeros_like(kwargs["hidden_states"])
+
+
+def _args() -> SimpleNamespace:
+    return SimpleNamespace(
+        disable_autocast=False,
+        model_loaded={"transformer": True},
+        model_paths={},
+    )
+
+
+def test_distilled_latent_preparation_preserves_official_fp32_noise(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(
+        "fastvideo.pipelines.stages.latent_preparation.get_local_torch_device",
+        lambda: torch.device("cpu"),
+    )
+    transformer = _RecordingTransformer()
+    stage = Cosmos25DistilledT2WLatentPreparationStage(Cosmos25DistilledScheduler(), transformer)
+    batch = ForwardBatch(
+        data_type="video",
+        prompt="test",
+        prompt_embeds=[torch.zeros(1, 2, 4)],
+        height=16,
+        width=24,
+        num_frames=5,
+        seed=7,
+        seeds=[7],
+    )
+
+    stage.forward(batch, _args())
+
+    assert batch.latents is not None
+    assert batch.latents.dtype is torch.float32
+    assert batch.latents.shape == (1, 2, 2, 2, 3)
+    expected = torch.randn((2, 2, 2, 3), generator=torch.Generator().manual_seed(7))
+    assert_close(batch.latents[0], expected, rtol=0, atol=0)
+
+
+def test_distilled_denoising_uses_per_frame_timesteps_and_official_rollout(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(
+        "fastvideo.pipelines.stages.denoising.set_forward_context",
+        lambda **_kwargs: nullcontext(),
+    )
+    scheduler = Cosmos25DistilledScheduler()
+    scheduler.set_timesteps(2)
+    transformer = _RecordingTransformer()
+    stage = Cosmos25DistilledT2WDenoisingStage.__new__(Cosmos25DistilledT2WDenoisingStage)
+    stage.transformer = transformer
+    stage.scheduler = scheduler
+    stage.pipeline = None
+    stage.progress_bar = lambda **_kwargs: _Progress()
+
+    initial = torch.arange(48, dtype=torch.float32).reshape(1, 2, 2, 3, 4) / 48
+    batch = ForwardBatch(
+        data_type="video",
+        prompt_embeds=[torch.zeros(1, 3, 4)],
+        latents=initial.clone(),
+        timesteps=scheduler.timesteps.clone(),
+        num_inference_steps=2,
+        fps=24,
+    )
+    stage.forward(batch, _args())
+
+    expected_scheduler = Cosmos25DistilledScheduler()
+    expected_scheduler.set_timesteps(2)
+    expected = initial.clone()
+    for timestep in expected_scheduler.timesteps:
+        expected_fp32 = expected.float()
+        expected_scheduler.scale_model_input(expected_fp32, timestep)
+        expected = expected_scheduler.step(torch.zeros_like(expected_fp32), timestep, expected_fp32).prev_sample
+
+    assert batch.latents is not None
+    assert batch.latents.dtype is torch.float32
+    assert_close(batch.latents, expected.float(), rtol=0, atol=0)
+    assert len(transformer.calls) == 2
+    assert transformer.calls[0]["timestep"].shape == (1, 2)
+    assert_close(transformer.calls[0]["timestep"].float(), torch.ones(1, 2))
+    assert_close(transformer.calls[0]["condition_mask"], torch.zeros(1, 1, 2, 3, 4))
+    assert_close(transformer.calls[0]["padding_mask"], torch.ones(1, 1, 3, 4))
+
+
+@pytest.mark.parametrize("field", ["image_path", "pil_image", "preprocessed_image", "video_path", "video_latent"])
+def test_distilled_validation_rejects_conditioning(field: str) -> None:
+    batch = ForwardBatch(data_type="video", prompt="test", height=16, width=16, seed=0)
+    setattr(batch, field, "provided")
+    with pytest.raises(ValueError, match="text-to-world"):
+        Cosmos25DistilledInputValidationStage().forward(batch, _args())
+
+
+def test_distilled_validation_rejects_classic_cfg() -> None:
+    batch = ForwardBatch(
+        data_type="video",
+        prompt="test",
+        negative_prompt="bad",
+        height=16,
+        width=16,
+        seed=0,
+        guidance_scale=2,
+    )
+    with pytest.raises(ValueError, match="guidance_scale=1"):
+        Cosmos25DistilledInputValidationStage().forward(batch, _args())
+
+
+def test_distilled_validation_rejects_non_distilled_step_count() -> None:
+    batch = ForwardBatch(
+        data_type="video",
+        prompt="test",
+        height=16,
+        width=16,
+        seed=0,
+        num_inference_steps=50,
+    )
+    with pytest.raises(ValueError, match="1 to 4"):
+        Cosmos25DistilledInputValidationStage().forward(batch, _args())

From b02390103a8b99bff9a8f4d45cf249c8186c8028 Mon Sep 17 00:00:00 2001
From: Raghav <ragg04@gmail.com>
Date: Thu, 27 Aug 2026 12:20:51 -0700
Subject: [PATCH 12/16] [Test] match Cosmos distilled mask dtype

---
 .../cosmos25/test_cosmos25_distilled_pipeline.py       | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/tests/local_tests/cosmos25/test_cosmos25_distilled_pipeline.py b/tests/local_tests/cosmos25/test_cosmos25_distilled_pipeline.py
index c0a85b4ace..e14342dedd 100644
--- a/tests/local_tests/cosmos25/test_cosmos25_distilled_pipeline.py
+++ b/tests/local_tests/cosmos25/test_cosmos25_distilled_pipeline.py
@@ -116,8 +116,14 @@ def test_distilled_denoising_uses_per_frame_timesteps_and_official_rollout(
     assert len(transformer.calls) == 2
     assert transformer.calls[0]["timestep"].shape == (1, 2)
     assert_close(transformer.calls[0]["timestep"].float(), torch.ones(1, 2))
-    assert_close(transformer.calls[0]["condition_mask"], torch.zeros(1, 1, 2, 3, 4))
-    assert_close(transformer.calls[0]["padding_mask"], torch.ones(1, 1, 3, 4))
+    assert_close(
+        transformer.calls[0]["condition_mask"],
+        torch.zeros(1, 1, 2, 3, 4, dtype=torch.bfloat16),
+    )
+    assert_close(
+        transformer.calls[0]["padding_mask"],
+        torch.ones(1, 1, 3, 4, dtype=torch.bfloat16),
+    )
 
 
 @pytest.mark.parametrize("field", ["image_path", "pil_image", "preprocessed_image", "video_path", "video_latent"])

From c5e56529d37ffd5d6d29d312a4eab8712af16298 Mon Sep 17 00:00:00 2001
From: Raghav <ragg04@gmail.com>
Date: Wed, 15 Jul 2026 14:13:14 -0700
Subject: [PATCH 13/16] [bugfix] reason1: handle BatchEncoding from
 apply_chat_template

Cosmos-Predict2.5 2B text encoding crashes on load because
`apply_chat_template(tokenize=True)` returns a transformers `BatchEncoding`
-- a `collections.UserDict`, i.e. a `Mapping` but NOT a `dict` -- which the
`isinstance(tokenizer_output, dict)` check missed, so the id extraction fell
through to `raise RuntimeError("Unexpected chat_template output type")`.

Match on `Mapping` instead (every `dict` is a `Mapping`, so existing
plain-dict and list outputs are unchanged) and lift the normalization into a
tested `_normalize_chat_template_ids` helper. Add weight-free parametrized
regression tests over every shape apply_chat_template can return.
---
 fastvideo/models/encoders/reason1.py          | 42 +++++++----
 .../encoders/test_reason1_chat_template.py    | 70 +++++++++++++++++++
 2 files changed, 99 insertions(+), 13 deletions(-)
 create mode 100644 fastvideo/tests/encoders/test_reason1_chat_template.py

diff --git a/fastvideo/models/encoders/reason1.py b/fastvideo/models/encoders/reason1.py
index 51f9f711c7..0b3e9fcbc2 100644
--- a/fastvideo/models/encoders/reason1.py
+++ b/fastvideo/models/encoders/reason1.py
@@ -3,7 +3,7 @@
 
 import os
 from dataclasses import dataclass
-from collections.abc import Iterable
+from collections.abc import Iterable, Mapping
 
 import torch
 from transformers import AutoProcessor
@@ -23,6 +23,33 @@
 logger = init_logger(__name__)
 
 
+def _normalize_chat_template_ids(tokenizer_output) -> list[int]:
+    """Normalize `apply_chat_template` output to a flat ``list[int]`` of token ids.
+
+    Depending on the transformers version and tokenizer, `apply_chat_template`
+    (with ``tokenize=True``) may return a ``list[int]``, a nested
+    ``list[list[int]]`` (batch dim), a tensor, or a mapping carrying an
+    ``"input_ids"`` field. transformers returns a ``BatchEncoding``, which
+    subclasses ``collections.UserDict`` -- a ``Mapping`` but **not** a ``dict``
+    -- so an ``isinstance(_, dict)`` check silently misses it and the caller
+    would otherwise raise. Match on ``Mapping`` instead: every ``dict`` is a
+    ``Mapping``, so plain-dict and list outputs behave exactly as before.
+    """
+    if isinstance(tokenizer_output, Mapping) and "input_ids" in tokenizer_output:
+        input_ids = tokenizer_output["input_ids"]
+    else:
+        input_ids = tokenizer_output
+    if hasattr(input_ids, "tolist"):
+        input_ids = input_ids.tolist()
+    if (isinstance(input_ids, list) and len(input_ids) == 1
+            and isinstance(input_ids[0], list)):
+        input_ids = input_ids[0]
+    if not isinstance(input_ids, list):
+        raise RuntimeError(
+            f"Unexpected chat_template output type: {type(tokenizer_output)}")
+    return input_ids
+
+
 @dataclass(frozen=True)
 class _WeightsSource:
     """Mimic `TextEncoderLoader.Source` (avoid import cycles)."""
@@ -258,18 +285,7 @@ def compute_text_embeddings(
                     add_generation_prompt=False,
                 )
 
-            if isinstance(tokenizer_output, dict) and "input_ids" in tokenizer_output:
-                input_ids = tokenizer_output["input_ids"]
-                if hasattr(input_ids, "tolist"):
-                    input_ids = input_ids.tolist()
-            else:
-                input_ids = tokenizer_output
-                if hasattr(input_ids, "tolist"):
-                    input_ids = input_ids.tolist()
-                if isinstance(input_ids, list) and len(input_ids) == 1 and isinstance(input_ids[0], list):
-                    input_ids = input_ids[0]
-                if not isinstance(input_ids, list):
-                    raise RuntimeError(f"Unexpected chat_template output type: {type(tokenizer_output)}")
+            input_ids = _normalize_chat_template_ids(tokenizer_output)
 
             if self.num_embedding_padding_tokens > len(input_ids):
                 pad_len = self.num_embedding_padding_tokens - len(input_ids)
diff --git a/fastvideo/tests/encoders/test_reason1_chat_template.py b/fastvideo/tests/encoders/test_reason1_chat_template.py
new file mode 100644
index 0000000000..d76523dbb6
--- /dev/null
+++ b/fastvideo/tests/encoders/test_reason1_chat_template.py
@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Weight-free regression tests for reason1's chat-template id normalization.
+
+Guards the fix for the Cosmos-Predict2.5 text-encoder crash: `apply_chat_template`
+returns a `BatchEncoding` (a `collections.UserDict`, i.e. a `Mapping` but NOT a
+`dict`), so an `isinstance(_, dict)` check missed it and raised. These cover every
+shape `apply_chat_template` can return, so a future transformers change can't
+silently reintroduce the crash.
+"""
+
+from collections import UserDict
+
+import pytest
+
+from fastvideo.models.encoders.reason1 import _normalize_chat_template_ids
+
+IDS = [10, 11, 12]
+
+
+class _FakeTensor:
+    """Stand-in for a torch tensor: exposes ``.tolist()`` like real output does."""
+
+    def __init__(self, value):
+        self._value = value
+
+    def tolist(self):
+        return self._value
+
+
+class _BatchEncodingLike(UserDict):
+    """Mirrors ``transformers.BatchEncoding`` (``class BatchEncoding(UserDict)``):
+    a Mapping that is deliberately NOT a ``dict`` subclass."""
+
+
+@pytest.mark.parametrize(
+    "tokenizer_output",
+    [
+        IDS,                                  # list[int]  (legacy list path)
+        [IDS],                                # nested list[list[int]] (batch dim)
+        _FakeTensor(IDS),                     # tensor
+        _FakeTensor([IDS]),                   # batched tensor
+        {"input_ids": IDS},                   # plain dict, flat
+        {"input_ids": [IDS]},                 # plain dict, batched
+        _BatchEncodingLike({"input_ids": IDS}),            # BatchEncoding, flat
+        _BatchEncodingLike({"input_ids": [IDS]}),          # BatchEncoding, batched
+        _BatchEncodingLike({"input_ids": _FakeTensor([IDS])}),  # the crash case
+    ],
+)
+def test_normalizes_every_shape_to_flat_ids(tokenizer_output):
+    assert _normalize_chat_template_ids(tokenizer_output) == IDS
+
+
+def test_batch_encoding_is_not_a_dict_but_is_handled():
+    # Documents the root cause: the old `isinstance(_, dict)` guard was False here.
+    be = _BatchEncodingLike({"input_ids": IDS})
+    assert not isinstance(be, dict)
+    assert _normalize_chat_template_ids(be) == IDS
+
+
+def test_unexpected_type_raises():
+    with pytest.raises(RuntimeError, match="Unexpected chat_template output type"):
+        _normalize_chat_template_ids(object())
+
+
+def test_real_batch_encoding_if_available():
+    # Belt-and-suspenders: exercise the actual transformers BatchEncoding when present.
+    pytest.importorskip("transformers")
+    from transformers.tokenization_utils_base import BatchEncoding
+    assert not isinstance(BatchEncoding(), dict)  # the property the fix relies on
+    assert _normalize_chat_template_ids(BatchEncoding({"input_ids": IDS})) == IDS

From 84ab9de6ddd0f406853ac982b0b136bdb2807c35 Mon Sep 17 00:00:00 2001
From: Raghav <ragg04@gmail.com>
Date: Wed, 15 Jul 2026 14:57:56 -0700
Subject: [PATCH 14/16] reason1: fail loudly on batch>1 chat_template output

Address review: the helper documents a flat list[int] contract, but a nested
list with batch>1 slipped through the len==1 unwrap guard and was returned
as-is, violating the contract and breaking downstream padding. Raise on
batch>1 instead of silently returning a nested list; add tests.
---
 fastvideo/models/encoders/reason1.py              | 11 +++++++++--
 .../tests/encoders/test_reason1_chat_template.py  | 15 +++++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/fastvideo/models/encoders/reason1.py b/fastvideo/models/encoders/reason1.py
index 0b3e9fcbc2..a6cfabd512 100644
--- a/fastvideo/models/encoders/reason1.py
+++ b/fastvideo/models/encoders/reason1.py
@@ -41,8 +41,15 @@ def _normalize_chat_template_ids(tokenizer_output) -> list[int]:
         input_ids = tokenizer_output
     if hasattr(input_ids, "tolist"):
         input_ids = input_ids.tolist()
-    if (isinstance(input_ids, list) and len(input_ids) == 1
-            and isinstance(input_ids[0], list)):
+    if isinstance(input_ids, list) and input_ids and isinstance(input_ids[0], list):
+        # A nested list is a batch dimension. We tokenize one conversation at a
+        # time, so batch size 1 is the only shape we can unambiguously flatten;
+        # fail loudly on anything else rather than return a nested list that
+        # violates this helper's flat-``list[int]`` contract downstream.
+        if len(input_ids) != 1:
+            raise RuntimeError(
+                f"Unexpected batched chat_template output: batch={len(input_ids)} "
+                f"type={type(tokenizer_output)}")
         input_ids = input_ids[0]
     if not isinstance(input_ids, list):
         raise RuntimeError(
diff --git a/fastvideo/tests/encoders/test_reason1_chat_template.py b/fastvideo/tests/encoders/test_reason1_chat_template.py
index d76523dbb6..d332ce4af4 100644
--- a/fastvideo/tests/encoders/test_reason1_chat_template.py
+++ b/fastvideo/tests/encoders/test_reason1_chat_template.py
@@ -62,6 +62,21 @@ def test_unexpected_type_raises():
         _normalize_chat_template_ids(object())
 
 
+@pytest.mark.parametrize(
+    "batched",
+    [
+        [[10, 11], [12, 13]],                          # nested list, batch 2
+        {"input_ids": [[10, 11], [12, 13]]},           # mapping w/ batched ids
+        _BatchEncodingLike({"input_ids": _FakeTensor([[10, 11], [12, 13]])}),
+    ],
+)
+def test_batch_gt_one_raises_instead_of_returning_nested(batched):
+    # The helper only unwraps batch size 1; a real batch must fail loudly rather
+    # than silently return a nested list that breaks downstream padding.
+    with pytest.raises(RuntimeError, match="Unexpected batched chat_template output"):
+        _normalize_chat_template_ids(batched)
+
+
 def test_real_batch_encoding_if_available():
     # Belt-and-suspenders: exercise the actual transformers BatchEncoding when present.
     pytest.importorskip("transformers")

From dbf65bf3180655e1a9062e1fed907f0baaae60e5 Mon Sep 17 00:00:00 2001
From: Raghav <ragg04@gmail.com>
Date: Thu, 27 Aug 2026 12:34:49 -0700
Subject: [PATCH 15/16] [Test] add Cosmos distilled frame return gate

---
 .../basic/basic_cosmos2_5_distilled_t2w.py       | 16 ++++++++++++++--
 tests/local_tests/cosmos25/PORT_STATUS.md        | 16 ++++++++++------
 tests/local_tests/cosmos25/README.md             | 10 ++++++++++
 3 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/examples/inference/basic/basic_cosmos2_5_distilled_t2w.py b/examples/inference/basic/basic_cosmos2_5_distilled_t2w.py
index d066c4614a..8d9f231693 100644
--- a/examples/inference/basic/basic_cosmos2_5_distilled_t2w.py
+++ b/examples/inference/basic/basic_cosmos2_5_distilled_t2w.py
@@ -16,6 +16,11 @@ def main() -> None:
     parser.add_argument("--width", type=int, default=1280)
     parser.add_argument("--fps", type=int, default=16)
     parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument(
+        "--return-frames",
+        action="store_true",
+        help="Return decoded frames without writing an MP4 (DreamVerse contract smoke)",
+    )
     args = parser.parse_args()
 
     generator = VideoGenerator.from_pretrained(
@@ -41,12 +46,19 @@ def main() -> None:
         "Bright blue-white sparks scatter over the metal while smoke rises, "
         "cinematic lighting, steady camera, realistic motion."
     )
-    generator.generate_video(
+    result = generator.generate_video(
         prompt,
         sampling_param=sampling,
         output_path=args.output,
-        save_video=True,
+        save_video=not args.return_frames,
+        return_frames=args.return_frames,
     )
+    if args.return_frames:
+        frames = result.get("frames") if isinstance(result, dict) else None
+        if not isinstance(frames, list) or not frames:
+            raise RuntimeError("DreamVerse contract failed: generation did not return a nonempty frames list")
+        first_shape = getattr(frames[0], "shape", None)
+        print(f"COSMOS25_DREAMVERSE_FRAMES: PASS count={len(frames)} first_shape={first_shape}")
     generator.shutdown()
 
 
diff --git a/tests/local_tests/cosmos25/PORT_STATUS.md b/tests/local_tests/cosmos25/PORT_STATUS.md
index 5c633904fb..2acd063d8c 100644
--- a/tests/local_tests/cosmos25/PORT_STATUS.md
+++ b/tests/local_tests/cosmos25/PORT_STATUS.md
@@ -11,7 +11,7 @@
 - local_tests_readme: `tests/local_tests/cosmos25/README.md`
 
 ## Current Phase
-- phase: pipeline integration
+- phase: DreamVerse frame-return contract
 - status: in_progress
 - owner: pipeline
 - last_updated: 2026-08-27
@@ -23,7 +23,7 @@
 | student DiT | transformer | reuse Cosmos25 architecture with distilled weights | `MinimalV1LVGDiT` through distillation model | `get_x0_fn_from_batch` / `denoise_edm` | `Cosmos25Transformer3DModel` | existing | complete | non-skip BF16 pass | none |
 | Reason1 encoder | text encoder | reuse | Predict2.5 Video2World config | distilled inference CLI | existing Cosmos25 encoder | existing | packaged passthrough | production loader pass | none |
 | tokenizer VAE | VAE | reuse | Predict2.5 tokenizer | distilled inference CLI | existing Cosmos25 VAE | existing | packaged passthrough | production loader pass | none |
-| T2W pipeline | pipeline | isolated scheduler-selected route | `generate_samples_from_batch` | distilled inference CLI | Cosmos2_5 staged pipeline | in progress | complete | pending end-to-end latent/video smoke | I003 |
+| T2W pipeline | pipeline | isolated scheduler-selected route | `generate_samples_from_batch` | distilled inference CLI | Cosmos2_5 staged pipeline | complete | complete | full-resolution video and eye gate pass | none |
 
 ## Conversion State
 - conversion_script: `scripts/checkpoint_conversion/cosmos25_distilled_to_diffusers.py`
@@ -40,8 +40,10 @@
 | official scheduler | `COSMOS25_OFFICIAL_REF_DIR=/path/to/Cosmos-Predict2.5 pytest tests/local_tests/cosmos25/test_cosmos25_distilled_scheduler_parity.py -v -s` | 2 passed, non-skip | CPU-only; pinned source; 2026-08-27 |
 | conversion contracts | `pytest tests/local_tests/cosmos25/test_cosmos25_distilled_conversion.py -q` | 7 passed | Synthetic checkpoints/layout; 2026-08-27 |
 | student DiT | `COSMOS25_OFFICIAL_REF_DIR=/path/to/Cosmos-Predict2.5 COSMOS25_DISTILLED_CHECKPOINT=/path/to/distilled.pt pytest tests/local_tests/cosmos25/test_cosmos25_distilled_transformer_parity.py -v -s` | passed, non-skip | Spark BF16: first-block relative mean 0.000655; final relative mean 0.038397; 2026-08-27 |
-| pipeline contracts | `pytest tests/local_tests/cosmos25/test_cosmos25_distilled_pipeline.py -q` | pending | CPU-only sampler/stage isolation checks |
-| pipeline smoke | pending command | not run | requires converted package and CUDA |
+| pipeline contracts | `pytest tests/local_tests/cosmos25/test_cosmos25_distilled_pipeline.py -q` | 9 passed | Spark; CPU-only sampler/stage isolation checks |
+| pipeline smoke | `python examples/inference/basic/basic_cosmos2_5_distilled_t2w.py --model /path/to/converted-model --steps 1 --frames 9 --height 256 --width 448` | passed | 2.19 s end-to-end after load; 2026-08-27 |
+| full T2W quality | `python examples/inference/basic/basic_cosmos2_5_distilled_t2w.py --model /path/to/converted-model` | passed + eye gate | 704x1280x77, 4 steps; 143.53 s end-to-end after load; visually coherent |
+| DreamVerse frames | example command with `--return-frames` | pending | requires nonempty decoded `frames` list without MP4 save |
 
 ## Open Questions
 | ID | Question | Owner | Needed By Phase | Status | Resolution |
@@ -54,7 +56,7 @@
 |---|---|---|---|---|---|---|---|---|
 | I001 | parity | student DiT | high | No non-skip real-weight distilled forward comparison yet | Spark official-vs-FastVideo BF16 comparison | parity | closed | passed at final relative mean 0.038397 |
 | I002 | conversion | packaged model | high | Released official checkpoint is not yet isolated in a FastVideo-loadable component layout | Converted package and strict production load | conversion | closed | 685 clean student tensors; load pass |
-| I003 | pipeline | T2W | high | End-to-end distilled generation is not yet validated | component gates are now complete | pipeline | open | isolated route implemented; run latent/video smoke |
+| I003 | pipeline | T2W | high | End-to-end distilled generation is not yet validated | small and full-resolution Spark runs plus visual inspection | pipeline | closed | full T2W quality gate passed |
 
 ## Escape Hatches
 | ID | Phase | Decision Type | Question | Recommended Option | Status | Resolution |
@@ -70,9 +72,11 @@
 | 2026-08-27 | Preserve native `net.*` student keys during conversion | Existing Cosmos25 loader owns the authoritative mapping | Converter only isolates student tensors and emits package metadata |
 | 2026-08-27 | Accept calibrated BF16 DiT parity | Preprocess is exact, first-block drift is 0.000655 relative, and drift grows smoothly to 0.038397 final relative | Clears the component gate without claiming bitwise equality |
 | 2026-08-27 | Select distilled stages from the packaged scheduler class | The package already carries authoritative inference semantics | Existing full Cosmos2.5 packages remain on their unchanged path |
+| 2026-08-27 | Accept the full-resolution T2W quality gate | The four-step 704x1280x77 run completed without runtime faults and passed visual inspection | Clears basic FastVideo T2W support; does not claim continuation or real-time latency |
 
 ## Handoff Notes
 - CPU scheduler unit and pinned-reference parity tests pass locally without skips.
 - Released checkpoint conversion, production strict load, and official-vs-FastVideo DiT parity pass on Spark.
-- Next gate is one isolated end-to-end T2W latent/video smoke from the converted package.
+- Small and full-resolution T2W generation pass on Spark; the full video passed visual inspection.
+- Next gate is DreamVerse's `save_video=False`, `return_frames=True` result contract.
 - Do not use the prior FlowUniPC/Karras Spark run as distilled parity evidence.
diff --git a/tests/local_tests/cosmos25/README.md b/tests/local_tests/cosmos25/README.md
index 7d8ea22209..9e9e5ca6d9 100644
--- a/tests/local_tests/cosmos25/README.md
+++ b/tests/local_tests/cosmos25/README.md
@@ -77,6 +77,16 @@ python examples/inference/basic/basic_cosmos2_5_distilled_t2w.py \
   --model /path/to/converted-model
 ```
 
+For the DreamVerse frame-return contract, rerun the small smoke without MP4
+output and require a nonempty decoded frame list:
+
+```bash
+FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA \
+python examples/inference/basic/basic_cosmos2_5_distilled_t2w.py \
+  --model /path/to/converted-model \
+  --steps 1 --frames 9 --height 256 --width 448 --return-frames
+```
+
 ## Real student DiT parity
 
 This gate loads the raw NVIDIA student into the official and FastVideo DiTs,

From 72f25e407f2076afa0c57b02f4a882f19f13c86b Mon Sep 17 00:00:00 2001
From: Raghav <ragg04@gmail.com>
Date: Thu, 27 Aug 2026 13:33:20 -0700
Subject: [PATCH 16/16] [Docs] record Cosmos 2.5 distilled validation

---
 docs/inference/support_matrix.md          |  6 ++++++
 tests/local_tests/cosmos25/PORT_STATUS.md | 10 ++++++----
 tests/local_tests/cosmos25/README.md      | 12 ++++++++----
 3 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/docs/inference/support_matrix.md b/docs/inference/support_matrix.md
index eef549aac4..e0aebf7976 100644
--- a/docs/inference/support_matrix.md
+++ b/docs/inference/support_matrix.md
@@ -99,6 +99,12 @@ under the generic T2V workload option in the registry.
 yet public. Follow the [MMAudio inference guide](https://github.com/hao-ai-lab/FastVideo/blob/main/fastvideo/pipelines/basic/mmaudio/README.md)
 to convert the official weights locally and set `MMAUDIO_MODEL_PATH`.
 
+**Note (Cosmos Predict2.5 distilled)**: the released 2B Text2World student is
+supported through local checkpoint conversion, but no public converted model ID
+is registered yet. Follow the validation guide in
+`tests/local_tests/cosmos25/README.md`, then pass the converted directory to
+`basic_cosmos2_5_distilled_t2w.py --model`.
+
 **Note (MiniMax H3)**: T2VA, FL2VA, and Ref2VA all generate video with stereo
 audio. Use the Ref2VA example when passing ordered image, video, or audio
 references.
diff --git a/tests/local_tests/cosmos25/PORT_STATUS.md b/tests/local_tests/cosmos25/PORT_STATUS.md
index 2acd063d8c..c61d77df32 100644
--- a/tests/local_tests/cosmos25/PORT_STATUS.md
+++ b/tests/local_tests/cosmos25/PORT_STATUS.md
@@ -11,8 +11,8 @@
 - local_tests_readme: `tests/local_tests/cosmos25/README.md`
 
 ## Current Phase
-- phase: DreamVerse frame-return contract
-- status: in_progress
+- phase: initial T2W support
+- status: complete
 - owner: pipeline
 - last_updated: 2026-08-27
 
@@ -43,7 +43,7 @@
 | pipeline contracts | `pytest tests/local_tests/cosmos25/test_cosmos25_distilled_pipeline.py -q` | 9 passed | Spark; CPU-only sampler/stage isolation checks |
 | pipeline smoke | `python examples/inference/basic/basic_cosmos2_5_distilled_t2w.py --model /path/to/converted-model --steps 1 --frames 9 --height 256 --width 448` | passed | 2.19 s end-to-end after load; 2026-08-27 |
 | full T2W quality | `python examples/inference/basic/basic_cosmos2_5_distilled_t2w.py --model /path/to/converted-model` | passed + eye gate | 704x1280x77, 4 steps; 143.53 s end-to-end after load; visually coherent |
-| DreamVerse frames | example command with `--return-frames` | pending | requires nonempty decoded `frames` list without MP4 save |
+| DreamVerse frames | example command with `--return-frames` | passed | 9 RGB frames; first shape `(256, 448, 3)`; 2026-08-27 |
 
 ## Open Questions
 | ID | Question | Owner | Needed By Phase | Status | Resolution |
@@ -73,10 +73,12 @@
 | 2026-08-27 | Accept calibrated BF16 DiT parity | Preprocess is exact, first-block drift is 0.000655 relative, and drift grows smoothly to 0.038397 final relative | Clears the component gate without claiming bitwise equality |
 | 2026-08-27 | Select distilled stages from the packaged scheduler class | The package already carries authoritative inference semantics | Existing full Cosmos2.5 packages remain on their unchanged path |
 | 2026-08-27 | Accept the full-resolution T2W quality gate | The four-step 704x1280x77 run completed without runtime faults and passed visual inspection | Clears basic FastVideo T2W support; does not claim continuation or real-time latency |
+| 2026-08-27 | Accept the decoded-frame return contract | The small Spark run returned 9 RGB frames with shape `(256, 448, 3)` without writing an MP4 | Clears the downstream frame-consumer contract without claiming DreamVerse integration |
 
 ## Handoff Notes
 - CPU scheduler unit and pinned-reference parity tests pass locally without skips.
 - Released checkpoint conversion, production strict load, and official-vs-FastVideo DiT parity pass on Spark.
 - Small and full-resolution T2W generation pass on Spark; the full video passed visual inspection.
-- Next gate is DreamVerse's `save_video=False`, `return_frames=True` result contract.
+- The `save_video=False`, `return_frames=True` result contract passes on Spark.
+- A public converted package ID remains open; until then, use the documented local conversion flow.
 - Do not use the prior FlowUniPC/Karras Spark run as distilled parity evidence.
diff --git a/tests/local_tests/cosmos25/README.md b/tests/local_tests/cosmos25/README.md
index 9e9e5ca6d9..fa621905b1 100644
--- a/tests/local_tests/cosmos25/README.md
+++ b/tests/local_tests/cosmos25/README.md
@@ -54,11 +54,11 @@ pytest tests/local_tests/cosmos25/test_cosmos25_distilled_conversion.py -q
 The released checkpoint conversion and production FastVideo strict load passed
 on the Spark validation host: 685 student tensors and no training counters.
 
-## Remaining GPU gate
+## Validated GPU gates
 
-Conversion, strict load, and the real-weight DiT comparison pass. The remaining
-gate is an end-to-end T2W generation from the converted package. Distilled
-V2W/rolling remains explicitly outside the initial support claim.
+Conversion, strict load, the real-weight DiT comparison, end-to-end T2W
+generation, and decoded-frame return all pass on the Spark validation host.
+Distilled V2W/rolling remains explicitly outside the initial support claim.
 
 Run the cheap pipeline contracts, then a small wiring smoke before the full
 four-step quality gate:
@@ -87,6 +87,10 @@ python examples/inference/basic/basic_cosmos2_5_distilled_t2w.py \
   --steps 1 --frames 9 --height 256 --width 448 --return-frames
 ```
 
+The Spark frame-return gate produced 9 RGB frames with shape `(256, 448, 3)`.
+The full four-step `704x1280x77` run completed in 143.53 seconds after model
+load and passed visual inspection.
+
 ## Real student DiT parity
 
 This gate loads the raw NVIDIA student into the official and FastVideo DiTs,