[feat] Add Cosmos 2.5 I2W/V2W support (staged pipeline + examples) (#1021)

Co-authored-by: Will Lin <wlsaidhi@gmail.com>

examples/inference/basic/basic_cosmos2_5_i2w.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+from fastvideo import VideoGenerator
+from fastvideo.configs.sample import SamplingParam
+
+
+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,
+    )
+
+    sampling_param = SamplingParam.from_pretrained(model_path)
+
+    # image2world example from official repo
+    image_path = "images/bus_terminal.jpg"
+
+    prompt = (
+        "A nighttime city bus terminal gradually shifts from stillness to subtle movement. "
+        "At first, multiple double-decker buses are parked under the glow of overhead lights, "
+        "with a central bus labeled '87D' facing forward and stationary. "
+        "As the video progresses, the bus in the middle moves ahead slowly, its headlights brightening the surrounding area "
+        "and casting reflections onto adjacent vehicles. "
+        "The motion creates space in the lineup, signaling activity within the otherwise quiet station. "
+        "It then comes to a smooth stop, resuming its position in line. "
+        "Overhead signage in Chinese characters remains illuminated, enhancing the vibrant, urban night scene."
+    )
+
+    generator.generate_video(
+        prompt,
+        sampling_param=sampling_param,
+        image_path=str(image_path),
+        num_cond_frames=1,
+        output_path="outputs_video/cosmos2_5_i2w.mp4",
+        save_video=True,
+    )
+
+    generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()
+
+

examples/inference/basic/basic_cosmos2_5_t2w.py

@@ -1,4 +1,5 @@
 from fastvideo import VideoGenerator
+from fastvideo.configs.sample import SamplingParam
 
 
 def main():
@@ -15,19 +16,27 @@ def main():
         pin_cpu_memory=True,
     )
 
+    # Load default sampling parameters (negative_prompt, resolution, steps, etc.)
+    sampling_param = SamplingParam.from_pretrained(model_path)
+
     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."
+        "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(
+    generator.generate_video(
         prompt,
-        negative_prompt="",
-        height=704,
-        width=1280,
-        num_frames=77,
-        num_inference_steps=35,
-        guidance_scale=7.0,
-        fps=24,
+        sampling_param=sampling_param,
         output_path="outputs_video/cosmos2_5_t2w.mp4",
         save_video=True,
     )

examples/inference/basic/basic_cosmos2_5_v2w.py

@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+from fastvideo import VideoGenerator
+from fastvideo.configs.sample import SamplingParam
+
+
+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,
+    )
+
+    sampling_param = SamplingParam.from_pretrained(model_path)
+
+    # video2world example from official repo
+    video_path = "videos/robot_pouring.mp4"
+
+    prompt = (
+        "A robotic arm, primarily white with black joints and cables, is shown in a clean, modern indoor setting with a white tabletop. "
+        "The arm, equipped with a gripper holding a small, light green pitcher, is positioned above a clear glass containing a reddish-brown liquid and a spoon. "
+        "The robotic arm is in the process of pouring a transparent liquid into the glass. "
+        "To the left of the pitcher, there is an opened jar with a similar reddish-brown substance visible through its transparent body. "
+        "In the background, a vase with white flowers and a brown couch are partially visible, adding to the contemporary ambiance. "
+        "The lighting is bright, casting soft shadows on the table. "
+        "The robotic arm's movements are smooth and controlled, demonstrating precision in its task. "
+        "As the video progresses, the robotic arm completes the pour, leaving the glass half-filled with the reddish-brown liquid. "
+        "The jar remains untouched throughout the sequence, and the spoon inside the glass remains stationary. "
+        "The other robotic arm on the right side also stays stationary throughout the video. "
+        "The final frame captures the robotic arm with the pitcher finishing the pour, with the glass now filled to a higher level, while the pitcher is slightly tilted but still held securely by the gripper."
+    )
+
+    generator.generate_video(
+        prompt,
+        sampling_param=sampling_param,
+        video_path=str(video_path),
+        num_cond_frames=1,
+        output_path="outputs_video/cosmos2_5_v2w.mp4",
+        save_video=True,
+    )
+
+    generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()
+
+

