[feat] Introduce Cosmos 2.5 Text2World pipeline (#974)

examples/inference/basic/basic_cosmos2_5_t2w.py

@@ -0,0 +1,42 @@
+from fastvideo import VideoGenerator
+
+
+def main():
+    # Point this to your local diffusers model dir (or replace with a HF model ID).
+    model_path = "KyleShao/Cosmos-Predict2.5-2B-Diffusers"
+
+    generator = VideoGenerator.from_pretrained(
+        model_path,
+        num_gpus=1,
+        use_fsdp_inference=False,  # set True if GPU is out of memory
+        dit_cpu_offload=False,
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=True,
+    )
+
+    prompt = (
+       "A high-definition video captures the precision of robotic welding in an industrial setting. The first frame showcases a robotic arm, equipped with a welding torch, positioned over a large metal structure. The welding process is in full swing, with bright sparks and intense light illuminating the scene, creating a vivid display of blue and white hues. A significant amount of smoke billows around the welding area, partially obscuring the view but emphasizing the heat and activity. The background reveals parts of the workshop environment, including a ventilation system and various pieces of machinery, indicating a busy and functional industrial workspace. As the video progresses, the robotic arm maintains its steady position, continuing the welding process and moving to its left. The welding torch consistently emits sparks and light, and the smoke continues to rise, diffusing slightly as it moves upward. The metal surface beneath the torch shows ongoing signs of heating and melting. The scene retains its industrial ambiance, with the welding sparks and smoke dominating the visual field, underscoring the ongoing nature of the welding operation."
+    )
+
+    video = generator.generate_video(
+        prompt,
+        negative_prompt="",
+        height=704,
+        width=1280,
+        num_frames=77,
+        num_inference_steps=35,
+        guidance_scale=7.0,
+        fps=24,
+        output_path="outputs_video/cosmos2_5_t2w.mp4",
+        save_video=True,
+    )
+
+    generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()
+
+
+

fastvideo/configs/models/encoders/__init__.py

@@ -7,10 +7,11 @@ from fastvideo.configs.models.encoders.clip import (
 from fastvideo.configs.models.encoders.llama import LlamaConfig
 from fastvideo.configs.models.encoders.t5 import T5Config, T5LargeConfig
 from fastvideo.configs.models.encoders.qwen2_5 import Qwen2_5_VLConfig
+from fastvideo.configs.models.encoders.reason1 import Reason1ArchConfig, Reason1Config
 
 __all__ = [
     "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig",
     "BaseEncoderOutput", "CLIPTextConfig", "CLIPVisionConfig",
     "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig",
-    "Qwen2_5_VLConfig"
+    "Qwen2_5_VLConfig", "Reason1ArchConfig", "Reason1Config"
 ]

fastvideo/configs/models/encoders/reason1.py

@@ -0,0 +1,72 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Config for Reason1 (Qwen2.5-VL) text encoder."""
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from fastvideo.configs.models.encoders.base import TextEncoderArchConfig, TextEncoderConfig
+
+
+@dataclass
+class Reason1ArchConfig(TextEncoderArchConfig):
+    """Architecture settings (defaults match Qwen2.5-VL-7B-Instruct)."""
+
+    architectures: list[str] = field(
+        default_factory=lambda: ["Qwen2_5_VLForConditionalGeneration"])
+    model_type: str = "qwen2_5_vl"
+
+    vocab_size: int = 152064
+    hidden_size: int = 3584
+    num_hidden_layers: int = 28
+    num_attention_heads: int = 28
+    num_key_value_heads: int = 4
+    intermediate_size: int = 18944
+
+    text_len: int = 512
+    hidden_state_skip_layer: int = 0
+    bos_token_id: int = 151643
+    pad_token_id: int = 151643
+    eos_token_id: int = 151645
+
+    image_token_id: int = 151655
+    video_token_id: int = 151656
+    vision_token_id: int = 151654
+    vision_start_token_id: int = 151652
+    vision_end_token_id: int = 151653
+
+    vision_config: dict[str, Any] | None = None
+
+    rope_theta: float = 1000000.0
+    rope_scaling: dict[str, Any] | None = field(default_factory=lambda: {
+        "type": "mrope",
+        "mrope_section": [16, 24, 24]
+    })
+    max_position_embeddings: int = 128000
+    max_window_layers: int = 28
+
+    embedding_concat_strategy: str = "mean_pooling"
+    n_layers_per_group: int = 5
+    num_embedding_padding_tokens: int = 512
+
+    attention_dropout: float = 0.0
+    hidden_act: str = "silu"
+    initializer_range: float = 0.02
+    rms_norm_eps: float = 1e-6
+
+    use_sliding_window: bool = False
+    sliding_window: int = 32768
+
+    tie_word_embeddings: bool = False
+    use_cache: bool = False
+    output_hidden_states: bool = True
+
+    torch_dtype: str = "bfloat16"
+    _attn_implementation: str = "flash_attention_2"
+
+
+@dataclass
+class Reason1Config(TextEncoderConfig):
+    """Reason1 text encoder config."""
+
+    arch_config: Reason1ArchConfig = field(default_factory=Reason1ArchConfig)
+    tokenizer_type: str = "Qwen/Qwen2.5-VL-7B-Instruct"

fastvideo/configs/models/vaes/__init__.py

@@ -1,4 +1,5 @@
 from fastvideo.configs.models.vaes.cosmosvae import CosmosVAEConfig
+from fastvideo.configs.models.vaes.cosmos2_5vae import Cosmos25VAEConfig
 from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
 from fastvideo.configs.models.vaes.hunyuan15vae import Hunyuan15VAEConfig
 from fastvideo.configs.models.vaes.stepvideovae import StepVideoVAEConfig
@@ -9,5 +10,6 @@ __all__ = [
     "WanVAEConfig",
     "StepVideoVAEConfig",
     "CosmosVAEConfig",
+    "Cosmos25VAEConfig",
     "Hunyuan15VAEConfig",
 ]

fastvideo/configs/models/vaes/cosmos2_5vae.py

@@ -0,0 +1,223 @@
+"""Cosmos 2.5 (Wan2.1-style) VAE config and checkpoint-key mapping."""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+
+import torch
+
+from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig
+
+
+@dataclass
+class Cosmos25VAEArchConfig(VAEArchConfig):
+    _name_or_path: str = ""
+    base_dim: int = 96
+    decoder_base_dim: int | None = None
+    z_dim: int = 16
+    dim_mult: tuple[int, ...] = (1, 2, 4, 4)
+    num_res_blocks: int = 2
+    attn_scales: tuple[float, ...] = ()
+    temperal_downsample: tuple[bool, ...] = (False, True, True)
+    dropout: float = 0.0
+    is_residual: bool = False
+    in_channels: int = 3
+    out_channels: int = 3
+    patch_size: int | None = None
+    scale_factor_temporal: int = 4
+    scale_factor_spatial: int = 8
+    clip_output: bool = True
+
+    latents_mean: tuple[float, ...] = (
+        -0.7571,
+        -0.7089,
+        -0.9113,
+        0.1075,
+        -0.1745,
+        0.9653,
+        -0.1517,
+        1.5508,
+        0.4134,
+        -0.0715,
+        0.5517,
+        -0.3632,
+        -0.1922,
+        -0.9497,
+        0.2503,
+        -0.2921,
+    )
+    latents_std: tuple[float, ...] = (
+        2.8184,
+        1.4541,
+        2.3275,
+        2.6558,
+        1.2196,
+        1.7708,
+        2.6052,
+        2.0743,
+        3.2687,
+        2.1526,
+        2.8652,
+        1.5579,
+        1.6382,
+        1.1253,
+        2.8251,
+        1.9160,
+    )
+
+    # Simple 1:1 renames. More complex decoder remapping is handled by
+    # `map_official_key()`.
+    param_names_mapping: dict[str, str] = field(
+        default_factory=lambda: {
+            r"^conv1\.(.*)$": r"quant_conv.\1",
+            r"^conv2\.(.*)$": r"post_quant_conv.\1",
+            r"^encoder\.conv1\.(.*)$": r"encoder.conv_in.\1",
+            r"^decoder\.conv1\.(.*)$": r"decoder.conv_in.\1",
+            r"^encoder\.head\.0\.gamma$": r"encoder.norm_out.gamma",
+            r"^encoder\.head\.2\.(.*)$": r"encoder.conv_out.\1",
+            r"^decoder\.head\.0\.gamma$": r"decoder.norm_out.gamma",
+            r"^decoder\.head\.2\.(.*)$": r"decoder.conv_out.\1",
+        })
+
+    @staticmethod
+    def map_official_key(key: str) -> str | None:
+        """Map a single official checkpoint key into FastVideo key space."""
+
+        def map_residual_subkey(prefix: str, sub: str) -> str | None:
+            if re.match(r"^residual\.0\.gamma$", sub):
+                return f"{prefix}.norm1.gamma"
+            m = re.match(r"^residual\.2\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.conv1.{m.group(1)}"
+            if re.match(r"^residual\.3\.gamma$", sub):
+                return f"{prefix}.norm2.gamma"
+            m = re.match(r"^residual\.6\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.conv2.{m.group(1)}"
+            m = re.match(r"^shortcut\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.conv_shortcut.{m.group(1)}"
+            return None
+
+        def map_attn_subkey(prefix: str, sub: str) -> str | None:
+            if re.match(r"^norm\.gamma$", sub):
+                return f"{prefix}.norm.gamma"
+            m = re.match(r"^to_qkv\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.to_qkv.{m.group(1)}"
+            m = re.match(r"^proj\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.proj.{m.group(1)}"
+            return None
+
+        def map_resample_subkey(prefix: str, sub: str) -> str | None:
+            m = re.match(r"^resample\.1\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.resample.1.{m.group(1)}"
+            m = re.match(r"^time_conv\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.time_conv.{m.group(1)}"
+            return None
+
+        m = re.match(r"^conv1\.(weight|bias)$", key)
+        if m:
+            return f"quant_conv.{m.group(1)}"
+        m = re.match(r"^conv2\.(weight|bias)$", key)
+        if m:
+            return f"post_quant_conv.{m.group(1)}"
+        m = re.match(r"^(encoder|decoder)\.conv1\.(weight|bias)$", key)
+        if m:
+            return f"{m.group(1)}.conv_in.{m.group(2)}"
+        m = re.match(r"^(encoder|decoder)\.head\.0\.gamma$", key)
+        if m:
+            return f"{m.group(1)}.norm_out.gamma"
+        m = re.match(r"^(encoder|decoder)\.head\.2\.(weight|bias)$", key)
+        if m:
+            return f"{m.group(1)}.conv_out.{m.group(2)}"
+
+        m = re.match(r"^(encoder|decoder)\.middle\.0\.(.*)$", key)
+        if m:
+            return map_residual_subkey(f"{m.group(1)}.mid_block.resnets.0",
+                                       m.group(2))
+        m = re.match(r"^(encoder|decoder)\.middle\.1\.(.*)$", key)
+        if m:
+            return map_attn_subkey(f"{m.group(1)}.mid_block.attentions.0",
+                                   m.group(2))
+        m = re.match(r"^(encoder|decoder)\.middle\.2\.(.*)$", key)
+        if m:
+            return map_residual_subkey(f"{m.group(1)}.mid_block.resnets.1",
+                                       m.group(2))
+
+        m = re.match(r"^encoder\.downsamples\.(\d+)\.(.*)$", key)
+        if m:
+            idx = int(m.group(1))
+            sub = m.group(2)
+            if sub.startswith("residual.") or sub.startswith("shortcut."):
+                return map_residual_subkey(f"encoder.down_blocks.{idx}", sub)
+            if sub.startswith("resample.") or sub.startswith("time_conv."):
+                return map_resample_subkey(f"encoder.down_blocks.{idx}", sub)
+            return None
+
+        m = re.match(r"^decoder\.upsamples\.(\d+)\.(.*)$", key)
+        if m:
+            uidx = int(m.group(1))
+            sub = m.group(2)
+
+            if uidx in (0, 1, 2):
+                block_i, res_i = 0, uidx
+            elif uidx == 3:
+                block_i, res_i = 0, None
+            elif uidx in (4, 5, 6):
+                block_i, res_i = 1, uidx - 4
+            elif uidx == 7:
+                block_i, res_i = 1, None
+            elif uidx in (8, 9, 10):
+                block_i, res_i = 2, uidx - 8
+            elif uidx == 11:
+                block_i, res_i = 2, None
+            elif uidx in (12, 13, 14):
+                block_i, res_i = 3, uidx - 12
+            else:
+                return None
+
+            if res_i is None:
+                return map_resample_subkey(
+                    f"decoder.up_blocks.{block_i}.upsamplers.0",
+                    sub,
+                )
+
+            return map_residual_subkey(
+                f"decoder.up_blocks.{block_i}.resnets.{res_i}",
+                sub,
+            )
+
+        return None
+
+    temporal_compression_ratio: int = 4
+    spatial_compression_ratio: int = 8
+
+    def __post_init__(self):
+        self.scaling_factor: torch.Tensor = 1.0 / torch.tensor(
+            self.latents_std).view(1, self.z_dim, 1, 1, 1)
+        self.shift_factor: torch.Tensor = torch.tensor(self.latents_mean).view(
+            1, self.z_dim, 1, 1, 1)
+        self.temporal_compression_ratio = self.scale_factor_temporal
+        self.spatial_compression_ratio = self.scale_factor_spatial
+
+
+@dataclass
+class Cosmos25VAEConfig(VAEConfig):
+    """Cosmos2.5 VAE config."""
+
+    arch_config: Cosmos25VAEArchConfig = field(
+        default_factory=Cosmos25VAEArchConfig)
+
+    use_feature_cache: bool = True
+    use_tiling: bool = False
+    use_temporal_tiling: bool = False
+    use_parallel_tiling: bool = False
+
+    def __post_init__(self):
+        self.blend_num_frames = (self.tile_sample_min_num_frames -
+                                 self.tile_sample_stride_num_frames) * 2

fastvideo/configs/pipelines/__init__.py

@@ -1,6 +1,7 @@
 from fastvideo.configs.pipelines.base import (PipelineConfig,
                                               SlidingTileAttnConfig)
 from fastvideo.configs.pipelines.cosmos import CosmosConfig
+from fastvideo.configs.pipelines.cosmos2_5 import Cosmos25Config
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig
 from fastvideo.configs.pipelines.registry import (
@@ -15,5 +16,5 @@ __all__ = [
     "Hunyuan15T2V480PConfig", "Hunyuan15T2V720PConfig", "SlidingTileAttnConfig",
     "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig",
     "WanI2V720PConfig", "StepVideoT2VConfig", "SelfForcingWanT2V480PConfig",
-    "CosmosConfig", "get_pipeline_config_cls_from_name"
+    "CosmosConfig", "Cosmos25Config", "get_pipeline_config_cls_from_name"
 ]

fastvideo/configs/pipelines/cosmos2_5.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import torch
+
+from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
+from fastvideo.configs.models.dits import Cosmos25VideoConfig
+from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25ArchConfig
+from fastvideo.configs.models.encoders import BaseEncoderOutput
+from fastvideo.configs.models.encoders.reason1 import Reason1Config, Reason1ArchConfig
+from fastvideo.configs.models.vaes import Cosmos25VAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig, STA_Mode
+
+
+def _identity_preprocess_text(prompt: str) -> str:
+    return prompt
+
+
+def reason1_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
+    hidden_states = getattr(outputs, "hidden_states", None)
+    if hidden_states is None:
+        raise ValueError("Reason1 postprocess requires outputs.hidden_states")
+
+    hs = list(hidden_states)[1:]
+    normed = []
+    for h in hs:
+        h = h.float()
+        h = (h - h.mean(dim=-1, keepdim=True)) / (h.std(dim=-1, keepdim=True) +
+                                                  1e-8)
+        normed.append(h)
+    return torch.cat(normed, dim=-1).to(hidden_states[0].dtype)
+
+
+@dataclass
+class Cosmos25Config(PipelineConfig):
+    """Configuration for Cosmos 2.5 (Predict2.5) video generation pipeline."""
+
+    dit_config: DiTConfig = field(default_factory=lambda: Cosmos25VideoConfig(
+        arch_config=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,
+            rope_enable_fps_modulation=False,
+            qk_norm="rms_norm",
+        )))
+
+    vae_config: VAEConfig = field(default_factory=Cosmos25VAEConfig)
+
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
+        default_factory=lambda: (Reason1Config(arch_config=Reason1ArchConfig(
+            embedding_concat_strategy="full_concat")), ))
+
+    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
+        default_factory=lambda: (_identity_preprocess_text, ))
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
+                                  ...] = field(default_factory=lambda:
+                                               (reason1_postprocess_text, ))
+
+    dit_precision: str = "bf16"
+    vae_precision: str = "bf16"
+    text_encoder_precisions: tuple[str, ...] = field(
+        default_factory=lambda: ("bf16", ))
+
+    embedded_cfg_scale: float = 0.0
+    flow_shift: float = 5.0
+
+    vae_tiling: bool = False
+    vae_sp: bool = False
+
+    STA_mode: STA_Mode = STA_Mode.NONE
+    skip_time_steps: int = 0
+
+    def __post_init__(self):
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True
+        self._vae_latent_dim = 16

fastvideo/configs/pipelines/registry.py

@@ -6,6 +6,7 @@ from collections.abc import Callable
 
 from fastvideo.configs.pipelines.base import PipelineConfig
 from fastvideo.configs.pipelines.cosmos import CosmosConfig
+from fastvideo.configs.pipelines.cosmos2_5 import Cosmos25Config
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig
 from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
@@ -55,6 +56,7 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
     "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_Config,
     "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_Config,
     "nvidia/Cosmos-Predict2-2B-Video2World": CosmosConfig,
+    "KyleShao/Cosmos-Predict2.5-2B-Diffusers": Cosmos25Config,
     "FastVideo/Matrix-Game-2.0-Base-Diffusers": MatrixGameI2V480PConfig,
     "FastVideo/Matrix-Game-2.0-GTA-Diffusers": MatrixGameI2V480PConfig,
     "FastVideo/Matrix-Game-2.0-TempleRun-Diffusers": MatrixGameI2V480PConfig,
@@ -94,7 +96,10 @@ PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = {
     "stepvideo":
     lambda id: "stepvideo" in id.lower(),
     "cosmos":
-    lambda id: "cosmos" in id.lower(),
+    lambda id: "cosmos" in id.lower() and ("2.5" not in id.lower(
+    ) and "2_5" not in id.lower() and "25" not in id.lower()),
+    "cosmos25":
+    lambda id: "cosmos25" in id.lower(),
     "turbodiffusion":
     lambda id: "turbodiffusion" in id.lower() or "turbowan" in id.lower(),
     # Add other pipeline architecture detectors
@@ -105,6 +110,7 @@ PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = {
     "longcatimagetovideo": LongCatT2V480PConfig,
     "longcatvideocontinuation": LongCatT2V480PConfig,
     "longcat": LongCatT2V480PConfig,
+    "cosmos25": Cosmos25Config,
     "hunyuan":
     HunyuanConfig,  # Base Hunyuan config as fallback for any Hunyuan variant
     "matrixgame": MatrixGameI2V480PConfig,

fastvideo/configs/sample/cosmos2_5.py

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass
+
+from fastvideo.configs.sample.base import SamplingParam
+
+
+@dataclass
+class Cosmos_Predict2_5_2B_Diffusers_SamplingParam(SamplingParam):
+    """Defaults for Cosmos 2.5 (Predict2.5) text-to-video diffusers-format model."""
+
+    height: int = 480
+    width: int = 832
+    num_frames: int = 121
+    fps: int = 24
+
+    guidance_scale: float = 7.0
+    # Official Cosmos2.5 sampling uses empty string as unconditional.
+    negative_prompt: str = ""
+    num_inference_steps: int = 35

fastvideo/configs/sample/registry.py

@@ -9,6 +9,7 @@ from fastvideo.configs.sample.hunyuan15 import Hunyuan15_480P_SamplingParam, Hun
 from fastvideo.configs.sample.stepvideo import StepVideoT2VSamplingParam
 
 from fastvideo.configs.sample.cosmos import Cosmos_Predict2_2B_Video2World_SamplingParam
+from fastvideo.configs.sample.cosmos2_5 import Cosmos_Predict2_5_2B_Diffusers_SamplingParam
 
 # isort: off
 from fastvideo.configs.sample.wan import (
@@ -96,6 +97,10 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
     "nvidia/Cosmos-Predict2-2B-Video2World":
     Cosmos_Predict2_2B_Video2World_SamplingParam,
 
+    # Cosmos2.5
+    "KyleShao/Cosmos-Predict2.5-2B-Diffusers":
+    Cosmos_Predict2_5_2B_Diffusers_SamplingParam,
+
     # MatrixGame2.0 models
     "FastVideo/Matrix-Game-2.0-Base-Diffusers":
     MatrixGame2_SamplingParam,
@@ -135,6 +140,10 @@ SAMPLING_PARAM_DETECTOR: dict[str, Callable[[str], bool]] = {
     lambda id: "matrixgame" in id.lower() or "matrix-game" in id.lower(),
     "turbodiffusion":
     lambda id: "turbodiffusion" in id.lower() or "turbowan" in id.lower(),
+    "cosmos25":
+    lambda id: "cosmos2_5" in id.lower(),
+    "cosmos":
+    lambda id: "cosmos" in id.lower() and "2_5" not in id.lower(),
     # Add other pipeline architecture detectors
 }
 
@@ -153,6 +162,8 @@ SAMPLING_FALLBACK_PARAM: dict[str, Any] = {
     "matrixgame": MatrixGame2_SamplingParam,
     "turbodiffusion":
     TurboDiffusionT2V_1_3B_SamplingParam,  # Default to T2V for fallback
+    "cosmos25": Cosmos_Predict2_5_2B_Diffusers_SamplingParam,
+    "cosmos": Cosmos_Predict2_2B_Video2World_SamplingParam,
     # Other fallbacks by architecture
 }
 
@@ -176,9 +187,6 @@ def get_sampling_param_cls_for_name(pipeline_name_or_path: str) -> Any | None:
 
     if os.path.exists(pipeline_name_or_path):
         config = verify_model_config_and_directory(pipeline_name_or_path)
-        logger.warning(
-            "FastVideo may not correctly identify the optimal sampling param for this model, as the local directory may have been renamed."
-        )
     else:
         config = maybe_download_model_index(pipeline_name_or_path)
 

fastvideo/models/encoders/reason1.py

@@ -0,0 +1,353 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Reason1 (Qwen2.5-VL) text encoder."""
+
+import os
+from dataclasses import dataclass
+from collections.abc import Iterable
+
+import torch
+from transformers import AutoProcessor
+
+from fastvideo.configs.models.encoders import BaseEncoderOutput, Reason1Config
+from fastvideo.logger import init_logger
+from fastvideo.models.encoders.base import TextEncoder
+from fastvideo.models.loader.weight_utils import default_weight_loader
+from fastvideo.platforms import AttentionBackendEnum
+
+from fastvideo.models.encoders.qwen2_5_vl_custom import (
+    Qwen2_5_VLForConditionalGenerationSimple,
+    Qwen2_5_VLConfig,
+    get_rope_index,
+)
+
+logger = init_logger(__name__)
+
+
+@dataclass(frozen=True)
+class _WeightsSource:
+    """Mimic `TextEncoderLoader.Source` (avoid import cycles)."""
+
+    model_or_path: str
+    prefix: str = ""
+    fall_back_to_pt: bool = True
+    allow_patterns_overrides: list[str] | None = None
+
+
+
+
+class Reason1TextEncoder(TextEncoder):
+    """Reason1 (Qwen2.5-VL) text encoder."""
+
+    _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
+        AttentionBackendEnum.FLASH_ATTN,
+        AttentionBackendEnum.TORCH_SDPA,
+    )
+    
+    def __init__(self, config: Reason1Config, prefix: str = "", checkpoint_path: str | None = None):
+        super().__init__(config)
+        
+        self.prefix = prefix
+        self.quant_config = None  # For future quantization support
+        
+        self.embedding_concat_strategy = config.arch_config.embedding_concat_strategy
+        self.n_layers_per_group = config.arch_config.n_layers_per_group
+        self.num_embedding_padding_tokens = config.arch_config.num_embedding_padding_tokens
+        
+        config_path = checkpoint_path if checkpoint_path else config.tokenizer_type
+        
+        logger.info("Initializing Reason1TextEncoder (Qwen2.5-VL) from %s", config_path)
+        try:
+            from transformers import AutoConfig as HFAutoConfig
+            hf_config = HFAutoConfig.from_pretrained(
+                config_path,
+                trust_remote_code=True,
+            )
+        except Exception as e:
+            logger.warning("Failed to load HF config from %s (%s). Using default Qwen2.5-VL-7B config.",
+                           config_path, e)
+            hf_config = Qwen2_5_VLConfig(
+                    hidden_size=3584,
+                    intermediate_size=18944,
+                    max_window_layers=28,
+                    num_attention_heads=28,
+                    num_hidden_layers=28,
+                    num_key_value_heads=4,
+                    tie_word_embeddings=False,
+                    vocab_size=152064,
+            )
+        
+        hf_config.output_hidden_states = True
+        
+        if hasattr(config.arch_config, '_attn_implementation') and config.arch_config._attn_implementation:
+            hf_config._attn_implementation = config.arch_config._attn_implementation
+        else:
+            hf_config._attn_implementation = "flash_attention_2"
+        logger.info("Reason1 attention implementation: %s", getattr(hf_config, "_attn_implementation", None))
+        
+        with torch.device("meta"):
+            self.model = Qwen2_5_VLForConditionalGenerationSimple(hf_config)
+        
+        self.processor = AutoProcessor.from_pretrained(
+            config_path,
+            trust_remote_code=True,
+        )
+
+        weights_override = os.getenv("FASTVIDEO_REASON1_WEIGHTS_PATH")
+        if weights_override:
+            self.secondary_weights = (
+                _WeightsSource(
+                    model_or_path=weights_override,
+                    prefix="",
+                    fall_back_to_pt=True,
+                    allow_patterns_overrides=None,
+                ),
+            )
+            logger.info("Reason1TextEncoder: overlaying weights from %s", weights_override)
+        
+        self._weights_loaded = False
+    
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        **kwargs,
+    ) -> BaseEncoderOutput:
+        # Cosmos2.5 alignment: keep attention_mask=None.
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=None,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            output_hidden_states=True,
+            return_dict=True,
+            pixel_values=kwargs.get('pixel_values', None),
+            pixel_values_videos=kwargs.get('pixel_values_videos', None),
+            image_grid_thw=kwargs.get('image_grid_thw', None),
+            video_grid_thw=kwargs.get('video_grid_thw', None),
+        )
+        
+        hidden_states = outputs.hidden_states
+        last_hidden_state = hidden_states[-1]
+        
+        return BaseEncoderOutput(
+            last_hidden_state=last_hidden_state,
+            hidden_states=hidden_states if output_hidden_states else None,
+            attention_mask=None,
+        )
+    
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        first_weight = None
+        weights_list = []
+        for name, weight in weights:
+            if first_weight is None:
+                first_weight = weight
+                self.model = self.model.to_empty(device=weight.device)
+                self.model.init_weights(buffer_device=weight.device)
+            weights_list.append((name, weight))
+        
+        params_dict = dict(self.model.named_parameters())
+        loaded_params: set[str] = set()
+        skipped_weights = {"lm_head": 0, "visual": 0, "decoder": 0}
+        
+        for name, loaded_weight in weights_list:
+            if "lm_head" in name:
+                skipped_weights["lm_head"] += 1
+                continue
+            if "visual" in name:
+                skipped_weights["visual"] += 1
+                continue
+            if "decoder" in name:
+                skipped_weights["decoder"] += 1
+                continue
+            
+            # Handle stacked params mapping (for quantized models)
+            for param_name, weight_name, shard_id in self.config.arch_config.stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                
+                if name not in params_dict:
+                    continue
+                
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                param_name_with_prefix = f"model.{name}" if self.prefix == "" else f"{self.prefix}.{name}"
+                loaded_params.add(param_name_with_prefix)
+                break
+            else:
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                
+                if name not in params_dict:
+                    continue
+                
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+                param_name_with_prefix = f"model.{name}" if self.prefix == "" else f"{self.prefix}.{name}"
+                loaded_params.add(param_name_with_prefix)
+        if first_weight is not None:
+            self.model = self.model.to(first_weight.device)
+        
+        all_params = set(f"model.{name}" if self.prefix == "" else f"{self.prefix}.{name}" 
+                        for name in params_dict.keys())
+        loaded_params.update(all_params)
+        
+        # Mark weights as loaded
+        self._weights_loaded = True
+        return loaded_params
+    
+    def compute_text_embeddings_online(
+        self,
+        data_batch: dict[str, list[str]],
+        input_caption_key: str,
+    ) -> torch.Tensor:
+        prompts = data_batch[input_caption_key]
+        return self.compute_text_embeddings(prompts)
+    
+    def compute_text_embeddings(
+        self,
+        prompts: list[str],
+        device: str | torch.device = "cuda",
+    ) -> torch.Tensor:
+        """Compute embeddings for a list of prompts."""
+        input_ids_batch = []
+        
+        tok = getattr(self.processor, "tokenizer", None)
+        if tok is None:
+            raise RuntimeError("Reason1TextEncoder requires processor.tokenizer")
+        pad_id = getattr(tok, "pad_id", None)
+        if pad_id is None:
+            pad_id = getattr(tok, "pad_token_id", None)
+        if pad_id is None:
+            pad_id = getattr(self.model.config, "pad_token_id", None)
+        if pad_id is None:
+            pad_id = 0
+        
+        for prompt in prompts:
+            conversations = [
+                {
+                    "role": "system",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": "You are a helpful assistant who will provide prompts to an image generator.",
+                        }
+                    ],
+                },
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": prompt,
+                        }
+                    ],
+                },
+            ]
+            
+            try:
+                tokenizer_output = tok.apply_chat_template(
+                    conversations,
+                    tokenize=True,
+                    add_generation_prompt=False,
+                    add_vision_id=False,
+                )
+            except TypeError:
+                tokenizer_output = tok.apply_chat_template(
+                conversations,
+                tokenize=True,
+                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)}"
+                    )
+            
+            if self.num_embedding_padding_tokens > len(input_ids):
+                pad_len = self.num_embedding_padding_tokens - len(input_ids)
+                input_ids = input_ids + [pad_id] * pad_len
+            else:
+                input_ids = input_ids[:self.num_embedding_padding_tokens]
+            
+            input_ids = torch.LongTensor(input_ids).to(device=device)
+            input_ids_batch.append(input_ids)
+        
+        input_ids_batch = torch.stack(input_ids_batch, dim=0)
+
+        # Cosmos2.5 alignment: keep attention_mask=None.
+        target_device = input_ids_batch.device
+        try:
+            embed_device = self.model.model.embed_tokens.weight.device  # type: ignore[attr-defined]
+        except Exception:
+            embed_device = None
+        if embed_device is not None and embed_device != target_device:
+            self.model = self.model.to(target_device)
+        
+        with torch.no_grad():
+            position_ids, _ = get_rope_index(
+                self.model.config,
+                input_ids_batch,
+                image_grid_thw=None,
+                video_grid_thw=None,
+                second_per_grid_ts=None,
+                attention_mask=None,
+            )
+            position_ids = position_ids.to(target_device)
+                
+            outputs = self.model.model(
+                input_ids=input_ids_batch,
+                position_ids=position_ids,
+                attention_mask=None,
+                output_hidden_states=True,
+                return_dict=True,
+                use_cache=False,
+            )
+            hidden_states = outputs.hidden_states
+        
+        normalized_hidden_states = []
+        for layer_idx in range(1, len(hidden_states)):
+            normalized_state = self._mean_normalize(hidden_states[layer_idx])
+            normalized_hidden_states.append(normalized_state)
+        
+        if self.embedding_concat_strategy == "full_concat":
+            text_embeddings = torch.cat(normalized_hidden_states, dim=-1)
+        elif self.embedding_concat_strategy == "mean_pooling":
+            text_embeddings = torch.stack(normalized_hidden_states).mean(dim=0)
+        elif self.embedding_concat_strategy == "pool_every_n_layers_and_concat":
+            pooled_embeddings = []
+            for i in range(0, len(normalized_hidden_states), self.n_layers_per_group):
+                group = normalized_hidden_states[i : i + self.n_layers_per_group]
+                pooled = torch.stack(group).mean(dim=0)
+                pooled_embeddings.append(pooled)
+            text_embeddings = torch.cat(pooled_embeddings, dim=-1)
+        else:
+            raise ValueError(
+                f"Unknown embedding_concat_strategy: {self.embedding_concat_strategy}"
+            )
+        
+        return text_embeddings
+    @staticmethod
+    def _mean_normalize(tensor: torch.Tensor) -> torch.Tensor:
+        return (tensor - tensor.mean(dim=-1, keepdim=True)) / (
+            tensor.std(dim=-1, keepdim=True) + 1e-8
+        )
+

