scripts/sample_text2video.py

import os
import time
import argparse
import yaml, math
from tqdm import trange
import torch
import numpy as np
from omegaconf import OmegaConf
import torch.distributed as dist
from pytorch_lightning import seed_everything

from lvdm.samplers.ddim import DDIMSampler
from lvdm.utils.common_utils import str2bool
from lvdm.utils.dist_utils import setup_dist, gather_data
from lvdm.utils.saving_utils import npz_to_video_grid, npz_to_imgsheet_5d
from scripts.sample_utils import load_model, get_conditions, make_model_input_shape, torch_to_np


# ------------------------------------------------------------------------------------------
def get_parser():
    parser = argparse.ArgumentParser()
    # basic args
    parser.add_argument("--ckpt_path", type=str, help="model checkpoint path")
    parser.add_argument("--config_path", type=str, help="model config path (a yaml file)")
    parser.add_argument("--prompt", type=str, help="input text prompts for text2video (a sentence OR a txt file).")
    parser.add_argument("--save_dir", type=str, help="results saving dir", default="results/")
    # device args
    parser.add_argument("--ddp", action='store_true', help="whether use pytorch ddp mode for parallel sampling (recommend for multi-gpu case)", default=False)
    parser.add_argument("--local_rank", type=int, help="is used for pytorch ddp mode", default=0)
    parser.add_argument("--gpu_id", type=int, help="choose a specific gpu", default=0)
    # sampling args
    parser.add_argument("--n_samples", type=int, help="how many samples for each text prompt", default=2)
    parser.add_argument("--batch_size", type=int, help="video batch size for sampling", default=1)
    parser.add_argument("--decode_frame_bs", type=int, help="frame batch size for framewise decoding", default=1)
    parser.add_argument("--sample_type", type=str, help="ddpm or ddim", default="ddim", choices=["ddpm", "ddim"])
    parser.add_argument("--ddim_steps", type=int, help="ddim sampling -- number of ddim denoising timesteps", default=50)
    parser.add_argument("--eta", type=float, help="ddim sampling -- eta (0.0 yields deterministic sampling, 1.0 yields random sampling)", default=1.0)
    parser.add_argument("--cfg_scale", type=float, default=15.0, help="classifier-free guidance scale")
    parser.add_argument("--seed", type=int, default=None, help="fix a seed for randomness (If you want to reproduce the sample results)")
    parser.add_argument("--show_denoising_progress", action='store_true', default=False, help="whether show denoising progress during sampling one batch",)
    # lora args
    parser.add_argument("--lora_path", type=str, help="lora checkpoint path")
    parser.add_argument("--inject_lora", action='store_true', default=False, help="",)
    parser.add_argument("--lora_scale", type=float, default=None, help="scale for lora weight")
    parser.add_argument("--lora_trigger_word", type=str, default="", help="",)
    # saving args
    parser.add_argument("--save_mp4", type=str2bool, default=True, help="whether save samples in separate mp4 files", choices=["True", "true", "False", "false"])
    parser.add_argument("--save_mp4_sheet", action='store_true', default=False, help="whether save samples in mp4 file",)
    parser.add_argument("--save_npz", action='store_true', default=False, help="whether save samples in npz file",)
    parser.add_argument("--save_jpg", action='store_true', default=False, help="whether save samples in jpg file",)
    parser.add_argument("--save_fps", type=int, default=8, help="fps of saved mp4 videos",)
    return parser

# ------------------------------------------------------------------------------------------
def sample_denoising_batch(model, noise_shape, condition, *args,
                           sample_type="ddim", sampler=None, 
                           ddim_steps=None, eta=None,
                           unconditional_guidance_scale=1.0, uc=None,
                           denoising_progress=False,
                           **kwargs,
                           ):
    
    if sample_type == "ddpm":
        samples = model.p_sample_loop(cond=condition, shape=noise_shape,
                                      return_intermediates=False, 
                                      verbose=denoising_progress,
                                      **kwargs,
                                      )
    elif sample_type == "ddim":
        assert(sampler is not None)
        assert(ddim_steps is not None)
        assert(eta is not None)
        ddim_sampler = sampler
        samples, _ = ddim_sampler.sample(S=ddim_steps,
                                         conditioning=condition,
                                         batch_size=noise_shape[0],
                                         shape=noise_shape[1:],
                                         verbose=denoising_progress,
                                         unconditional_guidance_scale=unconditional_guidance_scale,
                                         unconditional_conditioning=uc,
                                         eta=eta,
                                         **kwargs,
                                        )
    else:
        raise ValueError
    return samples


# ------------------------------------------------------------------------------------------
@torch.no_grad()
def sample_text2video(model, prompt, n_samples, batch_size,
                      sample_type="ddim", sampler=None, 
                      ddim_steps=50, eta=1.0, cfg_scale=7.5, 
                      decode_frame_bs=1,
                      ddp=False, all_gather=True, 
                      batch_progress=True, show_denoising_progress=False,
                      ):
    # get cond vector
    assert(model.cond_stage_model is not None)
    cond_embd = get_conditions(prompt, model, batch_size)
    uncond_embd = get_conditions("", model, batch_size) if cfg_scale != 1.0 else None

    # sample batches
    all_videos = []
    n_iter = math.ceil(n_samples / batch_size)
    iterator  = trange(n_iter, desc="Sampling Batches (text-to-video)") if batch_progress else range(n_iter)
    for _ in iterator:
        noise_shape = make_model_input_shape(model, batch_size)
        samples_latent = sample_denoising_batch(model, noise_shape, cond_embd,
                                            sample_type=sample_type,
                                            sampler=sampler,
                                            ddim_steps=ddim_steps,
                                            eta=eta,
                                            unconditional_guidance_scale=cfg_scale, 
                                            uc=uncond_embd,
                                            denoising_progress=show_denoising_progress,
                                            )
        samples = model.decode_first_stage(samples_latent, decode_bs=decode_frame_bs, return_cpu=False)
        
        # gather samples from multiple gpus
        if ddp and all_gather:
            data_list = gather_data(samples, return_np=False)
            all_videos.extend([torch_to_np(data) for data in data_list])
        else:
            all_videos.append(torch_to_np(samples))
    
    all_videos = np.concatenate(all_videos, axis=0)
    assert(all_videos.shape[0] >= n_samples)
    return all_videos


