interactive_inference.py

# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES
#
# Licensed under the Apache License, Version 2.0 (the "License").
# You may not use this file except in compliance with the License.
# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
#
# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
#
# SPDX-License-Identifier: Apache-2.0
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  # noqa: F401
from einops import rearrange

from utils.misc import set_seed
from utils.distributed import barrier  
from utils.memory import gpu, get_cuda_free_memory_gb, DynamicSwapInstaller

from pipeline.interactive_causal_inference import (
    InteractiveCausalInferencePipeline,
)
from utils.dataset import MultiTextDataset


def initialize_vae_decoder(use_taehv: bool, taehv_checkpoint_path: str, device: torch.device):
    if not use_taehv:
        return None

    from taehv import TAEHV

    if not os.path.exists(taehv_checkpoint_path):
        fallback_candidates = [
            "taew2_1.pth",
            "checkpoints/taew2_1.pth",
        ]
        resolved = None
        for candidate in fallback_candidates:
            if os.path.exists(candidate):
                resolved = candidate
                break
        if resolved is None:
            raise FileNotFoundError(
                f"TAEHV checkpoint not found at {taehv_checkpoint_path}. "
                "Please provide taew2_1.pth from "
                "local workspace (e.g., project root or checkpoints/)."
            )
        taehv_checkpoint_path = resolved

    class TinyVAEWrapper(torch.nn.Module):
        def __init__(self, checkpoint_path: str):
            super().__init__()
            self.taehv = TAEHV(checkpoint_path=checkpoint_path).to(torch.float16)

        def decode_to_pixel(self, latent: torch.Tensor, use_cache: bool = False) -> torch.Tensor:
            del use_cache
            latents_fp16 = latent.to(dtype=torch.float16)
            return self.taehv.decode_video(latents_fp16, parallel=False).mul_(2).sub_(1)

    tiny_vae = TinyVAEWrapper(taehv_checkpoint_path).eval().requires_grad_(False)
    tiny_vae.to(device)
    print("TAEHV tiny decoder enabled for interactive inference")
    return tiny_vae


# ----------------------------- Argument parsing -----------------------------
parser = argparse.ArgumentParser("Interactive causal inference")
parser.add_argument("--config_path", type=str, help="Path to the config file")
args = parser.parse_args()

config = OmegaConf.load(args.config_path)

# ----------------------------- Distributed setup -----------------------------
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)))
    
    # Set device first
    torch.cuda.set_device(local_rank)
    device = torch.device(f"cuda:{local_rank}")
    
    # Initialize process group with backend and timeout
    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:
        pipeline.generator.load_state_dict(raw_gen_state_dict)

# --------------------------- LoRA support (optional) ---------------------------
from utils.lora_utils import configure_lora_for_model
import peft

pipeline.is_lora_enabled = False
if getattr(config, "adapter", None) and configure_lora_for_model is not None:
    if local_rank == 0:
        print(f"LoRA enabled with config: {config.adapter}")
        print("Applying LoRA to generator (inference)...")
    # After loading base weights, apply LoRA wrapper to the generator's transformer model
    pipeline.generator.model = configure_lora_for_model(
        pipeline.generator.model,
        model_name="generator",
        lora_config=config.adapter,
        is_main_process=(local_rank == 0),
    )

    # Load LoRA weights (if lora_ckpt is provided)
    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")
        # Support both formats: containing the `generator_lora` key or a raw LoRA state dict
        if isinstance(lora_checkpoint, dict) and "generator_lora" in lora_checkpoint:
            peft.set_peft_model_state_dict(pipeline.generator.model, lora_checkpoint["generator_lora"])  # type: ignore
        else:
            peft.set_peft_model_state_dict(pipeline.generator.model, lora_checkpoint)  # type: ignore
        if local_rank == 0:
            print("LoRA weights loaded for generator")
    else:
        if local_rank == 0:
            print("No LoRA checkpoint specified; using base weights with LoRA adapters initialized")

    pipeline.is_lora_enabled = True

