inference/inference.py

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 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 pipeline import CausalInferencePipeline
from utils.dataset import TextDataset
from utils.misc import set_seed
from utils.memory import get_cuda_free_memory_gb, DynamicSwapInstaller
import json
import glob
parser = argparse.ArgumentParser()
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')
parser.add_argument('--num_samples', type=int, default=None, help='Override config num_samples')
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 args.num_samples is not None:
    config.num_samples = args.num_samples
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)
    config.distributed = True
    if rank == 0:
        print(f'[Rank {rank}] Initialized distributed processing on device {device}')
else:
    local_rank = 0
    rank = 0
    device = torch.device('cuda')
    set_seed(config.seed)
    config.distributed = False
    print(f'Single GPU mode on device {device}')
print(f'Free VRAM {get_cuda_free_memory_gb(device)} GB')
low_memory = get_cuda_free_memory_gb(device) < 40
torch.set_grad_enabled(False)
pipeline = CausalInferencePipeline(config, device=device)
if config.generator_ckpt:
    state_dict = torch.load(config.generator_ckpt, map_location='cpu')
    if 'generator' in state_dict or 'generator_ema' in state_dict:
        if config.use_ema and 'generator_ema' in state_dict:
            raw_gen_state_dict = state_dict['generator_ema']
            if 'generator' in state_dict:
                enc_keys = {k: v for k, v in state_dict['generator'].items() if 'query_memory_encoder' in k}
                if enc_keys:
                    raw_gen_state_dict = dict(raw_gen_state_dict)
                    raw_gen_state_dict.update(enc_keys)
        else:
            raw_gen_state_dict = state_dict.get('generator', state_dict.get('generator_ema'))
    elif 'model' in state_dict:
        raw_gen_state_dict = state_dict['model']
    else:
        raise ValueError(f'Generator state dict not found in {config.generator_ckpt}')

    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:
        enc_loaded = sum((1 for k in cleaned_state_dict if 'query_memory_encoder' in k))
        if len(missing) > 0:
            print(f'[Warning] {len(missing)} parameters missing: {missing[:8]} ...')
        if len(unexpected) > 0:
            print(f'[Warning] {len(unexpected)} unexpected parameters: {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
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)
extended_prompt_path = config.data_path
dataset = TextDataset(prompt_path=config.data_path, extended_prompt_path=extended_prompt_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)
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()
    if isinstance(batch_data, dict):
        batch = batch_data
    elif isinstance(batch_data, list):
        batch = batch_data[0]
    prompt = batch['prompts'][0]
    extended_prompt = batch['extended_prompts'][0] if 'extended_prompts' in batch else None
    index_str = f'{idx:05d}'
    manifest[index_str] = prompt
    first_output = os.path.join(config.output_folder, f'{index_str}-0.mp4')
    if idx < num_prompts and os.path.exists(first_output):
        print(f'Video already exists: {first_output}, skipping')
        continue
    if extended_prompt is not None:
        prompts = [extended_prompt] * config.num_samples
    else:
        prompts = [prompt] * config.num_samples
    sampled_noise = torch.randn([config.num_samples, config.num_output_frames, 16, 60, 104], device=device, dtype=torch.bfloat16)
    video, latents = pipeline.inference(noise=sampled_noise, text_prompts=prompts, return_latents=True, low_memory=low_memory, profile=False)
    current_video = rearrange(video, 'b t c h w -> b t h w c').cpu()
    video = 255.0 * current_video
    pipeline.vae.model.clear_cache()
    if idx < num_prompts:
        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, video[sample_idx], 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()