fastvideo/models/loader/component_loader.py

@@ -15,8 +15,7 @@ import torch.distributed as dist
 import torch.nn as nn
 from safetensors.torch import load_file as safetensors_load_file
 from torch.distributed import init_device_mesh
-from transformers import AutoImageProcessor, AutoTokenizer
-from transformers import UMT5EncoderModel
+from transformers import AutoImageProcessor, AutoProcessor, AutoTokenizer
 from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
 
 from fastvideo.configs.models import EncoderConfig
@@ -91,10 +90,12 @@ class ComponentLoader(ABC):
 
         if module_type in module_loaders:
             loader_cls, expected_library = module_loaders[module_type]
-            # Assert that the library matches what's expected for this module type
-            assert transformers_or_diffusers == expected_library, (
-                f"{module_type} must be loaded from {expected_library}, got {transformers_or_diffusers}"
-            )
+            # Allow fastvideo.* libraries for custom implementations (e.g. Cosmos2_5Pipeline)
+            # that aren't available in diffusers/transformers yet
+            is_fastvideo_module = transformers_or_diffusers.startswith("fastvideo.")
+            if not is_fastvideo_module:
+                # Assert that the library matches what's expected for this module type
+                assert transformers_or_diffusers == expected_library, f"{module_type} must be loaded from {expected_library}, got {transformers_or_diffusers}"
             return loader_cls()
 
         # For unknown module types, use a generic loader
