[new-model] Add Cosmos Predict2.5 distilled T2W and DFD V2W inference

.gitignore

@@ -142,3 +142,4 @@ fastvideo/tests/ssim/reference_videos/**
 *.nvimlog
 .nvimlog
 .python-version
+/DFDReference/

docs/inference/support_matrix.md

@@ -99,6 +99,17 @@ 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`.
+
+The public DFD Video2World student is also supported through local DCP
+conversion. It uses exactly one conditioning image, 4 steps, 81 frames at
+704x1280 and 24 FPS. Use `basic_cosmos2_5_dfd_i2w.py`; a public converted model
+ID is not registered yet.
+
 **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_dfd_i2w.py

@@ -0,0 +1,59 @@
+"""Run the public four-step Cosmos Predict2.5 DFD Video2World student."""
+
+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 DFD model directory")
+    parser.add_argument("--image", required=True, help="Conditioning image")
+    parser.add_argument("--prompt", required=True)
+    parser.add_argument("--output", default="outputs_video/cosmos2_5_dfd_i2w.mp4")
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument(
+        "--return-frames",
+        action="store_true",
+        help="Return decoded frames without writing an MP4",
+    )
+    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=4,
+        num_frames=81,
+        height=704,
+        width=1280,
+        fps=24,
+        seed=args.seed,
+        guidance_scale=1.0,
+    )
+    result = generator.generate_video(
+        args.prompt,
+        sampling_param=sampling,
+        image_path=args.image,
+        num_cond_frames=1,
+        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 len(frames) != 81:
+            raise RuntimeError("DFD frame-return contract failed: expected 81 decoded frames")
+        print(f"COSMOS25_DFD_FRAMES: PASS count={len(frames)}")
+    generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

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/qwen2_5_vl_custom.py

@@ -74,6 +74,27 @@
 
 logger = logging.get_logger(__name__)
 
+_SUPPORTED_ATTENTION_IMPLEMENTATIONS = {"flash_attention_2", "sdpa"}
+
+
+def _normalize_attention_implementation(config):
+    """Map HF attention defaults unsupported by this compact model to SDPA.
+
+    Recent Transformers releases may resolve the nested Qwen vision config to
+    ``eager`` even when the parent model explicitly requests FlashAttention 2.
+    This implementation intentionally provides only FlashAttention 2 and SDPA;
+    SDPA is the correct portable fallback for any other resolved value.
+    """
+    implementation = getattr(config, "_attn_implementation", None)
+    if implementation not in _SUPPORTED_ATTENTION_IMPLEMENTATIONS:
+        logger.warning_once(
+            "Qwen2.5-VL attention implementation %r is unsupported by the FastVideo "
+            "Reason1 encoder; using SDPA.",
+            implementation,
+        )
+        config._attn_implementation = "sdpa"
+    return config
+
 
 def _resolve_torch_dtype(dtype, default: torch.dtype = torch.float32) -> torch.dtype:
     if isinstance(dtype, torch.dtype):
@@ -329,6 +350,7 @@ class Qwen2_5_VisionTransformerPretrainedModel(nn.Module):
 
     def __init__(self, config, parent_torch_dtype=None) -> None:
         super().__init__()
+        config = _normalize_attention_implementation(config)
 
         config_torch_dtype = getattr(config, "torch_dtype", None)
         self.dtype = _resolve_torch_dtype(config_torch_dtype if config_torch_dtype is not None else parent_torch_dtype)
@@ -1453,6 +1475,8 @@ class Qwen2_5_VLForConditionalGenerationSimple(nn.Module):
     def __init__(self, config):
         super().__init__()
         config = _flatten_text_config(config)
+        config = _normalize_attention_implementation(config)
+        config.vision_config = _normalize_attention_implementation(config.vision_config)
         self.config = config
         self.visual = Qwen2_5_VisionTransformerPretrainedModel(
             config.vision_config,

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,16 @@
     "SelfForcingFlowMatchScheduler":
     ("schedulers", "scheduling_self_forcing_flow_match", "SelfForcingFlowMatchScheduler"),
     "RCMScheduler": ("schedulers", "scheduling_rcm", "RCMScheduler"),
+    "Cosmos25DistilledScheduler": (
+        "schedulers",
+        "scheduling_cosmos25_distilled",
+        "Cosmos25DistilledScheduler",
+    ),
+    "Cosmos25DFDScheduler": (
+        "schedulers",
+        "scheduling_cosmos25_dfd",
+        "Cosmos25DFDScheduler",
+    ),
 }
 
 _UPSAMPLERS = {