# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import math import os from abc import abstractmethod from contextlib import contextmanager from typing import Any, Callable, Dict, List, Tuple, Union import numpy as np import torch import torchvision from einops import rearrange from megatron.core import parallel_state from PIL import Image from torch.distributed.device_mesh import DeviceMesh from torch.distributed.fsdp import FSDPModule, fully_shard from tqdm import tqdm from cosmos_predict2.auxiliary.cosmos_reason1 import CosmosReason1 from cosmos_predict2.auxiliary.text_encoder import CosmosT5TextEncoder from cosmos_predict2.conditioner import DataType, T2VCondition from cosmos_predict2.configs.base.config_video2world import ConditioningStrategy, Video2WorldPipelineConfig from cosmos_predict2.datasets.utils import VIDEO_RES_SIZE_INFO from cosmos_predict2.models.utils import init_weights_on_device, load_state_dict from cosmos_predict2.module.denoise_prediction import DenoisePrediction from cosmos_predict2.module.denoiser_scaling import RectifiedFlowScaling from cosmos_predict2.pipelines.base import BasePipeline from cosmos_predict2.schedulers.rectified_flow_scheduler import RectifiedFlowAB2Scheduler from cosmos_predict2.tokenizers.tokenizer import TokenizerInterface from cosmos_predict2.utils.context_parallel import broadcast, broadcast_split_tensor, cat_outputs_cp, split_inputs_cp from cosmos_predict2.utils.dtensor_helper import DTensorFastEmaModelUpdater, broadcast_dtensor_model_states # from custom.any4d.a4d_config import Any4DConfig # avoid due to circular import # from custom.any4d.a4d_pipe import Any4DPipeline # avoid due to circular import from custom.utils.utils import get_checkpoint, get_folder from imaginaire.lazy_config import instantiate from imaginaire.utils import log, misc, distributed from imaginaire.utils.easy_io import easy_io from imaginaire.utils.ema import FastEmaModelUpdater IS_PREPROCESSED_KEY = "is_preprocessed" _IMAGE_EXTENSIONS = [".png", ".jpg", ".jpeg", "webp"] _VIDEO_EXTENSIONS = [".mp4"] NUM_CONDITIONAL_FRAMES_KEY: str = "num_conditional_frames" def resize_input(video: torch.Tensor, resolution: list[int]) -> torch.Tensor: r""" Resizes and crops the input video tensor while preserving aspect ratio. The video is first resized so that the smaller dimension matches the target resolution, preserving the aspect ratio. Then, it's center-cropped to the target resolution. Args: video (torch.Tensor): Input video tensor of shape (T, C, H, W). resolution (list[int]): Target resolution [H, W]. Returns: torch.Tensor: Resized and cropped video tensor of shape (T, C, target_H, target_W). """ orig_h, orig_w = video.shape[2], video.shape[3] target_h, target_w = resolution scaling_ratio = max((target_w / orig_w), (target_h / orig_h)) resizing_shape = (int(math.ceil(scaling_ratio * orig_h)), int(math.ceil(scaling_ratio * orig_w))) video_resized = torchvision.transforms.functional.resize(video, resizing_shape) video_cropped = torchvision.transforms.functional.center_crop(video_resized, resolution) return video_cropped def read_and_process_image( img_path: str, resolution: list[int], num_video_frames: int, resize: bool = True ) -> torch.Tensor: """ Reads an image, converts it to a video tensor, and processes it for model input. The image is loaded, converted to a tensor, and replicated to match the `num_video_frames`. It's then optionally resized and permuted to the standard video format (B, C, T, H, W). Args: img_path (str): Path to the input image file. resolution (list[int]): Target resolution [H, W] for resizing. num_video_frames (int): The number of frames the output video tensor should have. resize (bool, optional): Whether to resize the image to the target resolution. Defaults to True. Returns: torch.Tensor: Processed video tensor of shape (1, C, T, H, W). Raises: ValueError: If the image extension is not one of the supported types. """ ext = os.path.splitext(img_path)[1] if ext not in _IMAGE_EXTENSIONS: raise ValueError(f"Invalid image extension: {ext}") # Read the image and convert to tensor img = Image.open(img_path) img = torchvision.transforms.functional.to_tensor(img) # Resize the single image first if needed (more efficient than resizing full tensor later) if resize: orig_h, orig_w = img.shape[1], img.shape[2] target_h, target_w = resolution # Calculate scaling based on aspect ratio scaling_ratio = max((target_w / orig_w), (target_h / orig_h)) resizing_shape = (int(math.ceil(scaling_ratio * orig_h)), int(math.ceil(scaling_ratio * orig_w))) # Resize and crop the single image img_resized = torchvision.transforms.functional.resize(img.unsqueeze(0), resizing_shape).squeeze(0) img = torchvision.transforms.functional.center_crop(img_resized, resolution) # Get final dimensions C, H, W = img.shape # Pre-allocate tensor with target dimensions and directly in uint8 format # This avoids allocating a large float tensor that would be converted later vid_input = torch.zeros((num_video_frames, C, H, W), dtype=torch.uint8, device=img.device) # Set first frame and convert to uint8 (only the first frame needs conversion) vid_input[0] = (img * 255.0).to(torch.uint8) # Add batch dimension and permute in one step # [T, C, H, W] -> [1, C, T, H, W] vid_input = vid_input.unsqueeze(0).permute(0, 2, 1, 3, 4) return vid_input def read_and_process_video( video_path: str, resolution: list[int], num_video_frames: int, num_latent_conditional_frames: int = 2, resize: bool = True, ): """ Reads a video, processes it for model input. The video is loaded using easy_io, and uses the last 4x(num_latent_conditional_frames - 1) + 1 from the video. If the video is shorter than num_video_frames, it pads with the last frame repeated. The first num_latent_conditional_frames are marked as conditioning frames. Args: video_path (str): Path to the input video file. resolution (list[int]): Target resolution [H, W] for resizing. num_video_frames (int): Number of frames needed by the model (should equal model.tokenizer.get_pixel_num_frames(model.config.state_t)). num_latent_conditional_frames (int): Number of latent conditional frames from the input video (1 or 2). resize (bool, optional): Whether to resize the video to the target resolution. Defaults to True. Returns: torch.Tensor: Processed video tensor of shape (1, C, T, H, W) where T equals num_video_frames. Raises: ValueError: If the video extension is not supported or other validation errors. Note: Uses the last 4x(num_latent_conditional_frames - 1) + 1 frames from the video. If video is shorter, pads with last frame repeated. """ ext = os.path.splitext(video_path)[1] if ext.lower() not in _VIDEO_EXTENSIONS: raise ValueError(f"Invalid video extension: {ext}") # Load video using easy_io try: video_frames, video_metadata = easy_io.load(video_path) # Returns (T, H, W, C) numpy array log.info(f"Loaded video with shape {video_frames.shape}, metadata: {video_metadata}") except Exception as e: raise ValueError(f"Failed to load video {video_path}: {e}") # Convert numpy array to tensor and rearrange dimensions video_tensor = torch.from_numpy(video_frames).float() / 255.0 # Convert to [0, 1] range video_tensor = video_tensor.permute(3, 0, 1, 2) # (T, H, W, C) -> (C, T, H, W) available_frames = video_tensor.shape[1] # Calculate how many frames to extract from input video frames_to_extract = 4 * (num_latent_conditional_frames - 1) + 1 log.info(f"Will extract the last {frames_to_extract} frames from input video and pad to {num_video_frames}") # Validate num_latent_conditional_frames if num_latent_conditional_frames not in [1, 2]: raise ValueError(f"num_latent_conditional_frames must be 1 or 2, but got {num_latent_conditional_frames}") if available_frames < frames_to_extract: raise ValueError( f"Video has only {available_frames} frames but needs at least {frames_to_extract} frames for num_latent_conditional_frames={num_latent_conditional_frames}" ) # Extract the last frames_to_extract from input video start_idx = available_frames - frames_to_extract extracted_frames = video_tensor[:, start_idx:, :, :] # (C, frames_to_extract, H, W) # Convert to (frames_to_extract, C, H, W) for resize extracted_frames = extracted_frames.permute(1, 0, 2, 3) # (frames_to_extract, C, H, W) # Get last frame for padding last_frame = extracted_frames[-1] # (C, H, W) # Resize the extracted frames if needed (more efficient than resizing full tensor later) if resize: C, H, W = extracted_frames.shape[1], extracted_frames.shape[2], extracted_frames.shape[3] target_h, target_w = resolution # Calculate scaling based on aspect ratio scaling_ratio = max((target_w / W), (target_h / H)) resizing_shape = (int(math.ceil(scaling_ratio * H)), int(math.ceil(scaling_ratio * W))) # Resize and crop the extracted frames extracted_frames = torchvision.transforms.functional.resize(extracted_frames, resizing_shape) extracted_frames = torchvision.transforms.functional.center_crop(extracted_frames, resolution) # Resize and crop the last frame separately (for padding) last_frame = last_frame.unsqueeze(0) # Add batch dim for resize last_frame = torchvision.transforms.functional.resize(last_frame, resizing_shape) last_frame = torchvision.transforms.functional.center_crop(last_frame, resolution) last_frame = last_frame.squeeze(0) # Remove batch dim # Get final dimensions C, H, W = extracted_frames.shape[1], extracted_frames.shape[2], extracted_frames.shape[3] # Pre-allocate tensor with target dimensions and directly in uint8 format # This avoids allocating a large float tensor that would be converted later full_video = torch.zeros((num_video_frames, C, H, W), dtype=torch.uint8, device=extracted_frames.device) # Set extracted frames and convert to uint8 full_video[:frames_to_extract] = (extracted_frames * 255.0).to(torch.uint8) # Pad remaining frames with the last frame (already resized if needed) if frames_to_extract < num_video_frames: last_frame_uint8 = (last_frame * 255.0).to(torch.uint8) for i in range(num_video_frames - frames_to_extract): full_video[frames_to_extract + i] = last_frame_uint8 # Add batch dimension and permute in one operation to final format # [T, C, H, W] -> [1, C, T, H, W] full_video = full_video.unsqueeze(0).permute(0, 2, 1, 3, 4) return full_video class Video2WorldPipeline(BasePipeline): def __init__(self, device: str = "cuda", torch_dtype: torch.dtype = torch.bfloat16): super().__init__(device=device, torch_dtype=torch_dtype) self.text_encoder: CosmosT5TextEncoder = None self.dit: torch.nn.Module = None self.dit_ema: torch.nn.Module = None self.tokenizer: TokenizerInterface = None self.conditioner = None self.prompt_refiner = None self.text_guardrail_runner = None self.video_guardrail_runner = None self.model_names = ["text_encoder", "dit", "tokenizer"] self.height_division_factor = 16 self.width_division_factor = 16 self.use_unified_sequence_parallel = False # Any4D: self.any4d_active = False # since not accessible from pipeline config # Extra networks: self.extra_nets: torch.nn.ModuleDict = dict() # Overridden by Any4DPipeline for surgery def load_checkpoint(self, ckpt: dict, any4d_config=None): (missing_keys, surplus_keys) = self.dit.load_state_dict(ckpt, strict=False, assign=True) log.warning(f'all missing keys: {missing_keys}') log.warning(f'all surplus keys: {surplus_keys}') @staticmethod def from_config( config: Video2WorldPipelineConfig, dit_path: str = "", text_encoder_path: str = "", vae_path: str = "", device: str = "cuda", torch_dtype: torch.dtype = torch.bfloat16, load_ema_to_reg: bool = False, load_prompt_refiner: bool = False, any4d_config=None, pipe_cls=None, local_root: str = "", ) -> Any: any4d_active = (any4d_config is not None and hasattr(any4d_config, 'any4d_active') and any4d_config.any4d_active) log.info(f'from_config() called with pipe_cls: {pipe_cls}') log.info(f'any4d_active: {any4d_active}') # data_only_disable_model is the modern flag; remove_dit is the older one. Unify # by forcing remove_dit=True when data_only_disable_model is set, so the existing # `if config.remove_dit:` early-null block + `if not config.remove_dit:` gates below # (tokenizer / conditioner / DiT) also fire. This must happen BEFORE the early-null # block since that block nulls the local text_encoder_path variable, which gates T5. if any4d_config is not None and getattr(any4d_config, 'data_only_disable_model', False): config.remove_dit = True if config.remove_dit: log.warning("*** DiT will !not! be instantiated ***") dit_path = None text_encoder_path = None vae_path = None # Any4D: Prepare surgery from custom.any4d import a4d_surgery any4d_config = a4d_surgery.validate_config(any4d_config) # Create a pipe # pipe = Video2WorldPipeline(device=device, torch_dtype=torch_dtype) # NOTE(bvh): could now become either Video2WorldPipeline or Any4DPipeline: pipe_cls = Video2WorldPipeline if pipe_cls is None else pipe_cls pipe = pipe_cls(device=device, torch_dtype=torch_dtype) if any4d_active: pipe.any4d_config = any4d_config pipe.config = config pipe.precision = { "float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16, }[config.precision] pipe.tensor_kwargs = {"device": "cuda", "dtype": pipe.precision} log.warning(f"precision {pipe.precision}") # 1. set data keys and data information pipe.sigma_data = config.sigma_data pipe.setup_data_key() # 2. setup up diffusion processing and scaling~(pre-condition) pipe.scheduler = RectifiedFlowAB2Scheduler( sigma_min=config.timestamps.t_min, sigma_max=config.timestamps.t_max, order=config.timestamps.order, t_scaling_factor=config.rectified_flow_t_scaling_factor, ) pipe.scaling = RectifiedFlowScaling(pipe.sigma_data, config.rectified_flow_t_scaling_factor) # 3. Set up tokenizer # NOTE(bvh): This later gets passed to / reused by Any4D VAE. config.tokenizer.vae_pth = get_checkpoint(vae_path, local_root=local_root) # NOTE(bvh): This takes almost no time: if not config.remove_dit: pipe.tokenizer = instantiate(config.tokenizer) assert ( pipe.tokenizer.latent_ch == pipe.config.state_ch ), f"latent_ch {pipe.tokenizer.latent_ch} != state_shape {pipe.config.state_ch}" # 4. Load text encoder # NOTE(bvh): This later gets passed to / reused by Any4D VAE. # NOTE(bvh): Short strings like 'disable' or 'none' skip loading (useful for fast debugging). if text_encoder_path and len(text_encoder_path) > 10: # inference text_encoder_path = get_folder(text_encoder_path, local_root=local_root) # NOTE(bvh): For some reason, this takes about 1 to 3 minutes: with misc.timer('instantiate CosmosT5TextEncoder'): pipe.text_encoder = CosmosT5TextEncoder(device=device, cache_dir=text_encoder_path) pipe.text_encoder.to(device) else: # no # training log.warning('Any4D: Skipping loading of text_encoder because text_encoder_path not provided, ' 'will apply default embeddings at every batch instead...') pipe.text_encoder = None if not config.remove_dit: # 5. Initialize conditioner pipe.conditioner = instantiate(config.conditioner) assert ( sum(p.numel() for p in pipe.conditioner.parameters() if p.requires_grad) == 0 ), "conditioner should not have learnable parameters" if load_prompt_refiner: # no pipe.prompt_refiner = CosmosReason1( checkpoint_dir=config.prompt_refiner_config.checkpoint_dir, offload_model_to_cpu=config.prompt_refiner_config.offload_model_to_cpu, enabled=config.prompt_refiner_config.enabled, ) if config.guardrail_config.enabled: # no from cosmos_predict2.auxiliary.guardrail.common import presets as guardrail_presets pipe.text_guardrail_runner = guardrail_presets.create_text_guardrail_runner( config.guardrail_config.checkpoint_dir, config.guardrail_config.offload_model_to_cpu ) pipe.video_guardrail_runner = guardrail_presets.create_video_guardrail_runner( config.guardrail_config.checkpoint_dir, config.guardrail_config.offload_model_to_cpu ) else: # yes pipe.text_guardrail_runner = None pipe.video_guardrail_runner = None # 6. Set up DiT if dit_path: dit_path = get_checkpoint(dit_path, local_root=local_root) log.info(f"Loading DiT from {dit_path}") else: log.warning("dit_path not provided, initializing DiT with random weights") # NOTE(bvh): device here seems to be meta (= default), presumably to minimize RAM/VRAM, # but this causes trouble when new weights are involved (e.g. Any4D surgery). # Ideally, only the new ones should be initialized on CPU, whereas the rest can remain meta, # but this is probably too complicated to implement to be worth it for now. # with init_weights_on_device(): if not config.remove_dit: with init_weights_on_device(device=torch.device('cpu')): dit_config = config.net if any4d_active: dit_config.any4d_config = any4d_config with misc.timer('instantiate DiT'): pipe.dit = instantiate(dit_config).eval() if dit_path: # yes dit_path = get_checkpoint(dit_path, local_root=local_root) state_dict = load_state_dict(dit_path) # everything goes to CPU # drop net. prefix prefix_to_load = 'net.' state_dict_dit_compatible = dict() for k, v in state_dict.items(): if k.startswith(prefix_to_load): state_dict_dit_compatible[k[len(prefix_to_load) :]] = v else: state_dict_dit_compatible[k] = v # Overridden by Any4DPipeline for surgery with misc.timer('DiT load_checkpoint'): pipe.load_checkpoint(state_dict_dit_compatible, any4d_config=any4d_config) del state_dict, state_dict_dit_compatible log.success(f"Successfully loaded DiT from {dit_path}") # 6-2. Handle EMA if config.ema.enabled: # no pipe.dit_ema = instantiate(dit_config).eval() pipe.dit_ema.requires_grad_(False) pipe.dit_ema_worker = FastEmaModelUpdater() # default when not using FSDP s = config.ema.rate pipe.ema_exp_coefficient = np.roots([1, 7, 16 - s**-2, 12 - s**-2]).real.max() # copying is only necessary when starting the training at iteration 0. # Actual state_dict should be loaded after the pipe is created. pipe.dit_ema_worker.copy_to(src_model=pipe.dit, tgt_model=pipe.dit_ema) pipe.extra_nets = torch.nn.ModuleDict() for key, val in config.extra_nets.items(): if key == 'vggt': from custom.networks.cosmos_vggt import CosmosVGGT pipe.extra_nets[key] = CosmosVGGT(**val) if not config.remove_dit: pipe.dit = pipe.dit.to(device=device, dtype=torch_dtype) # TODO(bvh): ^ there are some meta tensors for some reason :( # pipe.dit = pipe.dit.to_empty(device=device) torch.cuda.empty_cache() # 7. training states if parallel_state.is_initialized(): pipe.data_parallel_size = parallel_state.get_data_parallel_world_size() else: pipe.data_parallel_size = 1 return pipe def setup_data_key(self) -> None: self.input_data_key = self.config.input_data_key # by default it is video key for Video diffusion model self.input_image_key = self.config.input_image_key def apply_fsdp(self, dp_mesh: DeviceMesh) -> None: # https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html # bookmark(bvh): FSDP sharding if self.dit is not None: self.dit.fully_shard(mesh=dp_mesh) self.dit = fully_shard(self.dit, mesh=dp_mesh, reshard_after_forward=True) broadcast_dtensor_model_states(self.dit, dp_mesh) if self.dit_ema: self.dit_ema.fully_shard(mesh=dp_mesh) self.dit_ema = fully_shard(self.dit_ema, mesh=dp_mesh, reshard_after_forward=True) broadcast_dtensor_model_states(self.dit_ema, dp_mesh) self.dit_ema_worker = DTensorFastEmaModelUpdater() # No need to copy weights to EMA when applying FSDP, it is already copied before applying FSDP. if len(self.extra_nets) > 0: for key in self.extra_nets.keys(): self.extra_nets[key].fully_shard(mesh=dp_mesh) self.extra_nets[key] = fully_shard(self.extra_nets[key], mesh=dp_mesh, reshard_after_forward=True) broadcast_dtensor_model_states(self.extra_nets[key], dp_mesh) def apply_cp(self) -> None: pass # used @ inference only def _get_data_batch_input( self, video: torch.Tensor, prompt: str, negative_prompt: str = "", num_latent_conditional_frames: int = 1 ): """ Prepares the input data batch for the diffusion model. Constructs a dictionary containing the video tensor, text embeddings, and other necessary metadata required by the model's forward pass. Optionally includes negative text embeddings. Args: video (torch.Tensor): The input video tensor (B, C, T, H, W). prompt (str): The text prompt for conditioning. negative_prompt (str): Negative prompt. num_latent_conditional_frames (int, optional): The number of latent conditional frames. Defaults to 1. Returns: dict: A dictionary containing the prepared data batch, moved to the correct device and dtype. """ B, C, T, H, W = video.shape self.batch_size = 1 data_batch = { "dataset_name": "video_data", "video": video, "t5_text_embeddings": self.encode_prompt(prompt).to(dtype=self.torch_dtype), "fps": torch.randint(16, 32, (self.batch_size,)), # Random FPS (might be used by model) "padding_mask": torch.zeros(self.batch_size, 1, H, W), # Padding mask (assumed no padding here) "num_conditional_frames": num_latent_conditional_frames, # Specify number of conditional frames } # Handle negative prompts for classifier-free guidance if negative_prompt: data_batch["neg_t5_text_embeddings"] = self.encode_prompt(negative_prompt).to(dtype=self.torch_dtype) # Move tensors to GPU and convert to bfloat16 if they are floating point for k, v in data_batch.items(): if isinstance(v, torch.Tensor) and torch.is_floating_point(data_batch[k]): data_batch[k] = v.cuda().to(dtype=torch.bfloat16) return data_batch def denoising_model(self) -> torch.nn.Module: return self.dit def encode_prompt( self, prompts: Union[str, List[str]], max_length: int = 512, return_mask: bool = False ) -> torch.Tensor: if isinstance(prompts, str): prompts = [prompts] return self.text_encoder.encode_prompts(prompts, max_length=max_length, return_mask=return_mask) # type: ignore # unused @torch.no_grad() def decode(self, latent: torch.Tensor) -> torch.Tensor: return self.tokenizer.decode(latent / self.sigma_data) # used @ train & inference def _normalize_video_databatch_inplace(self, data_batch: dict[str, torch.Tensor], input_key: str = None) -> None: """ Normalizes video data in-place on a CUDA device to reduce data loading overhead. This function modifies the video data tensor within the provided data_batch dictionary in-place, scaling the uint8 data from the range [0, 255] to the normalized range [-1, 1]. Warning: A warning is issued if the data has not been previously normalized. Args: data_batch (dict[str, Tensor]): A dictionary containing the video data under a specific key. This tensor is expected to be on a CUDA device and have dtype of torch.uint8. Side Effects: Modifies the 'input_data_key' tensor within the 'data_batch' dictionary in-place. Note: This operation is performed directly on the CUDA device to avoid the overhead associated with moving data to/from the GPU. Ensure that the tensor is already on the appropriate device and has the correct dtype (torch.uint8) to avoid unexpected behaviors. """ input_key = self.input_data_key if input_key is None else input_key # only handle video batch if input_key in data_batch: # Check if the data has already been normalized and avoid re-normalizing if IS_PREPROCESSED_KEY in data_batch and data_batch[IS_PREPROCESSED_KEY] is True: assert torch.is_floating_point(data_batch[input_key]), "Video data is not in float format." assert torch.all( (data_batch[input_key] >= -1.0001) & (data_batch[input_key] <= 1.0001) ), f"Video data is not in the range [-1, 1]. get data range [{data_batch[input_key].min()}, {data_batch[input_key].max()}]" else: assert data_batch[input_key].dtype == torch.uint8, "Video data is not in uint8 format." data_batch[input_key] = data_batch[input_key].to(**self.tensor_kwargs) / 127.5 - 1.0 data_batch[IS_PREPROCESSED_KEY] = True # if self.config.resize_online: # from torchvision.transforms.v2 import UniformTemporalSubsample # expected_length = self.tokenizer.get_pixel_num_frames(self.config.state_t) # original_length = data_batch[input_key].shape[2] # if original_length != expected_length: # video = rearrange(data_batch[input_key], "b c t h w -> b t c h w") # video = UniformTemporalSubsample(expected_length)(video) # data_batch[input_key] = rearrange(video, "b t c h w -> b c t h w") # used @ train & inference def _augment_image_dim_inplace(self, data_batch: dict[str, torch.Tensor], input_key: str = None) -> None: input_key = self.input_image_key if input_key is None else input_key if input_key in data_batch: # Check if the data has already been augmented and avoid re-augmenting if IS_PREPROCESSED_KEY in data_batch and data_batch[IS_PREPROCESSED_KEY] is True: assert ( data_batch[input_key].shape[2] == 1 ), f"Image data is claimed be augmented while its shape is {data_batch[input_key].shape}" return else: data_batch[input_key] = rearrange(data_batch[input_key], "b c h w -> b c 1 h w").contiguous() data_batch[IS_PREPROCESSED_KEY] = True # unused @torch.no_grad() def encode(self, state: torch.Tensor) -> torch.Tensor: return self.tokenizer.encode(state) * self.sigma_data @staticmethod def get_context_parallel_group() -> Any | None: if parallel_state.is_initialized(): return parallel_state.get_context_parallel_group() return None # used @ train & inference (but NOP for 2B) def broadcast_split_for_model_parallelsim( self, x0_B_C_T_H_W: torch.Tensor, condition: torch.Tensor, epsilon_B_C_T_H_W: torch.Tensor, sigma_B_T: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """ Broadcast and split the input data and condition for model parallelism. Currently, we only support context parallelism. """ cp_group = self.get_context_parallel_group() cp_size = 1 if cp_group is None else cp_group.size() if condition.is_video and cp_size > 1: # no x0_B_C_T_H_W = broadcast_split_tensor(x0_B_C_T_H_W, seq_dim=2, process_group=cp_group) epsilon_B_C_T_H_W = broadcast_split_tensor(epsilon_B_C_T_H_W, seq_dim=2, process_group=cp_group) if sigma_B_T is not None: assert sigma_B_T.ndim == 2, "sigma_B_T should be 2D tensor" if sigma_B_T.shape[-1] == 1: # single sigma is shared across all frames sigma_B_T = broadcast(sigma_B_T, cp_group) else: # different sigma for each frame sigma_B_T = broadcast_split_tensor(sigma_B_T, seq_dim=1, process_group=cp_group) if condition is not None: condition = condition.broadcast(cp_group) self.dit.enable_context_parallel(cp_group) else: # yes self.dit.disable_context_parallel() return x0_B_C_T_H_W, condition, epsilon_B_C_T_H_W, sigma_B_T # used @ train & inference def get_data_and_condition( self, data_batch: dict[str, torch.Tensor] ) -> Tuple[torch.Tensor, torch.Tensor, T2VCondition]: self._normalize_video_databatch_inplace(data_batch) self._augment_image_dim_inplace(data_batch) is_image_batch = self.is_image_batch(data_batch) # Obtain latent state if 'video_latent' in data_batch: # Vidar / Any4D: already processed & encoded etc by a4d_model.py raw_state = None # seems always unused anyway latent_state = data_batch['video_latent'] else: raw_state = data_batch[self.input_image_key if is_image_batch else self.input_data_key] latent_state = self.encode(raw_state).contiguous().float() # Condition condition = self.conditioner(data_batch) condition = condition.edit_data_type(DataType.IMAGE if is_image_batch else DataType.VIDEO) condition = condition.set_video_condition( gt_frames=latent_state.to(**self.tensor_kwargs), random_min_num_conditional_frames=self.config.min_num_conditional_frames, random_max_num_conditional_frames=self.config.max_num_conditional_frames, num_conditional_frames=data_batch.get(NUM_CONDITIONAL_FRAMES_KEY, None), ) return raw_state, latent_state, condition def is_image_batch(self, data_batch: dict[str, torch.Tensor]) -> bool: """We hanlde two types of data_batch. One comes from a joint_dataloader where "dataset_name" can be used to differenciate image_batch and video_batch. Another comes from a dataloader which we by default assumes as video_data for video model training. """ is_image = self.input_image_key in data_batch is_video = self.input_data_key in data_batch or 'video_latent' in data_batch assert ( is_image != is_video ), "Only one of the input_image_key or input_data_key should be present in the data_batch." return is_image # used @ train & inference def denoise(self, xt_B_C_T_H_W: torch.Tensor, sigma: torch.Tensor, condition: T2VCondition) -> DenoisePrediction: """ Performs denoising on the input noise data, noise level, and condition Args: xt (torch.Tensor): The input noise data. sigma (torch.Tensor): The noise level. condition (T2VCondition): conditional information, generated from self.conditioner Returns: DenoisePrediction: The denoised prediction, it includes clean data predicton (x0), \ noise prediction (eps_pred). """ if sigma.ndim == 1: sigma_B_T = rearrange(sigma, "b -> b 1") elif sigma.ndim == 2: sigma_B_T = sigma else: raise ValueError(f"sigma shape {sigma.shape} is not supported") sigma_B_1_T_1_1 = rearrange(sigma_B_T, "b t -> b 1 t 1 1") # get precondition for the network c_skip_B_1_T_1_1, c_out_B_1_T_1_1, c_in_B_1_T_1_1, c_noise_B_1_T_1_1 = self.scaling(sigma=sigma_B_1_T_1_1) # BVH shape notes (vanilla): # xt_B_C_T_H_W: tensor[2, 16, 7, 16, 20] n=71680 (0.3Mb) x∈[-80.249, 83.158] μ=-0.014 σ=12.790 cuda:1 # sigma: tensor[2, 1] μ=9.821 σ=11.563 cuda:1 [[17.997], [1.644]] # condition: Vid2VidCondition( # _is_broadcasted=False, crossattn_emb=tensor[2, 512, 1024] bf16 n=1048576 (2Mb) x∈[-0.656, 0.555] μ=3.052e-05 σ=0.018 cuda:1, # data_type=, # padding_mask=tensor[2, 1, 128, 160] bf16 n=40960 (80Kb) [38;2;127;127;127mall_zeros [0m cuda:1, # fps=tensor[2] i64 μ=10.000 σ=0. cuda:1 [10, 10], # use_video_condition=tensor[1] bool cuda:1 [True], # gt_frames=tensor[2, 16, 7, 16, 20] bf16 n=71680 (0.1Mb) x∈[-3.031, 2.594] μ=0.001 σ=0.645 cuda:1, # condition_video_input_mask_B_C_T_H_W=tensor[2, 1, 7, 16, 20] bf16 n=4480 (8.8Kb) x∈[0., 1.000] μ=0.428 σ=0.494 cuda:1 # ) # sigma_B_T: tensor[2, 1] μ=9.821 σ=11.563 cuda:1 [[17.997], [1.644]] # sigma_B_1_T_1_1: tensor[2, 1, 1, 1, 1] μ=9.821 σ=11.563 cuda:1 [[[[[17.997]]]], [[[[1.644]]]]] # c_skip_B_1_T_1_1: tensor[2, 1, 1, 1, 1] μ=0.215 σ=0.230 cuda:1 [[[[[0.053]]]], [[[[0.378]]]]] # c_out_B_1_T_1_1: tensor[2, 1, 1, 1, 1] μ=-0.785 σ=0.230 cuda:1 [[[[[-0.947]]]], [[[[-0.622]]]]] # c_in_B_1_T_1_1: tensor[2, 1, 1, 1, 1] μ=0.215 σ=0.230 cuda:1 [[[[[0.053]]]], [[[[0.378]]]]] # c_noise_B_1_T_1_1: tensor[2, 1, 7, 1, 1] n=14 x∈[9.999e-05, 0.947] μ=0.448 σ=0.423 cuda:1 # ^ TODO(bvh) revisit - I expect (2, 1, 1, 1, 1) instead here net_state_in_B_C_T_H_W = xt_B_C_T_H_W * c_in_B_1_T_1_1 if condition.is_video: # yes condition_state_in_B_C_T_H_W = condition.gt_frames.type_as(net_state_in_B_C_T_H_W) / self.config.sigma_data if not condition.use_video_condition: # no # When using random dropout, we zero out the ground truth frames condition_state_in_B_C_T_H_W = condition_state_in_B_C_T_H_W * 0 _, C, _, _, _ = xt_B_C_T_H_W.shape condition_video_mask = condition.condition_video_input_mask_B_C_T_H_W.repeat(1, C, 1, 1, 1).type_as( net_state_in_B_C_T_H_W ) if self.config.conditioning_strategy == str(ConditioningStrategy.FRAME_REPLACE): # yes # In case of frame replacement strategy, replace the first few frames of the video with the conditional frames # Update the c_noise as the conditional frames are clean and have very low noise # NOTE(bvh): New vs old Cosmos repo work very differently in this regard! # Notably, there is no explicit conditioning augmentation anymore: # https://github.com/nvidia-cosmos/cosmos-predict2/issues/128 # Make the first few frames of x_t be the ground truth frames net_state_in_B_C_T_H_W = ( condition_state_in_B_C_T_H_W * condition_video_mask + net_state_in_B_C_T_H_W * (1 - condition_video_mask) ) # Adjust c_noise for the conditional frames sigma_cond_B_1_T_1_1 = torch.ones_like(sigma_B_1_T_1_1) * self.config.sigma_conditional _, _, _, c_noise_cond_B_1_T_1_1 = self.scaling(sigma=sigma_cond_B_1_T_1_1) condition_video_mask_B_1_T_1_1 = condition_video_mask.mean(dim=[1, 3, 4], keepdim=True) c_noise_B_1_T_1_1 = c_noise_cond_B_1_T_1_1 * condition_video_mask_B_1_T_1_1 + c_noise_B_1_T_1_1 * ( 1 - condition_video_mask_B_1_T_1_1 ) # sigma_cond_B_1_T_1_1: tensor[2, 1, 1, 1, 1] μ=1.000e-04 σ=0. cuda:1 [[[[[1.000e-04]]]], [[[[1.000e-04]]]]] # condition_video_mask_B_1_T_1_1: tensor[2, 1, 7, 1, 1] n=14 x∈[0., 1.000] μ=0.429 σ=0.514 cuda:1 # c_noise_B_1_T_1_1: tensor[2, 1, 7, 1, 1] n=14 x∈[9.999e-05, 0.947] μ=0.448 σ=0.423 cuda:1 elif self.config.conditioning_strategy == str(ConditioningStrategy.CHANNEL_CONCAT): # no # In case of channel concatenation strategy, concatenate the conditional frames in the channel dimension condition_state_in_masked_B_C_T_H_W = condition_state_in_B_C_T_H_W * condition_video_mask net_state_in_B_C_T_H_W = torch.cat([net_state_in_B_C_T_H_W, condition_state_in_masked_B_C_T_H_W], dim=1) else: # no # In case of image batch, simply concatenate the 0 frames when channel concat strategy is used if self.config.conditioning_strategy == str(ConditioningStrategy.CHANNEL_CONCAT): # no net_state_in_B_C_T_H_W = torch.cat( [net_state_in_B_C_T_H_W, torch.zeros_like(net_state_in_B_C_T_H_W)], dim=1 ) # BVH shape notes (vanilla): # condition_state_in_B_C_T_H_W: tensor[2, 16, 7, 16, 20] n=71680 (0.3Mb) x∈[-2.547, 2.812] μ=3.401e-05 σ=0.638 cuda:0 # condition_video_mask: tensor[2, 16, 7, 16, 20] n=71680 (0.3Mb) x∈[0., 1.000] μ=0.429 σ=0.495 cuda:0 # net_state_in_B_C_T_H_W: tensor[2, 16, 7, 16, 20] n=71680 (0.3Mb) x∈[-3.730, 3.518] μ=-0.006 σ=0.760 cuda:0 # forward pass through the network net_output_B_C_T_H_W = self.dit( x_B_C_T_H_W=net_state_in_B_C_T_H_W.to(**self.tensor_kwargs), timesteps_B_T=c_noise_B_1_T_1_1.squeeze(dim=[1, 3, 4]).to(**self.tensor_kwargs), **condition.to_dict(), ).float() x0_pred_B_C_T_H_W = c_skip_B_1_T_1_1 * xt_B_C_T_H_W + c_out_B_1_T_1_1 * net_output_B_C_T_H_W if condition.is_video: # Set the first few frames to the ground truth frames. This will ensure that the loss is not computed for the first few frames. x0_pred_B_C_T_H_W = condition.gt_frames.type_as( x0_pred_B_C_T_H_W ) * condition_video_mask + x0_pred_B_C_T_H_W * (1 - condition_video_mask) # get noise prediction eps_pred_B_C_T_H_W = (xt_B_C_T_H_W - x0_pred_B_C_T_H_W) / sigma_B_1_T_1_1 # BVH shape notes (vanilla): # net_output_B_C_T_H_W: tensor[2, 16, 7, 16, 20] n=71680 (0.3Mb) x∈[-4.719, 4.438] μ=-0.036 σ=0.989 grad ToCopyBackward0 cuda:0 # x0_pred_B_C_T_H_W: tensor[2, 16, 7, 16, 20] n=71680 (0.3Mb) x∈[-2.575, 3.543] μ=0.018 σ=0.642 grad AddBackward0 cuda:0 # eps_pred_B_C_T_H_W: tensor[2, 16, 7, 16, 20] n=71680 (0.3Mb) x∈[-5.116, 4.268] μ=-0.006 σ=0.994 grad DivBackward0 cuda:0 return DenoisePrediction(x0_pred_B_C_T_H_W, eps_pred_B_C_T_H_W, None) # used @ inference only def get_x0_fn_from_batch( self, data_batch: Dict, guidance: float = 1.5, is_negative_prompt: bool = False, ) -> Callable: """ Generates a callable function `x0_fn` based on the provided data batch and guidance factor. This function first processes the input data batch through a conditioning workflow (`conditioner`) to obtain conditioned and unconditioned states. It then defines a nested function `x0_fn` which applies a denoising operation on an input `noise_x` at a given noise level `sigma` using both the conditioned and unconditioned states. Args: - data_batch (Dict): A batch of data used for conditioning. The format and content of this dictionary should align with the expectations of the `self.conditioner` - guidance (float, optional): A scalar value that modulates the influence of the conditioned state relative to the unconditioned state in the output. Defaults to 1.5. - is_negative_prompt (bool): use negative prompt t5 in uncondition if true Returns: - Callable: A function `x0_fn(noise_x, sigma)` that takes two arguments, `noise_x` and `sigma`, and return x0 predictoin The returned function is suitable for use in scenarios where a denoised state is required based on both conditioned and unconditioned inputs, with an adjustable level of guidance influence. """ if NUM_CONDITIONAL_FRAMES_KEY in data_batch: num_conditional_frames = data_batch[NUM_CONDITIONAL_FRAMES_KEY] else: num_conditional_frames = 1 if is_negative_prompt: condition, uncondition = self.conditioner.get_condition_with_negative_prompt(data_batch) else: condition, uncondition = self.conditioner.get_condition_uncondition(data_batch) is_image_batch = self.is_image_batch(data_batch) condition = condition.edit_data_type(DataType.IMAGE if is_image_batch else DataType.VIDEO) uncondition = uncondition.edit_data_type(DataType.IMAGE if is_image_batch else DataType.VIDEO) _, x0, _ = self.get_data_and_condition(data_batch) # override condition with inference mode; num_conditional_frames used Here! condition = condition.set_video_condition( gt_frames=x0, random_min_num_conditional_frames=self.config.min_num_conditional_frames, random_max_num_conditional_frames=self.config.max_num_conditional_frames, num_conditional_frames=num_conditional_frames, ) uncondition = uncondition.set_video_condition( gt_frames=x0, random_min_num_conditional_frames=self.config.min_num_conditional_frames, random_max_num_conditional_frames=self.config.max_num_conditional_frames, num_conditional_frames=num_conditional_frames, ) condition = condition.edit_for_inference( is_cfg_conditional=True, num_conditional_frames=num_conditional_frames) uncondition = uncondition.edit_for_inference( is_cfg_conditional=False, num_conditional_frames=num_conditional_frames ) # NOTE(bvh): NOP if context_parallel_size == 1 which it usually is _, condition, _, _ = self.broadcast_split_for_model_parallelsim(x0, condition, None, None) _, uncondition, _, _ = self.broadcast_split_for_model_parallelsim(x0, uncondition, None, None) if not parallel_state.is_initialized(): assert ( not self.dit.is_context_parallel_enabled ), "parallel_state is not initialized, context parallel should be turned off." def x0_fn(noise_x: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor: cond_x0 = self.denoise(noise_x, sigma, condition).x0 uncond_x0 = self.denoise(noise_x, sigma, uncondition).x0 raw_x0 = cond_x0 + guidance * (cond_x0 - uncond_x0) if "guided_image" in data_batch: # replacement trick that enables inpainting with base model assert "guided_mask" in data_batch, "guided_mask should be in data_batch if guided_image is present" guide_image = data_batch["guided_image"] guide_mask = data_batch["guided_mask"] raw_x0 = guide_mask * guide_image + (1 - guide_mask) * raw_x0 return raw_x0 return x0_fn # used @ inference only @torch.no_grad() def __call__( self, input_path: str, prompt: str, negative_prompt: str = "", num_conditional_frames: int = 1, guidance: float = 7.0, num_sampling_step: int = 35, seed: int = 0, solver_option: str = "2ab", ) -> torch.Tensor | None: # Parameter check height, width = VIDEO_RES_SIZE_INFO[self.config.resolution]["9,16"] # type: ignore height, width = self.check_resize_height_width(height, width) assert num_conditional_frames in [1, 5], "num_conditional_frames must be 1 or 5" num_latent_conditional_frames = self.tokenizer.get_latent_num_frames(num_conditional_frames) # Run text guardrail on the prompt if self.text_guardrail_runner is not None: from cosmos_predict2.auxiliary.guardrail.common import presets as guardrail_presets log.info("Running guardrail check on prompt...") if not guardrail_presets.run_text_guardrail(prompt, self.text_guardrail_runner): return None else: log.success("Passed guardrail on prompt") elif self.text_guardrail_runner is None: log.warning("Guardrail checks on prompt are disabled") # refine prompt only if prompt refiner is enabled if ( hasattr(self, "prompt_refiner") and self.prompt_refiner is not None and getattr(self.config, "prompt_refiner_config", None) and getattr(self.config.prompt_refiner_config, "enabled", False) ): log.info("Starting prompt refinement...") prompt = self.prompt_refiner.refine_prompt(input_path, prompt) log.info("Finished prompt refinement") # Run text guardrail on the refined prompt if self.text_guardrail_runner is not None: log.info("Running guardrail check on refined prompt...") if not guardrail_presets.run_text_guardrail(prompt, self.text_guardrail_runner): return None else: log.success("Passed guardrail on refined prompt") elif self.text_guardrail_runner is None: log.warning("Guardrail checks on refined prompt are disabled") elif ( hasattr(self, "config") and hasattr(self.config, "prompt_refiner_config") and not self.config.prompt_refiner_config.enabled ): log.warning("Prompt refinement is disabled") num_video_frames = self.tokenizer.get_pixel_num_frames(self.config.state_t) # Detect file extension to determine appropriate reading function ext = os.path.splitext(input_path)[1].lower() if ext in _VIDEO_EXTENSIONS: # Always use video reading for video files, regardless of num_latent_conditional_frames vid_input = read_and_process_video( input_path, [height, width], num_video_frames, num_latent_conditional_frames, resize=True ) elif ext in _IMAGE_EXTENSIONS: if num_latent_conditional_frames == 1: # Use image reading for single frame conditioning with image files vid_input = read_and_process_image(input_path, [height, width], num_video_frames, resize=True) else: raise ValueError( f"Cannot use multi-frame conditioning (num_conditional_frames={num_conditional_frames}) with image input. Please provide a video file." ) else: raise ValueError( f"Unsupported file extension: {ext}. Supported extensions are {_IMAGE_EXTENSIONS + _VIDEO_EXTENSIONS}" ) # Prepare the data batch with text embeddings data_batch = self._get_data_batch_input( vid_input, prompt, negative_prompt, num_latent_conditional_frames=num_latent_conditional_frames ) return self.process(data_batch, seed=seed, guidance=guidance, num_sampling_step=num_sampling_step) def process( self, data_batch, seed=0, guidance=7.0, num_sampling_step=10, return_latents=False, is_negative_prompt=True, ): if 'video_latent' in data_batch: # Vidar / Any4D: for vanilla validation, in case we have cached latents only (B, C,T, H, W) = data_batch['video_latent'].shape n_sample = B state_shape = [C, T, H, W] else: # preprocess self._normalize_video_databatch_inplace(data_batch) self._augment_image_dim_inplace(data_batch) is_image_batch = self.is_image_batch(data_batch) input_key = self.input_image_key if is_image_batch else self.input_data_key n_sample = data_batch[input_key].shape[0] _T, _H, _W = data_batch[input_key].shape[-3:] state_shape = [ self.config.state_ch, self.tokenizer.get_latent_num_frames(_T), _H // self.tokenizer.spatial_compression_factor, _W // self.tokenizer.spatial_compression_factor, ] x0_fn = self.get_x0_fn_from_batch( data_batch, guidance, is_negative_prompt=is_negative_prompt) # log.info("Starting video generation...") x_sigma_max = ( misc.arch_invariant_rand( (n_sample,) + tuple(state_shape), torch.float32, self.tensor_kwargs["device"], seed, ) * self.scheduler.config.sigma_max ) # Split the input data and condition for model parallelism, if context parallelism is enabled. if self.dit.is_context_parallel_enabled: x_sigma_max = split_inputs_cp(x=x_sigma_max, seq_dim=2, cp_group=self.get_context_parallel_group()) # ------------------------------------------------------------------ # # Sampling loop driven by `RectifiedFlowAB2Scheduler` # ------------------------------------------------------------------ # scheduler = self.scheduler # Construct sigma schedule (L + 1 entries including simga_min) and timesteps scheduler.set_timesteps(num_sampling_step, device=x_sigma_max.device) # Bring the initial latent into the precision expected by the scheduler sample = x_sigma_max.to(dtype=torch.float32) x0_prev: torch.Tensor | None = None for i, _ in enumerate(scheduler.timesteps): # Current noise level (sigma_t). sigma_t = scheduler.sigmas[i].to(sample.device, dtype=torch.float32) # `x0_fn` expects `sigma` as a tensor of shape [B] or [B, T]. We # pass a 1-D tensor broadcastable to any later shape handling. sigma_in = sigma_t.repeat(sample.shape[0]) # x0 prediction with conditional and unconditional branches x0_pred = x0_fn(sample, sigma_in) # Scheduler step updates the noisy sample and returns the cached x0. sample, x0_prev = scheduler.step( x0_pred=x0_pred, i=i, sample=sample, x0_prev=x0_prev, ) # Final clean pass at sigma_min. sigma_min = scheduler.sigmas[-1].to(sample.device, dtype=torch.float32) sigma_in = sigma_min.repeat(sample.shape[0]) samples = x0_fn(sample, sigma_in) # Merge context-parallel chunks back together if needed. if self.dit.is_context_parallel_enabled: # no samples = cat_outputs_cp(samples, seq_dim=2, cp_group=self.get_context_parallel_group()) # BVH shape notes (vanilla): # x_sigma_max: tensor[1, 16, 7, 16, 20] n=35840 (0.1Mb) x∈[-372.796, 306.543] μ=-0.438 σ=79.689 cuda:0 # scheduler: # # scheduler.timesteps: tensor[10] i64 x∈[0, 9] μ=4.500 σ=3.028 cuda:0 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] # scheduler.sigmas: tensor[11] f64 x∈[0.002, 80.000] μ=15.622 σ=25.611 cuda:0 # tensor([ 80.000, 45.314, 24.408, 12.382, 5.839, 2.515, # 0.965, 0.318, 0.085, 0.017, 0.002]) # sigma_t: tensor cuda:0 0.017 # x0_pred: tensor[1, 16, 7, 16, 20] n=35840 (0.1Mb) x∈[-4.334, 4.966] μ=0.022 σ=1.031 cuda:0 # sample: tensor[1, 16, 7, 16, 20] n=35840 (0.1Mb) x∈[-4.341, 4.957] μ=0.022 σ=1.035 cuda:0 # x0_prev: tensor[1, 16, 7, 16, 20] n=35840 (0.1Mb) x∈[-4.334, 4.966] μ=0.022 σ=1.031 cuda:0 # sigma_min: tensor cuda:0 0.002 # sigma_in: tensor[1] cuda:0 [0.002] # samples: tensor[1, 16, 7, 16, 20] n=35840 (0.1Mb) x∈[-4.339, 4.954] μ=0.022 σ=1.034 cuda:0 if return_latents: # yes return samples # Decode video = self.decode(samples) # shape: (B, C, T, H, W), possibly out of [-1, 1] # Run video guardrail on the generated video and apply postprocessing if self.video_guardrail_runner is not None: from cosmos_predict2.auxiliary.guardrail.common import presets as guardrail_presets # Clamp to safe range before normalization video = video.clamp(-1.0, 1.0) video_normalized = (video + 1) / 2 # [0, 1] # Convert tensor to NumPy frames for guardrail processing video_squeezed = video_normalized.squeeze(0) # (C, T, H, W) frames = (video_squeezed * 255).clamp(0, 255).to(torch.uint8) frames = frames.permute(1, 2, 3, 0).cpu().numpy() # (T, H, W, C) # Run guardrail processed_frames = guardrail_presets.run_video_guardrail(frames, self.video_guardrail_runner) if processed_frames is None: return None else: log.success("Passed guardrail on generated video") # Convert processed frames back to tensor format processed_video = torch.from_numpy(processed_frames).float().permute(3, 0, 1, 2) / 255.0 processed_video = processed_video * 2 - 1 # back to [-1, 1] processed_video = processed_video.unsqueeze(0) video = processed_video.to(video.device, dtype=video.dtype) log.success("Video generation completed successfully") return video # Taken from cosmos-predict2 PR # https://github.com/nvidia-cosmos/cosmos-predict2/pull/95 @contextmanager def ema_scope(self, context: Any = None, is_cpu: bool = False): if self.config.ema.enabled: # https://github.com/pytorch/pytorch/issues/144289 for module in self.dit.modules(): if isinstance(module, FSDPModule): module.reshard() self.dit_ema_worker.cache(self.dit.parameters(), is_cpu=is_cpu) self.dit_ema_worker.copy_to(src_model=self.dit_ema, tgt_model=self.dit) if context is not None: log.info(f"{context}: Switched to EMA weights") try: yield None finally: if self.config.ema.enabled: for module in self.dit.modules(): if isinstance(module, FSDPModule): module.reshard() self.dit_ema_worker.restore(self.dit.parameters()) if context is not None: log.info(f"{context}: Restored training weights")