# ------------------------------------------------------------------------------------------
def save_results(videos, save_dir, 
                 save_name="results", save_fps=8, save_mp4=True, 
                 save_npz=False, save_mp4_sheet=False, save_jpg=False
                 ):
    if save_mp4:
        save_subdir = os.path.join(save_dir, "videos")
        os.makedirs(save_subdir, exist_ok=True)
        for i in range(videos.shape[0]):
            npz_to_video_grid(videos[i:i+1,...], 
                              os.path.join(save_subdir, f"{save_name}_{i:03d}.mp4"), 
                              fps=save_fps)
        print(f'Successfully saved videos in {save_subdir}')
    
    if save_npz:
        save_path = os.path.join(save_dir, f"{save_name}.npz")
        np.savez(save_path, videos)
        print(f'Successfully saved npz in {save_path}')
    
    if save_mp4_sheet:
        save_path = os.path.join(save_dir, f"{save_name}.mp4")
        npz_to_video_grid(videos, save_path, fps=save_fps)
        print(f'Successfully saved mp4 sheet in {save_path}')

    if save_jpg:
        save_path = os.path.join(save_dir, f"{save_name}.jpg")
        npz_to_imgsheet_5d(videos, save_path, nrow=videos.shape[1])
        print(f'Successfully saved jpg sheet in {save_path}')


# ------------------------------------------------------------------------------------------
def main():
    """
    text-to-video generation
    """
    parser = get_parser()
    opt, unknown = parser.parse_known_args()
    os.makedirs(opt.save_dir, exist_ok=True)
    
    # set device
    if opt.ddp:
        setup_dist(opt.local_rank)
        opt.n_samples = math.ceil(opt.n_samples / dist.get_world_size())
        gpu_id = None
    else:
        gpu_id = opt.gpu_id
        os.environ["CUDA_VISIBLE_DEVICES"] = f"{gpu_id}"
    
    # set random seed
    if opt.seed is not None:
        if opt.ddp:
            seed = opt.local_rank + opt.seed
        else:
            seed = opt.seed
        seed_everything(seed)

    # dump args
    fpath = os.path.join(opt.save_dir, "sampling_args.yaml")
    with open(fpath, 'w') as f:
        yaml.dump(vars(opt), f, default_flow_style=False)

    # load & merge config
    config = OmegaConf.load(opt.config_path)
    cli = OmegaConf.from_dotlist(unknown)
    config = OmegaConf.merge(config, cli)
    print("config: \n", config)

    # get model & sampler
    model, _, _ = load_model(config, opt.ckpt_path, 
                             inject_lora=opt.inject_lora, 
                             lora_scale=opt.lora_scale, 
                             lora_path=opt.lora_path
                             )
    ddim_sampler = DDIMSampler(model) if opt.sample_type == "ddim" else None

    # prepare prompt
    if opt.prompt.endswith(".txt"):
        opt.prompt_file = opt.prompt
        opt.prompt = None
    else:
        opt.prompt_file = None

    if opt.prompt_file is not None:
        f = open(opt.prompt_file, 'r')
        prompts, line_idx = [], []
        for idx, line in enumerate(f.readlines()):
            l = line.strip()
            if len(l) != 0:
                prompts.append(l)
                line_idx.append(idx)
        f.close()
        cmd = f"cp {opt.prompt_file} {opt.save_dir}"
        os.system(cmd)
    else:
        prompts = [opt.prompt]
        line_idx = [None]

    if opt.inject_lora:
        assert(opt.lora_trigger_word != '')
        prompts = [p + opt.lora_trigger_word for p in prompts]
    
    # go
    start = time.time()  
    for prompt in prompts:
        # sample
        samples = sample_text2video(model, prompt, opt.n_samples, opt.batch_size,
                          sample_type=opt.sample_type, sampler=ddim_sampler,
                          ddim_steps=opt.ddim_steps, eta=opt.eta, 
                          cfg_scale=opt.cfg_scale,
                          decode_frame_bs=opt.decode_frame_bs,
                          ddp=opt.ddp, show_denoising_progress=opt.show_denoising_progress,
                          )
        # save
        if (opt.ddp and dist.get_rank() == 0) or (not opt.ddp):
            prompt_str = prompt.replace("/", "_slash_") if "/" in prompt else prompt
            save_name = prompt_str.replace(" ", "_") if " " in prompt else prompt_str
            if opt.seed is not None:
                save_name = save_name + f"_seed{seed:05d}"
            save_results(samples, opt.save_dir, save_name=save_name, save_fps=opt.save_fps)
    print("Finish sampling!")
    print(f"Run time = {(time.time() - start):.2f} seconds")

    if opt.ddp:
        dist.destroy_process_group()


# ------------------------------------------------------------------------------------------
if __name__ == "__main__":
    main()