inference/stream_long/inference_long_stream.py

import sys
from pathlib import Path
_PROJECT_ROOT = str(Path(__file__).resolve().parents[2])
if _PROJECT_ROOT in sys.path:
    sys.path.remove(_PROJECT_ROOT)
sys.path.insert(0, _PROJECT_ROOT)

import argparse
import os
import json
import glob
import time
import torch
import numpy as np
from omegaconf import OmegaConf
from tqdm import tqdm
from torch.utils.data import DataLoader, SequentialSampler
import torch.distributed as dist
try:
    import imageio.v2 as imageio
    import imageio_ffmpeg
except ImportError as e:
    sys.stderr.write(f'[inference_long_stream] imageio / imageio-ffmpeg not available: {e}\n  pip install imageio imageio-ffmpeg\n')
    sys.exit(1)
_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))
if _REPO_ROOT not in sys.path:
    sys.path.insert(0, _REPO_ROOT)
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

def _clean_key(name: str) -> str:
    return name.replace('_fsdp_wrapped_module.', '')

def build_pipeline(config, device, use_ema_cli: bool):
    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 use_ema_cli and 'generator_ema' in state_dict:
                raw = state_dict['generator_ema']
                if 'generator' in state_dict:
                    enc = {k: v for k, v in state_dict['generator'].items() if 'query_memory_encoder' in k}
                    if enc:
                        raw = dict(raw)
                        raw.update(enc)
            else:
                raw = state_dict.get('generator', state_dict.get('generator_ema'))
        elif 'model' in state_dict:
            raw = state_dict['model']
        else:
            raise ValueError(f'Generator state dict not found in {config.generator_ckpt}')
        cleaned = {_clean_key(k): v for k, v in raw.items()}
        missing, unexpected = pipeline.generator.load_state_dict(cleaned, strict=False)
        enc_loaded = sum((1 for k in cleaned if 'query_memory_encoder' in k))
        if missing:
            print(f'[Warning] {len(missing)} parameters missing: {missing[:8]} ...')
        if unexpected:
            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
        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=True)
        lora_ckpt_path = getattr(config, 'lora_ckpt', None)
        if lora_ckpt_path:
            print(f'Loading LoRA checkpoint from {lora_ckpt_path}')
            lora_ckpt = torch.load(lora_ckpt_path, map_location='cpu')
            if isinstance(lora_ckpt, dict) and 'generator_lora' in lora_ckpt:
                peft.set_peft_model_state_dict(pipeline.generator.model, lora_ckpt['generator_lora'])
            else:
                peft.set_peft_model_state_dict(pipeline.generator.model, lora_ckpt)
            print('LoRA weights loaded for generator')
            if isinstance(lora_ckpt, dict) and 'query_memory_encoder' in lora_ckpt:
                inner = pipeline.generator.model
                if inner.query_memory_encoder is not None:
                    inner.query_memory_encoder.load_state_dict(lora_ckpt['query_memory_encoder'], strict=False)
        else:
            print('No LoRA checkpoint specified; using base weights with LoRA adapters initialized')
        pipeline.is_lora_enabled = True
    return pipeline