@@ -449,6 +450,33 @@ class TokenizerLoader(ComponentLoader):
         """Load the tokenizer based on the model path, and inference args."""
         logger.info("Loading tokenizer from %s", model_path)
 
+        # Cosmos2.5 stores an AutoProcessor config in `tokenizer/config.json` (not a tokenizer
+        # config). Use its `_name_or_path` (e.g. Qwen/Qwen2.5-VL-7B-Instruct) as the source.
+        tokenizer_cfg_path = os.path.join(model_path, "config.json")
+        if os.path.exists(tokenizer_cfg_path):
+            try:
+                with open(tokenizer_cfg_path, "r") as f:
+                    tokenizer_cfg = json.load(f)
+                if isinstance(tokenizer_cfg, dict) and (
+                    tokenizer_cfg.get("_class_name") == "AutoProcessor"
+                    or "processor_type" in tokenizer_cfg
+                ):
+                    src = tokenizer_cfg.get("_name_or_path", "")
+                    if isinstance(src, str) and src.strip():
+                        processor = AutoProcessor.from_pretrained(
+                            src.strip(),
+                            trust_remote_code=True,
+                        )
+                        logger.info(
+                            "Loaded tokenizer/processor from %s: %s",
+                            src,
+                            processor.__class__.__name__,
+                        )
+                        return processor
+            except Exception:
+                # If parsing fails, fall through to AutoTokenizer below.
+                pass
+
         tokenizer = AutoTokenizer.from_pretrained(
             model_path,  # "<path to model>/tokenizer"
             # in v0, this was same string as encoder_name "ClipTextModel"
@@ -491,19 +519,40 @@ class VAELoader(ComponentLoader):
             if fastvideo_args.pipeline_config.vae_precision
             else torch.bfloat16
         ):