if (
    pipeline.is_lora_enabled
    and bool(getattr(config, "enable_torch_compile", False))
    and bool(getattr(config, "merge_lora_before_compile", True))
):
    if local_rank == 0:
        print("Merging LoRA adapters into base model before torch.compile")
    if not hasattr(pipeline.generator.model, "merge_and_unload"):
        raise RuntimeError(
            "LoRA model does not support merge_and_unload, cannot proceed with strict torch.compile mode."
        )
    pipeline.generator.model = pipeline.generator.model.merge_and_unload(
        progressbar=False,
        safe_merge=True,
    )
    if local_rank == 0:
        print("LoRA merge completed")

# Move pipeline to appropriate dtype and device
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)

use_taehv = bool(getattr(config, "use_taehv", False))
taehv_checkpoint_path = str(getattr(config, "taehv_checkpoint_path", "checkpoints/taew2_1.pth"))
if use_taehv:
    tiny_vae = initialize_vae_decoder(use_taehv=True, taehv_checkpoint_path=taehv_checkpoint_path, device=device)
    if tiny_vae is not None:
        pipeline.vae = tiny_vae

# Optional torch.compile for inference (config-controlled)
if bool(getattr(config, "enable_torch_compile", False)):
    compile_mode = str(getattr(config, "torch_compile_mode", "max-autotune-no-cudagraphs"))
    suppress_compile_errors = bool(getattr(config, "torch_compile_suppress_errors", True))
    require_compile_success = bool(getattr(config, "torch_compile_require_success", True))
    if suppress_compile_errors:
        torch._dynamo.config.suppress_errors = True
    if local_rank == 0:
        print(
            f"Enabling torch.compile with mode={compile_mode}, "
            f"suppress_errors={suppress_compile_errors}, "
            f"require_success={require_compile_success}"
        )
    try:
        compiled_gen = torch.compile(pipeline.generator, mode=compile_mode)
        if compiled_gen is not None:
            pipeline.generator = compiled_gen
    except Exception as e:
        if require_compile_success:
            raise RuntimeError(f"torch.compile failed for generator in strict mode: {e}") from e
        if local_rank == 0:
            print(f"[Warning] torch.compile failed for generator: {e}")
    try:
        compiled_vae = torch.compile(pipeline.vae, mode=compile_mode)
        if compiled_vae is not None:
            pipeline.vae = compiled_vae
    except Exception as e:
        if require_compile_success:
            raise RuntimeError(f"torch.compile failed for VAE in strict mode: {e}") from e
        if local_rank == 0:
            print(f"[Warning] torch.compile failed for VAE: {e}")

# ----------------------------- Build dataset -----------------------------
# Parse switch_frame_indices
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()
    ]

# Create dataset
dataset = MultiTextDataset(config.data_path)

# Validate number of segments & switch_frame_indices length
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)

# Create output directory
if local_rank == 0:
    os.makedirs(config.output_folder, exist_ok=True)

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

# ----------------------------- Inference loop -----------------------------
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"]  # type: ignore

    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 dist.is_initialized():
        rank = dist.get_rank()
    else:
        rank = 0

    # Determine model type for filename
    if hasattr(pipeline, 'is_lora_enabled') and pipeline.is_lora_enabled:
        model_type = "lora"
    elif getattr(config, 'use_ema', False):
        model_type = "ema"
    else:
        model_type = "regular"

    for seed_idx in range(config.num_samples):
        if config.save_with_index:
            output_path = os.path.join(config.output_folder, f"rank{rank}-{idx}-{seed_idx}_{model_type}.mp4")
        else:
            # Use the first prompt segment as the filename prefix to avoid overly long names
            short_name = prompts_list[0][0][:100].replace("/", "_")
            output_path = os.path.join(config.output_folder, f"rank{rank}-{short_name}-{seed_idx}_{model_type}.mp4")
        write_video(output_path, current_video[seed_idx].to(torch.uint8), fps=16)

    if config.inference_iter != -1 and i >= config.inference_iter:
        break

if dist.is_initialized():
    dist.destroy_process_group()