[new-model] Add Cosmos Predict2.5 2B distilled T2W inference

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.

examples/inference/basic/basic_cosmos2_5_distilled_t2w.py

@@ -0,0 +1,66 @@
+"""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)
+    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(
+        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."
+    )
+    result = generator.generate_video(
+        prompt,
+        sampling_param=sampling,
+        output_path=args.output,
+        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()
+
+
+if __name__ == "__main__":
+    main()

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,40 @@
 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 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(
+            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 +292,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)

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 = {

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