+            # Cosmos2.5 uses a Wan2.1 VAE stored as `tokenizer.safetensors` under the VAE folder.
+            is_cosmos25 = fastvideo_args.pipeline_config.__class__.__name__ == "Cosmos25Config"
+            if class_name == "AutoencoderKLWan" and is_cosmos25:
+                from fastvideo.models.vaes.cosmos25wanvae import Cosmos25WanVAE
+
+                dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
+                vae = Cosmos25WanVAE(device=target_device, dtype=dtype)
+
+                weight_path = os.path.join(model_path, "tokenizer.safetensors")
+                if not os.path.exists(weight_path):
+                    raise FileNotFoundError(
+                        f"Missing Cosmos2.5 VAE weights: {weight_path}"
+                    )
+                sd = safetensors_load_file(weight_path)
+                vae.load_state_dict(sd, strict=False)
+                return vae.eval()
+
             vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
             vae = vae_cls(vae_config).to(target_device)
 
         # Find all safetensors files
         safetensors_list = glob.glob(
-            os.path.join(str(model_path), "*.safetensors")
-        )
-        loaded = {}
-        for sf_file in safetensors_list:
-            loaded.update(safetensors_load_file(sf_file))
-        vae.load_state_dict(
-            loaded, strict=False
-        )  # We might only load encoder or decoder
+            os.path.join(str(model_path), "*.safetensors"))
+        if not safetensors_list:
+            raise ValueError(f"No safetensors files found in {model_path}")
+        # Common case: a single `.safetensors` checkpoint file.
+        # Some models may be sharded into multiple files; in that case we merge.
+        if len(safetensors_list) == 1:
+            loaded = safetensors_load_file(safetensors_list[0])
+        else:
+            loaded = {}
+            for sf_file in safetensors_list:
+                loaded.update(safetensors_load_file(sf_file))
+        vae.load_state_dict(loaded, strict=False)
 
         return vae.eval()
 
@@ -581,7 +630,17 @@ class TransformerLoader(ComponentLoader):
         ]
 
         # Load the model using FSDP loader
+        logger.info("Loading model from %s, default_dtype: %s", cls_name,
+                    default_dtype)
         assert fastvideo_args.hsdp_shard_dim is not None
