# 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. """ A variant of predict2_video2world.py for GR00T models that: 1. Supports prompt prefix for robot task descriptions 2. Turns off the guardrail and prompt refiner 3. Supports two GR00T variants: GR1 and DROID """ import argparse import json import os # Set TOKENIZERS_PARALLELISM environment variable to avoid deadlocks with multiprocessing os.environ["TOKENIZERS_PARALLELISM"] = "false" import torch from tqdm import tqdm from cosmos_predict2.configs.base.config_video2world import PREDICT2_VIDEO2WORLD_PIPELINE_14B from cosmos_predict2.pipelines.video2world import Video2WorldPipeline from cosmos_predict2.models.utils import load_lora_weights from examples.video2world import _DEFAULT_NEGATIVE_PROMPT, validate_input_file from imaginaire.utils.io import save_image_or_video from megatron.core import parallel_state from imaginaire.utils import distributed, log, misc def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="GR00T Video-to-World Generation with Cosmos Predict2") parser.add_argument( "--model_size", choices=["14B"], default="14B", help="Size of the model to use for GR00T video-to-world generation", ) parser.add_argument( "--dit_path", type=str, default="", help="Custom path to the DiT model checkpoint for post-trained models.", ) parser.add_argument( "--is_lora_trained", action="store_true", help="Whether the model is trained with lora weights", ) parser.add_argument( "--training_config_path", type=str, default="", help="Path to the training config, mandatory if is_lora_trained is True as we need to find lora configs from the training config", ) parser.add_argument( "--prompt", type=str, default="", help="Text prompt for video generation", ) parser.add_argument( "--input_path", type=str, default="assets/video2world/input0.jpg", help="Path to input image or video for conditioning (include file extension)", ) parser.add_argument( "--negative_prompt", type=str, default=_DEFAULT_NEGATIVE_PROMPT, help="Negative text prompt for video-to-world generation", ) parser.add_argument( "--num_conditional_frames", type=int, default=1, choices=[1, 5], help="Number of frames to condition on (1 for single frame, 5 for multi-frame conditioning)", ) parser.add_argument( "--batch_input_json", type=str, default=None, help="Path to JSON file containing batch inputs. Each entry should have 'input_video', 'prompt', and 'output_video' fields.", ) parser.add_argument("--guidance", type=float, default=7, help="Guidance value") parser.add_argument("--seed", type=int, default=0, help="Random seed for reproducibility") parser.add_argument( "--save_path", type=str, default="output/generated_video.mp4", help="Path to save the generated video (include file extension)", ) parser.add_argument( "--num_gpus", type=int, default=1, help="Number of GPUs to use for context parallel inference (should be a divisor of the total frames)", ) parser.add_argument( "--disable_guardrail", action="store_true", help="Disable guardrail checks on prompts", ) parser.add_argument( "--gr00t_variant", type=str, required=True, help="GR00T variant to use", choices=["gr1", "droid"] ) parser.add_argument( "--prompt_prefix", type=str, default="The robot arm is performing a task. ", help="Prefix to add to all prompts" ) return parser.parse_args() def setup_pipeline(args: argparse.Namespace): log.info(f"Using model size: {args.model_size} with GR00T variant: {args.gr00t_variant}") # Only 14B models are supported for GR00T if args.model_size == "14B": config = PREDICT2_VIDEO2WORLD_PIPELINE_14B config.resolution = "480" # GR00T models use 480p resolution config.prompt_refiner_config.enabled = False if args.gr00t_variant == "gr1": dit_path = "checkpoints/nvidia/Cosmos-Predict2-14B-Sample-GR00T-Dreams-GR1/model-480p-16fps.pt" elif args.gr00t_variant == "droid": dit_path = ( "checkpoints/nvidia/Cosmos-Predict2-14B-Sample-GR00T-Dreams-DROID/model-480p-16fps.pt" ) else: raise ValueError("Only 14B model size is supported for GR00T variants") if args.dit_path: dit_path = args.dit_path text_encoder_path = "checkpoints/google-t5/t5-11b" log.info(f"Loading model from: {dit_path}") misc.set_random_seed(seed=args.seed, by_rank=True) # Initialize cuDNN. torch.backends.cudnn.deterministic = False torch.backends.cudnn.benchmark = True # Floating-point precision settings. torch.backends.cudnn.allow_tf32 = True torch.backends.cuda.matmul.allow_tf32 = True # Initialize distributed environment for multi-GPU inference if args.num_gpus > 1: log.info(f"Initializing distributed environment with {args.num_gpus} GPUs for context parallelism") distributed.init() parallel_state.initialize_model_parallel(context_parallel_size=args.num_gpus) log.info(f"Context parallel group initialized with {args.num_gpus} GPUs") # Disable guardrail if requested if args.disable_guardrail: log.warning("Guardrail checks are disabled") config.guardrail_config.enabled = False # Load models log.info(f"Initializing Video2WorldPipeline with GR00T variant: {args.gr00t_variant}") pipe = Video2WorldPipeline.from_config( config=config, dit_path=None if args.is_lora_trained else dit_path, text_encoder_path=text_encoder_path, device="cuda", torch_dtype=torch.bfloat16, load_prompt_refiner=False, # Disable prompt refiner for GR00T ) if args.is_lora_trained: log.info(f"Loading DiT (with lora weights) from {dit_path}") assert os.path.exists(args.training_config_path), f"Training config file does not exist: {args.training_config_path}" load_lora_weights(pipe, dit_path, args.training_config_path) log.success(f"Successfully loaded DiT (with lora weights) from {dit_path}") return pipe def process_single_generation( pipe, input_path, prompt, output_path, negative_prompt, num_conditional_frames, guidance, seed, prompt_prefix ): # Validate input file if not validate_input_file(input_path, num_conditional_frames): log.warning(f"Input file validation failed: {input_path}") return False # Add prefix to prompt full_prompt = prompt_prefix + prompt log.info(f"Running Video2WorldPipeline\ninput: {input_path}\nprompt: {full_prompt}") video = pipe( prompt=full_prompt, negative_prompt=negative_prompt, input_path=input_path, num_conditional_frames=num_conditional_frames, guidance=guidance, seed=seed, ) if video is not None: # save the generated video output_dir = os.path.dirname(output_path) if output_dir: os.makedirs(output_dir, exist_ok=True) log.info(f"Saving generated video to: {output_path}") save_image_or_video(video, output_path, fps=16) log.success(f"Successfully saved video to: {output_path}") return True return False def generate_video(args: argparse.Namespace, pipe: Video2WorldPipeline) -> None: # Video-to-World if args.batch_input_json is not None: # Process batch inputs from JSON file log.info(f"Loading batch inputs from JSON file: {args.batch_input_json}") with open(args.batch_input_json, "r") as f: batch_inputs = json.load(f) for idx, item in enumerate(tqdm(batch_inputs)): input_video = item.get("input_video", "") prompt = item.get("prompt", "") output_video = item.get("output_video", f"output_{idx}.mp4") if not input_video or not prompt: log.warning(f"Skipping item {idx}: Missing input_video or prompt") continue process_single_generation( pipe=pipe, input_path=input_video, prompt=prompt, output_path=output_video, negative_prompt=args.negative_prompt, num_conditional_frames=args.num_conditional_frames, guidance=args.guidance, seed=args.seed, prompt_prefix=args.prompt_prefix, ) else: process_single_generation( pipe=pipe, input_path=args.input_path, prompt=args.prompt, output_path=args.save_path, negative_prompt=args.negative_prompt, num_conditional_frames=args.num_conditional_frames, guidance=args.guidance, seed=args.seed, prompt_prefix=args.prompt_prefix, ) return def cleanup_distributed(): """Clean up the distributed environment if initialized.""" if parallel_state.is_initialized(): parallel_state.destroy_model_parallel() if torch.distributed.is_initialized(): torch.distributed.destroy_process_group() if __name__ == "__main__": args = parse_args() try: pipe = setup_pipeline(args) generate_video(args, pipe) finally: # Make sure to clean up the distributed environment cleanup_distributed()