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()