+        # Cosmos2.5 checkpoints can include extra entries not present in the
+        # instantiated model (e.g. pos_embedder ranges / *_extra_state). Load
+        # non-strictly for Cosmos2.5 only; keep upstream strict behavior for others.
+        strict_load = not (
+            cls_name.startswith("Cosmos25")
+            or cls_name == "Cosmos25Transformer3DModel"
+            or getattr(fastvideo_args.pipeline_config, "prefix", "") == "Cosmos25"
+        )
         model = maybe_load_fsdp_model(
             model_cls=model_cls,
             init_params={"config": dit_config, "hf_config": hf_config},
@@ -589,6 +648,7 @@ class TransformerLoader(ComponentLoader):
             device=get_local_torch_device(),
             hsdp_replicate_dim=fastvideo_args.hsdp_replicate_dim,
             hsdp_shard_dim=fastvideo_args.hsdp_shard_dim,
+            strict=strict_load,
             cpu_offload=fastvideo_args.dit_cpu_offload,
             pin_cpu_memory=fastvideo_args.pin_cpu_memory,
             fsdp_inference=fastvideo_args.use_fsdp_inference,

fastvideo/models/loader/fsdp_load.py

@@ -67,6 +67,7 @@ def maybe_load_fsdp_model(
     default_dtype: torch.dtype,
     param_dtype: torch.dtype,
     reduce_dtype: torch.dtype,
+    strict: bool = True,
     cpu_offload: bool = False,
     fsdp_inference: bool = False,
     output_dtype: torch.dtype | None = None,
@@ -141,7 +142,7 @@ def maybe_load_fsdp_model(
         weight_iterator,
         device,
         default_dtype,
-        strict=True,
+        strict=strict,
         cpu_offload=cpu_offload,
         param_names_mapping=param_names_mapping_fn,
     )
@@ -293,6 +294,26 @@ def load_model_from_full_model_state_dict(
     for target_param_name, full_tensor in custom_param_sd.items():
         meta_sharded_param = meta_sd.get(target_param_name)
         if meta_sharded_param is None:
+            # Some checkpoints include extra entries that are not part of the
+            # instantiated model's state_dict (e.g. `_extra_state` keys from
+            # some FSDP checkpoint formats). These can be safely skipped.
+            if (target_param_name.endswith("._extra_state")
+                    or target_param_name.endswith("_extra_state")):
+                logger.warning(
+                    "Skipping non-parameter checkpoint key: %s",
+                    target_param_name,
+                )
+                continue
+
+            # For non-strict loads, treat this as an "unexpected key" and skip it
+            # (mirrors torch.nn.Module.load_state_dict(strict=False)).
+            if not strict:
+                logger.warning(
+                    "Skipping unexpected checkpoint key (not present in model): %s",
+                    target_param_name,
+                )
+                continue
+
             raise ValueError(
                 f"Parameter {target_param_name} not found in custom model state dict. The hf to custom mapping may be incorrect."
             )

fastvideo/models/registry.py

@@ -30,8 +30,9 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
     "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
     "StepVideoModel": ("dits", "stepvideo", "StepVideoModel"),
     "CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel"),
-    "LongCatVideoTransformer3DModel": ("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"),
-    "LongCatTransformer3DModel": ("dits", "longcat", "LongCatTransformer3DModel"),
+    "Cosmos25Transformer3DModel": ("dits", "cosmos2_5", "Cosmos25Transformer3DModel"),
+    "LongCatVideoTransformer3DModel": ("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"),  # Wrapper (Phase 1)
+    "LongCatTransformer3DModel": ("dits", "longcat", "LongCatTransformer3DModel"),  # Native (Phase 2)
 }
 
 _IMAGE_TO_VIDEO_DIT_MODELS = {
@@ -50,6 +51,9 @@ _TEXT_ENCODER_MODELS = {
     "STEP1TextEncoder": ("encoders", "stepllm", "STEP1TextEncoder"),
     "BertModel": ("encoders", "clip", "CLIPTextModel"),
     "Qwen2_5_VLTextModel": ("encoders", "qwen2_5", "Qwen2_5_VLTextModel"),
+    "Reason1TextEncoder": ("encoders", "reason1", "Reason1TextEncoder"),
+    "Qwen2_5_VLForConditionalGeneration":
+    ("encoders", "reason1", "Reason1TextEncoder"),
 }
 
 _IMAGE_ENCODER_MODELS: dict[str, tuple] = {
@@ -72,6 +76,8 @@ _SCHEDULERS = {
      "FlowMatchEulerDiscreteScheduler"),
     "UniPCMultistepScheduler":
     ("schedulers", "scheduling_unipc_multistep", "UniPCMultistepScheduler"),
+    "FlowUniPCMultistepScheduler":
+    ("schedulers", "scheduling_flow_unipc_multistep", "FlowUniPCMultistepScheduler"),
     "SelfForcingFlowMatchScheduler":
     ("schedulers", "scheduling_self_forcing_flow_match",
      "SelfForcingFlowMatchScheduler"),

fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py

@@ -109,6 +109,9 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         sigmas = 1.0 - alphas
         sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32)
 
+        # Needed when final_sigmas_type == "sigma_min" (kept for compatibility).
+        self.alphas_cumprod = torch.from_numpy(alphas).to(dtype=torch.float32)
+
         if not use_dynamic_shifting:
             # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
             assert shift is not None
@@ -171,6 +174,8 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         sigmas: list[float] | None = None,
         mu: float | None | None = None,
         shift: float | None | None = None,
+        use_karras_sigmas: bool | None = None,
+        use_kerras_sigma: bool | None = None,
     ):
         """
         Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -186,21 +191,44 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
                 " you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`"
             )
 
-        if sigmas is None:
-            assert num_inference_steps is not None
-            sigmas = np.linspace(self.sigma_max, self.sigma_min,
-                                 num_inference_steps +
-                                 1).copy()[:-1]  # pyright: ignore
+
+        # Cosmos official uses `use_kerras_sigma=True` and a specific EDM sigma schedule.
+        # Some external code uses the misspelling `use_kerras_sigma`; support both.
+        if use_karras_sigmas is None and use_kerras_sigma is not None:
+            use_karras_sigmas = use_kerras_sigma
+
+        if use_karras_sigmas:
+            # Force to use the exact sigma used in official EDM sampler:
+            # sigma_max=200, sigma_min=0.01, rho=7
+            sigma_max = 200.0
+            sigma_min = 0.01
+            rho = 7.0
+            # Match the official Cosmos implementation: Karras/EDM schedule with
+            # `num_inference_steps + 1` points, then `final_sigmas_type="zero"`
+            # appends the terminal sigma (0.0).
+            ramp = np.arange(num_inference_steps + 1,
+                             dtype=np.float32) / float(num_inference_steps)
+            min_inv_rho = sigma_min**(1 / rho)
+            max_inv_rho = sigma_max**(1 / rho)
+            sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho))**rho
+            # Convert EDM sigma to flow-matching sigma in [0, 1).
+            sigmas = sigmas / (1.0 + sigmas)
+        else:
+            if sigmas is None:
+                assert num_inference_steps is not None
+                sigmas = np.linspace(self.sigma_max, self.sigma_min,
+                                     num_inference_steps +
+                                     1).copy()[:-1]  # pyright: ignore
 
         if self.config.use_dynamic_shifting:
             assert mu is not None
             sigmas = self.time_shift(mu, 1.0, sigmas)  # pyright: ignore
         else:
-            if shift is None:
-                shift = self.config.shift
-            assert isinstance(sigmas, np.ndarray)
-            sigmas = shift * sigmas / (1 +
-                                       (shift - 1) * sigmas)  # pyright: ignore
+            if not use_karras_sigmas:
+                if shift is None:
+                    shift = self.config.shift
+                assert isinstance(sigmas, np.ndarray)
+                sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)  # pyright: ignore
 
         if self.config.final_sigmas_type == "sigma_min":
             sigma_last = ((1 - self.alphas_cumprod[0]) /
@@ -418,8 +446,12 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
         alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
 
-        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
-        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+        # Numerical safety
+        eps = 1e-12
+        lambda_t = torch.log(torch.clamp(alpha_t, min=eps)) - torch.log(
+            torch.clamp(sigma_t, min=eps))
+        lambda_s0 = torch.log(torch.clamp(alpha_s0, min=eps)) - torch.log(
+            torch.clamp(sigma_s0, min=eps))
 
         h = lambda_t - lambda_s0
         device = sample.device
@@ -430,7 +462,8 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             si = self.step_index - i  # pyright: ignore
             mi = model_output_list[-(i + 1)]
             alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
-            lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
+            lambda_si = torch.log(torch.clamp(alpha_si, min=eps)) - torch.log(
+                torch.clamp(sigma_si, min=eps))
             rk = (lambda_si - lambda_s0) / h
             rks.append(rk)
             assert mi is not None
@@ -563,8 +596,11 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
         alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
 
-        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
-        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+        eps = 1e-12
+        lambda_t = torch.log(torch.clamp(alpha_t, min=eps)) - torch.log(
+            torch.clamp(sigma_t, min=eps))
+        lambda_s0 = torch.log(torch.clamp(alpha_s0, min=eps)) - torch.log(
+            torch.clamp(sigma_s0, min=eps))
 
         h = lambda_t - lambda_s0
         device = this_sample.device
@@ -575,7 +611,8 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             si = self.step_index - (i + 1)  # pyright: ignore
             mi = model_output_list[-(i + 1)]
             alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
-            lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
+            lambda_si = torch.log(torch.clamp(alpha_si, min=eps)) - torch.log(
+                torch.clamp(sigma_si, min=eps))
             rk = (lambda_si - lambda_s0) / h
             rks.append(rk)
             assert mi is not None

fastvideo/models/vaes/cosmos25wanvae.py

@@ -0,0 +1,735 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+"""
+Cosmos 2.5 / Wan2.1 VAE adapter.
+
+Why this exists:
+- Cosmos2.5 uses a Wan2.1-style VAE, but the *diffusion model* operates in a
+  **normalized latent space**:
+    z_norm = (z - mean) / std
+
+  Meanwhile, FastVideo's `AutoencoderKLWan` operates in the VAE's native latent
+  space (denormalized):
+    z = z_norm * std + mean
+
+This adapter provides a single, stable interface for FastVideo pipelines:
+- `encode(x)` returns an object with `.mean` / `.sample()` / `.mode()`
+- `decode(z)` returns a tensor in pixel space
+
+It also exposes flags used by pipeline stages to avoid double (de)normalization:
+- `handles_latent_norm = True`   -> stages should NOT normalize encoder latents
+- `handles_latent_denorm = True` -> stages should NOT denormalize before decode
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+
+@dataclass
+class _TensorLatentDist:
+    """Minimal distribution-like wrapper used by pipeline stages."""
+
+    mean: torch.Tensor
+
+    def mode(self) -> torch.Tensor:
+        return self.mean
+
+    def sample(self, generator: Any | None = None) -> torch.Tensor:  # generator for API compatibility
+        # The official interface encodes deterministically; for compatibility we
+        # return the mean. (Stochastic posterior sampling isn't required for
+        # Cosmos2.5 inference.)
+        _ = generator
+        return self.mean
+
+
+class Cosmos25WanVAEAdapter(nn.Module):
+    """
+    Adapter that makes a Wan2.1-style VAE follow Cosmos2.5's latent contract:
+    - `encode()` returns **normalized** latents
+    - `decode()` expects **normalized** latents
+    """
+
+    # Pipeline stage hints (see latent_preparation.py / decoding.py / image_encoding.py)
+    handles_latent_norm: bool = True
+    handles_latent_denorm: bool = True
+    latent_norm_mode: str = "internal"  # informational
+
+    def __init__(
+        self,
+        inner: Any,
+        *,
+        latents_mean: Optional[torch.Tensor] = None,
+        latents_std: Optional[torch.Tensor] = None,
+    ) -> None:
+        super().__init__()
+        self.inner = inner
+
+        # Preserve `config` when available; some pipeline utilities expect it.
+        self.config = getattr(inner, "config", None)
+
+        # If not provided, try to derive from `config.latents_mean/std`.
+        cfg = self.config
+        if latents_mean is None and cfg is not None and hasattr(cfg, "latents_mean"):
+            latents_mean = torch.tensor(cfg.latents_mean, dtype=torch.float32).view(1, -1, 1, 1, 1)
+        if latents_std is None and cfg is not None and hasattr(cfg, "latents_std"):
+            latents_std = torch.tensor(cfg.latents_std, dtype=torch.float32).view(1, -1, 1, 1, 1)
+
+        if latents_mean is None or latents_std is None:
+            raise RuntimeError(
+                "Cosmos25WanVAEAdapter requires latents_mean/latents_std (either passed explicitly or available on inner.config)."
+            )
+
+        # Register as buffers so `.to(...)` moves them with the module.
+        self.register_buffer("_latents_mean", latents_mean, persistent=False)
+        self.register_buffer("_latents_std", latents_std, persistent=False)
+
+    def _to_latent_stats(self, like: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        mean = self._latents_mean.to(device=like.device, dtype=like.dtype)
+        std = self._latents_std.to(device=like.device, dtype=like.dtype)
+        return mean, std
+
+    def get_latent_num_frames(self, num_pixel_frames: int) -> int:
+        # Keep parity with official interface.
+        if hasattr(self.inner, "get_latent_num_frames"):
+            return int(self.inner.get_latent_num_frames(num_pixel_frames))
+        return 1 + (num_pixel_frames - 1) // 4
+
+    def encode(self, x: torch.Tensor) -> _TensorLatentDist:
+        """
+        Returns *normalized* latents (Cosmos contract).
+        """
+        enc_out = self.inner.encode(x)
+
+        # Support common encoder output shapes:
+        # - DiagonalGaussianDistribution (FastVideo VAE): has `.mean` / `.sample()` / `.mode()`
+        # - diffusers EncoderOutput: has `.latent_dist`
+        # - raw tensor
+        if hasattr(enc_out, "latent_dist"):
+            dist = enc_out.latent_dist
+            z_mean = dist.mode() if hasattr(dist, "mode") else dist.mean
+        elif hasattr(enc_out, "mode") and hasattr(enc_out, "mean"):
+            z_mean = enc_out.mode()
+        elif isinstance(enc_out, torch.Tensor):
+            z_mean = enc_out
+        else:
+            attrs = [a for a in dir(enc_out) if not a.startswith("_")]
+            raise RuntimeError(
+                f"Unsupported VAE encoder output type: {type(enc_out)}. attrs={attrs}"
+            )
+
+        mean, std = self._to_latent_stats(z_mean)
+        z_norm = (z_mean - mean) / std
+        return _TensorLatentDist(z_norm)
+
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        """
+        Expects *normalized* latents (Cosmos contract).
+        """
+        mean, std = self._to_latent_stats(z)
+        z_denorm = z * std + mean
+        out = self.inner.decode(z_denorm)
+        return out.sample if hasattr(out, "sample") else out
+
+
+#
+# Official-like Wan2.1 VAE implementation (ported from cosmos_predict2 wan2pt1.py)
+# -------------------------------------------------------------------------------
+# Motivation:
+# - We already solved checkpoint *key mapping* and can load official weights.
+# - Remaining output drift vs the official tokenizer is largely decoder-side.
+# - FastVideo's `AutoencoderKLWan` uses a different temporal upsample path
+#   (`DupUp3D` + `first_chunk` slicing), while the official tokenizer uses
+#   `Resample(mode="upsample3d")` with a time-conv + interleave reshape.
+#
+# This section ports the core modules (CausalConv3d/Resample/etc.) so we can run
+# a VAE that is behaviorally closer to the official implementation WITHOUT
+# importing any official repo classes at runtime.
+#
+
+CACHE_T = 2
+
+
+class Cosmos25CausalConv3d(nn.Conv3d):
+    """
+    Official-like causal 3D convolution.
+
+    Matches `CausalConv3d` in the official tokenizer: uses explicit F.pad and
+    supports a `cache_x` prefix for causal chunking.
+    """
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        # padding order for F.pad: (W_left, W_right, H_left, H_right, T_left, T_right)
+        self._padding: tuple[int, ...] = (
+            self.padding[2],
+            self.padding[2],
+            self.padding[1],
+            self.padding[1],
+            2 * self.padding[0],
+            0,
+        )
+        self.padding = (0, 0, 0)
+
+    def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
+        padding = list(self._padding)
+        if cache_x is not None and self._padding[4] > 0:
+            cache_x = cache_x.to(x.device)
+            x = torch.cat([cache_x, x], dim=2)
+            padding[4] -= cache_x.shape[2]
+        x = F.pad(x, padding)
+        return super().forward(x)
+
+
+class Cosmos25RMSNorm(nn.Module):
+    """Official-like RMS_norm (uses learnable gamma and optional bias)."""
+
+    def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False) -> None:
+        super().__init__()
+        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+
+        self.channel_first = channel_first
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(shape))
+        self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dim = 1 if self.channel_first else -1
+        return F.normalize(x, dim=dim) * self.scale * self.gamma + self.bias
+
+
+class Cosmos25Upsample(nn.Upsample):
+    """Official-like Upsample that is safe for bf16 (casts to fp32 internally)."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:  # type: ignore[override]
+        return super().forward(x.float()).type_as(x)
+
+
+class Cosmos25Resample(nn.Module):
+    """
+    Official-like Resample used for both spatial and temporal up/downsampling.
+    """
+
+    def __init__(self, dim: int, mode: str) -> None:
+        assert mode in ("none", "upsample2d", "upsample3d", "downsample2d", "downsample3d")
+        super().__init__()
+        self.dim = dim
+        self.mode = mode
+
+        if mode == "upsample2d":
+            self.resample = nn.Sequential(
+                Cosmos25Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim // 2, 3, padding=1),
+            )
+        elif mode == "upsample3d":
+            self.resample = nn.Sequential(
+                Cosmos25Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim // 2, 3, padding=1),
+            )
+            self.time_conv = Cosmos25CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
+        elif mode == "downsample2d":
+            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+        elif mode == "downsample3d":
+            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+            self.time_conv = Cosmos25CausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
+        else:
+            self.resample = nn.Identity()
+
+    def forward(self, x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor:
+        b, c, t, h, w = x.size()
+
+        # Temporal upsample uses a time-conv and then interleaves frames.
+        if self.mode == "upsample3d" and feat_cache is not None:
+            idx = feat_idx[0]
+            if feat_cache[idx] is None:
+                feat_cache[idx] = "Rep"
+                feat_idx[0] += 1
+            else:
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep":
+                    cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep":
+                    cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2)
+
+                if feat_cache[idx] == "Rep":
+                    x = self.time_conv(x)
+                else:
+                    x = self.time_conv(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+
+                x = x.reshape(b, 2, c, t, h, w)
+                x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3)
+                x = x.reshape(b, c, t * 2, h, w)
+
+        t = x.shape[2]
+        x = rearrange(x, "b c t h w -> (b t) c h w")
+        x = self.resample(x)
+        x = rearrange(x, "(b t) c h w -> b c t h w", t=t)
+
+        # Temporal downsample: time_conv consumes last-frame cache.
+        if self.mode == "downsample3d" and feat_cache is not None:
+            idx = feat_idx[0]
+            if feat_cache[idx] is None:
+                feat_cache[idx] = x.clone()
+                feat_idx[0] += 1
+            else:
+                cache_x = x[:, :, -1:, :, :].clone()
+                x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+
+        return x
+
+
+class Cosmos25ResidualBlock(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, dropout: float = 0.0) -> None:
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.residual = nn.Sequential(
+            Cosmos25RMSNorm(in_dim, images=False),
+            nn.SiLU(),
+            Cosmos25CausalConv3d(in_dim, out_dim, 3, padding=1),
+            Cosmos25RMSNorm(out_dim, images=False),
+            nn.SiLU(),
+            nn.Dropout(dropout),
+            Cosmos25CausalConv3d(out_dim, out_dim, 3, padding=1),
+        )
+        self.shortcut = Cosmos25CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity()
+
+    def forward(self, x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor:
+        h = self.shortcut(x)
+        for layer in self.residual:
+            if isinstance(layer, Cosmos25CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x + h
+
+
+class Cosmos25AttentionBlock(nn.Module):
+    """Official-like causal self-attention with a single head."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.dim = dim
+        self.norm = Cosmos25RMSNorm(dim)
+        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
+        self.proj = nn.Conv2d(dim, dim, 1)
+        nn.init.zeros_(self.proj.weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        identity = x
+        b, c, t, h, w = x.size()
+        x2 = rearrange(x, "b c t h w -> (b t) c h w")
+        x2 = self.norm(x2)
+        q, k, v = (
+            self.to_qkv(x2)
+            .reshape(b * t, 1, c * 3, -1)
+            .permute(0, 1, 3, 2)
+            .contiguous()
+            .chunk(3, dim=-1)
+        )
+        x2 = F.scaled_dot_product_attention(q, k, v)
+        x2 = x2.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
+        x2 = self.proj(x2)
+        x2 = rearrange(x2, "(b t) c h w-> b c t h w", t=t)
+        return x2 + identity
+
+
+class Cosmos25Encoder3d(nn.Module):
+    def __init__(
+        self,
+        dim: int = 96,
+        z_dim: int = 32,
+        dim_mult: list[int] = [1, 2, 4, 4],
+        num_res_blocks: int = 2,
+        attn_scales: list[float] = [],
+        temperal_downsample: list[bool] = [False, True, True],
+        dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+        dims = [dim * u for u in [1] + dim_mult]
+        scale = 1.0
+
+        self.conv1 = Cosmos25CausalConv3d(3, dims[0], 3, padding=1)
+
+        downsamples: list[nn.Module] = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            for _ in range(num_res_blocks):
+                downsamples.append(Cosmos25ResidualBlock(in_dim, out_dim, dropout))
+                if scale in attn_scales:
+                    downsamples.append(Cosmos25AttentionBlock(out_dim))
+                in_dim = out_dim
+
+            if i != len(dim_mult) - 1:
+                mode = "downsample3d" if temperal_downsample[i] else "downsample2d"
+                downsamples.append(Cosmos25Resample(out_dim, mode=mode))
+                scale /= 2.0
+        self.downsamples = nn.Sequential(*downsamples)
+
+        self.middle = nn.Sequential(
+            Cosmos25ResidualBlock(out_dim, out_dim, dropout),
+            Cosmos25AttentionBlock(out_dim),
+            Cosmos25ResidualBlock(out_dim, out_dim, dropout),
+        )
+
+        self.head = nn.Sequential(
+            Cosmos25RMSNorm(out_dim, images=False),
+            nn.SiLU(),
+            Cosmos25CausalConv3d(out_dim, z_dim, 3, padding=1),
+        )
+
+    def forward(self, x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor:
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        for layer in self.downsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)  # type: ignore[misc]
+            else:
+                x = layer(x)  # type: ignore[misc]
+
+        for layer in self.middle:
+            if isinstance(layer, Cosmos25ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)  # type: ignore[misc]
+
+        for layer in self.head:
+            if isinstance(layer, Cosmos25CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)  # type: ignore[misc]
+        return x
+
+
+class Cosmos25Decoder3d(nn.Module):
+    def __init__(
+        self,
+        dim: int = 96,
+        z_dim: int = 16,
+        dim_mult: list[int] = [1, 2, 4, 4],
+        num_res_blocks: int = 2,
+        attn_scales: list[float] = [],
+        temperal_upsample: list[bool] = [False, True, True],
+        dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+        dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
+        scale = 1.0 / 2 ** (len(dim_mult) - 2)
+
+        self.conv1 = Cosmos25CausalConv3d(z_dim, dims[0], 3, padding=1)
+        self.middle = nn.Sequential(
+            Cosmos25ResidualBlock(dims[0], dims[0], dropout),
+            Cosmos25AttentionBlock(dims[0]),
+            Cosmos25ResidualBlock(dims[0], dims[0], dropout),
+        )
+
+        upsamples: list[nn.Module] = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            if i in (1, 2, 3):
+                in_dim = in_dim // 2
+            for _ in range(num_res_blocks + 1):
+                upsamples.append(Cosmos25ResidualBlock(in_dim, out_dim, dropout))
+                if scale in attn_scales:
+                    upsamples.append(Cosmos25AttentionBlock(out_dim))
+                in_dim = out_dim
+            if i != len(dim_mult) - 1:
+                mode = "upsample3d" if temperal_upsample[i] else "upsample2d"
+                upsamples.append(Cosmos25Resample(out_dim, mode=mode))
+                scale *= 2.0
+        self.upsamples = nn.Sequential(*upsamples)
+
+        self.head = nn.Sequential(
+            Cosmos25RMSNorm(out_dim, images=False),
+            nn.SiLU(),
+            Cosmos25CausalConv3d(out_dim, 3, 3, padding=1),
+        )
+
+    def forward(self, x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor:
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        for layer in self.middle:
+            if isinstance(layer, Cosmos25ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)  # type: ignore[misc]
+
+        for layer in self.upsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)  # type: ignore[misc]
+            else:
+                x = layer(x)  # type: ignore[misc]
+
+        for layer in self.head:
+            if isinstance(layer, Cosmos25CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)  # type: ignore[misc]
+        return x
+
+
+def _count_cosmos25_conv3d(model: nn.Module) -> int:
+    return sum(1 for m in model.modules() if isinstance(m, Cosmos25CausalConv3d))
+
+
+class Cosmos25WanVAE(nn.Module):
+    """
+    A FastVideo-native copy of the *official-like* Wan2.1 VAE core.
+
+    Key properties:
+    - Module naming matches official tokenizer (`encoder`, `decoder`, `conv1`, `conv2`)
+      so it can consume `tokenizer.pth` keys directly.
+    - `encode()` returns **normalized** latents and `decode()` expects **normalized**
+      latents (Cosmos2.5 contract), matching `Wan2pt1VAEInterface`.
+    """
+
+    handles_latent_norm: bool = True
+    handles_latent_denorm: bool = True
+
+    def __init__(
+        self,
+        *,
+        device: torch.device | str = "cpu",
+        dtype: torch.dtype = torch.float32,
+        temporal_window: int = 4,
+        latents_mean: Optional[torch.Tensor] = None,
+        latents_std: Optional[torch.Tensor] = None,
+    ) -> None:
+        super().__init__()
+
+        # Official hyperparams for Cosmos2.5 tokenizer (Wan2.1 VAE).
+        cfg = dict(
+            dim=96,
+            z_dim=16,
+            dim_mult=[1, 2, 4, 4],
+            num_res_blocks=2,
+            attn_scales=[],
+            temperal_downsample=[False, True, True],
+            dropout=0.0,
+            temporal_window=temporal_window,
+        )
+        self.z_dim = 16
+        self.temporal_window = temporal_window
+
+        self.encoder = Cosmos25Encoder3d(
+            dim=cfg["dim"],
+            z_dim=cfg["z_dim"] * 2,
+            dim_mult=cfg["dim_mult"],
+            num_res_blocks=cfg["num_res_blocks"],
+            attn_scales=cfg["attn_scales"],
+            temperal_downsample=cfg["temperal_downsample"],
+            dropout=cfg["dropout"],
+        )
+        self.conv1 = Cosmos25CausalConv3d(self.z_dim * 2, self.z_dim * 2, 1)
+        self.conv2 = Cosmos25CausalConv3d(self.z_dim, self.z_dim, 1)
+        self.decoder = Cosmos25Decoder3d(
+            dim=cfg["dim"],
+            z_dim=cfg["z_dim"],
+            dim_mult=cfg["dim_mult"],
+            num_res_blocks=cfg["num_res_blocks"],
+            attn_scales=cfg["attn_scales"],
+            temperal_upsample=list(cfg["temperal_downsample"])[::-1],
+            dropout=cfg["dropout"],
+        )
+
+        # Default Cosmos2.5 latent stats (shared with configs).
+        if latents_mean is None:
+            latents_mean = torch.tensor(
+                [
+                    -0.7571,
+                    -0.7089,
+                    -0.9113,
+                    0.1075,
+                    -0.1745,
+                    0.9653,
+                    -0.1517,
+                    1.5508,
+                    0.4134,
+                    -0.0715,
+                    0.5517,
+                    -0.3632,
+                    -0.1922,
+                    -0.9497,
+                    0.2503,
+                    -0.2921,
+                ],
+                dtype=torch.float32,
+            ).view(1, 16, 1, 1, 1)
+        if latents_std is None:
+            latents_std = torch.tensor(
+                [
+                    2.8184,
+                    1.4541,
+                    2.3275,
+                    2.6558,
+                    1.2196,
+                    1.7708,
+                    2.6052,
+                    2.0743,
+                    3.2687,
+                    2.1526,
+                    2.8652,
+                    1.5579,
+                    1.6382,
+                    1.1253,
+                    2.8251,
+                    1.9160,
+                ],
+                dtype=torch.float32,
+            ).view(1, 16, 1, 1, 1)
+
+        self.register_buffer("_latents_mean", latents_mean, persistent=False)
+        self.register_buffer("_latents_std", latents_std, persistent=False)
+
+        self.to(device=device, dtype=dtype)
+        self.clear_cache()
+
+    def clear_cache(self) -> None:
+        # Decoder cache
+        self._conv_num = _count_cosmos25_conv3d(self.decoder)
+        self._conv_idx = [0]
+        self._feat_map: list[Any] = [None] * self._conv_num
+        # Encoder cache
+        self._enc_conv_num = _count_cosmos25_conv3d(self.encoder)
+        self._enc_conv_idx = [0]
+        self._enc_feat_map: list[Any] = [None] * self._enc_conv_num
+
+    def _scale(self, like: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        mean = self._latents_mean.to(device=like.device, dtype=like.dtype)
+        std = self._latents_std.to(device=like.device, dtype=like.dtype)
+        return mean, 1.0 / std
+
+    def _i0_encode(self, x: torch.Tensor) -> torch.Tensor:
+        return self.encoder(x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx)
+
+    def _i0_decode(self, x: torch.Tensor) -> torch.Tensor:
+        return self.decoder(x[:, :, 0:1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx)
+
+    def encode(self, x: torch.Tensor) -> _TensorLatentDist:
+        """
+        Encode to *normalized* latents (Cosmos contract).
+        """
+        self.clear_cache()
+        t = x.shape[2]
+        iters = 1 + (t - 1) // self.temporal_window
+
+        for i in range(iters):
+            self._enc_conv_idx = [0]
+            if i == 0:
+                out = self._i0_encode(x)
+            else:
+                out_ = self.encoder(
+                    x[:, :, 1 + self.temporal_window * (i - 1) : 1 + self.temporal_window * i, :, :],
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx,
+                )
+                out = torch.cat([out, out_], 2)
+
+        if (t - 1) % self.temporal_window:
+            self._enc_conv_idx = [0]
+            out_ = self.encoder(
+                x[:, :, 1 + self.temporal_window * (iters - 1) :, :, :],
+                feat_cache=self._enc_feat_map,
+                feat_idx=self._enc_conv_idx,
+            )
+            out = torch.cat([out, out_], 2)
+
+        mu, _log_var = self.conv1(out).chunk(2, dim=1)
+        mean, inv_std = self._scale(mu)
+        z_norm = (mu - mean) * inv_std
+        self.clear_cache()
+        return _TensorLatentDist(z_norm)
+
+    def decode(self, latent: torch.Tensor) -> torch.Tensor:
+        """
+        Decode from *normalized* latents (Cosmos contract).
+        """
+        self.clear_cache()
+        mean, inv_std = self._scale(latent)
+        z = latent / inv_std + mean  # z = z_norm * std + mean
+
+        iter_ = z.shape[2]
+        x = self.conv2(z)
+        for i in range(iter_):
+            self._conv_idx = [0]
+            if i == 0:
+                out = self._i0_decode(x)
+            else:
+                out_ = self.decoder(x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx)
+                out = torch.cat([out, out_], 2)
+
+        self.clear_cache()
+        return out
+
+    # --- Interface helpers (match official Wan2pt1VAEInterface) ---
+    def get_latent_num_frames(self, num_pixel_frames: int) -> int:
+        return 1 + (int(num_pixel_frames) - 1) // 4
+
+    def get_pixel_num_frames(self, num_latent_frames: int) -> int:
+        return (int(num_latent_frames) - 1) * 4 + 1
+
+    @property
+    def spatial_compression_factor(self) -> int:
+        return 8
+
+    @property
+    def temporal_compression_factor(self) -> int:
+        return 4
+
+    @property
+    def latent_ch(self) -> int:
+        return 16
+
+
+
+
+
+
+
+
+

fastvideo/pipelines/basic/cosmos/cosmos2_5_pipeline.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos 2.5 pipeline entry (staged pipeline)."""
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.stages import (ConditioningStage,
+                                        Cosmos25DenoisingStage,
+                                        Cosmos25LatentPreparationStage,
+                                        DecodingStage, InputValidationStage,
+                                        Cosmos25TextEncodingStage,
+                                        Cosmos25TimestepPreparationStage)
+
+logger = init_logger(__name__)
+
+
+class Cosmos2_5Pipeline(ComposedPipelineBase):
+    """Cosmos 2.5 video generation pipeline."""
+
+    _required_config_modules = [
+        "text_encoder", "tokenizer", "vae", "transformer", "scheduler",
+        "safety_checker"
+    ]
+
+    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())
+
+        self.add_stage(
+            stage_name="prompt_encoding_stage",
+            stage=Cosmos25TextEncodingStage(
+                text_encoder=self.get_module("text_encoder"), ),
+        )
+
+        self.add_stage(stage_name="conditioning_stage",
+                       stage=ConditioningStage())
+
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=Cosmos25TimestepPreparationStage(
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=Cosmos25LatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer"),
+                           vae=self.get_module("vae")))
+
+        self.add_stage(stage_name="denoising_stage",
+                       stage=Cosmos25DenoisingStage(
+                           transformer=self.get_module("transformer"),
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae")))
+        logger.info("Cosmos 2.5 pipeline stages created")
+
+
+# Entry point for pipeline registry
+EntryClass = Cosmos2_5Pipeline

fastvideo/pipelines/pipeline_registry.py

@@ -29,6 +29,7 @@ _PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = {
     "HunyuanVideoPipeline": "hunyuan",
     "HunyuanVideo15Pipeline": "hunyuan15",
     "Cosmos2VideoToWorldPipeline": "cosmos",
+    "Cosmos2_5Pipeline": "cosmos",
     "MatrixGamePipeline": "matrixgame",
     "MatrixGameCausalDMDPipeline": "matrixgame",
     "LongCatPipeline": "longcat",

fastvideo/pipelines/stages/__init__.py

@@ -10,7 +10,8 @@ from fastvideo.pipelines.stages.base import PipelineStage
 from fastvideo.pipelines.stages.causal_denoising import CausalDMDDenosingStage
 from fastvideo.pipelines.stages.conditioning import ConditioningStage
 from fastvideo.pipelines.stages.decoding import DecodingStage
-from fastvideo.pipelines.stages.denoising import (CosmosDenoisingStage,
+from fastvideo.pipelines.stages.denoising import (Cosmos25DenoisingStage,
+                                                  CosmosDenoisingStage,
                                                   DenoisingStage,
                                                   DmdDenoisingStage)
 from fastvideo.pipelines.stages.encoding import EncodingStage
@@ -19,14 +20,16 @@ from fastvideo.pipelines.stages.image_encoding import (
     ImageVAEEncodingStage, VideoVAEEncodingStage, Hy15ImageEncodingStage)
 from fastvideo.pipelines.stages.input_validation import InputValidationStage
 from fastvideo.pipelines.stages.latent_preparation import (
-    CosmosLatentPreparationStage, LatentPreparationStage)
+    Cosmos25LatentPreparationStage, CosmosLatentPreparationStage,
+    LatentPreparationStage)
 from fastvideo.pipelines.stages.matrixgame_denoising import (
     MatrixGameCausalDenoisingStage)
 from fastvideo.pipelines.stages.stepvideo_encoding import (
     StepvideoPromptEncodingStage)
-from fastvideo.pipelines.stages.text_encoding import TextEncodingStage
+from fastvideo.pipelines.stages.text_encoding import (Cosmos25TextEncodingStage,
+                                                      TextEncodingStage)
 from fastvideo.pipelines.stages.timestep_preparation import (
-    TimestepPreparationStage)
+    Cosmos25TimestepPreparationStage, TimestepPreparationStage)
 
 # LongCat stages
 from fastvideo.pipelines.stages.longcat_video_vae_encoding import LongCatVideoVAEEncodingStage
@@ -37,14 +40,17 @@ __all__ = [
     "PipelineStage",
     "InputValidationStage",
     "TimestepPreparationStage",
+    "Cosmos25TimestepPreparationStage",
     "LatentPreparationStage",
     "CosmosLatentPreparationStage",
+    "Cosmos25LatentPreparationStage",
     "ConditioningStage",
     "DenoisingStage",
     "DmdDenoisingStage",
     "CausalDMDDenosingStage",
     "MatrixGameCausalDenoisingStage",
     "CosmosDenoisingStage",
+    "Cosmos25DenoisingStage",
     "EncodingStage",
     "DecodingStage",
     "ImageEncodingStage",
@@ -54,6 +60,7 @@ __all__ = [
     "ImageVAEEncodingStage",
     "VideoVAEEncodingStage",
     "TextEncodingStage",
+    "Cosmos25TextEncodingStage",
     "StepvideoPromptEncodingStage",
     # LongCat stages
     "LongCatVideoVAEEncodingStage",

fastvideo/pipelines/stages/decoding.py

@@ -51,39 +51,42 @@ class DecodingStage(PipelineStage):
         return result
 
     def _denormalize_latents(self, latents: torch.Tensor) -> torch.Tensor:
-        # denormalization for MatrixGame VAE
-        # z = z * std + mean during decode
-        if (hasattr(self.vae.config, 'latents_mean')
-                and hasattr(self.vae.config, 'latents_std')):
-            # Convert config values to tensors
-            latents_mean = torch.tensor(self.vae.config.latents_mean,
+        """Convert normalized latents into the VAE's expected latent space."""
+        # Some VAEs handle latent (de)normalization internally.
+        if bool(getattr(self.vae, "handles_latent_denorm", False)):
+            return latents
+
+        cfg = getattr(self.vae, "config", None)
+
+        # MatrixGame-style: z = z * std + mean
+        if (cfg is not None and hasattr(cfg, "latents_mean")
+                and hasattr(cfg, "latents_std")):
+            latents_mean = torch.tensor(cfg.latents_mean,
                                         device=latents.device,
                                         dtype=latents.dtype).view(
                                             1, -1, 1, 1, 1)
-
-            latents_std = torch.tensor(self.vae.config.latents_std,
+            latents_std = torch.tensor(cfg.latents_std,
                                        device=latents.device,
                                        dtype=latents.dtype).view(
                                            1, -1, 1, 1, 1)
+            return latents * latents_std + latents_mean
 
-            # Apply denormalization: z = z * std + mean
-            latents = latents * latents_std + latents_mean
-        elif hasattr(self.vae, 'scaling_factor'):
-            # Standard VAE scaling
+        # Diffusers-style: scaling_factor (+ optional shift_factor)
+        if hasattr(self.vae, "scaling_factor"):
             if isinstance(self.vae.scaling_factor, torch.Tensor):
                 latents = latents / self.vae.scaling_factor.to(
                     latents.device, latents.dtype)
             else:
                 latents = latents / self.vae.scaling_factor
 
-            # Apply shifting if needed
-            if (hasattr(self.vae, "shift_factor")
-                    and self.vae.shift_factor is not None):
+            if hasattr(self.vae,
+                       "shift_factor") and self.vae.shift_factor is not None:
                 if isinstance(self.vae.shift_factor, torch.Tensor):
-                    latents += self.vae.shift_factor.to(latents.device,
-                                                        latents.dtype)
+                    latents = latents + self.vae.shift_factor.to(
+                        latents.device, latents.dtype)
                 else:
-                    latents += self.vae.shift_factor
+                    latents = latents + self.vae.shift_factor
+
         return latents
 
     @torch.no_grad()
@@ -273,4 +276,4 @@ class DecodingStage(PipelineStage):
                 del pipeline.modules["vae"]
             fastvideo_args.model_loaded["vae"] = False
 
-        return batch
+        return batch

fastvideo/pipelines/stages/denoising.py

@@ -1018,6 +1018,145 @@ class CosmosDenoisingStage(DenoisingStage):
         return result
 
 
+class Cosmos25DenoisingStage(CosmosDenoisingStage):
+    """Denoising stage for Cosmos 2.5 DiT (expects 1D/2D timestep, not 5D)."""
+
+    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
+
+        extra_step_kwargs = self.prepare_extra_func_kwargs(
+            self.scheduler.step,
+            {
+                "generator": batch.generator,
+                "eta": batch.eta
+            },
+        )
+
+        if hasattr(self.transformer, 'module'):
+            transformer_dtype = next(self.transformer.module.parameters()).dtype
+        else:
+            transformer_dtype = next(self.transformer.parameters()).dtype
+        target_dtype = transformer_dtype
+        autocast_enabled = (target_dtype != torch.float32
+                            ) and not fastvideo_args.disable_autocast
+
+        latents = batch.latents
+        if latents is None:
+            raise ValueError(
+                "latents must be provided for Cosmos25DenoisingStage")
+        guidance_scale = batch.guidance_scale
+
+        # Use timesteps prepared by Cosmos25TimestepPreparationStage when available.
+        if batch.timesteps is None:
+            self.scheduler.set_timesteps(batch.num_inference_steps,
+                                         device=latents.device)
+            timesteps = self.scheduler.timesteps
+        else:
+            timesteps = batch.timesteps.to(latents.device)
+
+        # Match official behavior: pass fps as a tensor.
+        fps_val = batch.fps if isinstance(batch.fps, int | float) else 24
+        fps_tensor = torch.tensor([fps_val],
+                                  device=latents.device,
+                                  dtype=target_dtype)
+
+        # Cosmos2.5 denoises a 4D latent (C,T,H,W) and the scheduler.step expects (B,C,T,H,W).
+        latents_4d = latents[0]
+
+        # Masks from latent prep stage
+        condition_mask = batch.cond_mask.to(target_dtype) if hasattr(
+            batch, 'cond_mask') else None
+        padding_mask = batch.padding_mask.to(target_dtype) if hasattr(
+            batch, 'padding_mask') else None
+        if condition_mask is None:
+            _, t, h, w = latents_4d.shape
+            condition_mask = torch.zeros(1,
+                                         1,
+                                         t,
+                                         h,
+                                         w,
+                                         device=latents.device,
+                                         dtype=target_dtype)
+        if padding_mask is None:
+            _, _, h, w = latents_4d.shape
+            padding_mask = torch.ones(1,
+                                      1,
+                                      h,
+                                      w,
+                                      device=latents.device,
+                                      dtype=target_dtype)
+
+        # Cosmos2.5 timestep scaling (see compare_pipelines.py): t * 0.001
+        timestep_scale = 0.001
+
+        with self.progress_bar(total=len(timesteps)) as progress_bar:
+            for i, t in enumerate(timesteps):
+                t_val = float(t)
+                timestep_val = t_val * timestep_scale
+                timestep = torch.tensor([[timestep_val]],
+                                        device=latents.device,
+                                        dtype=target_dtype)
+
+                model_hidden_states = latents_4d.unsqueeze(0)
+
+                with (
+                        set_forward_context(current_timestep=int(t_val),
+                                            attn_metadata=None,
+                                            forward_batch=batch),
+                        torch.autocast(device_type="cuda",
+                                       dtype=target_dtype,
+                                       enabled=autocast_enabled),
+                ):
+                    cond_v = self.transformer(
+                        hidden_states=model_hidden_states.to(target_dtype),
+                        encoder_hidden_states=batch.prompt_embeds[0].to(
+                            target_dtype),
+                        timestep=timestep,
+                        fps=fps_tensor,
+                        condition_mask=condition_mask,
+                        padding_mask=padding_mask,
+                        return_dict=False,
+                    )[0]
+
+                    if batch.do_classifier_free_guidance and batch.negative_prompt_embeds:
+                        uncond_v = self.transformer(
+                            hidden_states=model_hidden_states.to(target_dtype),
+                            encoder_hidden_states=batch.
+                            negative_prompt_embeds[0].to(target_dtype),
+                            timestep=timestep,
+                            fps=fps_tensor,
+                            condition_mask=condition_mask,
+                            padding_mask=padding_mask,
+                            return_dict=False,
+                        )[0]
+                        v = uncond_v + guidance_scale * (cond_v - uncond_v)
+                    else:
+                        v = cond_v
+
+                prev = self.scheduler.step(v.unsqueeze(0),
+                                           t,
+                                           latents_4d.unsqueeze(0),
+                                           **extra_step_kwargs,
+                                           return_dict=False)[0]
+                latents_4d = prev.squeeze(0)
+
+                progress_bar.update()
+
+        batch.latents = latents_4d.unsqueeze(0)
+        return batch
+
+
 class DmdDenoisingStage(DenoisingStage):
     """
     Denoising stage for DMD.

fastvideo/pipelines/stages/latent_preparation.py

@@ -5,6 +5,7 @@ Latent preparation stage for diffusion pipelines.
 
 from typing import Any
 
+import numpy as np
 import torch
 from diffusers.utils.torch_utils import randn_tensor
 
@@ -427,6 +428,210 @@ class CosmosLatentPreparationStage(PipelineStage):
 
         return batch
 
+
+class Cosmos25LatentPreparationStage(CosmosLatentPreparationStage):
+    """Latent preparation for Cosmos 2.5 DiT input conventions."""
+
+    @staticmethod
+    def _arch_invariant_randn(
+        shape: tuple[int, ...],
+        *,
+        seed: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Architecture-invariant RNG (matches cosmos_predict2.misc.arch_invariant_rand)."""
+        rng = np.random.RandomState(seed)
+        arr = rng.standard_normal(shape).astype(np.float32)
+        return torch.from_numpy(arr).to(device=device, dtype=dtype)
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        # Differences vs `CosmosLatentPreparationStage`: channel convention, seed usage,
+        # and `padding_mask` for concat_padding_mask=True.
+
+        # Determine batch size
+        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
+
+        # Match `compare_pipelines.py`: initialize noise in fp32, then run the
+        # denoising computation in bf16.
+        dtype = torch.float32
+        device = get_local_torch_device()
+        generator = batch.generator
+        latents = batch.latents
+        num_frames = batch.num_frames
+        height = batch.height
+        width = batch.width
+
+        if height is None or width is None:
+            raise ValueError("Height and width must be provided")
+
+        vae_scale_factor_spatial = 8
+        vae_scale_factor_temporal = 4
+
+        latent_height = height // 8
+        latent_width = width // vae_scale_factor_spatial
+        num_latent_frames = (num_frames - 1) // vae_scale_factor_temporal + 1
+
+        # Cosmos 2.5 convention: transformer in_channels == latent channels
+        num_channels_latents = self.transformer.config.in_channels
+
+        shape = (batch_size, num_channels_latents, num_latent_frames,
+                 latent_height, latent_width)
+
+        init_latents = None
+        conditioning_latents = None
+        video = None
+
+        if hasattr(batch, 'video') and batch.video is not None:
+            video = batch.video
+        elif hasattr(batch, 'pil_image') and batch.pil_image is not None:
+            vae_scale_factor_spatial = 8
+            image_processor = ImageProcessor(
+                vae_scale_factor=vae_scale_factor_spatial)
+            processed_image = image_processor.preprocess(
+                batch.pil_image, height, width)
+            video = processed_image.unsqueeze(2)
+            video = video.to(device=device, dtype=torch.bfloat16)
+        elif hasattr(
+                batch,
+                'preprocessed_image') and batch.preprocessed_image is not None:
+            if isinstance(batch.preprocessed_image, torch.Tensor):
+                if batch.preprocessed_image.dim() == 4:
+                    video = batch.preprocessed_image.unsqueeze(2)
+                elif batch.preprocessed_image.dim() == 5:
+                    video = batch.preprocessed_image
+        else:
+            logger.info(
+                "CosmosLatentPreparationStage - No video input sources found")
+
+        if video is not None:
+            num_cond_frames = video.size(2)
+            if num_cond_frames >= num_frames:
+                num_cond_latent_frames = (num_frames -
+                                          1) // vae_scale_factor_temporal + 1
+                video = video[:, :, -num_frames:]
+            else:
+                num_cond_latent_frames = (num_cond_frames -
+                                          1) // vae_scale_factor_temporal + 1
+                num_padding_frames = num_frames - num_cond_frames
+                last_frame = video[:, :, -1:]
+                padding = last_frame.repeat(1, 1, num_padding_frames, 1, 1)
+                video = torch.cat([video, padding], dim=2)
+
+            if self.vae is not None:
+                self.vae = self.vae.to(device)
+                self.vae = self.vae.to(dtype=video.dtype)
+
+                def retrieve_latents(
+                        encoder_output: Any,
+                        generator: Any | None = None) -> torch.Tensor:
+                    if hasattr(encoder_output, "latent_dist"):
+                        return encoder_output.latent_dist.sample(generator)
+                    elif hasattr(encoder_output, "latents"):
+                        return encoder_output.latents
+                    elif hasattr(encoder_output, "sample"):
+                        return encoder_output.sample(generator)
+                    elif isinstance(encoder_output, torch.Tensor):
+                        return encoder_output
+                    else:
+                        attrs = [
+                            attr for attr in dir(encoder_output)
+                            if not attr.startswith('_')
+                        ]
+                        raise AttributeError(
+                            f"Could not access latents of provided encoder_output. Available attributes: {attrs}"
+                        )
+
+                if isinstance(generator, list):
+                    init_latents = [
+                        retrieve_latents(self.vae.encode(video[i].unsqueeze(0)),
+                                         generator=torch.Generator(
+                                             device="cpu").manual_seed(100))
+                        for i in range(batch_size)
+                    ]
+                else:
+                    init_latents = [
+                        retrieve_latents(
+                            self.vae.encode(vid.unsqueeze(0)),
+                            torch.Generator(device="cpu").manual_seed(100))
+                        for vid in video
+                    ]
+
+                init_latents = torch.cat(init_latents, dim=0).to(dtype)
+
+                cfg = getattr(self.vae, "config", None)
+                if (not bool(getattr(self.vae, "handles_latent_norm", False))
+                        and cfg is not None and hasattr(cfg, 'latents_mean')
+                        and hasattr(cfg, 'latents_std')):
+                    latents_mean = torch.tensor(cfg.latents_mean).view(
+                        1, cfg.z_dim, 1, 1, 1).to(device, dtype)
+                    latents_std = torch.tensor(cfg.latents_std).view(
+                        1, cfg.z_dim, 1, 1, 1).to(device, dtype)
+                    init_latents = (init_latents - latents_mean
+                                    ) / latents_std * self.scheduler.sigma_data
+
+                conditioning_latents = init_latents
+                self.vae.to("cpu")
+        else:
+            num_cond_latent_frames = 0
+
+        if latents is None:
+            seed = int(batch.seed if batch.seed is not None else 0)
+            # Use arch-invariant RNG to match Cosmos2.5 reference sampling.
+            latents_fp32 = self._arch_invariant_randn(shape,
+                                                      seed=seed,
+                                                      device=device,
+                                                      dtype=torch.float32)
+            latents = latents_fp32.to(torch.bfloat16)
+        else:
+            # If latents are supplied, keep compute dtype consistent with Cosmos sampling.
+            latents = latents.to(device=device, dtype=torch.bfloat16)
+
+        # Cosmos2.5 starts from unit Gaussian noise (no extra sigma_max scaling).
+
+        padding_shape = (batch_size, 1, num_latent_frames, latent_height,
+                         latent_width)
+        ones_padding = latents.new_ones(padding_shape)
+        zeros_padding = latents.new_zeros(padding_shape)
+
+        cond_indicator = latents.new_zeros(1, 1, latents.size(2), 1, 1)
+        cond_indicator[:, :, :num_cond_latent_frames] = 1.0
+        cond_mask = cond_indicator * ones_padding + (
+            1 - cond_indicator) * zeros_padding
+
+        uncond_indicator = None
+        uncond_mask = None
+        if batch.do_classifier_free_guidance:
+            uncond_indicator = latents.new_zeros(1, 1, latents.size(2), 1, 1)
+            uncond_indicator[:, :, :num_cond_latent_frames] = 1.0
+            uncond_mask = uncond_indicator * ones_padding + (
+                1 - uncond_indicator) * zeros_padding
+
+        # Cosmos 2.5 requires a spatial padding mask when concat_padding_mask=True
+        padding_mask = latents.new_ones(batch_size, 1, latent_height,
+                                        latent_width)
+
+        batch.latents = latents
+        batch.raw_latent_shape = latents.shape
+        batch.conditioning_latents = conditioning_latents
+        batch.cond_indicator = cond_indicator
+        batch.uncond_indicator = uncond_indicator
+        batch.cond_mask = cond_mask
+        batch.uncond_mask = uncond_mask
+        batch.padding_mask = padding_mask
+        return batch
+
     def adjust_video_length(self, batch: ForwardBatch,
                             fastvideo_args: FastVideoArgs) -> int:
         """

fastvideo/pipelines/stages/text_encoding.py

@@ -328,3 +328,69 @@ class TextEncodingStage(PipelineStage):
             lambda x: not batch.do_classifier_free_guidance or V.
             list_of_tensors_with_min_dims(x, 2))
         return result
+
+
+class Cosmos25TextEncodingStage(PipelineStage):
+    """Cosmos 2.5 text encoding stage.
+
+    Cosmos 2.5 uses Reason1 (Qwen2.5-VL) and relies on the encoder's
+    `compute_text_embeddings_online()`.
+    """
+
+    def __init__(self, text_encoder) -> None:
+        super().__init__()
+        self.text_encoder = text_encoder
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch,
+                fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        assert batch.prompt is not None
+        prompts = [batch.prompt] if isinstance(batch.prompt,
+                                               str) else batch.prompt
+
+        encoder = self.text_encoder
+        if not hasattr(encoder, "compute_text_embeddings_online"):
+            raise RuntimeError(
+                "Cosmos25TextEncodingStage requires text_encoder.compute_text_embeddings_online()"
+            )
+
+        with set_forward_context(current_timestep=0, attn_metadata=None):
+            prompt_embeds = encoder.compute_text_embeddings_online(
+                {"text": prompts}, "text")
+
+        batch.prompt_embeds = [prompt_embeds]
+
+        if batch.do_classifier_free_guidance:
+            neg = batch.negative_prompt
+            neg_prompts = ([neg] *
+                           len(prompts)) if isinstance(neg, str) else neg
+            with set_forward_context(current_timestep=0, attn_metadata=None):
+                neg_embeds = encoder.compute_text_embeddings_online(
+                    {"text": neg_prompts}, "text")
+            batch.negative_prompt_embeds = [neg_embeds]
+        else:
+            batch.negative_prompt_embeds = []
+
+        return batch
+
+    def verify_input(self, batch: ForwardBatch,
+                     fastvideo_args: FastVideoArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("prompt", batch.prompt, V.string_or_list_strings)
+        result.add_check(
+            "negative_prompt",
+            batch.negative_prompt,
+            lambda x:
+            (not batch.do_classifier_free_guidance) or isinstance(x, str),
+        )
+        return result
+
+    def verify_output(self, batch: ForwardBatch,
+                      fastvideo_args: FastVideoArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("prompt_embeds", batch.prompt_embeds,
+                         V.list_of_tensors_min_dims(2))
+        result.add_check(
+            "negative_prompt_embeds", batch.negative_prompt_embeds, lambda x:
+            not batch.do_classifier_free_guidance or V.list_not_empty(x))
+        return result

fastvideo/pipelines/stages/timestep_preparation.py

@@ -123,3 +123,32 @@ class TimestepPreparationStage(PipelineStage):
         result.add_check("timesteps", batch.timesteps,
                          [V.is_tensor, V.with_dims(1)])
         return result
+
+
+class Cosmos25TimestepPreparationStage(TimestepPreparationStage):
+    """Cosmos 2.5 timestep preparation with scheduler-specific kwargs."""
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        scheduler = self.scheduler
+        device = get_local_torch_device()
+        num_inference_steps = batch.num_inference_steps
+
+        extra_kwargs: dict = {}
+        sig = inspect.signature(scheduler.set_timesteps)
+        if "shift" in sig.parameters:
+            extra_kwargs["shift"] = fastvideo_args.pipeline_config.flow_shift
+        # Prefer the canonical diffusers kwarg name if available.
+        if "use_karras_sigmas" in sig.parameters:
+            extra_kwargs["use_karras_sigmas"] = True
+        elif "use_kerras_sigma" in sig.parameters:
+            extra_kwargs["use_kerras_sigma"] = True
+
+        scheduler.set_timesteps(num_inference_steps,
+                                device=device,
+                                **extra_kwargs)
+        batch.timesteps = scheduler.timesteps
+        return batch

scripts/checkpoint_conversion/pt_to_safetensors.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+"""Convert a PyTorch checkpoint (.pt) to a safetensors file."""
+
+import argparse
+from pathlib import Path
+
+import torch
+from safetensors.torch import save_file
+
+
+def convert_pt_to_safetensors(
+    input_path: str,
+    output_path: str,
+    key: str | None = None,
+    force: bool = False,
+    skip_patterns: list[str] | None = None,
+):
+    input_path = Path(input_path)
+    output_path = Path(output_path)
+
+    if not input_path.exists():
+        raise FileNotFoundError(f"Input file not found: {input_path}")
+
+    if output_path.exists() and not force:
+        raise FileExistsError(
+            f"Output file already exists: {output_path}. Use --force to overwrite."
+        )
+
+    checkpoint = torch.load(input_path, map_location="cpu")
+
+    state_dict: dict[str, torch.Tensor]
+    if isinstance(checkpoint, dict):
+        if key is not None:
+            if key not in checkpoint:
+                raise KeyError(f"Key {key!r} not found in checkpoint.")
+            state_dict = checkpoint[key]
+        else:
+            for k in ("state_dict", "model_state_dict", "model", "ema"):
+                if k in checkpoint:
+                    state_dict = checkpoint[k]
+                    break
+            else:
+                state_dict = checkpoint
+    else:
+        state_dict = checkpoint
+
+    if not isinstance(state_dict, dict):
+        raise TypeError(f"Expected a dict state_dict, got {type(state_dict)}")
+
+    if skip_patterns:
+        state_dict = {
+            k: v
+            for k, v in state_dict.items()
+            if not any(pat in k for pat in skip_patterns)
+        }
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    save_file(state_dict, str(output_path))
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Convert a PyTorch checkpoint (.pt) to safetensors."
+    )
+    parser.add_argument(
+        "input",
+        type=str,
+        help="Path to input .pt checkpoint file"
+    )
+    parser.add_argument(
+        "output",
+        type=str,
+        help="Path to output .safetensors file"
+    )
+    parser.add_argument(
+        "--key",
+        type=str,
+        default=None,
+        help="Optional key to extract from checkpoint dict (e.g. 'model_state_dict')"
+    )
+    parser.add_argument(
+        "--force",
+        action="store_true",
+        help="Overwrite output file if it exists"
+    )
+    parser.add_argument(
+        "--skip-pattern",
+        action="append",
+        dest="skip_patterns",
+        help="Parameter name patterns to skip (can be used multiple times)"
+    )
+
+    args = parser.parse_args()
+
+    convert_pt_to_safetensors(
+        args.input,
+        args.output,
+        args.key,
+        args.force,
+        args.skip_patterns
+    )
+
+
+if __name__ == "__main__":
+    main()
+