|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import argparse |
|
import copy |
|
import json |
|
import os |
|
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false" |
|
|
|
import sys |
|
from io import BytesIO |
|
|
|
import torch |
|
|
|
from cosmos_transfer1.checkpoints import ( |
|
BASE_7B_CHECKPOINT_AV_SAMPLE_PATH, |
|
BASE_7B_CHECKPOINT_PATH, |
|
EDGE2WORLD_CONTROLNET_DISTILLED_CHECKPOINT_PATH, |
|
) |
|
from cosmos_transfer1.diffusion.inference.inference_utils import load_controlnet_specs, validate_controlnet_specs |
|
from cosmos_transfer1.diffusion.inference.preprocessors import Preprocessors |
|
from cosmos_transfer1.diffusion.inference.world_generation_pipeline import ( |
|
DiffusionControl2WorldGenerationPipeline, |
|
DistilledControl2WorldGenerationPipeline, |
|
) |
|
from cosmos_transfer1.utils import log, misc |
|
from cosmos_transfer1.utils.io import read_prompts_from_file, save_video |
|
|
|
torch.enable_grad(False) |
|
torch.serialization.add_safe_globals([BytesIO]) |
|
|
|
|
|
def parse_arguments() -> argparse.Namespace: |
|
parser = argparse.ArgumentParser(description="Control to world generation demo script", conflict_handler="resolve") |
|
|
|
|
|
parser.add_argument( |
|
"--prompt", |
|
type=str, |
|
default="The video captures a stunning, photorealistic scene with remarkable attention to detail, giving it a lifelike appearance that is almost indistinguishable from reality. It appears to be from a high-budget 4K movie, showcasing ultra-high-definition quality with impeccable resolution.", |
|
help="prompt which the sampled video condition on", |
|
) |
|
parser.add_argument( |
|
"--negative_prompt", |
|
type=str, |
|
default="The video captures a game playing, with bad crappy graphics and cartoonish frames. It represents a recording of old outdated games. The lighting looks very fake. The textures are very raw and basic. The geometries are very primitive. The images are very pixelated and of poor CG quality. There are many subtitles in the footage. Overall, the video is unrealistic at all.", |
|
help="negative prompt which the sampled video condition on", |
|
) |
|
parser.add_argument( |
|
"--input_video_path", |
|
type=str, |
|
default="", |
|
help="Optional input RGB video path", |
|
) |
|
parser.add_argument( |
|
"--num_input_frames", |
|
type=int, |
|
default=1, |
|
help="Number of conditional frames for long video generation", |
|
choices=[1], |
|
) |
|
parser.add_argument("--sigma_max", type=float, default=70.0, help="sigma_max for partial denoising") |
|
parser.add_argument( |
|
"--blur_strength", |
|
type=str, |
|
default="medium", |
|
choices=["very_low", "low", "medium", "high", "very_high"], |
|
help="blur strength.", |
|
) |
|
parser.add_argument( |
|
"--canny_threshold", |
|
type=str, |
|
default="medium", |
|
choices=["very_low", "low", "medium", "high", "very_high"], |
|
help="blur strength of canny threshold applied to input. Lower means less blur or more detected edges, which means higher fidelity to input.", |
|
) |
|
parser.add_argument( |
|
"--controlnet_specs", |
|
type=str, |
|
help="Path to JSON file specifying multicontrolnet configurations", |
|
required=True, |
|
) |
|
parser.add_argument( |
|
"--is_av_sample", action="store_true", help="Whether the model is an driving post-training model" |
|
) |
|
parser.add_argument( |
|
"--checkpoint_dir", type=str, default="checkpoints", help="Base directory containing model checkpoints" |
|
) |
|
parser.add_argument( |
|
"--tokenizer_dir", |
|
type=str, |
|
default="Cosmos-Tokenize1-CV8x8x8-720p", |
|
help="Tokenizer weights directory relative to checkpoint_dir", |
|
) |
|
parser.add_argument( |
|
"--video_save_name", |
|
type=str, |
|
default="output", |
|
help="Output filename for generating a single video", |
|
) |
|
parser.add_argument( |
|
"--video_save_folder", |
|
type=str, |
|
default="outputs/", |
|
help="Output folder for generating a batch of videos", |
|
) |
|
parser.add_argument( |
|
"--batch_input_path", |
|
type=str, |
|
help="Path to a JSONL file of input prompts for generating a batch of videos", |
|
) |
|
parser.add_argument("--batch_size", type=int, default=1, help="Batch size") |
|
parser.add_argument("--num_steps", type=int, default=35, help="Number of diffusion sampling steps") |
|
parser.add_argument("--guidance", type=float, default=5, help="Classifier-free guidance scale value") |
|
parser.add_argument("--fps", type=int, default=24, help="FPS of the output video") |
|
parser.add_argument("--seed", type=int, default=1, help="Random seed") |
|
parser.add_argument("--num_gpus", type=int, default=1, help="Number of GPUs used to run inference in parallel.") |
|
parser.add_argument( |
|
"--offload_diffusion_transformer", |
|
action="store_true", |
|
help="Offload DiT after inference", |
|
) |
|
parser.add_argument( |
|
"--offload_text_encoder_model", |
|
action="store_true", |
|
help="Offload text encoder model after inference", |
|
) |
|
parser.add_argument( |
|
"--offload_guardrail_models", |
|
action="store_true", |
|
help="Offload guardrail models after inference", |
|
) |
|
parser.add_argument( |
|
"--upsample_prompt", |
|
action="store_true", |
|
help="Upsample prompt using Pixtral upsampler model", |
|
) |
|
parser.add_argument( |
|
"--offload_prompt_upsampler", |
|
action="store_true", |
|
help="Offload prompt upsampler model after inference", |
|
) |
|
parser.add_argument("--use_distilled", action="store_true", help="Use distilled ControlNet model variant") |
|
|
|
cmd_args = parser.parse_args() |
|
|
|
|
|
control_inputs, json_args = load_controlnet_specs(cmd_args) |
|
|
|
log.info(f"control_inputs: {json.dumps(control_inputs, indent=4)}") |
|
log.info(f"args in json: {json.dumps(json_args, indent=4)}") |
|
|
|
|
|
for key in json_args: |
|
if f"--{key}" not in sys.argv: |
|
setattr(cmd_args, key, json_args[key]) |
|
|
|
log.info(f"final args: {json.dumps(vars(cmd_args), indent=4)}") |
|
|
|
return cmd_args, control_inputs |
|
|
|
|
|
def demo(cfg, control_inputs): |
|
"""Run control-to-world generation demo. |
|
|
|
This function handles the main control-to-world generation pipeline, including: |
|
- Setting up the random seed for reproducibility |
|
- Initializing the generation pipeline with the provided configuration |
|
- Processing single or multiple prompts/images/videos from input |
|
- Generating videos from prompts and images/videos |
|
- Saving the generated videos and corresponding prompts to disk |
|
|
|
Args: |
|
cfg (argparse.Namespace): Configuration namespace containing: |
|
- Model configuration (checkpoint paths, model settings) |
|
- Generation parameters (guidance, steps, dimensions) |
|
- Input/output settings (prompts/images/videos, save paths) |
|
- Performance options (model offloading settings) |
|
|
|
The function will save: |
|
- Generated MP4 video files |
|
- Text files containing the processed prompts |
|
|
|
If guardrails block the generation, a critical log message is displayed |
|
and the function continues to the next prompt if available. |
|
""" |
|
|
|
control_inputs = validate_controlnet_specs(cfg, control_inputs) |
|
misc.set_random_seed(cfg.seed) |
|
|
|
device_rank = 0 |
|
process_group = None |
|
if cfg.num_gpus > 1: |
|
from megatron.core import parallel_state |
|
|
|
from cosmos_transfer1.utils import distributed |
|
|
|
distributed.init() |
|
parallel_state.initialize_model_parallel(context_parallel_size=cfg.num_gpus) |
|
process_group = parallel_state.get_context_parallel_group() |
|
|
|
device_rank = distributed.get_rank(process_group) |
|
|
|
preprocessors = Preprocessors() |
|
|
|
if cfg.use_distilled: |
|
assert not cfg.is_av_sample |
|
checkpoint = EDGE2WORLD_CONTROLNET_DISTILLED_CHECKPOINT_PATH |
|
pipeline = DistilledControl2WorldGenerationPipeline( |
|
checkpoint_dir=cfg.checkpoint_dir, |
|
checkpoint_name=checkpoint, |
|
offload_network=cfg.offload_diffusion_transformer, |
|
offload_text_encoder_model=cfg.offload_text_encoder_model, |
|
offload_guardrail_models=cfg.offload_guardrail_models, |
|
guidance=cfg.guidance, |
|
num_steps=cfg.num_steps, |
|
fps=cfg.fps, |
|
seed=cfg.seed, |
|
num_input_frames=cfg.num_input_frames, |
|
control_inputs=control_inputs, |
|
sigma_max=cfg.sigma_max, |
|
blur_strength=cfg.blur_strength, |
|
canny_threshold=cfg.canny_threshold, |
|
upsample_prompt=cfg.upsample_prompt, |
|
offload_prompt_upsampler=cfg.offload_prompt_upsampler, |
|
process_group=process_group, |
|
) |
|
else: |
|
checkpoint = BASE_7B_CHECKPOINT_AV_SAMPLE_PATH if cfg.is_av_sample else BASE_7B_CHECKPOINT_PATH |
|
|
|
|
|
pipeline = DiffusionControl2WorldGenerationPipeline( |
|
checkpoint_dir=cfg.checkpoint_dir, |
|
checkpoint_name=checkpoint, |
|
offload_network=cfg.offload_diffusion_transformer, |
|
offload_text_encoder_model=cfg.offload_text_encoder_model, |
|
offload_guardrail_models=cfg.offload_guardrail_models, |
|
guidance=cfg.guidance, |
|
num_steps=cfg.num_steps, |
|
fps=cfg.fps, |
|
seed=cfg.seed, |
|
num_input_frames=cfg.num_input_frames, |
|
control_inputs=control_inputs, |
|
sigma_max=cfg.sigma_max, |
|
blur_strength=cfg.blur_strength, |
|
canny_threshold=cfg.canny_threshold, |
|
upsample_prompt=cfg.upsample_prompt, |
|
offload_prompt_upsampler=cfg.offload_prompt_upsampler, |
|
process_group=process_group, |
|
) |
|
|
|
if cfg.batch_input_path: |
|
log.info(f"Reading batch inputs from path: {cfg.batch_input_path}") |
|
prompts = read_prompts_from_file(cfg.batch_input_path) |
|
else: |
|
|
|
prompts = [{"prompt": cfg.prompt, "visual_input": cfg.input_video_path}] |
|
|
|
batch_size = cfg.batch_size if hasattr(cfg, "batch_size") else 1 |
|
if any("upscale" in control_input for control_input in control_inputs) and batch_size > 1: |
|
batch_size = 1 |
|
log.info("Setting batch_size=1 as upscale does not support batch generation") |
|
os.makedirs(cfg.video_save_folder, exist_ok=True) |
|
for batch_start in range(0, len(prompts), batch_size): |
|
|
|
batch_prompts = prompts[batch_start : batch_start + batch_size] |
|
actual_batch_size = len(batch_prompts) |
|
|
|
batch_prompt_texts = [p.get("prompt", None) for p in batch_prompts] |
|
batch_video_paths = [p.get("visual_input", None) for p in batch_prompts] |
|
|
|
batch_control_inputs = [] |
|
for i, input_dict in enumerate(batch_prompts): |
|
current_prompt = input_dict.get("prompt", None) |
|
current_video_path = input_dict.get("visual_input", None) |
|
|
|
if cfg.batch_input_path: |
|
video_save_subfolder = os.path.join(cfg.video_save_folder, f"video_{batch_start+i}") |
|
os.makedirs(video_save_subfolder, exist_ok=True) |
|
else: |
|
video_save_subfolder = cfg.video_save_folder |
|
|
|
current_control_inputs = copy.deepcopy(control_inputs) |
|
if "control_overrides" in input_dict: |
|
for hint_key, override in input_dict["control_overrides"].items(): |
|
if hint_key in current_control_inputs: |
|
current_control_inputs[hint_key].update(override) |
|
else: |
|
log.warning(f"Ignoring unknown control key in override: {hint_key}") |
|
|
|
|
|
log.info("running preprocessor") |
|
preprocessors( |
|
current_video_path, |
|
current_prompt, |
|
current_control_inputs, |
|
video_save_subfolder, |
|
cfg.regional_prompts if hasattr(cfg, "regional_prompts") else None, |
|
) |
|
batch_control_inputs.append(current_control_inputs) |
|
|
|
regional_prompts = [] |
|
region_definitions = [] |
|
if hasattr(cfg, "regional_prompts") and cfg.regional_prompts: |
|
log.info(f"regional_prompts: {cfg.regional_prompts}") |
|
for regional_prompt in cfg.regional_prompts: |
|
regional_prompts.append(regional_prompt["prompt"]) |
|
if "region_definitions_path" in regional_prompt: |
|
log.info(f"region_definitions_path: {regional_prompt['region_definitions_path']}") |
|
region_definition_path = regional_prompt["region_definitions_path"] |
|
if isinstance(region_definition_path, str) and region_definition_path.endswith(".json"): |
|
with open(region_definition_path, "r") as f: |
|
region_definitions_json = json.load(f) |
|
region_definitions.extend(region_definitions_json) |
|
else: |
|
region_definitions.append(region_definition_path) |
|
|
|
if hasattr(pipeline, "regional_prompts"): |
|
pipeline.regional_prompts = regional_prompts |
|
if hasattr(pipeline, "region_definitions"): |
|
pipeline.region_definitions = region_definitions |
|
|
|
|
|
batch_outputs = pipeline.generate( |
|
prompt=batch_prompt_texts, |
|
video_path=batch_video_paths, |
|
negative_prompt=cfg.negative_prompt, |
|
control_inputs=batch_control_inputs, |
|
save_folder=video_save_subfolder, |
|
batch_size=actual_batch_size, |
|
) |
|
if batch_outputs is None: |
|
log.critical("Guardrail blocked generation for entire batch.") |
|
continue |
|
|
|
videos, final_prompts = batch_outputs |
|
for i, (video, prompt) in enumerate(zip(videos, final_prompts)): |
|
if cfg.batch_input_path: |
|
video_save_subfolder = os.path.join(cfg.video_save_folder, f"video_{batch_start+i}") |
|
video_save_path = os.path.join(video_save_subfolder, "output.mp4") |
|
prompt_save_path = os.path.join(video_save_subfolder, "prompt.txt") |
|
else: |
|
video_save_path = os.path.join(cfg.video_save_folder, f"{cfg.video_save_name}.mp4") |
|
prompt_save_path = os.path.join(cfg.video_save_folder, f"{cfg.video_save_name}.txt") |
|
|
|
if device_rank == 0: |
|
os.makedirs(os.path.dirname(video_save_path), exist_ok=True) |
|
save_video( |
|
video=video, |
|
fps=cfg.fps, |
|
H=video.shape[1], |
|
W=video.shape[2], |
|
video_save_quality=5, |
|
video_save_path=video_save_path, |
|
) |
|
|
|
|
|
with open(prompt_save_path, "wb") as f: |
|
f.write(prompt.encode("utf-8")) |
|
|
|
log.info(f"Saved video to {video_save_path}") |
|
log.info(f"Saved prompt to {prompt_save_path}") |
|
|
|
|
|
if cfg.num_gpus > 1: |
|
parallel_state.destroy_model_parallel() |
|
import torch.distributed as dist |
|
|
|
dist.destroy_process_group() |
|
|
|
|
|
if __name__ == "__main__": |
|
args, control_inputs = parse_arguments() |
|
demo(args, control_inputs) |
|
|