# 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 importlib import json import os from contextlib import contextmanager from typing import Any, Dict, List, NamedTuple, Optional, Tuple import cv2 import einops import imageio import numpy as np import torch import torchvision.transforms.functional as transforms_F from einops import rearrange from cosmos_transfer1.auxiliary.guardrail.common.io_utils import save_video from cosmos_transfer1.checkpoints import ( DEPTH2WORLD_CONTROLNET_7B_CHECKPOINT_PATH, EDGE2WORLD_CONTROLNET_7B_CHECKPOINT_PATH, EDGE2WORLD_CONTROLNET_DISTILLED_CHECKPOINT_PATH, HDMAP2WORLD_CONTROLNET_7B_CHECKPOINT_PATH, KEYPOINT2WORLD_CONTROLNET_7B_CHECKPOINT_PATH, LIDAR2WORLD_CONTROLNET_7B_CHECKPOINT_PATH, SEG2WORLD_CONTROLNET_7B_CHECKPOINT_PATH, UPSCALER_CONTROLNET_7B_CHECKPOINT_PATH, VIS2WORLD_CONTROLNET_7B_CHECKPOINT_PATH, SV2MV_t2w_HDMAP2WORLD_CONTROLNET_7B_CHECKPOINT_PATH, SV2MV_t2w_LIDAR2WORLD_CONTROLNET_7B_CHECKPOINT_PATH, SV2MV_v2w_HDMAP2WORLD_CONTROLNET_7B_CHECKPOINT_PATH, SV2MV_v2w_LIDAR2WORLD_CONTROLNET_7B_CHECKPOINT_PATH, ) from cosmos_transfer1.diffusion.config.transfer.augmentors import BilateralOnlyBlurAugmentorConfig from cosmos_transfer1.diffusion.datasets.augmentors.control_input import get_augmentor_for_eval from cosmos_transfer1.diffusion.model.model_t2w import DiffusionT2WModel from cosmos_transfer1.diffusion.model.model_v2w import DiffusionV2WModel from cosmos_transfer1.diffusion.model.model_v2w_multiview import DiffusionV2WMultiviewModel from cosmos_transfer1.utils import log from cosmos_transfer1.utils.config_helper import get_config_module, override from cosmos_transfer1.utils.io import load_from_fileobj TORCH_VERSION: Tuple[int, ...] = tuple(int(x) for x in torch.__version__.split(".")[:2]) if TORCH_VERSION >= (1, 11): from torch.ao import quantization from torch.ao.quantization import FakeQuantizeBase, ObserverBase elif ( TORCH_VERSION >= (1, 8) and hasattr(torch.quantization, "FakeQuantizeBase") and hasattr(torch.quantization, "ObserverBase") ): from torch import quantization from torch.quantization import FakeQuantizeBase, ObserverBase DEFAULT_AUGMENT_SIGMA = 0.001 NUM_MAX_FRAMES = 5000 VIDEO_RES_SIZE_INFO = { "1,1": (960, 960), "4,3": (960, 704), "3,4": (704, 960), "16,9": (1280, 704), "9,16": (704, 1280), } # Default model names for each control type default_model_names = { "vis": VIS2WORLD_CONTROLNET_7B_CHECKPOINT_PATH, "seg": SEG2WORLD_CONTROLNET_7B_CHECKPOINT_PATH, "edge": EDGE2WORLD_CONTROLNET_7B_CHECKPOINT_PATH, "depth": DEPTH2WORLD_CONTROLNET_7B_CHECKPOINT_PATH, "keypoint": KEYPOINT2WORLD_CONTROLNET_7B_CHECKPOINT_PATH, "upscale": UPSCALER_CONTROLNET_7B_CHECKPOINT_PATH, "hdmap": HDMAP2WORLD_CONTROLNET_7B_CHECKPOINT_PATH, "lidar": LIDAR2WORLD_CONTROLNET_7B_CHECKPOINT_PATH, } default_distilled_model_names = { "edge": EDGE2WORLD_CONTROLNET_DISTILLED_CHECKPOINT_PATH, } class _IncompatibleKeys( NamedTuple( "IncompatibleKeys", [ ("missing_keys", List[str]), ("unexpected_keys", List[str]), ("incorrect_shapes", List[Tuple[str, Tuple[int], Tuple[int]]]), ], ) ): pass def non_strict_load_model(model: torch.nn.Module, checkpoint_state_dict: dict) -> _IncompatibleKeys: """Load a model checkpoint with non-strict matching, handling shape mismatches. Args: model (torch.nn.Module): Model to load weights into checkpoint_state_dict (dict): State dict from checkpoint Returns: _IncompatibleKeys: Named tuple containing: - missing_keys: Keys present in model but missing from checkpoint - unexpected_keys: Keys present in checkpoint but not in model - incorrect_shapes: Keys with mismatched tensor shapes The function handles special cases like: - Uninitialized parameters - Quantization observers - TransformerEngine FP8 states """ # workaround https://github.com/pytorch/pytorch/issues/24139 model_state_dict = model.state_dict() incorrect_shapes = [] for k in list(checkpoint_state_dict.keys()): if k in model_state_dict: if "_extra_state" in k: # Key introduced by TransformerEngine for FP8 log.debug(f"Skipping key {k} introduced by TransformerEngine for FP8 in the checkpoint.") continue model_param = model_state_dict[k] # Allow mismatch for uninitialized parameters if TORCH_VERSION >= (1, 8) and isinstance(model_param, torch.nn.parameter.UninitializedParameter): continue if not isinstance(model_param, torch.Tensor): raise ValueError( f"Find non-tensor parameter {k} in the model. type: {type(model_param)} {type(checkpoint_state_dict[k])}, please check if this key is safe to skip or not." ) shape_model = tuple(model_param.shape) shape_checkpoint = tuple(checkpoint_state_dict[k].shape) if shape_model != shape_checkpoint: has_observer_base_classes = ( TORCH_VERSION >= (1, 8) and hasattr(quantization, "ObserverBase") and hasattr(quantization, "FakeQuantizeBase") ) if has_observer_base_classes: # Handle the special case of quantization per channel observers, # where buffer shape mismatches are expected. def _get_module_for_key(model: torch.nn.Module, key: str) -> torch.nn.Module: # foo.bar.param_or_buffer_name -> [foo, bar] key_parts = key.split(".")[:-1] cur_module = model for key_part in key_parts: cur_module = getattr(cur_module, key_part) return cur_module cls_to_skip = ( ObserverBase, FakeQuantizeBase, ) target_module = _get_module_for_key(model, k) if isinstance(target_module, cls_to_skip): # Do not remove modules with expected shape mismatches # them from the state_dict loading. They have special logic # in _load_from_state_dict to handle the mismatches. continue incorrect_shapes.append((k, shape_checkpoint, shape_model)) checkpoint_state_dict.pop(k) incompatible = model.load_state_dict(checkpoint_state_dict, strict=False) # Remove keys with "_extra_state" suffix, which are non-parameter items introduced by TransformerEngine for FP8 handling missing_keys = [k for k in incompatible.missing_keys if "_extra_state" not in k] unexpected_keys = [k for k in incompatible.unexpected_keys if "_extra_state" not in k] return _IncompatibleKeys( missing_keys=missing_keys, unexpected_keys=unexpected_keys, incorrect_shapes=incorrect_shapes, ) @contextmanager def skip_init_linear(): # skip init of nn.Linear orig_reset_parameters = torch.nn.Linear.reset_parameters torch.nn.Linear.reset_parameters = lambda x: x xavier_uniform_ = torch.nn.init.xavier_uniform_ torch.nn.init.xavier_uniform_ = lambda x: x yield torch.nn.Linear.reset_parameters = orig_reset_parameters torch.nn.init.xavier_uniform_ = xavier_uniform_ def load_model_by_config( config_job_name, config_file="projects/cosmos_video/config/config.py", model_class=DiffusionT2WModel, base_checkpoint_dir="", ): config_module = get_config_module(config_file) config = importlib.import_module(config_module).make_config() config = override(config, ["--", f"experiment={config_job_name}"]) if base_checkpoint_dir != "" and hasattr(config.model, "base_load_from"): if hasattr(config.model.base_load_from, "load_path"): if config.model.base_load_from.load_path != "": config.model.base_load_from.load_path = config.model.base_load_from.load_path.replace( "checkpoints", base_checkpoint_dir ) log.info( f"Model need to load a base model weight, change the loading path from default folder to the {base_checkpoint_dir}" ) # Check that the config is valid config.validate() # Freeze the config so developers don't change it during training. config.freeze() # type: ignore # Initialize model with skip_init_linear(): model = model_class(config.model) return model def load_network_model(model: DiffusionT2WModel, ckpt_path: str): if ckpt_path: with skip_init_linear(): model.set_up_model() net_state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=False) # , weights_only=True) non_strict_load_model(model.model, net_state_dict) else: model.set_up_model() model.cuda() def load_tokenizer_model(model: DiffusionT2WModel, tokenizer_dir: str): with skip_init_linear(): model.set_up_tokenizer(tokenizer_dir) model.cuda() def prepare_data_batch( height: int, width: int, num_frames: int, fps: int, prompt_embedding: torch.Tensor, negative_prompt_embedding: Optional[torch.Tensor] = None, ): """Prepare input batch tensors for video generation. Args: height (int): Height of video frames width (int): Width of video frames num_frames (int): Number of frames to generate fps (int): Frames per second prompt_embedding (torch.Tensor): Encoded text prompt embeddings negative_prompt_embedding (torch.Tensor, optional): Encoded negative prompt embeddings Returns: dict: Batch dictionary containing: - video: Zero tensor of target video shape - t5_text_mask: Attention mask for text embeddings - image_size: Target frame dimensions - fps: Target frame rate - num_frames: Number of frames - padding_mask: Frame padding mask - t5_text_embeddings: Prompt embeddings - neg_t5_text_embeddings: Negative prompt embeddings (if provided) - neg_t5_text_mask: Mask for negative embeddings (if provided) """ # Create base data batch data_batch = { "video": torch.zeros((1, 3, num_frames, height, width), dtype=torch.uint8).cuda(), "t5_text_mask": torch.ones(1, 512, dtype=torch.bfloat16).cuda(), "image_size": torch.tensor([[height, width, height, width]] * 1, dtype=torch.bfloat16).cuda(), "fps": torch.tensor([fps] * 1, dtype=torch.bfloat16).cuda(), "num_frames": torch.tensor([num_frames] * 1, dtype=torch.bfloat16).cuda(), "padding_mask": torch.zeros((1, 1, height, width), dtype=torch.bfloat16).cuda(), } # Handle text embeddings t5_embed = prompt_embedding.to(dtype=torch.bfloat16).cuda() data_batch["t5_text_embeddings"] = t5_embed if negative_prompt_embedding is not None: neg_t5_embed = negative_prompt_embedding.to(dtype=torch.bfloat16).cuda() data_batch["neg_t5_text_embeddings"] = neg_t5_embed data_batch["neg_t5_text_mask"] = torch.ones(1, 512, dtype=torch.bfloat16).cuda() return data_batch def get_video_batch(model, prompt_embedding, negative_prompt_embedding, height, width, fps, num_video_frames): """Prepare complete input batch for video generation including latent dimensions. Args: model: Diffusion model instance prompt_embedding (torch.Tensor): Text prompt embeddings negative_prompt_embedding (torch.Tensor): Negative prompt embeddings height (int): Output video height width (int): Output video width fps (int): Output video frame rate num_video_frames (int): Number of frames to generate Returns: tuple: - data_batch (dict): Complete model input batch - state_shape (list): Shape of latent state [C,T,H,W] accounting for VAE compression """ raw_video_batch = prepare_data_batch( height=height, width=width, num_frames=num_video_frames, fps=fps, prompt_embedding=prompt_embedding, negative_prompt_embedding=negative_prompt_embedding, ) state_shape = [ model.tokenizer.channel, model.tokenizer.get_latent_num_frames(num_video_frames), height // model.tokenizer.spatial_compression_factor, width // model.tokenizer.spatial_compression_factor, ] return raw_video_batch, state_shape def resize_video(video_np, h, w, interpolation=cv2.INTER_AREA): """Resize video frames to the specified height and width.""" video_np = video_np[0].transpose((1, 2, 3, 0)) # Convert to T x H x W x C t = video_np.shape[0] resized_video = np.zeros((t, h, w, 3), dtype=np.uint8) for i in range(t): resized_video[i] = cv2.resize(video_np[i], (w, h), interpolation=interpolation) return resized_video.transpose((3, 0, 1, 2))[None] # Convert back to B x C x T x H x W def detect_aspect_ratio(img_size: tuple[int]): """Function for detecting the closest aspect ratio.""" _aspect_ratios = np.array([(16 / 9), (4 / 3), 1, (3 / 4), (9 / 16)]) _aspect_ratio_keys = ["16,9", "4,3", "1,1", "3,4", "9,16"] w, h = img_size current_ratio = w / h closest_aspect_ratio = np.argmin((_aspect_ratios - current_ratio) ** 2) return _aspect_ratio_keys[closest_aspect_ratio] def get_upscale_size(orig_size: tuple[int], aspect_ratio: str, upscale_factor: int = 3, patch_overlap: int = 256): patch_w, patch_h = orig_size if aspect_ratio == "16,9" or aspect_ratio == "4,3": ratio = int(aspect_ratio.split(",")[1]) / int(aspect_ratio.split(",")[0]) target_w = patch_w * upscale_factor - patch_overlap target_h = patch_h * upscale_factor - int(patch_overlap * ratio) elif aspect_ratio == "9,16" or aspect_ratio == "3,4": ratio = int(aspect_ratio.split(",")[0]) / int(aspect_ratio.split(",")[1]) target_h = patch_h * upscale_factor - patch_overlap target_w = patch_w * upscale_factor - int(patch_overlap * ratio) else: target_h = patch_h * upscale_factor - patch_overlap target_w = patch_w * upscale_factor - patch_overlap return target_w, target_h def read_and_resize_input(input_control_path, num_total_frames, interpolation): control_input, fps = read_video_or_image_into_frames_BCTHW( input_control_path, normalize=False, # s.t. output range is [0, 255] max_frames=num_total_frames, also_return_fps=True, ) # BCTHW aspect_ratio = detect_aspect_ratio((control_input.shape[-1], control_input.shape[-2])) w, h = VIDEO_RES_SIZE_INFO[aspect_ratio] control_input = resize_video(control_input, h, w, interpolation=interpolation) # BCTHW, range [0, 255] control_input = torch.from_numpy(control_input[0]) # CTHW, range [0, 255] return control_input, fps, aspect_ratio def get_video_batch_for_multiview_model( model, prompt_embedding, height, width, fps, num_video_frames, frame_repeat_negative_condition ): """Prepare complete input batch for video generation including latent dimensions. Args: model: Diffusion model instance prompt_embedding list(torch.Tensor): Text prompt embeddings height (int): Output video height width (int): Output video width fps (int): Output video frame rate num_video_frames (int): Number of frames to generate frame_repeat_negative_condition (int): Number of frames to generate Returns: tuple: - data_batch (dict): Complete model input batch - state_shape (list): Shape of latent state [C,T,H,W] accounting for VAE compression """ n_views = len(prompt_embedding) prompt_embedding = einops.rearrange(prompt_embedding, "n t d -> (n t) d").unsqueeze(0) raw_video_batch = prepare_data_batch( height=height, width=width, num_frames=num_video_frames, fps=fps, prompt_embedding=prompt_embedding, ) if frame_repeat_negative_condition != -1: frame_repeat = torch.zeros(n_views) frame_repeat[-1] = frame_repeat_negative_condition frame_repeat[-2] = frame_repeat_negative_condition raw_video_batch["frame_repeat"] = frame_repeat.unsqueeze(0).to(dtype=torch.bfloat16).cuda() state_shape = [ model.tokenizer.channel, model.tokenizer.get_latent_num_frames(int(num_video_frames / n_views)) * n_views, height // model.tokenizer.spatial_compression_factor, width // model.tokenizer.spatial_compression_factor, ] return raw_video_batch, state_shape def get_ctrl_batch_mv(H, W, data_batch, num_total_frames, control_inputs, num_views, num_video_frames): # Initialize control input dictionary control_input_dict = {k: v for k, v in data_batch.items()} control_weights = [] hint_keys = [] for hint_key, control_info in control_inputs.items(): if hint_key not in valid_hint_keys: continue if "input_control" in control_info: cond_videos = [] for in_file in control_info["input_control"]: log.info(f"reading control input {in_file} for hint {hint_key}") cond_vid, fps = read_video_or_image_into_frames_BCTHW( in_file, normalize=False, # s.t. output range is [0, 255] max_frames=num_total_frames, also_return_fps=True, ) cond_vid = resize_video(cond_vid, H, W, interpolation=cv2.INTER_LINEAR) cond_vid = torch.from_numpy(cond_vid[0]) cond_videos.append(cond_vid) input_frames = torch.cat(cond_videos, dim=1) control_input_dict[f"control_input_{hint_key}"] = input_frames hint_keys.append(hint_key) control_weights.append(control_info["control_weight"]) target_w, target_h = W, H hint_key = "control_input_" + "_".join(hint_keys) add_control_input = get_augmentor_for_eval(input_key="video", output_key=hint_key) if len(control_input_dict): control_input = add_control_input(control_input_dict)[hint_key] if control_input.ndim == 4: control_input = control_input[None] control_input = control_input.bfloat16() / 255 * 2 - 1 control_weights = load_spatial_temporal_weights( control_weights, B=1, T=num_total_frames, H=target_h, W=target_w, patch_h=H, patch_w=W ) data_batch["control_weight"] = control_weights if len(control_inputs) > 1: # Multicontrol enabled data_batch["hint_key"] = "control_input_multi" data_batch["control_input_multi"] = control_input else: # Single-control case data_batch["hint_key"] = hint_key data_batch[hint_key] = control_input data_batch["target_h"], data_batch["target_w"] = target_h // 8, target_w // 8 data_batch["video"] = torch.zeros((1, 3, 57, H, W), dtype=torch.uint8).cuda() # ????? data_batch["image_size"] = torch.tensor([[H, W, H, W]] * 1, dtype=torch.bfloat16).cuda() data_batch["padding_mask"] = torch.zeros((1, 1, H, W), dtype=torch.bfloat16).cuda() # add view indices for post-train model if num_views == 5: mapped_indices = [0, 1, 2, 4, 5] view_indices_conditioning = [] for v_index in mapped_indices: view_indices_conditioning.append(torch.ones(num_video_frames, device="cuda") * v_index) view_indices_conditioning = torch.cat(view_indices_conditioning, dim=0) data_batch["view_indices"] = view_indices_conditioning.unsqueeze(0).contiguous() return data_batch def get_batched_ctrl_batch( model, prompt_embeddings, # [B, ...] negative_prompt_embeddings, # [B, ...] or None height, width, fps, num_video_frames, input_video_paths, # List[str], length B control_inputs_list, # List[dict], length B blur_strength, canny_threshold, ): """ Create a fully batched data_batch for video generation, including all control and video inputs. Args: model: The diffusion model instance. prompt_embeddings: [B, ...] tensor of prompt embeddings. negative_prompt_embeddings: [B, ...] tensor of negative prompt embeddings or None. height, width, fps, num_video_frames: Video parameters. input_video_paths: List of input video paths, length B. control_inputs_list: List of control input dicts, length B. blur_strength, canny_threshold: ControlNet augmentation parameters. Returns: data_batch: Dict with all fields batched along dim 0 (batch dimension). state_shape: List describing the latent state shape. """ B = len(input_video_paths) def prepare_single_data_batch(b): data_batch = { "video": torch.zeros((1, 3, num_video_frames, height, width), dtype=torch.uint8).cuda(), "t5_text_mask": torch.ones(1, 512, dtype=torch.bfloat16).cuda(), "image_size": torch.tensor([[height, width, height, width]], dtype=torch.bfloat16).cuda(), "fps": torch.tensor([fps], dtype=torch.bfloat16).cuda(), "num_frames": torch.tensor([num_video_frames], dtype=torch.bfloat16).cuda(), "padding_mask": torch.zeros((1, 1, height, width), dtype=torch.bfloat16).cuda(), "t5_text_embeddings": prompt_embeddings[b : b + 1].to(dtype=torch.bfloat16).cuda(), } if negative_prompt_embeddings is not None: data_batch["neg_t5_text_embeddings"] = negative_prompt_embeddings[b : b + 1].to(dtype=torch.bfloat16).cuda() data_batch["neg_t5_text_mask"] = torch.ones(1, 512, dtype=torch.bfloat16).cuda() return data_batch # Prepare and process each sample single_batches = [] for b in range(B): single_data_batch = prepare_single_data_batch(b) processed = get_ctrl_batch( model, single_data_batch, num_video_frames, input_video_paths[b], control_inputs_list[b], blur_strength, canny_threshold, ) single_batches.append(processed) # Merge all single-sample batches into a batched data_batch batched_data_batch = {} for k in single_batches[0]: if isinstance(single_batches[0][k], torch.Tensor): if k == "control_weight" and single_batches[0][k].ndim == 6: # [num_controls, 1, 1, T, H, W] per sample # Stack along dim=1 to get [num_controls, B, 1, T, H, W] batched_data_batch[k] = torch.cat([d[k] for d in single_batches], dim=1) else: # Concatenate along batch dimension (dim=0) for other tensors batched_data_batch[k] = torch.cat([d[k] for d in single_batches], dim=0) else: batched_data_batch[k] = single_batches[0][k] # assume they're the same for now state_shape = [ model.tokenizer.channel, model.tokenizer.get_latent_num_frames(num_video_frames), height // model.tokenizer.spatial_compression_factor, width // model.tokenizer.spatial_compression_factor, ] return batched_data_batch, state_shape def get_ctrl_batch( model, data_batch, num_video_frames, input_video_path, control_inputs, blur_strength, canny_threshold ): """Prepare complete input batch for video generation including latent dimensions. Args: model: Diffusion model instance Returns: - data_batch (dict): Complete model input batch """ state_shape = model.state_shape H, W = ( state_shape[-2] * model.tokenizer.spatial_compression_factor, state_shape[-1] * model.tokenizer.spatial_compression_factor, ) # Initialize control input dictionary control_input_dict = {k: v for k, v in data_batch.items()} num_total_frames = NUM_MAX_FRAMES if input_video_path: input_frames, fps, aspect_ratio = read_and_resize_input( input_video_path, num_total_frames=num_total_frames, interpolation=cv2.INTER_AREA ) _, num_total_frames, H, W = input_frames.shape control_input_dict["video"] = input_frames.numpy() # CTHW data_batch["input_video"] = input_frames.bfloat16()[None] / 255 * 2 - 1 # BCTHW else: data_batch["input_video"] = None target_w, target_h = W, H control_weights = [] for hint_key, control_info in control_inputs.items(): if "input_control" in control_info: in_file = control_info["input_control"] interpolation = cv2.INTER_NEAREST if hint_key == "seg" else cv2.INTER_LINEAR log.info(f"reading control input {in_file} for hint {hint_key}") control_input_dict[f"control_input_{hint_key}"], fps, aspect_ratio = read_and_resize_input( in_file, num_total_frames=num_total_frames, interpolation=interpolation ) # CTHW num_total_frames = min(num_total_frames, control_input_dict[f"control_input_{hint_key}"].shape[1]) target_h, target_w = H, W = control_input_dict[f"control_input_{hint_key}"].shape[2:] if hint_key == "upscale": orig_size = (W, H) target_w, target_h = get_upscale_size(orig_size, aspect_ratio, upscale_factor=3) input_resized = resize_video( input_frames[None].numpy(), target_h, target_w, interpolation=cv2.INTER_LINEAR, ) # BCTHW control_input_dict["control_input_upscale"] = split_video_into_patches( torch.from_numpy(input_resized), H, W ) data_batch["input_video"] = control_input_dict["control_input_upscale"].bfloat16() / 255 * 2 - 1 control_weights.append(control_info["control_weight"]) # Trim all control videos and input video to be the same length. log.info(f"Making all control and input videos to be length of {num_total_frames} frames.") if len(control_inputs) > 1: for hint_key in control_inputs.keys(): cur_key = f"control_input_{hint_key}" if cur_key in control_input_dict: control_input_dict[cur_key] = control_input_dict[cur_key][:, :num_total_frames] if input_video_path: control_input_dict["video"] = control_input_dict["video"][:, :num_total_frames] data_batch["input_video"] = data_batch["input_video"][:, :, :num_total_frames] hint_key = "control_input_" + "_".join(control_inputs.keys()) add_control_input = get_augmentor_for_eval( input_key="video", output_key=hint_key, preset_blur_strength=blur_strength, preset_canny_threshold=canny_threshold, blur_config=BilateralOnlyBlurAugmentorConfig[blur_strength], ) if len(control_input_dict): control_input = add_control_input(control_input_dict)[hint_key] if control_input.ndim == 4: control_input = control_input[None] control_input = control_input.bfloat16() / 255 * 2 - 1 control_weights = load_spatial_temporal_weights( control_weights, B=1, T=num_video_frames, H=target_h, W=target_w, patch_h=H, patch_w=W ) data_batch["control_weight"] = control_weights if len(control_inputs) > 1: # Multicontrol enabled data_batch["hint_key"] = "control_input_multi" data_batch["control_input_multi"] = control_input else: # Single-control case data_batch["hint_key"] = hint_key data_batch[hint_key] = control_input data_batch["target_h"], data_batch["target_w"] = target_h // 8, target_w // 8 data_batch["video"] = torch.zeros((1, 3, 121, H, W), dtype=torch.uint8).cuda() data_batch["image_size"] = torch.tensor([[H, W, H, W]] * 1, dtype=torch.bfloat16).cuda() data_batch["padding_mask"] = torch.zeros((1, 1, H, W), dtype=torch.bfloat16).cuda() return data_batch def generate_control_input(input_file_path, save_folder, hint_key, blur_strength, canny_threshold, num_total_frames=10): log.info( f"Generating control input for {hint_key} with blur strength {blur_strength} and canny threshold {canny_threshold}" ) video_input = read_video_or_image_into_frames_BCTHW(input_file_path, normalize=False)[0, :, :num_total_frames] control_input = get_augmentor_for_eval( input_key="video", output_key=hint_key, preset_blur_strength=blur_strength, preset_canny_threshold=canny_threshold, blur_config=BilateralOnlyBlurAugmentorConfig[blur_strength], ) control_input = control_input({"video": video_input})[hint_key] control_input = control_input.numpy().transpose((1, 2, 3, 0)) output_file_path = f"{save_folder}/{hint_key}_upsampler.mp4" log.info(f"Saving control input to {output_file_path}") save_video(frames=control_input, fps=24, filepath=output_file_path) return output_file_path def generate_world_from_control( model: DiffusionV2WModel, state_shape: list[int], is_negative_prompt: bool, data_batch: dict, guidance: float, num_steps: int, seed: int, condition_latent: torch.Tensor, num_input_frames: int, sigma_max: float, x_sigma_max=None, augment_sigma=None, use_batch_processing: bool = True, chunking: Optional[int] = None, ) -> Tuple[np.array, list, list]: """Generate video using a conditioning video/image input. Args: model (DiffusionV2WModel): The diffusion model instance state_shape (list[int]): Shape of the latent state [C,T,H,W] is_negative_prompt (bool): Whether negative prompt is provided data_batch (dict): Batch containing model inputs including text embeddings guidance (float): Classifier-free guidance scale for sampling num_steps (int): Number of diffusion sampling steps seed (int): Random seed for generation condition_latent (torch.Tensor): Latent tensor from conditioning video/image file num_input_frames (int): Number of input frames chunking: Chunking size, if None, chunking is disabled Returns: np.array: Generated video frames in shape [T,H,W,C], range [0,255] """ assert not model.config.conditioner.video_cond_bool.sample_tokens_start_from_p_or_i, "not supported" if augment_sigma is None: augment_sigma = DEFAULT_AUGMENT_SIGMA b, c, t, h, w = condition_latent.shape if condition_latent.shape[2] < state_shape[1]: # Padding condition latent to state shape condition_latent = torch.cat( [ condition_latent, condition_latent.new_zeros(b, c, state_shape[1] - t, h, w), ], dim=2, ).contiguous() num_of_latent_condition = compute_num_latent_frames(model, num_input_frames) sample = model.generate_samples_from_batch( data_batch, guidance=guidance, state_shape=[c, t, h, w], num_steps=num_steps, is_negative_prompt=is_negative_prompt, seed=seed, condition_latent=condition_latent, num_condition_t=num_of_latent_condition, condition_video_augment_sigma_in_inference=augment_sigma, x_sigma_max=x_sigma_max, sigma_max=sigma_max, target_h=data_batch["target_h"], target_w=data_batch["target_w"], patch_h=h, patch_w=w, use_batch_processing=use_batch_processing, chunking=chunking, ) return sample def read_video_or_image_into_frames_BCTHW( input_path: str, input_path_format: str = "mp4", H: int = None, W: int = None, normalize: bool = True, max_frames: int = -1, also_return_fps: bool = False, ) -> torch.Tensor: """Read video or image file and convert to tensor format. Args: input_path (str): Path to input video/image file input_path_format (str): Format of input file (default: "mp4") H (int, optional): Height to resize frames to W (int, optional): Width to resize frames to normalize (bool): Whether to normalize pixel values to [-1,1] (default: True) max_frames (int): Maximum number of frames to read (-1 for all frames) also_return_fps (bool): Whether to return fps along with frames Returns: torch.Tensor | tuple: Video tensor in shape [B,C,T,H,W], optionally with fps if requested """ log.debug(f"Reading video from {input_path}") loaded_data = load_from_fileobj(input_path, format=input_path_format) frames, meta_data = loaded_data if input_path.endswith(".png") or input_path.endswith(".jpg") or input_path.endswith(".jpeg"): frames = np.array(frames[0]) # HWC, [0,255] if frames.shape[-1] > 3: # RGBA, set the transparent to white # Separate the RGB and Alpha channels rgb_channels = frames[..., :3] alpha_channel = frames[..., 3] / 255.0 # Normalize alpha channel to [0, 1] # Create a white background white_bg = np.ones_like(rgb_channels) * 255 # White background in RGB # Blend the RGB channels with the white background based on the alpha channel frames = (rgb_channels * alpha_channel[..., None] + white_bg * (1 - alpha_channel[..., None])).astype( np.uint8 ) frames = [frames] fps = 0 else: fps = int(meta_data.get("fps")) if max_frames != -1: frames = frames[:max_frames] input_tensor = np.stack(frames, axis=0) input_tensor = einops.rearrange(input_tensor, "t h w c -> t c h w") if normalize: input_tensor = input_tensor / 128.0 - 1.0 input_tensor = torch.from_numpy(input_tensor).bfloat16() # TCHW log.debug(f"Raw data shape: {input_tensor.shape}") if H is not None and W is not None: input_tensor = transforms_F.resize( input_tensor, size=(H, W), # type: ignore interpolation=transforms_F.InterpolationMode.BICUBIC, antialias=True, ) input_tensor = einops.rearrange(input_tensor, "(b t) c h w -> b c t h w", b=1) if normalize: input_tensor = input_tensor.to("cuda") log.debug(f"Load shape {input_tensor.shape} value {input_tensor.min()}, {input_tensor.max()}") if also_return_fps: return input_tensor, fps return input_tensor def compute_num_latent_frames(model: DiffusionV2WModel, num_input_frames: int, downsample_factor=8) -> int: """This function computes the number of latent frames given the number of input frames. Args: model (DiffusionV2WModel): video generation model num_input_frames (int): number of input frames downsample_factor (int): downsample factor for temporal reduce Returns: int: number of latent frames """ # First find how many vae chunks are contained with in num_input_frames num_latent_frames = ( num_input_frames // model.tokenizer.video_vae.pixel_chunk_duration * model.tokenizer.video_vae.latent_chunk_duration ) # Then handle the remainder if num_input_frames % model.tokenizer.video_vae.pixel_chunk_duration == 1: num_latent_frames += 1 elif num_input_frames % model.tokenizer.video_vae.pixel_chunk_duration > 1: assert ( num_input_frames % model.tokenizer.video_vae.pixel_chunk_duration - 1 ) % downsample_factor == 0, f"num_input_frames % model.tokenizer.video_vae.pixel_chunk_duration - 1 must be divisible by {downsample_factor}" num_latent_frames += ( 1 + (num_input_frames % model.tokenizer.video_vae.pixel_chunk_duration - 1) // downsample_factor ) return num_latent_frames def create_condition_latent_from_input_frames( model: DiffusionV2WModel, input_frames: torch.Tensor, num_frames_condition: int = 25, from_back: bool = True, ): """Create condition latent for video generation from input frames. Takes the last num_frames_condition frames from input as conditioning. Args: model (DiffusionV2WModel): Video generation model input_frames (torch.Tensor): Input video tensor [B,C,T,H,W], range [-1,1] num_frames_condition (int): Number of frames to use for conditioning Returns: tuple: (condition_latent, encode_input_frames) where: - condition_latent (torch.Tensor): Encoded latent condition [B,C,T,H,W] - encode_input_frames (torch.Tensor): Padded input frames used for encoding """ B, C, T, H, W = input_frames.shape num_frames_encode = ( model.tokenizer.pixel_chunk_duration ) # (model.state_shape[1] - 1) / model.vae.pixel_chunk_duration + 1 log.debug( f"num_frames_encode not set, set it based on pixel chunk duration and model state shape: {num_frames_encode}" ) log.debug( f"Create condition latent from input frames {input_frames.shape}, value {input_frames.min()}, {input_frames.max()}, dtype {input_frames.dtype}" ) assert ( input_frames.shape[2] >= num_frames_condition ), f"input_frames not enough for condition, require at least {num_frames_condition}, get {input_frames.shape[2]}, {input_frames.shape}" assert ( num_frames_encode >= num_frames_condition ), f"num_frames_encode should be larger than num_frames_condition, get {num_frames_encode}, {num_frames_condition}" # Put the conditioal frames to the begining of the video, and pad the end with zero if model.config.conditioner.video_cond_bool.condition_location == "first_and_last_1": condition_frames_first = input_frames[:, :, :num_frames_condition] condition_frames_last = input_frames[:, :, -num_frames_condition:] padding_frames = condition_frames_first.new_zeros(B, C, num_frames_encode + 1 - 2 * num_frames_condition, H, W) encode_input_frames = torch.cat([condition_frames_first, padding_frames, condition_frames_last], dim=2) elif not from_back: condition_frames = input_frames[:, :, :num_frames_condition] padding_frames = condition_frames.new_zeros(B, C, num_frames_encode - num_frames_condition, H, W) encode_input_frames = torch.cat([condition_frames, padding_frames], dim=2) else: condition_frames = input_frames[:, :, -num_frames_condition:] padding_frames = condition_frames.new_zeros(B, C, num_frames_encode - num_frames_condition, H, W) encode_input_frames = torch.cat([condition_frames, padding_frames], dim=2) log.info( f"create latent with input shape {encode_input_frames.shape} including padding {num_frames_encode - num_frames_condition} at the end" ) if hasattr(model, "n_views") and encode_input_frames.shape[0] == model.n_views: encode_input_frames = einops.rearrange(encode_input_frames, "(B V) C T H W -> B C (V T) H W", V=model.n_views) latent = model.encode(encode_input_frames) elif model.config.conditioner.video_cond_bool.condition_location == "first_and_last_1": latent1 = model.encode(encode_input_frames[:, :, :num_frames_encode]) # BCTHW latent2 = model.encode(encode_input_frames[:, :, num_frames_encode:]) latent = torch.cat([latent1, latent2], dim=2) # BCTHW elif encode_input_frames.shape[0] == 1: # treat as single view video latent = model.tokenizer.encode(encode_input_frames) * model.sigma_data else: raise ValueError( f"First dimension of encode_input_frames {encode_input_frames.shape[0]} does not match " f"model.n_views or model.n_views is not defined and first dimension is not 1" ) return latent, encode_input_frames def compute_num_frames_condition(model: DiffusionV2WModel, num_of_latent_overlap: int, downsample_factor=8) -> int: """This function computes the number of condition pixel frames given the number of latent frames to overlap. Args: model (ExtendDiffusionModel): video generation model num_of_latent_overlap (int): number of latent frames to overlap downsample_factor (int): downsample factor for temporal reduce Returns: int: number of condition frames in output space """ if getattr(model.tokenizer.video_vae, "is_casual", True): # For casual model num_frames_condition = ( num_of_latent_overlap // model.tokenizer.video_vae.latent_chunk_duration * model.tokenizer.video_vae.pixel_chunk_duration ) if num_of_latent_overlap % model.tokenizer.video_vae.latent_chunk_duration == 1: num_frames_condition += 1 elif num_of_latent_overlap % model.tokenizer.video_vae.latent_chunk_duration > 1: num_frames_condition += ( 1 + (num_of_latent_overlap % model.tokenizer.video_vae.latent_chunk_duration - 1) * downsample_factor ) else: num_frames_condition = num_of_latent_overlap * downsample_factor return num_frames_condition def get_condition_latent( model: DiffusionV2WModel, input_image_or_video_path: str, num_input_frames: int = 1, state_shape: list[int] = None, frame_index: int = 0, frame_stride: int = 1, from_back: bool = True, start_frame: int = 0, ) -> torch.Tensor: """Get condition latent from input image/video file. Args: model (DiffusionV2WModel): Video generation model input_image_or_video_path (str): Path to conditioning image/video num_input_frames (int): Number of input frames for video2world prediction Returns: tuple: (condition_latent, input_frames) where: - condition_latent (torch.Tensor): Encoded latent condition [B,C,T,H,W] - input_frames (torch.Tensor): Input frames tensor [B,C,T,H,W] """ if state_shape is None: state_shape = model.state_shape assert num_input_frames > 0, "num_input_frames must be greater than 0" H, W = ( state_shape[-2] * model.tokenizer.spatial_compression_factor, state_shape[-1] * model.tokenizer.spatial_compression_factor, ) input_path_format = input_image_or_video_path.split(".")[-1] input_frames = read_video_or_image_into_frames_BCTHW( input_image_or_video_path, input_path_format=input_path_format, H=H, W=W, ) if model.config.conditioner.video_cond_bool.condition_location == "first_and_last_1": start_frame = frame_index * frame_stride end_frame = (frame_index + 1) * frame_stride curr_input_frames = torch.cat( [input_frames[:, :, start_frame : start_frame + 1], input_frames[:, :, end_frame : end_frame + 1]], dim=2 ).contiguous() # BCTHW num_of_latent_condition = 1 num_frames_condition = compute_num_frames_condition( model, num_of_latent_condition, downsample_factor=model.tokenizer.temporal_compression_factor ) condition_latent, _ = create_condition_latent_from_input_frames(model, curr_input_frames, num_frames_condition) condition_latent = condition_latent.to(torch.bfloat16) return condition_latent input_frames = input_frames[:, :, start_frame:, :, :] condition_latent, _ = create_condition_latent_from_input_frames( model, input_frames, num_input_frames, from_back=from_back ) condition_latent = condition_latent.to(torch.bfloat16) return condition_latent def check_input_frames(input_path: str, required_frames: int) -> bool: """Check if input video/image has sufficient frames. Args: input_path: Path to input video or image required_frames: Number of required frames Returns: np.ndarray of frames if valid, None if invalid """ if input_path.endswith((".jpg", ".jpeg", ".png")): if required_frames > 1: log.error(f"Input ({input_path}) is an image but {required_frames} frames are required") return False return True # Let the pipeline handle image loading # For video input try: vid = imageio.get_reader(input_path, "ffmpeg") frame_count = vid.count_frames() if frame_count < required_frames: log.error(f"Input video has {frame_count} frames but {required_frames} frames are required") return False else: return True except Exception as e: log.error(f"Error reading video file {input_path}: {e}") return False def load_spatial_temporal_weights(weight_paths, B, T, H, W, patch_h, patch_w): """ Load and process spatial-temporal weight maps from .pt files Args: weight_paths: List of weights that can be scalars, paths to .pt files, or empty strings B, T, H, W: Desired tensor dimensions patch_h, patch_w: Patch dimensions for splitting Returns: For all scalar weights: tensor of shape [num_controls] For any spatial maps: tensor of shape [num_controls, B, 1, T, H, W] """ # Process each weight path weights = [] has_spatial_weights = False for path in weight_paths: if not path or (isinstance(path, str) and path.lower() == "none"): # Use default weight of 1.0 w = torch.ones((T, H, W), dtype=torch.bfloat16) else: try: # Try to parse as scalar scalar_value = float(path) w = torch.full((T, H, W), scalar_value, dtype=torch.bfloat16) except ValueError: # Not a scalar, must be a path to a weight map has_spatial_weights = True w = torch.load(path, weights_only=False).to(dtype=torch.bfloat16) # [T, H, W] if w.ndim == 2: # Spatial only w = w.unsqueeze(0).repeat(T, 1, 1) elif w.ndim != 3: raise ValueError(f"Weight map must be 2D or 3D, got shape {w.shape}") if w.shape != (T, H, W): w = ( torch.nn.functional.interpolate( w.unsqueeze(0).unsqueeze(0), size=(T, H, W), mode="trilinear", align_corners=False, ) .squeeze(0) .squeeze(0) ) w = torch.clamp(w, min=0) w = w.unsqueeze(0).unsqueeze(1) w = w.expand(B, 1, -1, -1, -1) weights.append(w) if not has_spatial_weights: scalar_weights = [float(w) for w in weight_paths] weights_tensor = torch.tensor(scalar_weights, dtype=torch.bfloat16) weights_tensor = weights_tensor / (weights_tensor.sum().clip(1)) return weights_tensor.cuda() weights = torch.stack(weights, dim=0).cuda() weights = weights / (weights.sum(dim=0, keepdim=True).clip(1)) # Split into patches if needed if patch_h != H or patch_w != W: num_controls = len(weights) weights = weights.reshape(num_controls * B, 1, T, H, W) weights = split_video_into_patches(weights, patch_h, patch_w) B_new = weights.shape[0] // num_controls weights = weights.reshape(num_controls, B_new, 1, T, H, W) return weights def resize_control_weight_map(control_weight_map, size): assert control_weight_map.shape[2] == 1 # [num_control, B, 1, T, H, W] weight_map = control_weight_map.squeeze(2) # [num_control, B, T, H, W] T, H, W = size if weight_map.shape[2:5] != (T, H, W): assert (weight_map.shape[2] == T) or (weight_map.shape[2] == 8 * (T - 1) + 1) weight_map_i = [ torch.nn.functional.interpolate( weight_map[:, :, :1], size=(1, H, W), mode="trilinear", align_corners=False, ) ] weight_map_i += [ torch.nn.functional.interpolate( weight_map[:, :, 1:], size=(T - 1, H, W), mode="trilinear", align_corners=False, ) ] weight_map = torch.cat(weight_map_i, dim=2) return weight_map.unsqueeze(2) def split_video_into_patches(tensor, patch_h, patch_w): h, w = tensor.shape[-2:] n_img_w = (w - 1) // patch_w + 1 n_img_h = (h - 1) // patch_h + 1 overlap_size_h = overlap_size_w = 0 if n_img_w > 1: overlap_size_w = (n_img_w * patch_w - w) // (n_img_w - 1) # 512 for n=2, 320 for n=4 assert n_img_w * patch_w - overlap_size_w * (n_img_w - 1) == w if n_img_h > 1: overlap_size_h = (n_img_h * patch_h - h) // (n_img_h - 1) assert n_img_h * patch_h - overlap_size_h * (n_img_h - 1) == h p_h = patch_h - overlap_size_h p_w = patch_w - overlap_size_w patches = [] for i in range(n_img_h): for j in range(n_img_w): patches += [tensor[:, :, :, p_h * i : (p_h * i + patch_h), p_w * j : (p_w * j + patch_w)]] return torch.cat(patches) def merge_patches_into_video(imgs, overlap_size_h, overlap_size_w, n_img_h, n_img_w): b, c, t, h, w = imgs.shape imgs = rearrange(imgs, "(b m n) c t h w -> m n b c t h w", m=n_img_h, n=n_img_w) H = n_img_h * h - (n_img_h - 1) * overlap_size_h W = n_img_w * w - (n_img_w - 1) * overlap_size_w img_sum = torch.zeros((b // (n_img_h * n_img_w), c, t, H, W)).to(imgs) mask_sum = torch.zeros((H, W)).to(imgs) # Create a linear mask for blending. def create_linear_gradient_tensor(H, W, overlap_size_h, overlap_size_w): y, x = torch.meshgrid( torch.minimum(torch.arange(H), H - torch.arange(H)) / (overlap_size_h + 1e-6), torch.minimum(torch.arange(W), W - torch.arange(W)) / (overlap_size_w + 1e-6), ) return torch.clamp(y, 0.01, 1) * torch.clamp(x, 0.01, 1) mask_ij = create_linear_gradient_tensor(h, w, overlap_size_h, overlap_size_w).to(imgs) for i in range(n_img_h): for j in range(n_img_w): h_start = i * (h - overlap_size_h) w_start = j * (w - overlap_size_w) img_sum[:, :, :, h_start : h_start + h, w_start : w_start + w] += ( imgs[i, j] * mask_ij[None, None, None, :, :] ) mask_sum[h_start : h_start + h, w_start : w_start + w] += mask_ij return img_sum / (mask_sum[None, None, None, :, :] + 1e-6) valid_hint_keys = {"vis", "seg", "edge", "depth", "keypoint", "upscale", "hdmap", "lidar"} def load_controlnet_specs(cfg) -> Dict[str, Any]: with open(cfg.controlnet_specs, "r") as f: controlnet_specs_in = json.load(f) controlnet_specs = {} args = {} for hint_key, config in controlnet_specs_in.items(): if hint_key in valid_hint_keys: controlnet_specs[hint_key] = config else: if type(config) == dict: raise ValueError(f"Invalid hint_key: {hint_key}. Must be one of {valid_hint_keys}") else: args[hint_key] = config continue return controlnet_specs, args def validate_controlnet_specs(cfg, controlnet_specs) -> Dict[str, Any]: """ Load and validate controlnet specifications from a JSON file. Args: json_path (str): Path to the JSON file containing controlnet specs. checkpoint_dir (str): Base directory for checkpoint files. Returns: Dict[str, Any]: Validated and processed controlnet specifications. """ checkpoint_dir = cfg.checkpoint_dir sigma_max = cfg.sigma_max input_video_path = cfg.input_video_path use_distilled = cfg.use_distilled for hint_key, config in controlnet_specs.items(): if hint_key not in valid_hint_keys: raise ValueError(f"Invalid hint_key: {hint_key}. Must be one of {valid_hint_keys}") if not input_video_path and sigma_max < 80: raise ValueError("Must have 'input_video' specified if sigma_max < 80") if not input_video_path and "input_control" not in config: raise ValueError( f"{hint_key} controlnet must have 'input_control' video specified if no 'input_video' specified." ) if "ckpt_path" not in config: log.info(f"No checkpoint path specified for {hint_key}. Using default.") ckpt_path = os.path.join(checkpoint_dir, default_model_names[hint_key]) if use_distilled: if hint_key in default_distilled_model_names: ckpt_path = os.path.join(checkpoint_dir, default_distilled_model_names[hint_key]) else: log.info(f"No default distilled checkpoint for {hint_key}. Using full checkpoint") config["ckpt_path"] = ckpt_path log.info(f"Using default checkpoint path: {config['ckpt_path']}") # Regardless whether "control_weight_prompt" is provided (i.e. whether we automatically # generate spatiotemporal control weight binary masks), control_weight is needed to. if "control_weight" not in config: log.warning(f"No control weight specified for {hint_key}. Setting to 0.5.") config["control_weight"] = "0.5" else: # Check if control weight is a path or a scalar weight = config["control_weight"] if not isinstance(weight, str) or not weight.endswith(".pt"): try: # Try converting to float scalar_value = float(weight) if scalar_value < 0: raise ValueError(f"Control weight for {hint_key} must be non-negative.") except ValueError: raise ValueError( f"Control weight for {hint_key} must be a valid non-negative float or a path to a .pt file." ) return controlnet_specs