Spaces:
Running on Zero
Running on Zero
File size: 9,547 Bytes
3e936b2 | 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 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 | import sys
from pathlib import Path
_PROJECT_ROOT = str(Path(__file__).resolve().parents[1])
if _PROJECT_ROOT in sys.path:
sys.path.remove(_PROJECT_ROOT)
sys.path.insert(0, _PROJECT_ROOT)
import argparse
import os
from typing import List
import torch
import torch.distributed as dist
from omegaconf import OmegaConf
from tqdm import tqdm
from torch.utils.data import DataLoader, SequentialSampler
from torch.utils.data.distributed import DistributedSampler
from torchvision.io import write_video
from torchvision import transforms
from einops import rearrange
from utils.misc import set_seed
from utils.distributed import barrier
from utils.memory import get_cuda_free_memory_gb, DynamicSwapInstaller
from pipeline.interactive_causal_inference import InteractiveCausalInferencePipeline
from utils.dataset import MultiTextDataset
import json
import glob
parser = argparse.ArgumentParser('Interactive causal inference')
parser.add_argument('--config_path', type=str, help='Path to the config file')
parser.add_argument('--checkpoint_path', type=str, default=None, help='Override config generator_ckpt')
parser.add_argument('--lora_ckpt', type=str, default=None, help='Override config lora_ckpt')
parser.add_argument('--data_path', type=str, default=None, help='Override config data_path')
parser.add_argument('--output_folder', type=str, default=None, help='Override config output_folder')
parser.add_argument('--use_ema', action='store_true', help='Override config use_ema')
parser.add_argument('--seed', type=int, default=None, help='Override config seed')
args = parser.parse_args()
config = OmegaConf.load(args.config_path)
if args.checkpoint_path:
config.generator_ckpt = args.checkpoint_path
if args.lora_ckpt:
config.lora_ckpt = args.lora_ckpt
if args.data_path:
config.data_path = args.data_path
if args.output_folder:
config.output_folder = args.output_folder
if args.use_ema:
config.use_ema = True
if args.seed is not None:
config.seed = args.seed
if 'LOCAL_RANK' in os.environ:
os.environ['NCCL_CROSS_NIC'] = '1'
os.environ['NCCL_DEBUG'] = os.environ.get('NCCL_DEBUG', 'INFO')
os.environ['NCCL_TIMEOUT'] = os.environ.get('NCCL_TIMEOUT', '1800')
local_rank = int(os.environ['LOCAL_RANK'])
world_size = int(os.environ.get('WORLD_SIZE', '1'))
rank = int(os.environ.get('RANK', str(local_rank)))
torch.cuda.set_device(local_rank)
device = torch.device(f'cuda:{local_rank}')
if not dist.is_initialized():
dist.init_process_group(backend='nccl', rank=rank, world_size=world_size, timeout=torch.distributed.constants.default_pg_timeout)
set_seed(config.seed + local_rank)
print(f'[Rank {rank}] Initialized distributed processing on device {device}')
else:
local_rank = 0
rank = 0
device = torch.device('cuda')
set_seed(config.seed)
print(f'Single GPU mode on device {device}')
low_memory = get_cuda_free_memory_gb(device) < 40
torch.set_grad_enabled(False)
pipeline = InteractiveCausalInferencePipeline(config, device=device)
if config.generator_ckpt:
state_dict = torch.load(config.generator_ckpt, map_location='cpu')
raw_gen_state_dict = state_dict['generator_ema' if config.use_ema else 'generator']
if config.use_ema:
def _clean_key(name: str) -> str:
return name.replace('_fsdp_wrapped_module.', '')
cleaned_state_dict = {_clean_key(k): v for k, v in raw_gen_state_dict.items()}
missing, unexpected = pipeline.generator.load_state_dict(cleaned_state_dict, strict=False)
if local_rank == 0:
if missing:
print(f'[Warning] {len(missing)} parameters missing: {missing[:8]} ...')
if unexpected:
print(f'[Warning] {len(unexpected)} unexpected params: {unexpected[:8]} ...')
else:
missing, unexpected = pipeline.generator.load_state_dict(raw_gen_state_dict, strict=False)
if local_rank == 0:
if missing:
print(f'[Warning] {len(missing)} parameters missing: {missing[:8]} ...')
if unexpected:
print(f'[Warning] {len(unexpected)} unexpected params: {unexpected[:8]} ...')
pipeline.is_lora_enabled = False
if getattr(config, 'adapter', None):
from utils.lora_utils import configure_lora_for_model
import peft
if local_rank == 0:
print(f'LoRA enabled with config: {config.adapter}')
print('Applying LoRA to generator (inference)...')
pipeline.generator.model = configure_lora_for_model(pipeline.generator.model, model_name='generator', lora_config=config.adapter, is_main_process=local_rank == 0)
lora_ckpt_path = getattr(config, 'lora_ckpt', None)
if lora_ckpt_path:
if local_rank == 0:
print(f'Loading LoRA checkpoint from {lora_ckpt_path}')
lora_checkpoint = torch.load(lora_ckpt_path, map_location='cpu')
if isinstance(lora_checkpoint, dict) and 'generator_lora' in lora_checkpoint:
peft.set_peft_model_state_dict(pipeline.generator.model, lora_checkpoint['generator_lora'])
else:
peft.set_peft_model_state_dict(pipeline.generator.model, lora_checkpoint)
if local_rank == 0:
print('LoRA weights loaded for generator')
if isinstance(lora_checkpoint, dict) and 'query_memory_encoder' in lora_checkpoint:
inner = pipeline.generator.model
if inner.query_memory_encoder is not None:
inner.query_memory_encoder.load_state_dict(lora_checkpoint['query_memory_encoder'], strict=False)
elif local_rank == 0:
print('No LoRA checkpoint specified; using base weights with LoRA adapters initialized')
pipeline.is_lora_enabled = True
print('dtype', pipeline.generator.model.dtype)
pipeline = pipeline.to(dtype=torch.bfloat16)
if low_memory:
DynamicSwapInstaller.install_model(pipeline.text_encoder, device=device)
pipeline.generator.to(device=device)
pipeline.vae.to(device=device)
if isinstance(config.switch_frame_indices, int):
switch_frame_indices: List[int] = [int(config.switch_frame_indices)]
else:
switch_frame_indices: List[int] = [int(x) for x in str(config.switch_frame_indices).split(',') if str(x).strip()]
dataset = MultiTextDataset(config.data_path)
num_segments = len(dataset[0]['prompts_list'])
assert len(switch_frame_indices) == num_segments - 1, 'The number of switch_frame_indices should be the number of prompt segments minus 1'
print('Number of segments:', num_segments)
print('Switch frame indices:', switch_frame_indices)
num_prompts_total = len(dataset)
print(f'Number of prompt lines: {num_prompts_total}')
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)
if local_rank == 0:
os.makedirs(config.output_folder, exist_ok=True)
if dist.is_initialized():
dist.barrier()
manifest = {}
for i, batch_data in tqdm(enumerate(dataloader), disable=local_rank != 0):
idx = batch_data['idx'].item()
prompts_list: List[str] = batch_data['prompts_list']
index_str = f'{idx:05d}'
def _unwrap(p):
while isinstance(p, list) and len(p) == 1:
p = p[0]
return p
manifest[index_str] = [_unwrap(p) for p in prompts_list]
first_output = os.path.join(config.output_folder, f'{index_str}-0.mp4')
if idx < num_prompts_total and os.path.exists(first_output):
print(f'Video already exists: {first_output}, skipping')
continue
sampled_noise = torch.randn([config.num_samples, config.num_output_frames, 16, 60, 104], device=device, dtype=torch.bfloat16)
video = pipeline.inference(noise=sampled_noise, text_prompts_list=prompts_list, switch_frame_indices=switch_frame_indices, return_latents=False)
current_video = rearrange(video, 'b t c h w -> b t h w c').cpu() * 255.0
if idx < num_prompts_total:
for sample_idx in range(config.num_samples):
output_path = os.path.join(config.output_folder, f'{index_str}-{sample_idx}.mp4')
write_video(output_path, current_video[sample_idx].to(torch.uint8), fps=16)
if config.inference_iter != -1 and i >= config.inference_iter:
break
if dist.is_initialized():
rank_manifest_path = os.path.join(config.output_folder, f'.manifest_rank{rank}.json')
with open(rank_manifest_path, 'w', encoding='utf-8') as f:
json.dump(manifest, f, indent=2, ensure_ascii=False)
dist.barrier()
if local_rank == 0:
merged = {}
manifest_path = os.path.join(config.output_folder, 'manifest.json')
if os.path.exists(manifest_path):
with open(manifest_path) as f:
merged = json.load(f)
for rfile in sorted(glob.glob(os.path.join(config.output_folder, '.manifest_rank*.json'))):
with open(rfile) as f:
merged.update(json.load(f))
os.remove(rfile)
with open(manifest_path, 'w', encoding='utf-8') as f:
json.dump(merged, f, indent=2, ensure_ascii=False)
else:
manifest_path = os.path.join(config.output_folder, 'manifest.json')
if os.path.exists(manifest_path):
with open(manifest_path) as f:
existing = json.load(f)
existing.update(manifest)
manifest = existing
with open(manifest_path, 'w', encoding='utf-8') as f:
json.dump(manifest, f, indent=2, ensure_ascii=False)
if dist.is_initialized():
dist.destroy_process_group()
|