ShotStream/Inference_Causal.py at main · KlingAIResearch/ShotStream · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
import argparse
import torch
import os
from omegaconf import OmegaConf
from tqdm import tqdm
from torchvision import transforms
from torchvision.io import write_video
from einops import rearrange
import torch.distributed as dist
from torch.utils.data import DataLoader, SequentialSampler
from torch.utils.data.distributed import DistributedSampler
from utils.dataset import MultiShots_FrameConcat_Dataset
import sys
from pipeline import (
    CausalInferenceArPipeline,
)
from utils.misc import set_seed

import logging

parser = argparse.ArgumentParser()
parser.add_argument("--config_path", type=str, default=None)
parser.add_argument("--resume_ckpt", type=str, default=None)
parser.add_argument("--resume_lora_ckpt", type=str, default=None)
parser.add_argument("--data_path", type=str, default=None)
parser.add_argument("--output_folder", type=str, default=None)
parser.add_argument("--seed", type=int, default=0, help="Random seed")
parser.add_argument("--multi_caption", type=bool, default=True)
parser.add_argument("--use_wo_rope_cache", type=bool, default=False)

args = parser.parse_args()

# Config setting
def _init_logging(rank):
    if rank == 0:
        logging.basicConfig(
            level=logging.INFO,
            format="[%(asctime)s] %(levelname)s: %(message)s",
            handlers=[logging.StreamHandler(stream=sys.stdout)])
    else:
        logging.basicConfig(level=logging.ERROR)

# Initialize distributed inference
rank = int(os.getenv("RANK", 0))
world_size = int(os.getenv("WORLD_SIZE", 1))
local_rank = int(os.getenv("LOCAL_RANK", 0))
device = local_rank
print(f"Device is {device}")
_init_logging(rank)

if world_size > 1:
    torch.cuda.set_device(local_rank)
    dist.init_process_group(
        backend="nccl",
        init_method="env://",
        rank=rank,
        world_size=world_size)

logging.info(f"Generation job args: {args}")

if dist.is_initialized():
    base_seed = [args.seed] if rank == 0 else [None]
    dist.broadcast_object_list(base_seed, src=0)
    args.seed = base_seed[0]

current_seed = args.seed + rank
set_seed(current_seed)
logging.info(f"Rank {rank} set seed to {current_seed}")

torch.set_grad_enabled(False)

config = OmegaConf.load(args.config_path)
default_config = OmegaConf.load("ckpts/default_config.yaml")
config = OmegaConf.merge(default_config, config)

config.multi_caption = args.multi_caption
config.use_wo_rope_cache = args.use_wo_rope_cache

# Initialize pipeline
if hasattr(config, 'denoising_step_list'):
    pipeline = CausalInferenceArPipeline(config, device=device)

# Load Ckpt
if args.resume_ckpt is not None:
    config.resume_ckpt = args.resume_ckpt

if config.resume_ckpt:
    state_dict = torch.load(config.resume_ckpt, map_location="cpu")
    print(f"resume generator's ckpt from {config.resume_ckpt}")
    pipeline.generator.load_state_dict(state_dict['generator'])

pipeline = pipeline.to(dtype=torch.bfloat16)

# Dataset
if args.data_path is not None:
    config.data_path = args.data_path
print(f"Dataset Path is {config.data_path}")

dataset = MultiShots_FrameConcat_Dataset(csv_path=config.data_path)
num_prompts = len(dataset)
print(f"Number of prompts: {num_prompts}")

if dist.is_initialized():
    sampler = DistributedSampler(dataset, shuffle=False, drop_last=True)
else:
    sampler = SequentialSampler(dataset)
dataloader = DataLoader(dataset, batch_size=1, sampler=sampler, num_workers=0, drop_last=False)

# Output Dir
if local_rank == 0:
    os.makedirs(args.output_folder, exist_ok=True)

if dist.is_initialized():
    dist.barrier()


for i, batch_data in tqdm(enumerate(dataloader), disable=(local_rank != 0)):
    idx = batch_data['idx'].item()

    if isinstance(batch_data, dict):
        batch = batch_data
    elif isinstance(batch_data, list):
        batch = batch_data[0]  # First (and only) item in the batch

    all_video = []
    num_generated_frames = 0  # Number of generated (latent) frames

    video = pipeline.inference(
        batch=batch,
        use_wo_rope_cache=config.use_wo_rope_cache,
    )
    current_video = rearrange(video, 'b t c h w -> b t h w c').cpu()
    all_video.append(current_video)

    video = 255.0 * torch.cat(all_video, dim=1)

    pipeline.vae.model.clear_cache()

    caption=batch['shots_captions'][0][-1][0][0]
    output_path = f"{args.output_folder}/{i:03d}_{caption[:50]}.mp4"
    write_video(output_path, video[0], fps=16)