diff --git a/examples/inference/basic/basic_cosmos2_5_i2w.py b/examples/inference/basic/basic_cosmos2_5_i2w.py new file mode 100644 index 000000000..5cdb631fb --- /dev/null +++ b/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() + + diff --git a/examples/inference/basic/basic_cosmos2_5_t2w.py b/examples/inference/basic/basic_cosmos2_5_t2w.py index d2d11a0ef..d3ed30c83 100644 --- a/examples/inference/basic/basic_cosmos2_5_t2w.py +++ b/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, ) diff --git a/examples/inference/basic/basic_cosmos2_5_v2w.py b/examples/inference/basic/basic_cosmos2_5_v2w.py new file mode 100644 index 000000000..c10218b01 --- /dev/null +++ b/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() + + diff --git a/fastvideo/configs/sample/cosmos2_5.py b/fastvideo/configs/sample/cosmos2_5.py index 48287b303..1bcb5573e 100644 --- a/fastvideo/configs/sample/cosmos2_5.py +++ b/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 diff --git a/fastvideo/configs/sample/registry.py b/fastvideo/configs/sample/registry.py index 22014b06b..f4443b91b 100644 --- a/fastvideo/configs/sample/registry.py +++ b/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 diff --git a/fastvideo/image_processor.py b/fastvideo/image_processor.py index 3483631f5..bef175cee 100644 --- a/fastvideo/image_processor.py +++ b/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) diff --git a/fastvideo/pipelines/basic/cosmos/cosmos2_5_pipeline.py b/fastvideo/pipelines/basic/cosmos/cosmos2_5_pipeline.py index c21720100..8724659e3 100644 --- a/fastvideo/pipelines/basic/cosmos/cosmos2_5_pipeline.py +++ b/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"))) diff --git a/fastvideo/pipelines/stages/__init__.py b/fastvideo/pipelines/stages/__init__.py index fe957ad2d..f0bc5b86d 100644 --- a/fastvideo/pipelines/stages/__init__.py +++ b/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", diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index a6c235b5d..c09254365 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/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. diff --git a/fastvideo/pipelines/stages/input_validation.py b/fastvideo/pipelines/stages/input_validation.py index 2768668fc..3559ae1db 100644 --- a/fastvideo/pipelines/stages/input_validation.py +++ b/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( diff --git a/fastvideo/pipelines/stages/latent_preparation.py b/fastvideo/pipelines/stages/latent_preparation.py index 4af23ed2d..bf23a454e 100644 --- a/fastvideo/pipelines/stages/latent_preparation.py +++ b/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) diff --git a/fastvideo/pipelines/stages/validators.py b/fastvideo/pipelines/stages/validators.py index dac57f001..9462a2ae0 100644 --- a/fastvideo/pipelines/stages/validators.py +++ b/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': diff --git a/images/bus_terminal.jpg b/images/bus_terminal.jpg new file mode 100644 index 000000000..17191028e Binary files /dev/null and b/images/bus_terminal.jpg differ diff --git a/videos/README.md b/videos/README.md new file mode 100644 index 000000000..d2615f4dd --- /dev/null +++ b/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. diff --git a/videos/robot_pouring.mp4 b/videos/robot_pouring.mp4 new file mode 100644 index 000000000..738eeb88b Binary files /dev/null and b/videos/robot_pouring.mp4 differ