def stream_decode_and_write(pipeline, latents, out_path, chunk_size, fps, codec, quality):
    assert latents.shape[0] == 1, 'streaming writer assumes batch size 1'
    B, T_latent, C, H, W = latents.shape
    zs = latents.permute(0, 2, 1, 3, 4)
    dtype = zs.dtype
    vae_wrapper = pipeline.vae
    vae = vae_wrapper.model
    device = next(vae.parameters()).device
    mean = vae_wrapper.mean.to(device=device, dtype=dtype)
    std_inv = (1.0 / vae_wrapper.std).to(device=device, dtype=dtype)
    scale = [mean, std_inv]
    writer = imageio.get_writer(out_path, fps=fps, codec=codec, quality=quality, macro_block_size=None)
    vae.clear_cache()
    frames_written = 0
    t0 = time.time()
    try:
        n_chunks = (T_latent + chunk_size - 1) // chunk_size
        for ci, start in enumerate(range(0, T_latent, chunk_size)):
            end = min(start + chunk_size, T_latent)
            chunk = zs[:, :, start:end].to(device, non_blocking=True)
            decoded = vae.cached_decode(chunk, scale)
            decoded = decoded.float().clamp_(-1, 1)
            decoded = (decoded * 0.5 + 0.5).mul_(255.0).to(torch.uint8).cpu()
            T_chunk_px = decoded.shape[2]
            for t in range(T_chunk_px):
                frame = decoded[0, :, t].permute(1, 2, 0).contiguous().numpy()
                writer.append_data(frame)
            frames_written += T_chunk_px
            del decoded
            torch.cuda.empty_cache()
            if ci % 8 == 0 or ci == n_chunks - 1:
                print(f'[stream] chunk {ci + 1}/{n_chunks} (latent {start}-{end}, +{T_chunk_px}px → total {frames_written}px, elapsed {time.time() - t0:.1f}s)', flush=True)
    finally:
        writer.close()
        vae.clear_cache()
    return frames_written

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--config_path', type=str, required=True)
    parser.add_argument('--checkpoint_path', type=str, default=None)
    parser.add_argument('--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('--use_ema', action='store_true')
    parser.add_argument('--seed', type=int, default=None)
    parser.add_argument('--num_samples', type=int, default=None)
    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
    device = torch.device('cuda')
    set_seed(config.seed)
    config.distributed = False
    print(f'Single GPU streaming mode on device {device}')
    print(f'Free VRAM {get_cuda_free_memory_gb(device)} GB')
    torch.set_grad_enabled(False)
    pipeline = build_pipeline(config, device, use_ema_cli=config.get('use_ema', False))
    pipeline = pipeline.to(dtype=torch.bfloat16)
    low_memory = get_cuda_free_memory_gb(device) < 40 or int(config.num_output_frames) > 100000
    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}')
    sampler = SequentialSampler(dataset)
    dataloader = DataLoader(dataset, batch_size=1, sampler=sampler, num_workers=0, drop_last=False)
    os.makedirs(config.output_folder, exist_ok=True)
    stream_cfg = OmegaConf.to_container(getattr(config, 'stream_decode', OmegaConf.create({})))
    chunk_size = int(stream_cfg.get('chunk_size', 120))
    codec = stream_cfg.get('codec', 'libx264')
    quality = int(stream_cfg.get('quality', 8))
    fps = 16
    manifest = {}
    for i, batch_data in tqdm(enumerate(dataloader)):
        idx = batch_data['idx'].item()
        batch = batch_data
        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
        all_exist = all((os.path.exists(os.path.join(config.output_folder, f'{index_str}-{s}.mp4')) for s in range(config.num_samples)))
        if idx < num_prompts and all_exist:
            print(f'All {config.num_samples} samples already exist for {index_str}, skipping')
            continue
        prompts = [extended_prompt] * config.num_samples if extended_prompt is not None else [prompt] * config.num_samples
        sampled_noise = torch.randn([config.num_samples, config.num_output_frames, 16, 60, 104], device=device, dtype=torch.bfloat16)
        print(f'[stream] generating latents: shape {list(sampled_noise.shape)}')
        t_gen0 = time.time()

        def _stub_decode(latent, use_cache=False, **kwargs):
            return torch.zeros([latent.shape[0], 1, 3, 480, 832], device=latent.device, dtype=torch.float32)
        _orig_decode = pipeline.vae.decode_to_pixel
        _orig_decode_chunk = pipeline.vae.decode_to_pixel_chunk
        pipeline.vae.decode_to_pixel = _stub_decode
        pipeline.vae.decode_to_pixel_chunk = _stub_decode
        try:
            _, latents = pipeline.inference(noise=sampled_noise, text_prompts=prompts, return_latents=True, low_memory=low_memory, profile=False)
        finally:
            pipeline.vae.decode_to_pixel = _orig_decode
            pipeline.vae.decode_to_pixel_chunk = _orig_decode_chunk
        torch.cuda.empty_cache()
        print(f'[stream] latents ready in {time.time() - t_gen0:.1f}s, shape {list(latents.shape)}')
        pipeline.vae.model.clear_cache()
        for sample_idx in range(config.num_samples):
            out_path = os.path.join(config.output_folder, f'{index_str}-{sample_idx}.mp4')
            if os.path.exists(out_path):
                print(f'[stream] {out_path} exists, skipping')
                continue
            print(f'[stream] writing sample {sample_idx} → {out_path}')
            t_dec0 = time.time()
            n_px = stream_decode_and_write(pipeline, latents[sample_idx:sample_idx + 1], out_path, chunk_size=chunk_size, fps=fps, codec=codec, quality=quality)
            expected_px = 4 * config.num_output_frames - 3
            print(f'[stream] wrote {n_px} pixel frames (expected {expected_px}) in {time.time() - t_dec0:.1f}s → {n_px / fps:.2f}s video')
        del latents, sampled_noise
        torch.cuda.empty_cache()
        if config.inference_iter != -1 and i >= config.inference_iter:
            break
    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 __name__ == '__main__':
    main()