# 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 gc import os from typing import Any, Optional import einops import numpy as np import torch from cosmos_predict1.diffusion.inference.inference_utils import ( generate_world_from_text, generate_world_from_video, get_condition_latent, get_condition_latent_multiview, get_video_batch, get_video_batch_for_multiview_model, load_model_by_config, load_network_model, load_tokenizer_model, read_video_or_image_into_frames_BCTHW, ) from cosmos_predict1.diffusion.model.model_t2w import DiffusionT2WModel from cosmos_predict1.diffusion.model.model_t2w_multiview import DiffusionMultiviewT2WModel from cosmos_predict1.diffusion.model.model_v2w import DiffusionV2WModel from cosmos_predict1.diffusion.model.model_v2w_multiview import DiffusionMultiviewV2WModel from cosmos_predict1.diffusion.model.model_world_interpolator import DiffusionWorldInterpolatorWModel from cosmos_predict1.diffusion.prompt_upsampler.text2world_prompt_upsampler_inference import ( create_prompt_upsampler, run_chat_completion, ) from cosmos_predict1.diffusion.prompt_upsampler.video2world_prompt_upsampler_inference import ( create_vlm_prompt_upsampler, prepare_dialog, ) from cosmos_predict1.diffusion.prompt_upsampler.video2world_prompt_upsampler_inference import ( run_chat_completion as run_chat_completion_vlm, ) from cosmos_predict1.diffusion.training.utils.inference_long_video import generate_video_from_batch_with_loop from cosmos_predict1.utils import log from cosmos_predict1.utils.base_world_generation_pipeline import BaseWorldGenerationPipeline MODEL_NAME_DICT = { # text2world "Cosmos-Predict1-7B-Text2World": "Cosmos_Predict1_Text2World_7B", "Cosmos-Predict1-14B-Text2World": "Cosmos_Predict1_Text2World_14B", "Cosmos-Predict1-7B-Text2World_post-trained": "Cosmos_Predict1_Text2World_7B_Post_trained", "Cosmos-Predict1-14B-Text2World_post-trained": "Cosmos_Predict1_Text2World_14B_Post_trained", # text2world low-memory "Cosmos-Predict1-7B-Text2World_post-trained-4gpu_80gb": "Cosmos_Predict1_Text2World_7B_Post_trained_4gpu_80gb", "Cosmos-Predict1-7B-Text2World_post-trained-8gpu_40gb": "Cosmos_Predict1_Text2World_7B_Post_trained_8gpu_40gb", "Cosmos-Predict1-7B-Text2World_post-trained-4gpu_40gb": "Cosmos_Predict1_Text2World_7B_Post_trained_4gpu_40gb", # text2world lora "Cosmos-Predict1-7B-Text2World_post-trained-lora": "Cosmos_Predict1_Text2World_7B_Post_trained_lora", # video2world "Cosmos-Predict1-7B-Video2World": "Cosmos_Predict1_Video2World_7B", "Cosmos-Predict1-14B-Video2World": "Cosmos_Predict1_Video2World_14B", "Cosmos-Predict1-7B-Video2World_post-trained": "Cosmos_Predict1_Video2World_7B_Post_trained", "Cosmos-Predict1-14B-Video2World_post-trained": "Cosmos_Predict1_Video2World_14B_Post_trained", # video2world low-memory "Cosmos-Predict1-7B-Video2World_post-trained-4gpu_80gb": "Cosmos_Predict1_Video2World_7B_Post_trained_4gpu_80gb", "Cosmos-Predict1-7B-Video2World_post-trained-8gpu_40gb": "Cosmos_Predict1_Video2World_7B_Post_trained_8gpu_40gb", "Cosmos-Predict1-7B-Video2World_post-trained-4gpu_40gb": "Cosmos_Predict1_Video2World_7B_Post_trained_4gpu_40gb", # video2world lora "Cosmos-Predict1-7B-Video2World_post-trained-lora": "Cosmos_Predict1_Video2World_7B_Post_trained_lora", "Cosmos-Predict1-7B-Text2World-Sample-AV-Multiview": "Cosmos_Predict1_Text2World_7B_Multiview", "Cosmos-Predict1-7B-Video2World-Sample-AV-Multiview": "Cosmos_Predict1_Video2World_7B_Multiview", "Cosmos-Predict1-7B-WorldInterpolator": "Cosmos_Predict1_WorldInterpolator_7B", # Gen3c "Gen3C-Cosmos-7B": "GEN3C_Cosmos_7B", } class DiffusionText2WorldGenerationPipeline(BaseWorldGenerationPipeline): def __init__( self, inference_type: str, checkpoint_dir: str, checkpoint_name: str, prompt_upsampler_dir: Optional[str] = None, enable_prompt_upsampler: bool = True, has_text_input: bool = True, offload_network: bool = False, offload_tokenizer: bool = False, offload_text_encoder_model: bool = False, offload_prompt_upsampler: bool = False, offload_guardrail_models: bool = False, disable_guardrail: bool = False, guidance: float = 7.0, num_steps: int = 35, height: int = 704, width: int = 1280, fps: int = 24, num_video_frames: int = 121, seed: int = 0, ): """Initialize the diffusion world generation pipeline. Args: inference_type: Type of world generation ('text2world' or 'video2world') checkpoint_dir: Base directory containing model checkpoints checkpoint_name: Name of the diffusion transformer checkpoint to use prompt_upsampler_dir: Directory containing prompt upsampler model weights enable_prompt_upsampler: Whether to use prompt upsampling has_text_input: Whether the pipeline takes text input for world generation offload_network: Whether to offload diffusion transformer after inference offload_tokenizer: Whether to offload tokenizer after inference offload_text_encoder_model: Whether to offload T5 model after inference offload_prompt_upsampler: Whether to offload prompt upsampler offload_guardrail_models: Whether to offload guardrail models disable_guardrail: Whether to disable guardrail guidance: Classifier-free guidance scale num_steps: Number of diffusion sampling steps height: Height of output video width: Width of output video fps: Frames per second of output video num_video_frames: Number of frames to generate seed: Random seed for sampling """ assert inference_type in [ "text2world", "video2world", "world_interpolator", ], "Invalid inference_type, must be 'text2world' or 'video2world'" self.model_name = MODEL_NAME_DICT[checkpoint_name] self.guidance = guidance self.num_steps = num_steps self.height = height self.width = width self.fps = fps self.num_video_frames = num_video_frames self.seed = seed super().__init__( inference_type=inference_type, checkpoint_dir=checkpoint_dir, checkpoint_name=checkpoint_name, has_text_input=has_text_input, offload_network=offload_network, offload_tokenizer=offload_tokenizer, offload_text_encoder_model=offload_text_encoder_model, offload_guardrail_models=offload_guardrail_models, disable_guardrail=disable_guardrail, ) self.prompt_upsampler_dir = prompt_upsampler_dir self.enable_prompt_upsampler = enable_prompt_upsampler self.offload_prompt_upsampler = offload_prompt_upsampler self.prompt_upsampler = None if enable_prompt_upsampler and not offload_prompt_upsampler: self._load_prompt_upsampler_model() def _load_prompt_upsampler_model(self): self.prompt_upsampler = create_prompt_upsampler( checkpoint_dir=os.path.join(self.checkpoint_dir, self.prompt_upsampler_dir), ) def _load_model(self): self.model = load_model_by_config( config_job_name=self.model_name, config_file="cosmos_predict1/diffusion/config/config.py", model_class=DiffusionT2WModel, ) def _load_network(self): load_network_model(self.model, f"{self.checkpoint_dir}/{self.checkpoint_name}/model.pt") def _load_tokenizer(self): load_tokenizer_model(self.model, f"{self.checkpoint_dir}/Cosmos-Tokenize1-CV8x8x8-720p") def _offload_prompt_upsampler_model(self): """Move prompt enhancement model to CPU/disk. Offloads prompt upsampling model after processing input to reduce GPU memory usage. """ if self.prompt_upsampler: del self.prompt_upsampler self.prompt_upsampler = None gc.collect() torch.cuda.empty_cache() def _run_prompt_upsampler_on_prompt(self, prompt: str) -> str: """Enhance the input prompt using the prompt upsampler model. Args: prompt: Raw text prompt to be enhanced Returns: str: Enhanced version of the input prompt with more descriptive details """ upsampled_prompt = run_chat_completion(self.prompt_upsampler, prompt) log.info(f"Upsampled prompt: {upsampled_prompt}") return upsampled_prompt def _run_prompt_upsampler_on_prompt_with_offload(self, *args: Any, **kwargs: Any) -> str: """Enhance prompt with prompt upsampler model. Args: *args: Positional arguments **kwargs: Keyword arguments Returns: Enhanced prompt string """ if self.offload_prompt_upsampler: self._load_prompt_upsampler_model() enhanced_prompt = self._run_prompt_upsampler_on_prompt(*args, **kwargs) if self.offload_prompt_upsampler: self._offload_prompt_upsampler_model() return enhanced_prompt def _run_tokenizer_decoding(self, sample: torch.Tensor) -> np.ndarray: """Decode latent samples to video frames using the tokenizer decoder. Args: sample: Latent tensor from diffusion model [B, C, T, H, W] Returns: np.ndarray: Decoded video frames as uint8 numpy array [T, H, W, C] with values in range [0, 255] """ # Decode video video = (1.0 + self.model.decode(sample)).clamp(0, 2) / 2 # [B, 3, T, H, W] video = (video[0].permute(1, 2, 3, 0) * 255).to(torch.uint8).cpu().numpy() return video def _run_model( self, embedding: torch.Tensor, negative_prompt_embedding: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Generate video latents using the diffusion model. Args: embedding: Text embedding tensor from text encoder negative_prompt_embedding: Optional embedding for negative prompt guidance Returns: torch.Tensor: Generated video latents before tokenizer decoding Note: The model and tokenizer are automatically offloaded after inference if offloading is enabled in the config. """ # Get video batch and state shape data_batch, state_shape = get_video_batch( model=self.model, prompt_embedding=embedding, negative_prompt_embedding=negative_prompt_embedding, height=self.height, width=self.width, fps=self.fps, num_video_frames=self.num_video_frames, ) # Generate video frames sample = generate_world_from_text( model=self.model, state_shape=state_shape, is_negative_prompt=True if negative_prompt_embedding is not None else False, data_batch=data_batch, guidance=self.guidance, num_steps=self.num_steps, seed=self.seed, ) return sample def _run_model_with_offload( self, prompt_embedding: torch.Tensor, negative_prompt_embedding: Optional[torch.Tensor] = None ) -> np.ndarray: """Generate world representation with automatic model offloading. Wraps the core generation process with model loading/offloading logic to minimize GPU memory usage during inference. Args: prompt_embedding: Text embedding tensor from text encoder negative_prompt_embedding: Optional embedding for negative prompt guidance Returns: np.ndarray: Generated world representation """ if self.offload_network: self._load_network() if self.offload_tokenizer: self._load_tokenizer() sample = self._run_model(prompt_embedding, negative_prompt_embedding) if self.offload_network: self._offload_network() if self.offload_tokenizer: self._load_tokenizer() sample = self._run_tokenizer_decoding(sample) if self.offload_tokenizer: self._offload_tokenizer() return sample def generate( self, prompt: str, negative_prompt: Optional[str] = None, word_limit_to_skip_upsampler: Optional[int] = None, ) -> tuple[np.ndarray, str] | None: """Generate video from text prompt with optional negative prompt guidance. Pipeline steps: 1. Run safety checks on input prompt 2. Enhance prompt using upsampler if enabled 3. Run safety checks on upsampled prompt if applicable 4. Convert prompt to embeddings 5. Generate video frames using diffusion 6. Run safety checks and apply face blur on generated video frames Args: prompt: Text description of desired video negative_prompt: Optional text to guide what not to generate word_limit_to_skip_upsampler: Skip prompt upsampler for better robustness if the number of words in the prompt is greater than this value Returns: tuple: ( Generated video frames as uint8 np.ndarray [T, H, W, C], Final prompt used for generation (may be enhanced) ), or None if content fails guardrail safety checks """ log.info(f"Run with prompt: {prompt}") log.info(f"Run with negative prompt: {negative_prompt}") log.info(f"Run with prompt upsampler: {self.enable_prompt_upsampler}") if not self.disable_guardrail: log.info("Run guardrail on prompt") is_safe = self._run_guardrail_on_prompt_with_offload(prompt) if not is_safe: log.critical("Input text prompt is not safe") return None log.info("Pass guardrail on prompt") # Enhance prompt if self.enable_prompt_upsampler: word_count = len(prompt.split()) if word_limit_to_skip_upsampler is None or word_count <= word_limit_to_skip_upsampler: log.info("Run prompt upsampler on prompt") prompt = self._run_prompt_upsampler_on_prompt_with_offload(prompt) if not self.disable_guardrail: log.info("Run guardrail on upsampled prompt") is_safe = self._run_guardrail_on_prompt_with_offload(prompt=prompt) if not is_safe: log.critical("Upsampled text prompt is not safe") return None log.info("Pass guardrail on upsampled prompt") else: log.info( f"Skip prompt upsampler for better robustness because the number of words ({word_count}) in the prompt is greater than {word_limit_to_skip_upsampler}" ) log.info("Run text embedding on prompt") if negative_prompt: prompts = [prompt, negative_prompt] else: prompts = [prompt] prompt_embeddings, _ = self._run_text_embedding_on_prompt_with_offload(prompts) prompt_embedding = prompt_embeddings[0] negative_prompt_embedding = prompt_embeddings[1] if negative_prompt else None log.info("Finish text embedding on prompt") # Generate video log.info("Run generation") video = self._run_model_with_offload( prompt_embedding, negative_prompt_embedding=negative_prompt_embedding, ) log.info("Finish generation") if not self.disable_guardrail: log.info("Run guardrail on generated video") video = self._run_guardrail_on_video_with_offload(video) if video is None: log.critical("Generated video is not safe") return None log.info("Pass guardrail on generated video") return video, prompt class DiffusionVideo2WorldGenerationPipeline(DiffusionText2WorldGenerationPipeline): def __init__( self, inference_type: str, checkpoint_dir: str, checkpoint_name: str, prompt_upsampler_dir: Optional[str] = None, enable_prompt_upsampler: bool = True, has_text_input: bool = True, offload_network: bool = False, offload_tokenizer: bool = False, offload_text_encoder_model: bool = False, offload_prompt_upsampler: bool = False, offload_guardrail_models: bool = False, disable_guardrail: bool = False, guidance: float = 7.0, num_steps: int = 35, height: int = 704, width: int = 1280, fps: int = 24, num_video_frames: int = 121, seed: int = 0, num_input_frames: int = 1, ): """Initialize diffusion world generation pipeline. Args: inference_type: Type of world generation ('text2world' or 'video2world') checkpoint_dir: Base directory containing model checkpoints checkpoint_name: Name of the diffusion transformer checkpoint to use prompt_upsampler_dir: Directory containing prompt upsampler model weights enable_prompt_upsampler: Whether to use prompt upsampling has_text_input: Whether the pipeline takes text input for world generation offload_network: Whether to offload diffusion transformer after inference offload_tokenizer: Whether to offload tokenizer after inference offload_text_encoder_model: Whether to offload T5 model after inference offload_prompt_upsampler: Whether to offload prompt upsampler offload_guardrail_models: Whether to offload guardrail models disable_guardrail: Whether to disable guardrail guidance: Classifier-free guidance scale num_steps: Number of diffusion sampling steps height: Height of output video width: Width of output video fps: Frames per second of output video num_video_frames: Number of frames to generate seed: Random seed for sampling num_input_frames: Number of latent conditions """ self.num_input_frames = num_input_frames super().__init__( inference_type=inference_type, checkpoint_dir=checkpoint_dir, checkpoint_name=checkpoint_name, prompt_upsampler_dir=prompt_upsampler_dir, enable_prompt_upsampler=enable_prompt_upsampler, has_text_input=has_text_input, offload_network=offload_network, offload_tokenizer=offload_tokenizer, offload_text_encoder_model=offload_text_encoder_model, offload_prompt_upsampler=offload_prompt_upsampler, offload_guardrail_models=offload_guardrail_models, disable_guardrail=disable_guardrail, guidance=guidance, num_steps=num_steps, height=height, width=width, fps=fps, num_video_frames=num_video_frames, seed=seed, ) def _run_prompt_upsampler_on_prompt(self, image_or_video_path: str) -> str: """Enhance the input prompt using visual context from the conditioning image. Args: image_or_video_path: Path to conditioning image or video used for visual context Returns: str: Enhanced prompt incorporating visual details from the image """ dialog = prepare_dialog(image_or_video_path) upsampled_prompt = run_chat_completion_vlm( self.prompt_upsampler, dialog, max_gen_len=400, temperature=0.01, top_p=0.9, logprobs=False ) log.info(f"Upsampled prompt: {upsampled_prompt}") return upsampled_prompt def _load_prompt_upsampler_model(self): self.prompt_upsampler = create_vlm_prompt_upsampler( checkpoint_dir=os.path.join(self.checkpoint_dir, self.prompt_upsampler_dir), ) def _load_model(self): self.model = load_model_by_config( config_job_name=self.model_name, config_file="cosmos_predict1/diffusion/config/config.py", model_class=DiffusionV2WModel, ) def _run_model( self, embedding: torch.Tensor, condition_latent: torch.Tensor, negative_prompt_embedding: torch.Tensor | None = None, ) -> torch.Tensor: """Generate video frames using the diffusion model. Args: embedding: Text embedding tensor from T5 encoder condition_latent: Latent tensor from conditioning image or video negative_prompt_embedding: Optional embedding for negative prompt guidance Returns: Tensor of generated video frames Note: Model and tokenizer are automatically offloaded after inference if offloading is enabled. """ # Get video batch and state shape data_batch, state_shape = get_video_batch( model=self.model, prompt_embedding=embedding, negative_prompt_embedding=negative_prompt_embedding, height=self.height, width=self.width, fps=self.fps, num_video_frames=self.num_video_frames, ) # Generate video frames video = generate_world_from_video( model=self.model, state_shape=self.model.state_shape, is_negative_prompt=True, data_batch=data_batch, guidance=self.guidance, num_steps=self.num_steps, seed=self.seed, condition_latent=condition_latent, num_input_frames=self.num_input_frames, ) return video def _run_tokenizer_encoding(self, image_or_video_path: str) -> torch.Tensor: """ Encode image to latent space Args: image_or_video_path: Path to conditioning image Returns: torch.Tensor: Latent tensor from tokenizer encoding """ condition_latent = get_condition_latent( model=self.model, input_image_or_video_path=image_or_video_path, num_input_frames=self.num_input_frames, state_shape=self.model.state_shape, ) return condition_latent def _run_model_with_offload( self, prompt_embedding: torch.Tensor, image_or_video_path: str, negative_prompt_embedding: Optional[torch.Tensor] = None, ) -> np.ndarray: """Generate world representation with automatic model offloading. Wraps the core generation process with model loading/offloading logic to minimize GPU memory usage during inference. Args: prompt_embedding: Text embedding tensor from T5 encoder image_or_video_path: Path to conditioning image or video negative_prompt_embedding: Optional embedding for negative prompt guidance Returns: np.ndarray: Generated world representation as numpy array """ if self.offload_tokenizer: self._load_tokenizer() condition_latent = self._run_tokenizer_encoding(image_or_video_path) if self.offload_network: self._load_network() sample = self._run_model(prompt_embedding, condition_latent, negative_prompt_embedding) if self.offload_network: self._offload_network() sample = self._run_tokenizer_decoding(sample) if self.offload_tokenizer: self._offload_tokenizer() return sample def generate( self, prompt: str, image_or_video_path: str, negative_prompt: Optional[str] = None, ) -> tuple[np.ndarray, str] | None: """Generate video from text prompt and optional image. Pipeline steps: 1. Run safety checks on input prompt 2. Enhance prompt using upsampler if enabled 3. Run safety checks on upsampled prompt if applicable 4. Convert prompt to embeddings 5. Generate video frames using diffusion 6. Run safety checks and apply face blur on generated video frames Args: prompt: Text description of desired video image_or_video_path: Path to conditioning image or video negative_prompt: Optional text to guide what not to generate Returns: tuple: ( Generated video frames as uint8 np.ndarray [T, H, W, C], Final prompt used for generation (may be enhanced) ), or None if content fails guardrail safety checks """ log.info(f"Run with image or video path: {image_or_video_path}") log.info(f"Run with negative prompt: {negative_prompt}") log.info(f"Run with prompt upsampler: {self.enable_prompt_upsampler}") if self.enable_prompt_upsampler: log.info("Run prompt upsampler on image or video, input prompt is not used") prompt = self._run_prompt_upsampler_on_prompt_with_offload(image_or_video_path=image_or_video_path) log.info(f"Run with prompt: {prompt}") if not self.disable_guardrail: log.info(f"Run guardrail on {'upsampled' if self.enable_prompt_upsampler else 'text'} prompt") is_safe = self._run_guardrail_on_prompt_with_offload(prompt) if not is_safe: log.critical(f"Input {'upsampled' if self.enable_prompt_upsampler else 'text'} prompt is not safe") return None log.info(f"Pass guardrail on {'upsampled' if self.enable_prompt_upsampler else 'text'} prompt") else: log.info("Not running guardrail") log.info("Run text embedding on prompt") if negative_prompt: prompts = [prompt, negative_prompt] else: prompts = [prompt] prompt_embeddings, _ = self._run_text_embedding_on_prompt_with_offload(prompts) prompt_embedding = prompt_embeddings[0] negative_prompt_embedding = prompt_embeddings[1] if negative_prompt else None log.info("Finish text embedding on prompt") # Generate video log.info("Run generation") video = self._run_model_with_offload( prompt_embedding, negative_prompt_embedding=negative_prompt_embedding, image_or_video_path=image_or_video_path, ) log.info("Finish generation") if not self.disable_guardrail: log.info("Run guardrail on generated video") video = self._run_guardrail_on_video_with_offload(video) if video is None: log.critical("Generated video is not safe") return None log.info("Pass guardrail on generated video") return video, prompt class DiffusionText2WorldMultiviewGenerationPipeline(DiffusionText2WorldGenerationPipeline): def __init__( self, inference_type: str, checkpoint_dir: str, checkpoint_name: str, prompt_upsampler_dir: Optional[str] = None, has_text_input: bool = True, offload_network: bool = False, offload_tokenizer: bool = False, offload_text_encoder_model: bool = False, offload_prompt_upsampler: bool = False, offload_guardrail_models: bool = False, disable_guardrail: bool = False, guidance: float = 7.0, num_steps: int = 35, height: int = 704, width: int = 1280, fps: int = 24, num_video_frames: int = 121, n_views: int = 6, frame_repeat_negative_condition: int = 10, seed: int = 0, ): """Initialize the diffusion multi-view world generation pipeline. Args: inference_type: Type of world generation ('text2world' or 'video2world') checkpoint_dir: Base directory containing model checkpoints checkpoint_name: Name of the diffusion transformer checkpoint to use prompt_upsampler_dir: Directory containing prompt upsampler model weights enable_prompt_upsampler: Whether to use prompt upsampling has_text_input: Whether the pipeline takes text input for world generation offload_network: Whether to offload diffusion transformer after inference offload_tokenizer: Whether to offload tokenizer after inference offload_text_encoder_model: Whether to offload T5 model after inference offload_prompt_upsampler: Whether to offload prompt upsampler offload_guardrail_models: Whether to offload guardrail models disable_guardrail: Whether to disable guardrail guidance: Classifier-free guidance scale num_steps: Number of diffusion sampling steps height: Height of output video width: Width of output video fps: Frames per second of output video num_video_frames: Number of frames to generate n_views: Number of views frame_repeat_negative_condition: Number of frames to repeat to be used as negative condition. seed: Random seed for sampling """ assert inference_type in [ "text2world", "video2world", ], "Invalid inference_type, must be 'text2world' or 'video2world'" self.n_views = n_views self.frame_repeat_negative_condition = frame_repeat_negative_condition super().__init__( inference_type=inference_type, checkpoint_dir=checkpoint_dir, checkpoint_name=checkpoint_name, prompt_upsampler_dir=prompt_upsampler_dir, enable_prompt_upsampler=False, has_text_input=has_text_input, offload_network=offload_network, offload_tokenizer=offload_tokenizer, offload_text_encoder_model=offload_text_encoder_model, offload_prompt_upsampler=offload_prompt_upsampler, offload_guardrail_models=offload_guardrail_models, disable_guardrail=disable_guardrail, guidance=guidance, num_steps=num_steps, height=height, width=width, fps=fps, num_video_frames=num_video_frames, seed=seed, ) def _load_model(self): self.model = load_model_by_config( config_job_name=self.model_name, config_file="cosmos_predict1/diffusion/config/config.py", model_class=DiffusionMultiviewT2WModel, ) def _run_tokenizer_decoding(self, sample: torch.Tensor) -> np.ndarray: """Decode latent samples to video frames using the tokenizer decoder. Args: sample: Latent tensor from diffusion model [B, C, T, H, W] Returns: np.ndarray: Decoded video frames as uint8 numpy array [T, H, W, C] with values in range [0, 255] """ # Decode video video = (1.0 + self.model.decode(sample)).clamp(0, 2) / 2 # [B, 3, T, H, W] video_segments = einops.rearrange(video, "b c (v t) h w -> b c v t h w", v=self.n_views) grid_video = torch.stack( [video_segments[:, :, i] for i in [1, 0, 2, 4, 3, 5]], dim=2, ) grid_video = einops.rearrange(grid_video, "b c (h w) t h1 w1 -> b c t (h h1) (w w1)", h=2, w=3) grid_video = (grid_video[0].permute(1, 2, 3, 0) * 255).to(torch.uint8).cpu().numpy() video = (video[0].permute(1, 2, 3, 0) * 255).to(torch.uint8).cpu().numpy() return [grid_video, video] def _run_model( self, embedding: torch.Tensor, negative_prompt_embedding: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Generate video latents using the diffusion model. Args: embedding: Text embedding tensor from text encoder negative_prompt_embedding: Optional embedding for negative prompt guidance Returns: torch.Tensor: Generated video latents before tokenizer decoding Note: The model and tokenizer are automatically offloaded after inference if offloading is enabled in the config. """ # Get video batch and state shape data_batch, state_shape = get_video_batch_for_multiview_model( model=self.model, prompt_embedding=embedding, height=self.height, width=self.width, fps=self.fps, num_video_frames=self.num_video_frames * len(embedding), # number of views frame_repeat_negative_condition=self.frame_repeat_negative_condition, ) # Generate video frames sample = generate_world_from_text( model=self.model, state_shape=state_shape, is_negative_prompt=False, data_batch=data_batch, guidance=self.guidance, num_steps=self.num_steps, seed=self.seed, ) return sample def generate( self, prompt: dict, ) -> tuple[np.ndarray, str] | None: """Generate video from text prompt with optional negative prompt guidance. Pipeline steps: 1. Convert prompt to embeddings 2. Generate video frames using diffusion Args: prompt: A dictionary of text description of desired video. Returns: tuple: ( Generated video frames as uint8 np.ndarray [T, H, W, C], Final prompt used for generation (may be enhanced) ), or None if content fails guardrail safety checks """ log.info(f"Run with prompt: {prompt}") prompts = [ prompt["prompt"], prompt["prompt_left"], prompt["prompt_right"], prompt["prompt_back"], prompt["prompt_back_left"], prompt["prompt_back_right"], ] prompt_embeddings, _ = self._run_text_embedding_on_prompt_with_offload(prompts) log.info("Finish text embedding on prompt") # Generate video log.info("Run generation") videos = self._run_model_with_offload( prompt_embeddings, ) log.info("Finish generation") return videos, prompt class DiffusionVideo2WorldMultiviewGenerationPipeline(DiffusionText2WorldMultiviewGenerationPipeline): def __init__( self, inference_type: str, checkpoint_dir: str, checkpoint_name: str, prompt_upsampler_dir: Optional[str] = None, enable_prompt_upsampler: bool = True, has_text_input: bool = True, offload_network: bool = False, offload_tokenizer: bool = False, offload_text_encoder_model: bool = False, offload_prompt_upsampler: bool = False, offload_guardrail_models: bool = False, disable_guardrail: bool = False, guidance: float = 7.0, num_steps: int = 35, height: int = 704, width: int = 1280, fps: int = 24, num_video_frames: int = 121, seed: int = 0, num_input_frames: int = 1, n_views: int = 6, frame_repeat_negative_condition: int = 10, ): """Initialize diffusion world multi-view generation pipeline. Args: inference_type: Type of world generation ('text2world' or 'video2world') checkpoint_dir: Base directory containing model checkpoints checkpoint_name: Name of the diffusion transformer checkpoint to use prompt_upsampler_dir: Directory containing prompt upsampler model weights enable_prompt_upsampler: Whether to use prompt upsampling has_text_input: Whether the pipeline takes text input for world generation offload_network: Whether to offload diffusion transformer after inference offload_tokenizer: Whether to offload tokenizer after inference offload_text_encoder_model: Whether to offload T5 model after inference offload_prompt_upsampler: Whether to offload prompt upsampler offload_guardrail_models: Whether to offload guardrail models disable_guardrail: Whether to disable guardrail guidance: Classifier-free guidance scale num_steps: Number of diffusion sampling steps height: Height of output video width: Width of output video fps: Frames per second of output video num_video_frames: Number of frames to generate seed: Random seed for sampling num_input_frames: Number of latent conditions """ self.num_input_frames = num_input_frames super().__init__( inference_type=inference_type, checkpoint_dir=checkpoint_dir, checkpoint_name=checkpoint_name, prompt_upsampler_dir=prompt_upsampler_dir, has_text_input=has_text_input, offload_network=offload_network, offload_tokenizer=offload_tokenizer, offload_text_encoder_model=offload_text_encoder_model, offload_prompt_upsampler=offload_prompt_upsampler, offload_guardrail_models=offload_guardrail_models, disable_guardrail=disable_guardrail, guidance=guidance, num_steps=num_steps, height=height, width=width, fps=fps, num_video_frames=num_video_frames, seed=seed, n_views=n_views, frame_repeat_negative_condition=frame_repeat_negative_condition, ) def _load_model(self): self.model = load_model_by_config( config_job_name=self.model_name, config_file="cosmos_predict1/diffusion/config/config.py", model_class=DiffusionMultiviewV2WModel, ) def _run_model( self, embedding: torch.Tensor, condition_latent: torch.Tensor, negative_prompt_embedding: torch.Tensor | None = None, data_batch: dict = None, state_shape: list = None, ) -> torch.Tensor: """Generate video frames using the diffusion model. Args: embedding: Text embedding tensor from T5 encoder condition_latent: Latent tensor from conditioning image or video negative_prompt_embedding: Optional embedding for negative prompt guidance Returns: Tensor of generated video frames Note: Model and tokenizer are automatically offloaded after inference if offloading is enabled. """ # Generate video frames video = generate_world_from_video( model=self.model, state_shape=state_shape, is_negative_prompt=False, data_batch=data_batch, guidance=self.guidance, num_steps=self.num_steps, seed=self.seed, condition_latent=condition_latent, num_input_frames=self.num_input_frames, ) return video def _run_tokenizer_encoding(self, image_or_video_path: str, state_shape: list) -> torch.Tensor: """ Encode image to latent space Args: image_or_video_path: Path to conditioning image Returns: torch.Tensor: Latent tensor from tokenizer encoding """ condition_latent, condition_frames = get_condition_latent_multiview( model=self.model, input_image_or_video_path=image_or_video_path, num_input_frames=self.num_input_frames, state_shape=state_shape, ) return condition_latent, condition_frames def _run_model_with_offload( self, prompt_embedding: torch.Tensor, image_or_video_path: str, negative_prompt_embedding: Optional[torch.Tensor] = None, ) -> np.ndarray: """Generate world representation with automatic model offloading. Wraps the core generation process with model loading/offloading logic to minimize GPU memory usage during inference. Args: prompt_embedding: Text embedding tensor from T5 encoder image_or_video_path: Path to conditioning image or video negative_prompt_embedding: Optional embedding for negative prompt guidance Returns: np.ndarray: Generated world representation as numpy array """ if self.offload_tokenizer: self._load_tokenizer() data_batch, state_shape = get_video_batch_for_multiview_model( model=self.model, prompt_embedding=prompt_embedding, height=self.height, width=self.width, fps=self.fps, num_video_frames=self.num_video_frames * len(prompt_embedding), # number of views frame_repeat_negative_condition=self.frame_repeat_negative_condition, ) condition_latent, condition_frames = self._run_tokenizer_encoding(image_or_video_path, state_shape) if self.offload_network: self._load_network() sample = self._run_model(prompt_embedding, condition_latent, negative_prompt_embedding, data_batch, state_shape) if self.offload_network: self._offload_network() sample = self._run_tokenizer_decoding(sample) if self.offload_tokenizer: self._offload_tokenizer() return sample def generate( self, prompt: dict, image_or_video_path: str, ) -> tuple[np.ndarray, str] | None: """Generate video from text prompt with optional negative prompt guidance. Pipeline steps: 1. Convert prompt to embeddings 2. Generate video frames using diffusion Args: prompt: A dictionary of text description of desired video. Returns: tuple: ( Generated video frames as uint8 np.ndarray [T, H, W, C], Final prompt used for generation (may be enhanced) ), or None if content fails guardrail safety checks """ log.info(f"Run with prompt: {prompt}") prompts = [ prompt["prompt"], prompt["prompt_left"], prompt["prompt_right"], prompt["prompt_back"], prompt["prompt_back_left"], prompt["prompt_back_right"], ] prompt_embeddings, _ = self._run_text_embedding_on_prompt_with_offload(prompts) log.info("Finish text embedding on prompt") # Generate video log.info("Run generation") video = self._run_model_with_offload( prompt_embeddings, image_or_video_path=image_or_video_path, ) log.info("Finish generation") return video, prompt class DiffusionWorldInterpolatorGenerationPipeline(DiffusionVideo2WorldGenerationPipeline): def __init__( self, inference_type: str, checkpoint_dir: str, checkpoint_name: str, prompt_upsampler_dir: Optional[str] = None, enable_prompt_upsampler: bool = True, has_text_input: bool = True, offload_network: bool = False, offload_tokenizer: bool = False, offload_text_encoder_model: bool = False, offload_prompt_upsampler: bool = False, offload_guardrail_models: bool = False, disable_guardrail: bool = False, guidance: float = -1.0, num_steps: int = 35, height: int = 704, width: int = 1280, fps: int = 24, num_video_frames: int = 121, seed: int = 11, num_input_frames: int = 1, num_frame_pairs: int = 1, frame_index_start: int = 0, frame_stride: int = 1, ): """Initialize diffusion world generation pipeline. Args: inference_type: Type of world generation ('text2world' or 'video2world') checkpoint_dir: Base directory containing model checkpoints checkpoint_name: Name of the diffusion transformer checkpoint to use prompt_upsampler_dir: Directory containing prompt upsampler model weights enable_prompt_upsampler: Whether to use prompt upsampling has_text_input: Whether the pipeline takes text input for world generation offload_network: Whether to offload diffusion transformer after inference offload_tokenizer: Whether to offload tokenizer after inference offload_text_encoder_model: Whether to offload T5 model after inference offload_prompt_upsampler: Whether to offload prompt upsampler offload_guardrail_models: Whether to offload guardrail models disable_guardrail: Whether to disable guardrail guidance: Classifier-free guidance scale num_steps: Number of diffusion sampling steps height: Height of output video width: Width of output video fps: Frames per second of output video num_video_frames: Number of frames to generate seed: Random seed for sampling num_input_frames: Number of latent conditions """ self.num_input_frames = num_input_frames self.num_frame_pairs = num_frame_pairs self.frame_index_start = frame_index_start self.frame_stride = frame_stride self.num_steps = num_steps self.height = height self.width = width self.fps = fps super().__init__( inference_type=inference_type, checkpoint_dir=checkpoint_dir, checkpoint_name=checkpoint_name, prompt_upsampler_dir=prompt_upsampler_dir, enable_prompt_upsampler=enable_prompt_upsampler, has_text_input=has_text_input, offload_network=offload_network, offload_tokenizer=offload_tokenizer, offload_text_encoder_model=offload_text_encoder_model, offload_prompt_upsampler=offload_prompt_upsampler, offload_guardrail_models=offload_guardrail_models, disable_guardrail=disable_guardrail, guidance=guidance, num_steps=num_steps, height=height, width=width, fps=fps, num_video_frames=num_video_frames, seed=seed, num_input_frames=num_input_frames, ) def _run_prompt_upsampler_on_prompt(self, image_or_video_path: str) -> str: """Enhance the input prompt using visual context from the conditioning image. Args: image_or_video_path: Path to conditioning image or video used for visual context Returns: str: Enhanced prompt incorporating visual details from the image """ dialog = prepare_dialog(image_or_video_path) upsampled_prompt = run_chat_completion_vlm( self.prompt_upsampler, dialog, max_gen_len=400, temperature=0.01, top_p=0.9, logprobs=False ) log.info(f"Upsampled prompt: {upsampled_prompt}") return upsampled_prompt def _load_prompt_upsampler_model(self): self.prompt_upsampler = create_vlm_prompt_upsampler( checkpoint_dir=os.path.join(self.checkpoint_dir, self.prompt_upsampler_dir), ) def _load_model(self): self.model = load_model_by_config( config_job_name=self.model_name, config_file="cosmos_predict1/diffusion/config/config.py", model_class=DiffusionWorldInterpolatorWModel, ) @torch.inference_mode() def _run_model( self, condition_latent: torch.Tensor | None = None, negative_prompt_embedding: torch.Tensor | None = None, num_of_loops: int = 1, num_of_latent_overlap_list: list[int] = [1], augment_sigma_list: list[float] = [0.001], add_input_frames_guidance: float = 0, skip_reencode: int = 0, state_shape: list = None, raw_data_batch: dict = None, ) -> np.ndarray: """Generate video frames using the diffusion model, supporting chunk processing for video extension. Args: condition_latent: Latent tensor from conditioning image or video (optional for video extension). negative_prompt_embedding: Optional embedding for negative prompt guidance. num_of_loops: Number of loops for generating video segments. num_of_latent_overlap_list: List of overlaps for latent conditions in each loop. augment_sigma_list: List of sigma values for augmentation. add_input_frames_guidance: Guidance strength for input frames. skip_reencode: Whether to skip reencoding. frame_index_start: Starting index for frame pairs. num_frame_pairs: Number of frame pairs to process. frame_stride: Stride between frame pairs. is_interpolator_model: Whether the model is an interpolator. input_frames: Input video frames for interpolation (optional). Returns: np.ndarray: Generated video frames in shape (T, H, W, C). """ video_np_THWC, _, _ = generate_video_from_batch_with_loop( model=self.model, data_batch=raw_data_batch, condition_latent=condition_latent, num_of_loops=num_of_loops, num_of_latent_overlap_list=num_of_latent_overlap_list, guidance=self.guidance, state_shape=state_shape, num_steps=self.num_steps, seed=self.seed, is_negative_prompt=True if negative_prompt_embedding is not None else False, visualize=False, save_fig_path=None, augment_sigma_list=augment_sigma_list, add_input_frames_guidance=add_input_frames_guidance, skip_reencode=skip_reencode, ) return video_np_THWC def _run_tokenizer_encoding( self, image_or_video_path: str, frame_index: int = 0, frame_stride: int = 1 ) -> torch.Tensor: """Encode image to latent space Args: image_or_video_path: Path to conditioning image frame_index: Starting frame index for encoding frame_stride: Stride between frames for encoding Returns: torch.Tensor: Latent tensor from tokenizer encoding """ condition_latent = get_condition_latent( model=self.model, input_image_or_video_path=image_or_video_path, num_input_frames=self.num_input_frames, state_shape=self.model.state_shape, frame_index=frame_index, frame_stride=frame_stride, ) return condition_latent def _run_model_with_offload( self, prompt_embedding: torch.Tensor, image_or_video_path: str, negative_prompt_embedding: Optional[torch.Tensor] = None, frame_index_start: int = 0, num_frame_pairs: int = 1, ) -> np.ndarray: """Generate world representation with automatic model offloading. Wraps the core generation process with model loading/offloading logic to minimize GPU memory usage during inference. Args: prompt_embedding: Text embedding tensor from T5 encoder image_or_video_path: Path to conditioning image or video negative_prompt_embedding: Optional embedding for negative prompt guidance frame_index_start: Starting index for frame pairs num_frame_pairs: Number of frame pairs to process Returns: np.ndarray: Generated world representation as numpy array """ if self.offload_tokenizer: self._load_tokenizer() # Prepare video batch and state shape raw_data_batch, state_shape = get_video_batch( model=self.model, prompt_embedding=prompt_embedding, negative_prompt_embedding=negative_prompt_embedding, height=self.height, width=self.width, fps=self.fps, num_video_frames=self.num_video_frames, ) H, W = ( state_shape[-2] * self.model.tokenizer.spatial_compression_factor, state_shape[-1] * self.model.tokenizer.spatial_compression_factor, ) input_path_format = image_or_video_path.split(".")[-1] input_frames = read_video_or_image_into_frames_BCTHW( image_or_video_path, input_path_format=input_path_format, H=H, W=W, ) num_frames = input_frames.shape[2] num_frame_pairs = num_frame_pairs or num_frames // self.frame_stride frame_stride = self.frame_stride video_output = [] for frame_index in range(frame_index_start, num_frame_pairs): print(f"Processing frame pair {frame_index + 1} / {num_frame_pairs}...") condition_latent = self._run_tokenizer_encoding(image_or_video_path, frame_index, frame_stride) video_np_THWC = self._run_model( condition_latent=condition_latent, negative_prompt_embedding=negative_prompt_embedding, raw_data_batch=raw_data_batch, state_shape=state_shape, ) # Convert to tensor, rearrange, and normalize to [0, 1] video_0_1 = einops.rearrange(torch.from_numpy(video_np_THWC), "t h w c -> c t h w") / 255.0 # Handle overlap by skipping the first frame of subsequent segments if len(video_output) == 0: video_output.append(video_0_1) else: video_output.append(video_0_1[:, 1:, :, :]) # Skip first frame to avoid duplication # Concatenate all segments video_tensor = torch.cat(video_output, dim=1) # Shape: (C, total_num_frames, H, W) # Convert to NumPy array for guardrail: [T, H, W, C], uint8, [0, 255] video_np = (video_tensor.permute(1, 2, 3, 0) * 255).to(torch.uint8).cpu().numpy() # Shape: (T, H, W, C) if self.offload_network: self._offload_network() if self.offload_tokenizer: self._offload_tokenizer() return video_np def generate( self, prompt: str, image_or_video_path: str, negative_prompt: Optional[str] = None, ) -> tuple[np.ndarray, str] | None: """Generate video from text prompt and optional image. Pipeline steps: 1. Run safety checks on input prompt 2. Enhance prompt using upsampler if enabled 3. Run safety checks on upsampled prompt if applicable 4. Convert prompt to embeddings 5. Generate video frames using diffusion 6. Run safety checks and apply face blur on generated video frames Args: prompt: Text description of desired video image_or_video_path: Path to conditioning image or video negative_prompt: Optional text to guide what not to generate Returns: tuple: ( Generated video frames as uint8 np.ndarray [T, H, W, C], Final prompt used for generation (may be enhanced) ), or None if content fails guardrail safety checks """ log.info(f"Run with prompt: {prompt}") log.info(f"Run with image or video path: {image_or_video_path}") log.info(f"Run with negative prompt: {negative_prompt}") log.info(f"Run with prompt upsampler: {self.enable_prompt_upsampler}") if not self.disable_guardrail and not self.enable_prompt_upsampler: log.info("Run guardrail on prompt") is_safe = self._run_guardrail_on_prompt_with_offload(prompt) if not is_safe: log.critical("Input text prompt is not safe") return None log.info("Pass guardrail on prompt") else: log.info("Run prompt upsampler on image or video, input prompt is not used") prompt = self._run_prompt_upsampler_on_prompt_with_offload(image_or_video_path=image_or_video_path) if not self.disable_guardrail: log.info("Run guardrail on upsampled prompt") is_safe = self._run_guardrail_on_prompt_with_offload(prompt) if not is_safe: log.critical("Upsampled text prompt is not safe") return None log.info("Pass guardrail on upsampled prompt") log.info("Run text embedding on prompt") if negative_prompt: prompts = [prompt, negative_prompt] else: prompts = [prompt] prompt_embeddings, _ = self._run_text_embedding_on_prompt_with_offload(prompts) prompt_embedding = prompt_embeddings[0] negative_prompt_embedding = prompt_embeddings[1] if negative_prompt else None log.info("Finish text embedding on prompt") # Generate video log.info("Run generation") video = self._run_model_with_offload( prompt_embedding, negative_prompt_embedding=negative_prompt_embedding, image_or_video_path=image_or_video_path, frame_index_start=self.frame_index_start, num_frame_pairs=self.num_frame_pairs, ) log.info("Finish generation") if not self.disable_guardrail: log.info("Run guardrail on generated video") video = self._run_guardrail_on_video_with_offload(video) if video is None: log.critical("Generated video is not safe") return None log.info("Pass guardrail on generated video") return video, prompt