import argparse import os from pathlib import Path import imageio import numpy as np import torch import torch.distributed as dist import torchvision from einops import rearrange from fastvideo.models.hunyuan.inference import HunyuanVideoSampler from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, nccl_info def initialize_distributed(): local_rank = int(os.getenv("RANK", 0)) world_size = int(os.getenv("WORLD_SIZE", 1)) print("world_size", world_size) torch.cuda.set_device(local_rank) dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank) initialize_sequence_parallel_state(world_size) def main(args): initialize_distributed() print(nccl_info.sp_size) print(args) models_root_path = Path(args.model_path) if not models_root_path.exists(): raise ValueError(f"`models_root` not exists: {models_root_path}") # Create save folder to save the samples save_path = args.output_path os.makedirs(os.path.dirname(save_path), exist_ok=True) # Load models hunyuan_video_sampler = HunyuanVideoSampler.from_pretrained(models_root_path, args=args) # Get the updated args args = hunyuan_video_sampler.args if args.prompt.endswith('.txt'): with open(args.prompt) as f: prompts = [line.strip() for line in f.readlines()] else: prompts = [args.prompt] for prompt in prompts: outputs = hunyuan_video_sampler.predict( prompt=prompt, height=args.height, width=args.width, video_length=args.num_frames, seed=args.seed, negative_prompt=args.neg_prompt, infer_steps=args.num_inference_steps, guidance_scale=args.guidance_scale, num_videos_per_prompt=args.num_videos, flow_shift=args.flow_shift, batch_size=args.batch_size, embedded_guidance_scale=args.embedded_cfg_scale, ) videos = rearrange(outputs["samples"], "b c t h w -> t b c h w") outputs = [] for x in videos: x = torchvision.utils.make_grid(x, nrow=6) x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) outputs.append((x * 255).numpy().astype(np.uint8)) os.makedirs(os.path.dirname(args.output_path), exist_ok=True) imageio.mimsave(os.path.join(args.output_path, f"{prompt[:100]}.mp4"), outputs, fps=args.fps) if __name__ == "__main__": parser = argparse.ArgumentParser() # Basic parameters parser.add_argument("--prompt", type=str, help="prompt file for inference") parser.add_argument("--num_frames", type=int, default=16) parser.add_argument("--height", type=int, default=256) parser.add_argument("--width", type=int, default=256) parser.add_argument("--num_inference_steps", type=int, default=50) parser.add_argument("--model_path", type=str, default="data/hunyuan") parser.add_argument("--output_path", type=str, default="./outputs/video") parser.add_argument("--fps", type=int, default=24) # Additional parameters parser.add_argument( "--denoise-type", type=str, default="flow", help="Denoise type for noised inputs.", ) parser.add_argument("--seed", type=int, default=None, help="Seed for evaluation.") parser.add_argument("--neg_prompt", type=str, default=None, help="Negative prompt for sampling.") parser.add_argument( "--guidance_scale", type=float, default=1.0, help="Classifier free guidance scale.", ) parser.add_argument( "--embedded_cfg_scale", type=float, default=6.0, help="Embedded classifier free guidance scale.", ) parser.add_argument("--flow_shift", type=int, default=7, help="Flow shift parameter.") parser.add_argument("--batch_size", type=int, default=1, help="Batch size for inference.") parser.add_argument( "--num_videos", type=int, default=1, help="Number of videos to generate per prompt.", ) parser.add_argument( "--load-key", type=str, default="module", help="Key to load the model states. 'module' for the main model, 'ema' for the EMA model.", ) parser.add_argument( "--use-cpu-offload", action="store_true", help="Use CPU offload for the model load.", ) parser.add_argument( "--dit-weight", type=str, default="data/hunyuan/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt", ) parser.add_argument( "--reproduce", action="store_true", help="Enable reproducibility by setting random seeds and deterministic algorithms.", ) parser.add_argument( "--disable-autocast", action="store_true", help="Disable autocast for denoising loop and vae decoding in pipeline sampling.", ) # Flow Matching parser.add_argument( "--flow-reverse", action="store_true", help="If reverse, learning/sampling from t=1 -> t=0.", ) parser.add_argument("--flow-solver", type=str, default="euler", help="Solver for flow matching.") parser.add_argument( "--use-linear-quadratic-schedule", action="store_true", help= "Use linear quadratic schedule for flow matching. Following MovieGen (https://ai.meta.com/static-resource/movie-gen-research-paper)", ) parser.add_argument( "--linear-schedule-end", type=int, default=25, help="End step for linear quadratic schedule for flow matching.", ) # Model parameters parser.add_argument("--model", type=str, default="HYVideo-T/2-cfgdistill") parser.add_argument("--latent-channels", type=int, default=16) parser.add_argument("--precision", type=str, default="bf16", choices=["fp32", "fp16", "bf16"]) parser.add_argument("--rope-theta", type=int, default=256, help="Theta used in RoPE.") parser.add_argument("--vae", type=str, default="884-16c-hy") parser.add_argument("--vae-precision", type=str, default="fp16", choices=["fp32", "fp16", "bf16"]) parser.add_argument("--vae-tiling", action="store_true", default=True) parser.add_argument("--vae-sp", action="store_true", default=False) parser.add_argument("--text-encoder", type=str, default="llm") parser.add_argument( "--text-encoder-precision", type=str, default="fp16", choices=["fp32", "fp16", "bf16"], ) parser.add_argument("--text-states-dim", type=int, default=4096) parser.add_argument("--text-len", type=int, default=256) parser.add_argument("--tokenizer", type=str, default="llm") parser.add_argument("--prompt-template", type=str, default="dit-llm-encode") parser.add_argument("--prompt-template-video", type=str, default="dit-llm-encode-video") parser.add_argument("--hidden-state-skip-layer", type=int, default=2) parser.add_argument("--apply-final-norm", action="store_true") parser.add_argument("--text-encoder-2", type=str, default="clipL") parser.add_argument( "--text-encoder-precision-2", type=str, default="fp16", choices=["fp32", "fp16", "bf16"], ) parser.add_argument( "--enable_torch_compile", action="store_true", help="Use torch.compile for speeding up STA inference without teacache", ) parser.add_argument("--text-states-dim-2", type=int, default=768) parser.add_argument("--tokenizer-2", type=str, default="clipL") parser.add_argument("--text-len-2", type=int, default=77) args = parser.parse_args() # process for vae sequence parallel if args.vae_sp and not args.vae_tiling: raise ValueError("Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True.") main(args)