()
| 21 | print(f"Saved {len(video)} images for camera {cam_name} to {images_dir}") |
| 22 | |
| 23 | def main(): |
| 24 | parser = argparse.ArgumentParser(description="Generate video from image sequence.") |
| 25 | parser.add_argument('--input_folder', type=str, required=True, help='Path to folder containing images') |
| 26 | parser.add_argument('--output_folder', type=str, required=True, help='Path to output video folder') |
| 27 | parser.add_argument('--model_path', type=str, default='./checkpoints/4DSloMo_LoRA.ckpt', help='Path to the LoRA model checkpoint') |
| 28 | parser.add_argument('--num_inference_steps', type=int, default=50, help='Number of inference steps') |
| 29 | args = parser.parse_args() |
| 30 | |
| 31 | dist.init_process_group(backend="nccl") |
| 32 | local_rank = int(os.environ["LOCAL_RANK"]) |
| 33 | rank = dist.get_rank() |
| 34 | world_size = dist.get_world_size() |
| 35 | torch.cuda.set_device(local_rank) |
| 36 | device = f"cuda:{local_rank}" |
| 37 | |
| 38 | if rank == 0: |
| 39 | print(f"Running distributed inference on {world_size} GPUs.") |
| 40 | |
| 41 | image_folder = args.input_folder |
| 42 | output_folder = args.output_folder |
| 43 | |
| 44 | model_manager = ModelManager(torch_dtype=torch.bfloat16, device="cpu") |
| 45 | model_manager.load_models( |
| 46 | ["checkpoints/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth"], |
| 47 | torch_dtype=torch.float32, |
| 48 | ) |
| 49 | |
| 50 | model_manager.load_models( |
| 51 | [ |
| 52 | [ |
| 53 | "checkpoints/diffusion_pytorch_model-00001-of-00007.safetensors", |
| 54 | "checkpoints/diffusion_pytorch_model-00002-of-00007.safetensors", |
| 55 | "checkpoints/diffusion_pytorch_model-00003-of-00007.safetensors", |
| 56 | "checkpoints/diffusion_pytorch_model-00004-of-00007.safetensors", |
| 57 | "checkpoints/diffusion_pytorch_model-00005-of-00007.safetensors", |
| 58 | "checkpoints/diffusion_pytorch_model-00006-of-00007.safetensors", |
| 59 | "checkpoints/diffusion_pytorch_model-00007-of-00007.safetensors", |
| 60 | ], |
| 61 | "checkpoints/models_t5_umt5-xxl-enc-bf16.pth", |
| 62 | "checkpoints/Wan2.1_VAE.pth", |
| 63 | ], |
| 64 | torch_dtype=torch.bfloat16, |
| 65 | ) |
| 66 | model_manager.load_lora(args.model_path, lora_alpha=1.0) |
| 67 | pipe = FixPipeline.from_model_manager(model_manager, device=device) |
| 68 | pipe.enable_vram_management(num_persistent_param_in_dit=None) |
| 69 | |
| 70 | cam_list = ["19305323","19305319","19305336","19305328","19305326","19305340","19305309","19305329","19224108","19305334","19305337","19305314"] |
| 71 | |
| 72 | cameras_for_this_rank = cam_list[rank::world_size] |
| 73 | print(f"Rank {rank} is assigned {len(cameras_for_this_rank)} cameras.") |
| 74 | |
| 75 | for cam_name in cameras_for_this_rank: |
| 76 | print(f"------ Rank {rank} processing {cam_name} ------") |
| 77 | image_path = f"{image_folder}/test/ours_None/gt_crop/{cam_name}_0000.png" |
| 78 | if not os.path.exists(image_path): |
| 79 | raise FileNotFoundError(f"Image not found for camera {cam_name}: {image_path}") |
| 80 | try: |
no test coverage detected