fastvideo/configs/sample/cosmos2_5.py

@@ -5,15 +5,19 @@ 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
+class Cosmos25SamplingParamBase(SamplingParam):
+    height: int = 704
+    width: int = 1280
+    num_frames: int = 77
     fps: int = 24
+    seed: int = 0
 
     guidance_scale: float = 7.0
-    # Official Cosmos2.5 sampling uses empty string as unconditional.
-    negative_prompt: str = ""
+    negative_prompt: str = (
+        "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, "
+        "over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, "
+        "underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, "
+        "low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, "
+        "unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. "
+        "Overall, the video is of poor quality.")
     num_inference_steps: int = 35

fastvideo/configs/sample/registry.py

@@ -10,7 +10,7 @@ from fastvideo.configs.sample.hyworld import HYWorld_SamplingParam
 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
+from fastvideo.configs.sample.cosmos2_5 import Cosmos25SamplingParamBase
 from fastvideo.configs.sample.ltx2 import LTX2SamplingParam
 
 # isort: off
@@ -89,8 +89,7 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
     Cosmos_Predict2_2B_Video2World_SamplingParam,
 
     # Cosmos2.5
-    "KyleShao/Cosmos-Predict2.5-2B-Diffusers":
-    Cosmos_Predict2_5_2B_Diffusers_SamplingParam,
+    "KyleShao/Cosmos-Predict2.5-2B-Diffusers": Cosmos25SamplingParamBase,
 
     # MatrixGame2.0 models
     "FastVideo/Matrix-Game-2.0-Base-Diffusers": MatrixGame2_SamplingParam,
@@ -160,7 +159,7 @@ 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,
+    "cosmos25": Cosmos25SamplingParamBase,
     "cosmos": Cosmos_Predict2_2B_Video2World_SamplingParam,
     "ltx2": LTX2SamplingParam,
     # Other fallbacks by architecture

fastvideo/image_processor.py

@@ -193,3 +193,46 @@ class ImageProcessor:
         tensor = 2.0 * tensor - 1.0
 
         return tensor
