ALSv/self-forcing

1import argparse2import torch3import os4from omegaconf import OmegaConf5from tqdm import tqdm6from torchvision import transforms7from torchvision.io import write_video8from einops import rearrange9import torch.distributed as dist10from torch.utils.data import DataLoader, SequentialSampler11from torch.utils.data.distributed import DistributedSampler12 13from pipeline import (14    CausalDiffusionInferencePipeline,15    CausalInferencePipeline16)17from utils.dataset import TextDataset, TextImagePairDataset18from utils.misc import set_seed19 20parser = argparse.ArgumentParser()21parser.add_argument("--config_path", type=str, help="Path to the config file")22parser.add_argument("--checkpoint_path", type=str, help="Path to the checkpoint folder")23parser.add_argument("--data_path", type=str, help="Path to the dataset")24parser.add_argument("--extended_prompt_path", type=str, help="Path to the extended prompt")25parser.add_argument("--output_folder", type=str, help="Output folder")26parser.add_argument("--num_output_frames", type=int, default=21,27                    help="Number of overlap frames between sliding windows")28parser.add_argument("--i2v", action="store_true", help="Whether to perform I2V (or T2V by default)")29parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA parameters")30parser.add_argument("--seed", type=int, default=0, help="Random seed")31parser.add_argument("--num_samples", type=int, default=1, help="Number of samples to generate per prompt")32parser.add_argument("--save_with_index", action="store_true",33                    help="Whether to save the video using the index or prompt as the filename")34args = parser.parse_args()35 36# Initialize distributed inference37if "LOCAL_RANK" in os.environ:38    dist.init_process_group(backend='nccl')39    local_rank = int(os.environ["LOCAL_RANK"])40    torch.cuda.set_device(local_rank)41    device = torch.device(f"cuda:{local_rank}")42    world_size = dist.get_world_size()43    set_seed(args.seed + local_rank)44else:45    device = torch.device("cuda")46    local_rank = 047    world_size = 148    set_seed(args.seed)49 50torch.set_grad_enabled(False)51 52config = OmegaConf.load(args.config_path)53default_config = OmegaConf.load("configs/default_config.yaml")54config = OmegaConf.merge(default_config, config)55 56# Initialize pipeline57if hasattr(config, 'denoising_step_list'):58    # Few-step inference59    pipeline = CausalInferencePipeline(config, device=device)60else:61    # Multi-step diffusion inference62    pipeline = CausalDiffusionInferencePipeline(config, device=device)63 64if args.checkpoint_path:65    state_dict = torch.load(args.checkpoint_path, map_location="cpu")66    pipeline.generator.load_state_dict(state_dict['generator' if not args.use_ema else 'generator_ema'])67 68pipeline = pipeline.to(device=device, dtype=torch.bfloat16)69 70# Create dataset71if args.i2v:72    assert not dist.is_initialized(), "I2V does not support distributed inference yet"73    transform = transforms.Compose([74        transforms.Resize((480, 832)),75        transforms.ToTensor(),76        transforms.Normalize([0.5], [0.5])77    ])78    dataset = TextImagePairDataset(args.data_path, transform=transform)79else:80    dataset = TextDataset(prompt_path=args.data_path, extended_prompt_path=args.extended_prompt_path)81num_prompts = len(dataset)82print(f"Number of prompts: {num_prompts}")83 84if dist.is_initialized():85    sampler = DistributedSampler(dataset, shuffle=False, drop_last=True)86else:87    sampler = SequentialSampler(dataset)88dataloader = DataLoader(dataset, batch_size=1, sampler=sampler, num_workers=0, drop_last=False)89 90# Create output directory (only on main process to avoid race conditions)91if local_rank == 0:92    os.makedirs(args.output_folder, exist_ok=True)93 94if dist.is_initialized():95    dist.barrier()96 97 98def encode(self, videos: torch.Tensor) -> torch.Tensor:99    device, dtype = videos[0].device, videos[0].dtype100    scale = [self.mean.to(device=device, dtype=dtype),101             1.0 / self.std.to(device=device, dtype=dtype)]102    output = [103        self.model.encode(u.unsqueeze(0), scale).float().squeeze(0)104        for u in videos105    ]106 107    output = torch.stack(output, dim=0)108    return output109 110 111for i, batch_data in tqdm(enumerate(dataloader), disable=(local_rank != 0)):112    idx = batch_data['idx'].item()113 114    # For DataLoader batch_size=1, the batch_data is already a single item, but in a batch container115    # Unpack the batch data for convenience116    if isinstance(batch_data, dict):117        batch = batch_data118    elif isinstance(batch_data, list):119        batch = batch_data[0]  # First (and only) item in the batch120 121    all_video = []122    num_generated_frames = 0  # Number of generated (latent) frames123 124    if args.i2v:125        # For image-to-video, batch contains image and caption126        prompt = batch['prompts'][0]  # Get caption from batch127        prompts = [prompt] * args.num_samples128 129        # Process the image130        image = batch['image'].squeeze(0).unsqueeze(0).unsqueeze(2).to(device=device, dtype=torch.bfloat16)131 132        # Encode the input image as the first latent133        initial_latent = pipeline.vae.encode_to_latent(image).to(device=device, dtype=torch.bfloat16)134        initial_latent = initial_latent.repeat(args.num_samples, 1, 1, 1, 1)135 136        sampled_noise = torch.randn(137            [args.num_samples, args.num_output_frames - 1, 16, 60, 104], device=device, dtype=torch.bfloat16138        )139    else:140        # For text-to-video, batch is just the text prompt141        prompt = batch['prompts'][0]142        extended_prompt = batch['extended_prompts'][0] if 'extended_prompts' in batch else None143        if extended_prompt is not None:144            prompts = [extended_prompt] * args.num_samples145        else:146            prompts = [prompt] * args.num_samples147        initial_latent = None148 149        sampled_noise = torch.randn(150            [args.num_samples, args.num_output_frames, 16, 60, 104], device=device, dtype=torch.bfloat16151        )152 153    # Generate 81 frames154    video, latents = pipeline.inference(155        noise=sampled_noise,156        text_prompts=prompts,157        return_latents=True,158        initial_latent=initial_latent,159    )160    current_video = rearrange(video, 'b t c h w -> b t h w c').cpu()161    all_video.append(current_video)162    num_generated_frames += latents.shape[1]163 164    # Final output video165    video = 255.0 * torch.cat(all_video, dim=1)166 167    # Clear VAE cache168    pipeline.vae.model.clear_cache()169 170    # Save the video if the current prompt is not a dummy prompt171    if idx < num_prompts:172        model = "regular" if not args.use_ema else "ema"173        for seed_idx in range(args.num_samples):174            # All processes save their videos175            if args.save_with_index:176                output_path = os.path.join(args.output_folder, f'{idx}-{seed_idx}_{model}.mp4')177            else:178                output_path = os.path.join(args.output_folder, f'{prompt[:100]}-{seed_idx}.mp4')179            write_video(output_path, video[seed_idx], fps=16)180