+
+    def _preprocess_cosmos25(
+        self,
+        image: PIL.Image.Image | np.ndarray | torch.Tensor,
+        height: int,
+        width: int,
+    ) -> torch.Tensor:
+        """
+        Cosmos-Predict2.5-style preprocessing for image2world:
+        - aspect-preserving resize (scale so both dims >= target)
+        - center crop to (height, width)
+        - normalize to [-1, 1]
+
+        Returns:
+            torch.Tensor: (1, 3, height, width) in [-1, 1]
+        """
+        if not isinstance(image, PIL.Image.Image):
+            # Reuse the generic preprocess for non-PIL inputs.
+            return self.preprocess(image, height=height, width=width)
+
+        # Ensure RGB (official uses torchvision.to_tensor on PIL, which yields C=3 for RGB inputs)
+        image = image.convert("RGB")
+
+        # Make sure target dims are multiples of VAE scale factor (FastVideo convention)
+        height = height - (height % self.vae_scale_factor)
+        width = width - (width % self.vae_scale_factor)
+
+        orig_w, orig_h = image.size
+        # Scale so that resized image fully covers the target box (like official resize_input()).
+        scale = max(width / float(orig_w), height / float(orig_h))
+        resized_w = int(np.ceil(scale * orig_w))
+        resized_h = int(np.ceil(scale * orig_h))
+
+        image = image.resize((resized_w, resized_h),
+                             resample=PIL.Image.Resampling.BILINEAR)
+
+        # Center crop
+        left = max(0, (resized_w - width) // 2)
+        top = max(0, (resized_h - height) // 2)
+        image = image.crop((left, top, left + width, top + height))
+
+        image_np = np.array(image, dtype=np.float32) / 255.0  # [0,1]
+        return self._normalize_to_tensor(image_np)  # -> [-1,1], (1,3,H,W)

fastvideo/pipelines/basic/cosmos/cosmos2_5_pipeline.py

@@ -5,8 +5,8 @@ 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,
+                                        Cosmos25AutoDenoisingStage,
+                                        Cosmos25AutoLatentPreparationStage,
                                         DecodingStage, InputValidationStage,
                                         Cosmos25TextEncodingStage,
                                         Cosmos25TimestepPreparationStage)
@@ -42,13 +42,13 @@ class Cosmos2_5Pipeline(ComposedPipelineBase):
                            scheduler=self.get_module("scheduler")))
 
         self.add_stage(stage_name="latent_preparation_stage",
-                       stage=Cosmos25LatentPreparationStage(
+                       stage=Cosmos25AutoLatentPreparationStage(
                            scheduler=self.get_module("scheduler"),
                            transformer=self.get_module("transformer"),
                            vae=self.get_module("vae")))
 
         self.add_stage(stage_name="denoising_stage",
-                       stage=Cosmos25DenoisingStage(
+                       stage=Cosmos25AutoDenoisingStage(
                            transformer=self.get_module("transformer"),
                            scheduler=self.get_module("scheduler")))
 

fastvideo/pipelines/stages/__init__.py

@@ -10,10 +10,10 @@ 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 (Cosmos25DenoisingStage,
-                                                  CosmosDenoisingStage,
-                                                  DenoisingStage,
-                                                  DmdDenoisingStage)
+from fastvideo.pipelines.stages.denoising import (
+    Cosmos25AutoDenoisingStage, Cosmos25DenoisingStage,
+    Cosmos25V2WDenoisingStage, Cosmos25T2WDenoisingStage, CosmosDenoisingStage,
+    DenoisingStage, DmdDenoisingStage)
 from fastvideo.pipelines.stages.encoding import EncodingStage
 from fastvideo.pipelines.stages.image_encoding import (
     ImageEncodingStage, MatrixGameImageEncodingStage, RefImageEncodingStage,
@@ -22,7 +22,8 @@ from fastvideo.pipelines.stages.image_encoding import (
 from fastvideo.pipelines.stages.input_validation import InputValidationStage
 from fastvideo.pipelines.stages.latent_preparation import (
     Cosmos25LatentPreparationStage, CosmosLatentPreparationStage,
-    LatentPreparationStage)
+    Cosmos25AutoLatentPreparationStage, Cosmos25T2WLatentPreparationStage,
+    Cosmos25V2WLatentPreparationStage, LatentPreparationStage)
 from fastvideo.pipelines.stages.ltx2_audio_decoding import LTX2AudioDecodingStage
 from fastvideo.pipelines.stages.ltx2_denoising import LTX2DenoisingStage
 from fastvideo.pipelines.stages.ltx2_latent_preparation import (
@@ -50,6 +51,9 @@ __all__ = [
     "LatentPreparationStage",
     "CosmosLatentPreparationStage",
     "Cosmos25LatentPreparationStage",
+    "Cosmos25T2WLatentPreparationStage",
+    "Cosmos25V2WLatentPreparationStage",
+    "Cosmos25AutoLatentPreparationStage",
     "LTX2LatentPreparationStage",
     "LTX2AudioDecodingStage",
     "ConditioningStage",
@@ -60,6 +64,9 @@ __all__ = [
     "HYWorldDenoisingStage",
     "CosmosDenoisingStage",
     "Cosmos25DenoisingStage",
+    "Cosmos25T2WDenoisingStage",
+    "Cosmos25V2WDenoisingStage",
+    "Cosmos25AutoDenoisingStage",
     "LTX2DenoisingStage",
     "EncodingStage",
     "DecodingStage",

fastvideo/pipelines/stages/denoising.py

@@ -1057,7 +1057,6 @@ class Cosmos25DenoisingStage(CosmosDenoisingStage):
                 "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)
@@ -1065,20 +1064,45 @@ class Cosmos25DenoisingStage(CosmosDenoisingStage):
         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)
+        cfg = fastvideo_args.pipeline_config
+
+        if batch.fps is None:
+            gen = batch.generator
+            if isinstance(gen, list) and len(gen) > 0:
+                gen = gen[0]
+            fps_tensor = torch.randint(
+                16,
+                32,
+                (1, ),
+                generator=gen if isinstance(gen, torch.Generator) else None,
+                device=latents.device,
+            ).float().to(dtype=target_dtype)
+        else:
+            fps_val = batch.fps
+            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
+        # Masks are optional for T2W.
+        cond_mask = getattr(batch, "cond_mask", None)
+        condition_mask = cond_mask.to(target_dtype) if isinstance(
+            cond_mask, torch.Tensor) else None
+        pad_mask = getattr(batch, "padding_mask", None)
+        padding_mask = pad_mask.to(target_dtype) if isinstance(
+            pad_mask, torch.Tensor) else None
+
+        # Conditioning fields are attached by latent preparation stage.
+        conditioning_latents = getattr(batch, "conditioning_latents", None)
+        cond_indicator = getattr(batch, "cond_indicator", None)
+        # Infer whether this is a conditioned run (V2W/I2W) purely from the presence
+        # of conditioning latents. Avoid carrying explicit mode flags on the batch.
+        is_conditioned = (conditioning_latents is not None)
+
+        init_noise_4d = latents_4d.clone()
         if condition_mask is None:
             _, t, h, w = latents_4d.shape
             condition_mask = torch.zeros(1,
@@ -1090,23 +1114,58 @@ class Cosmos25DenoisingStage(CosmosDenoisingStage):
                                          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)
+            padding_default = 0.0 if is_conditioned else 1.0
+            padding_mask = torch.full(
+                (1, 1, h, w),
+                float(padding_default),
+                device=latents.device,
+                dtype=target_dtype,
+            )
 
-        # Cosmos2.5 timestep scaling (see compare_pipelines.py): t * 0.001
         timestep_scale = 0.001
 
+        state_dtype = torch.float32
+
+        conditional_frame_timestep = 0.1
+        latents_4d = latents_4d.to(state_dtype)
+        init_noise_4d = init_noise_4d.to(state_dtype)
+
+        clamp_every_step = bool(getattr(cfg, "cosmos25_clamp_every_step",
+                                        True)) if is_conditioned else False
+
         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)
+                if is_conditioned:
+                    t_frames = int(latents_4d.shape[1])
+                    timestep = torch.full(
+                        (1, t_frames),
+                        float(t_val * timestep_scale),
+                        device=latents.device,
+                        dtype=torch.float32,
+                    )
+                    if cond_indicator is not None and t_frames > 0:
+                        cond_t = cond_indicator[0, 0, :t_frames, 0, 0]
+                        cond_mask_t = (cond_t > 0.5)
+                        if bool(cond_mask_t.any().item()):
+                            timestep[0, cond_mask_t] = float(
+                                conditional_frame_timestep)
+                else:
+                    timestep_val = t_val * timestep_scale
+                    timestep = torch.tensor(
+                        [[float(timestep_val)]],
+                        device=latents.device,
+                        dtype=target_dtype,
+                    )
+
+                # Conditioned runs: replace x_t with GT x0 on the conditioned frames.
+                if (is_conditioned and cond_indicator is not None
+                        and conditioning_latents is not None
+                        and (clamp_every_step or i == 0)):
+                    cond_ind_4d = cond_indicator[0].to(state_dtype)
+                    gt_x0 = conditioning_latents[0].to(state_dtype)
+                    latents_4d = gt_x0 * cond_ind_4d + latents_4d * (
+                        1 - cond_ind_4d)
 
                 model_hidden_states = latents_4d.unsqueeze(0)
 
@@ -1140,10 +1199,22 @@ class Cosmos25DenoisingStage(CosmosDenoisingStage):
                             padding_mask=padding_mask,
                             return_dict=False,
                         )[0]
-                        v = uncond_v + guidance_scale * (cond_v - uncond_v)
+                        if is_conditioned:
+                            v = cond_v + guidance_scale * (cond_v - uncond_v)
+                        else:
+                            v = uncond_v + guidance_scale * (cond_v - uncond_v)
                     else:
                         v = cond_v
 
+                # Conditioned runs: replace velocity on conditioned frames with GT velocity.
+                if (is_conditioned and cond_indicator is not None
+                        and conditioning_latents is not None):
+                    cond_ind_4d = cond_indicator[0].to(state_dtype)
+                    gt_x0 = conditioning_latents[0].to(state_dtype)
+                    gt_v = init_noise_4d.to(state_dtype) - gt_x0
+                    v = cond_ind_4d * gt_v + (1 -
+                                              cond_ind_4d) * v.to(state_dtype)
+
                 prev = self.scheduler.step(v.unsqueeze(0),
                                            t,
                                            latents_4d.unsqueeze(0),
@@ -1153,10 +1224,79 @@ class Cosmos25DenoisingStage(CosmosDenoisingStage):
 
                 progress_bar.update()
 
-        batch.latents = latents_4d.unsqueeze(0)
+        batch.latents = latents_4d.to(target_dtype).unsqueeze(0)
         return batch
 
 
+class Cosmos25T2WDenoisingStage(Cosmos25DenoisingStage):
+    """Cosmos 2.5 Text2World denoising stage."""
+
+    _CONDITIONING_FIELDS = (
+        "conditioning_latents",
+        "cond_indicator",
+        "uncond_indicator",
+    )
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        for name in self._CONDITIONING_FIELDS:
+            if hasattr(batch, name):
+                setattr(batch, name, None)
+        return super().forward(batch, fastvideo_args)
+
+
+class Cosmos25V2WDenoisingStage(Cosmos25DenoisingStage):
+    """Cosmos 2.5 Video2World denoising stage."""
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        return super().forward(batch, fastvideo_args)
+
+
+class Cosmos25AutoDenoisingStage(PipelineStage):
+    """Route Cosmos 2.5 denoising to T2W vs V2W/I2W."""
+
+    def __init__(self, transformer, scheduler) -> None:
+        super().__init__()
+        self._t2w = Cosmos25T2WDenoisingStage(transformer=transformer,
+                                              scheduler=scheduler)
+        self._v2w = Cosmos25V2WDenoisingStage(transformer=transformer,
+                                              scheduler=scheduler)
+
+    def pipeline(self):
+        return self._v2w.pipeline() if self._v2w.pipeline else None
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        conditioning_latents = getattr(batch, "conditioning_latents", None)
+        if conditioning_latents is not None:
+            return self._v2w.forward(batch, fastvideo_args)
+        return self._t2w.forward(batch, fastvideo_args)
+
+    def verify_input(self, batch: ForwardBatch,
+                     fastvideo_args: FastVideoArgs) -> VerificationResult:
+        conditioning_latents = getattr(batch, "conditioning_latents", None)
+        if conditioning_latents is not None:
+            return self._v2w.verify_input(batch, fastvideo_args)
+        return self._t2w.verify_input(batch, fastvideo_args)
+
+    def verify_output(self, batch: ForwardBatch,
+                      fastvideo_args: FastVideoArgs) -> VerificationResult:
+        conditioning_latents = getattr(batch, "conditioning_latents", None)
+        if conditioning_latents is not None:
+            return self._v2w.verify_output(batch, fastvideo_args)
+        return self._t2w.verify_output(batch, fastvideo_args)
+
+
 class DmdDenoisingStage(DenoisingStage):
     """
     Denoising stage for DMD.

fastvideo/pipelines/stages/input_validation.py

@@ -238,7 +238,8 @@ class InputValidationStage(PipelineStage):
                      fastvideo_args: FastVideoArgs) -> VerificationResult:
         """Verify input validation stage inputs."""
         result = VerificationResult()
-        result.add_check("seed", batch.seed, [V.not_none, V.positive_int])
+        # Cosmos-Predict2.5 default seed is 0; allow non-negative seeds here.
+        result.add_check("seed", batch.seed, [V.not_none, V.non_negative_int])
         result.add_check("num_videos_per_prompt", batch.num_videos_per_prompt,
                          V.positive_int)
         result.add_check(

fastvideo/pipelines/stages/latent_preparation.py

@@ -3,6 +3,7 @@
 Latent preparation stage for diffusion pipelines.
 """
 
+import os
 from typing import Any
 
 import numpy as np
@@ -450,10 +451,6 @@ class Cosmos25LatentPreparationStage(CosmosLatentPreparationStage):
         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:
@@ -463,8 +460,6 @@ class Cosmos25LatentPreparationStage(CosmosLatentPreparationStage):
 
         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
@@ -483,7 +478,6 @@ class Cosmos25LatentPreparationStage(CosmosLatentPreparationStage):
         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,
@@ -492,16 +486,16 @@ class Cosmos25LatentPreparationStage(CosmosLatentPreparationStage):
         init_latents = None
         conditioning_latents = None
         video = None
+        image_conditioning = False
 
-        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:
+        if 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(
+            processed_image = image_processor._preprocess_cosmos25(
                 batch.pil_image, height, width)
             video = processed_image.unsqueeze(2)
+            image_conditioning = True
             video = video.to(device=device, dtype=torch.bfloat16)
         elif hasattr(
                 batch,
@@ -509,25 +503,77 @@ class Cosmos25LatentPreparationStage(CosmosLatentPreparationStage):
             if isinstance(batch.preprocessed_image, torch.Tensor):
                 if batch.preprocessed_image.dim() == 4:
                     video = batch.preprocessed_image.unsqueeze(2)
+                    image_conditioning = True
                 elif batch.preprocessed_image.dim() == 5:
                     video = batch.preprocessed_image
+        elif hasattr(batch, "video_latent") and isinstance(
+                batch.video_latent, torch.Tensor):
+            if batch.video_latent.dim() == 5:
+                video = batch.video_latent
         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:
+            # Avoid missing CUDA fallback kernels in some builds by using fp16 contiguous input.
+            if isinstance(video, torch.Tensor):
+                video = video.to(device=device,
+                                 dtype=torch.float16).contiguous()
+            num_video_frames = int(video.size(2))
+
+            num_cond_frames_eff: int | None = None
+            try:
+                ncf = getattr(batch, "num_cond_frames", None)
+                if isinstance(ncf, int) and ncf > 0:
+                    num_cond_frames_eff = min(int(ncf), num_video_frames)
+            except Exception:
+                num_cond_frames_eff = None
+            if num_cond_frames_eff is None:
+                num_cond_frames_eff = 1 if (
+                    not image_conditioning) else num_video_frames
+
+            if num_video_frames >= num_cond_frames_eff:
+                cond_video = video[:, :, -num_cond_frames_eff:]
+            else:
+                pad = video[:, :,
+                            -1:].repeat(1, 1,
+                                        num_cond_frames_eff - num_video_frames,
+                                        1, 1)
+                cond_video = torch.cat([video, pad], dim=2)
+
+            if num_cond_frames_eff >= num_frames:
+                cond_video = cond_video[:, :, -num_frames:]
                 num_cond_latent_frames = (num_frames -
                                           1) // vae_scale_factor_temporal + 1
-                video = video[:, :, -num_frames:]
+                video = cond_video
             else:
-                num_cond_latent_frames = (num_cond_frames -
+                num_cond_latent_frames = (num_cond_frames_eff -
                                           1) // vae_scale_factor_temporal + 1
-                num_padding_frames = num_frames - num_cond_frames
-                last_frame = video[:, :, -1:]
+                num_padding_frames = num_frames - num_cond_frames_eff
+                last_frame = cond_video[:, :, -1:]
                 padding = last_frame.repeat(1, 1, num_padding_frames, 1, 1)
-                video = torch.cat([video, padding], dim=2)
+                if image_conditioning:
+                    padding = cond_video.new_full(
+                        (cond_video.size(0), cond_video.size(1),
+                         num_padding_frames, cond_video.size(3),
+                         cond_video.size(4)),
+                        -1.0,
+                    )
+                video = torch.cat([cond_video, padding], dim=2)
+
+            if os.environ.get("FASTVIDEO_COSMOS25_LOG_KNOBS",
+                              "0") in ("1", "true", "True"):
+                logger.info(
+                    "[Cosmos2.5 latent_prep] using conditioning input: image=%s video_latent=%s video_path=%s | "
+                    "num_video_frames=%s num_cond_frames=%s num_frames=%s",
+                    bool(image_conditioning),
+                    isinstance(getattr(batch, "video_latent", None),
+                               torch.Tensor),
+                    getattr(batch, "video_path", None),
+                    num_video_frames,
+                    num_cond_frames_eff,
+                    num_frames,
+                )
 
             if self.vae is not None:
                 self.vae = self.vae.to(device)
@@ -588,18 +634,14 @@ class Cosmos25LatentPreparationStage(CosmosLatentPreparationStage):
 
         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)
@@ -618,9 +660,7 @@ class Cosmos25LatentPreparationStage(CosmosLatentPreparationStage):
             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)
+        padding_mask = latents.new_zeros(batch_size, 1, height, width)
 
         batch.latents = latents
         batch.raw_latent_shape = latents.shape
@@ -681,3 +721,127 @@ class Cosmos25LatentPreparationStage(CosmosLatentPreparationStage):
                          [V.is_tensor, V.with_dims(5)])
         result.add_check("raw_latent_shape", batch.raw_latent_shape, V.is_tuple)
         return result
+
+
+class Cosmos25T2WLatentPreparationStage(PipelineStage):
+    """Cosmos 2.5 Text2World latent preparation."""
+
+    def __init__(self, scheduler, transformer) -> None:
+        super().__init__()
+        self.scheduler = scheduler
+        self.transformer = transformer
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        if isinstance(batch.prompt, list):
+            batch_size = len(batch.prompt)
+        elif batch.prompt is not None:
+            batch_size = 1
+        else:
+            batch_size = batch.prompt_embeds[0].shape[0]
+        batch_size *= batch.num_videos_per_prompt
+
+        device = get_local_torch_device()
+        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 // vae_scale_factor_spatial
+        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)
+
+        if latents is None:
+            seed = int(batch.seed if batch.seed is not None else 0)
+            latents_fp32 = Cosmos25LatentPreparationStage._arch_invariant_randn(
+                shape,
+                seed=seed,
+                device=device,
+                dtype=torch.float32,
+            )
+            latents = latents_fp32.to(torch.bfloat16)
+        else:
+            latents = latents.to(device=device, dtype=torch.bfloat16)
+
+        batch.latents = latents
+        batch.raw_latent_shape = latents.shape
+
+        batch.conditioning_latents = None
+        batch.cond_indicator = None
+        batch.uncond_indicator = None
+        batch.cond_mask = None
+        batch.uncond_mask = None
+        if hasattr(batch, "padding_mask"):
+            batch.padding_mask = None
+
+        return batch
+
+    def verify_input(self, batch: ForwardBatch,
+                     fastvideo_args: FastVideoArgs) -> VerificationResult:
+        return Cosmos25LatentPreparationStage.verify_input(
+            self, batch, fastvideo_args)  # type: ignore[misc]
+
+    def verify_output(self, batch: ForwardBatch,
+                      fastvideo_args: FastVideoArgs) -> VerificationResult:
+        return Cosmos25LatentPreparationStage.verify_output(
+            self, batch, fastvideo_args)  # type: ignore[misc]
+
+
+class Cosmos25V2WLatentPreparationStage(Cosmos25LatentPreparationStage):
+    """Cosmos 2.5 V2W/I2W latent preparation stage (conditioning-aware)."""
+
+
+class Cosmos25AutoLatentPreparationStage(PipelineStage):
+    """Route Cosmos 2.5 latent prep to T2W vs V2W/I2W."""
+
+    def __init__(self, scheduler, transformer, vae=None) -> None:
+        super().__init__()
+        self._t2w = Cosmos25T2WLatentPreparationStage(
+            scheduler=scheduler,
+            transformer=transformer,
+        )
+        self._v2w = Cosmos25V2WLatentPreparationStage(
+            scheduler=scheduler,
+            transformer=transformer,
+            vae=vae,
+        )
+
+    @staticmethod
+    def _has_conditioning_input(batch: ForwardBatch) -> bool:
+        return (getattr(batch, "pil_image", None) is not None
+                or getattr(batch, "preprocessed_image", None) is not None
+                or bool(getattr(batch, "video_path", None)) or isinstance(
+                    getattr(batch, "video_latent", None), torch.Tensor))
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        if self._has_conditioning_input(batch):
+            return self._v2w.forward(batch, fastvideo_args)
+        return self._t2w.forward(batch, fastvideo_args)
+
+    def verify_input(self, batch: ForwardBatch,
+                     fastvideo_args: FastVideoArgs) -> VerificationResult:
+        if self._has_conditioning_input(batch):
+            return self._v2w.verify_input(batch, fastvideo_args)
+        return self._t2w.verify_input(batch, fastvideo_args)
+
+    def verify_output(self, batch: ForwardBatch,
+                      fastvideo_args: FastVideoArgs) -> VerificationResult:
+        if self._has_conditioning_input(batch):
+            return self._v2w.verify_output(batch, fastvideo_args)
+        return self._t2w.verify_output(batch, fastvideo_args)

fastvideo/pipelines/stages/validators.py

@@ -25,6 +25,11 @@ class StageValidators:
         """Check if value is a positive integer."""
         return isinstance(value, int) and value > 0
 
+    @staticmethod
+    def non_negative_int(value: Any) -> bool:
+        """Check if value is a non-negative integer (allows 0)."""
+        return isinstance(value, int) and value >= 0
+
     @staticmethod
     def positive_float(value: Any) -> bool:
         """Check if value is a positive float."""
@@ -405,6 +410,8 @@ class VerificationResult:
                 error_msg = f"tensor contains {torch.isnan(value).sum().item()} NaN values"
         elif validator_name == 'positive_int':
             expected = "positive integer"
+        elif validator_name == 'non_negative_int':
+            expected = "non-negative integer"
         elif validator_name == 'not_none':
             expected = "non-None value"
         elif validator_name == 'list_not_empty':

videos/README.md

@@ -0,0 +1,7 @@
+# FastVideo/videos
+
+This folder is used to store **video assets for examples**, primarily **input videos** consumed by scripts under `FastVideo/examples/`.
+
+- **Typical contents**: short input clips for demos (e.g., video2world / image2video examples).
+- **Non-critical**: these assets are for convenience and are not required to use the FastVideo library.
+- **Large files**: avoid committing large videos to git; prefer shared storage or download-on-demand.