# Adopted from https://github.com/guandeh17/Self-Forcing
# SPDX-License-Identifier: Apache-2.0
import os
import sys

# torchrun no longer consistently prepends the script directory to sys.path,
# which breaks absolute project imports when launched from another cwd.
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

# torchvision 0.27+ removed write_video/read_video. Several modules import the
# symbols at module import time, so patch them before importing project code.
import torchvision.io as _tv_io
if not hasattr(_tv_io, "write_video"):
    import imageio.v2 as _imageio_v2

    def _shim_write_video(filename, video_array, fps, **_unused):
        if hasattr(video_array, "detach"):
            video_array = video_array.detach().cpu().numpy()
        _imageio_v2.mimwrite(filename, video_array, fps=fps, codec="libx264", quality=8)

    _tv_io.write_video = _shim_write_video
if not hasattr(_tv_io, "read_video"):
    import imageio.v3 as _imageio_v3
    import torch as _torch_for_shim

    def _shim_read_video(filename, pts_unit="sec", output_format="THWC", **_unused):
        frames = _imageio_v3.imread(filename, plugin="pyav")
        tensor = _torch_for_shim.from_numpy(frames)
        if output_format == "TCHW":
            tensor = tensor.permute(0, 3, 1, 2)
        return tensor, None, {}

    _tv_io.read_video = _shim_read_video

import argparse
import torch
from omegaconf import OmegaConf
from tqdm import tqdm
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 CausalDiffusionInferencePipeline
from utils.dataset import MultiTextConcatDataset, MultiVideoConcatDataset, eval_collate_fn, multi_video_collate_fn
from utils.misc import set_seed
from utils.config import normalize_config, section_get, wan_default_config
from utils.nvfp4_checkpoint import (
    clean_fsdp_state_dict_keys,
    drop_fouroversix_master_weights,
    is_nvfp4_state_dict,
    is_te_nvfp4_checkpoint,
    quantize_model_for_fouroversix_nvfp4,
    unwrap_generator_state_dict,
)

from utils.memory import get_cuda_free_memory_gb, DynamicSwapInstaller


def save_prompts_to_txt(prompts_for_sample, prompt_txt_path: str, is_main_process: bool):
    """Save per-block prompts alongside the video.

    Consecutive identical prompts are merged, e.g.:
        [0] a, [1] a, [2] b  =>  [0,1] a\\n[2] b\\n
    """
    try:
        with open(prompt_txt_path, "w", encoding="utf-8") as f:
            if len(prompts_for_sample) == 0:
                return
            current_prompt = prompts_for_sample[0]
            current_indices = [0]
            for seg_idx in range(1, len(prompts_for_sample)):
                p = prompts_for_sample[seg_idx]
                if p == current_prompt:
                    current_indices.append(seg_idx)
                else:
                    indices_str = ",".join(str(i) for i in current_indices)
                    f.write(f"[{indices_str}] {current_prompt}\n")
                    current_prompt = p
                    current_indices = [seg_idx]
            indices_str = ",".join(str(i) for i in current_indices)
            f.write(f"[{indices_str}] {current_prompt}\n")
    except Exception as e:
        if is_main_process:
            print(f"Warning: failed to save prompts to {prompt_txt_path}: {e}")

parser = argparse.ArgumentParser()
parser.add_argument("--config_path", type=str, help="Path to the config file")
te_quant_group = parser.add_mutually_exclusive_group()
te_quant_group.add_argument(
    "--use_te_quant",
    dest="use_te_quant",
    action="store_true",
    help="Override config and enable TransformerEngine quantization",
)
te_quant_group.add_argument(
    "--no_use_te_quant",
    dest="use_te_quant",
    action="store_false",
    help="Override config and disable TransformerEngine quantization",
)
parser.set_defaults(use_te_quant=None)
args = parser.parse_args()

config = normalize_config(OmegaConf.load(args.config_path))
if args.use_te_quant is not None:
    config.model_quant_use_transformer_engine = args.use_te_quant

if not hasattr(config, "sampling_steps") or config.sampling_steps is None:
    raise ValueError("sampling_steps must be defined in the inference config")

if not hasattr(config, "guidance_scale") or config.guidance_scale is None:
    config.guidance_scale = 1.0

config.use_ema = section_get(config, "inference", "use_ema", getattr(config, "use_ema", False))
config.output_folder = section_get(config, "inference", "output_folder", getattr(config, "output_folder", "videos/longlive2"))
config.num_samples = section_get(config, "inference", "num_samples", getattr(config, "num_samples", 1))
config.num_output_frames = getattr(config, "num_output_frames", config.image_or_video_shape[1])
config.save_with_index = getattr(config, "save_with_index", False)
config.inference_iter = getattr(config, "inference_iter", -1)

if bool(getattr(config, "fp8_quant", False)) and bool(
    getattr(config, "model_quant", False)
):
    raise ValueError("fp8_quant and model_quant (NVFP4) are mutually exclusive.")


def _maybe_to_dict(value):
    if value is None:
        return None
    if OmegaConf.is_config(value):
        value = OmegaConf.to_container(value, resolve=True)
    return dict(value)


def _config_bool(value, default=False):
    if value is None:
        return default
    if isinstance(value, str):
        return value.strip().lower() in {"1", "true", "yes", "y", "on"}
    return bool(value)


def _expected_inference_samples(config):
    inference_iter = int(getattr(config, "inference_iter", -1))
    if inference_iter >= 0:
        return inference_iter + 1
    return None


def _resolve_torch_compile(config):
    setting = getattr(config, "torch_compile", False)
    if isinstance(setting, str) and setting.strip().lower() == "auto":
        if not (
            bool(getattr(config, "model_quant", False))
            or bool(getattr(config, "fp8_quant", False))
        ):
            return False, "auto disabled because quantization is false"
        min_samples = int(getattr(config, "torch_compile_min_samples", 2))
        expected_samples = _expected_inference_samples(config)
        if expected_samples is not None and expected_samples < min_samples:
            return (
                False,
                "auto disabled because expected samples "
                f"({expected_samples}) < torch_compile_min_samples ({min_samples})",
            )
        return True, "auto enabled for repeated quantized inference"
    return _config_bool(setting, default=False), "explicit setting"


def quantize_generator_model(model, config, keep_master_weights):
    from utils.quant import (
        ModelQuantizationConfig,
        _materialize_mixed_quantized_weights_for_inference,
        _materialize_quantized_weights_for_inference,
        _materialize_transformer_engine_weights_for_inference,
        quantize_model_with_filter,
    )

    use_transformer_engine = bool(getattr(config, "model_quant_use_transformer_engine", False))
    te_inference_only = bool(getattr(config, "model_quant_te_inference_only", use_transformer_engine))
    te_low_precision_weights = bool(getattr(config, "model_quant_te_low_precision_weights", te_inference_only))
    te_fallback_to_fouroversix = bool(getattr(config, "model_quant_te_fallback_to_fouroversix", False))

    quant_cfg = ModelQuantizationConfig(
        scale_rule=getattr(config, "model_quant_scale_rule", "static_6"),
        quantize_backend=getattr(config, "model_quant_backend", None),
        activation_scale_rule=getattr(
            config,
            "model_quant_activation_scale_rule",
            getattr(config, "model_quant_scale_rule", "static_6"),
        ),
        weight_scale_rule=getattr(config, "model_quant_weight_scale_rule", None),
        gradient_scale_rule=getattr(config, "model_quant_gradient_scale_rule", None),
    )
    quant_cfg.keep_master_weights = keep_master_weights
    model, matched_modules = quantize_model_with_filter(
        model,
        quant_config=quant_cfg,
        filtered_modules=getattr(config, "model_quant_filtered_modules", None),
        use_default_filtered_modules=getattr(config, "model_quant_use_default_filtered_modules", True),
        cast_model_to_bf16=True,
        materialize_for_inference=False,
        use_transformer_engine=use_transformer_engine,
        te_inference_only=te_inference_only,
        te_low_precision_weights=te_low_precision_weights,
        te_recipe_kwargs=_maybe_to_dict(getattr(config, "model_quant_te_recipe_kwargs", None)),
        te_module_kwargs=_maybe_to_dict(getattr(config, "model_quant_te_module_kwargs", None)),
        te_fallback_to_fouroversix=te_fallback_to_fouroversix,
        verbose=True,
    )
    materialize_fn = _materialize_quantized_weights_for_inference
    if use_transformer_engine and te_fallback_to_fouroversix:
        materialize_fn = _materialize_mixed_quantized_weights_for_inference
    elif use_transformer_engine:
        materialize_fn = _materialize_transformer_engine_weights_for_inference
    if local_rank == 0:
        print(f"[NVFP4] Generator quantized; {len(matched_modules)} modules excluded")
    return model, materialize_fn


def materialize_quantized_generator(model, device, materialize_fn, stage_desc):
    mat_modules, master_bytes, quantized_bytes = materialize_fn(
        model,
        target_device=device,
    )
    if local_rank == 0:
        print(
            f"[NVFP4] Materialized quantized generator weights {stage_desc}: "
            f"{len(mat_modules)} modules, "
            f"master_weight={master_bytes / (1024 ** 3):.3f} GiB, "
            f"quantized_weight={quantized_bytes / (1024 ** 3):.3f} GiB"
        )


def configure_generator_torch_compile(pipeline, config):
    compile_enabled, reason = _resolve_torch_compile(config)
    if not compile_enabled:
        if local_rank == 0 and str(getattr(config, "torch_compile", "false")).lower() == "auto":
            print(f"[torch.compile] skipped: {reason}")
        return
    target = str(getattr(config, "torch_compile_target", "generator_model")).lower()
    if target not in {"generator_model", "model"}:
        if local_rank == 0:
            print(f"[torch.compile][warn] Unsupported target={target}; expected generator_model")
        return
    if not hasattr(pipeline.generator, "configure_torch_compile"):
        if local_rank == 0:
            print("[torch.compile][warn] Current generator does not expose configure_torch_compile; skipping")
        return
    compiled = pipeline.generator.configure_torch_compile(
        backend=str(getattr(config, "torch_compile_backend", "inductor")),
        mode=getattr(config, "torch_compile_mode", "max-autotune-no-cudagraphs"),
        fullgraph=_config_bool(getattr(config, "torch_compile_fullgraph", False)),
        dynamic=_config_bool(getattr(config, "torch_compile_dynamic", False)),
        options=_maybe_to_dict(getattr(config, "torch_compile_options", None)),
        suppress_errors=_config_bool(getattr(config, "torch_compile_suppress_errors", True), default=True),
    )
    if local_rank == 0:
        status = "enabled" if compiled else "not enabled"
        print(f"[torch.compile] {status}: target={target}")

# Initialize distributed inference
if "LOCAL_RANK" in os.environ:
    dist.init_process_group(backend='nccl')
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    device = torch.device(f"cuda:{local_rank}")
    set_seed(config.seed + local_rank)
    config.distributed = True  # Mark as distributed for pipeline
else:
    local_rank = 0
    device = torch.device("cuda")
    set_seed(config.seed)
    config.distributed = False  # Mark as non-distributed

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)


# Initialize pipeline
pipeline = CausalDiffusionInferencePipeline(config, device=device)

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

merge_lora = bool(getattr(config, "merge_lora", False))
has_lora_adapter = bool(getattr(config, "adapter", None) and configure_lora_for_model is not None)
if has_lora_adapter and (
    bool(getattr(config, "model_quant", False))
    or bool(getattr(config, "fp8_quant", False))
) and not merge_lora:
    if local_rank == 0:
        print(
            "[quant][LoRA] merge_lora=false is unsupported with quantization; "
            "forcing merge_lora=true so the LoRA is folded into the BF16 base before quantization."
        )
    merge_lora = True
    config.merge_lora = True
materialize_quantized_weights_for_inference = None
generator_checkpoint = None
generator_lora_state = None
generator_ckpt_path = getattr(config, "generator_ckpt", None)
loaded_prequantized_generator = False
prequantized_generator_backend = None
if generator_ckpt_path:
    generator_checkpoint = torch.load(generator_ckpt_path, map_location="cpu")
    is_lora_only_checkpoint = (
        isinstance(generator_checkpoint, dict)
        and "generator_lora" in generator_checkpoint
        and not any(key in generator_checkpoint for key in ("generator", "generator_ema", "model"))
    )
    if is_lora_only_checkpoint:
        generator_lora_state = generator_checkpoint["generator_lora"]
        if local_rank == 0:
            print(f"Found LoRA generator weights in {generator_ckpt_path}")
    else:
        raw_gen_state_dict = unwrap_generator_state_dict(generator_checkpoint, use_ema=config.use_ema)
        if config.use_ema:
            raw_gen_state_dict = clean_fsdp_state_dict_keys(raw_gen_state_dict)
        if is_te_nvfp4_checkpoint(generator_checkpoint):
            raise ValueError(
                "This checkpoint was saved as a TransformerEngine module state_dict, "
                "which is not packed NVFP4 and is no longer a supported export format. "
                "Regenerate with `--backend transformer_engine` to save merged bf16 weights "
                "for TE runtime quantization, or use `--backend fouroversix` for a compact "
                "materialized NVFP4 checkpoint."
            )
        elif is_nvfp4_state_dict(raw_gen_state_dict):
            if not getattr(config, "model_quant", False):
                raise ValueError(
                    "generator_ckpt is a materialized NVFP4 checkpoint, but model_quant is false. "
                    "Set model_quant: true in the inference yaml."
                )
            if getattr(config, "model_quant_use_transformer_engine", False):
                raise ValueError(
                    "Materialized NVFP4 generator checkpoints use FourOverSix modules. "
                    "Set model_quant_use_transformer_engine: false when loading this checkpoint."
                )
            if local_rank == 0:
                print(f"[NVFP4] Loading materialized generator checkpoint from {generator_ckpt_path}")
            pipeline.generator.model, matched_modules = quantize_model_for_fouroversix_nvfp4(
                pipeline.generator.model,
                config=config,
                keep_master_weights=False,
                verbose=(local_rank == 0),
            )
            dropped_modules = drop_fouroversix_master_weights(pipeline.generator.model)
            pipeline.generator.load_state_dict(raw_gen_state_dict, strict=True)
            loaded_prequantized_generator = True
            prequantized_generator_backend = "fouroversix"
            if local_rank == 0:
                print(
                    "[NVFP4] Prepared quantized generator architecture: "
                    f"{len(dropped_modules)} materialized modules, "
                    f"{len(matched_modules)} modules excluded"
                )
        elif config.use_ema:
            missing, unexpected = pipeline.generator.load_state_dict(raw_gen_state_dict, strict=False)
            if local_rank == 0:
                if len(missing) > 0:
                    print(f"[Warning] {len(missing)} parameters are missing when loading checkpoint: {missing[:8]} ...")
                if len(unexpected) > 0:
                    print(f"[Warning] {len(unexpected)} unexpected parameters encountered when loading checkpoint: {unexpected[:8]} ...")
        else:
            print(f"Loading generator from {generator_ckpt_path}")
            pipeline.generator.load_state_dict(raw_gen_state_dict, strict=True)

pipeline.is_lora_enabled = False
pipeline.is_lora_merged = False

if loaded_prequantized_generator:
    if has_lora_adapter or merge_lora or getattr(config, "lora_ckpt", None):
        if local_rank == 0:
            print("[NVFP4] Pre-quantized generator checkpoint is already saved with merged weights; skipping LoRA setup")
    has_lora_adapter = False
    merge_lora = False
    config.merge_lora = False

if getattr(config, "model_quant", False) and not merge_lora and not loaded_prequantized_generator:
    pipeline.generator.model, materialize_quantized_weights_for_inference = quantize_generator_model(
        pipeline.generator.model,
        config=config,
        keep_master_weights=has_lora_adapter,
    )

if has_lora_adapter:
    if local_rank == 0:
        print(f"LoRA enabled with config: {config.adapter}")
        print("Applying LoRA to generator (inference)...")
        if merge_lora:
            print("LoRA weights will be merged into the base model before inference")
    # Apply LoRA to the generator transformer after loading base weights.
    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 from lora_ckpt. If omitted, fall back to generator_ckpt
    # when it directly contains generator_lora.
    lora_ckpt_path = getattr(config, "lora_ckpt", None)
    if lora_ckpt_path:
        if local_rank == 0:
            print(f"Loading LoRA weights from lora_ckpt: {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"])  # 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")
    elif generator_lora_state is not None:
        if local_rank == 0:
            print(f"Loading LoRA weights from generator_ckpt: {generator_ckpt_path}")
        peft.set_peft_model_state_dict(pipeline.generator.model, generator_lora_state)  # type: ignore
        if local_rank == 0:
            print("LoRA weights loaded for generator")
    else:
        if local_rank == 0:
            print("No LoRA checkpoint configured; using initialized LoRA adapters")

    if merge_lora:
        if local_rank == 0:
            print("Merging LoRA weights into generator before quantization / inference...")
        pipeline.generator.model = pipeline.generator.model.merge_and_unload(safe_merge=True)
        pipeline.is_lora_merged = True
    else:
        pipeline.is_lora_enabled = True
elif merge_lora and local_rank == 0:
    print("merge_lora=True requested but no adapter config was found; continuing without LoRA merge")

del generator_checkpoint


# Move pipeline to appropriate dtype and device
if loaded_prequantized_generator:
    pipeline.text_encoder.to(dtype=torch.bfloat16)
    pipeline.vae.to(dtype=torch.bfloat16)
else:
    pipeline = pipeline.to(dtype=torch.bfloat16)
if low_memory:
    DynamicSwapInstaller.install_model(pipeline.text_encoder, device=device)
pipeline.generator.to(device=device)

if getattr(config, "model_quant", False) and not loaded_prequantized_generator:
    if merge_lora:
        pipeline.generator.model, materialize_quantized_weights_for_inference = quantize_generator_model(
            pipeline.generator.model,
            config=config,
            keep_master_weights=False,
        )
        stage_desc = "after LoRA merge" if pipeline.is_lora_merged else "for inference"
    else:
        stage_desc = "after LoRA wrapping" if pipeline.is_lora_enabled else "for inference"
    materialize_quantized_generator(
        pipeline.generator.model,
        device=device,
        materialize_fn=materialize_quantized_weights_for_inference,
        stage_desc=stage_desc,
    )
elif loaded_prequantized_generator and local_rank == 0:
    print(f"[NVFP4] Using pre-saved {prequantized_generator_backend} generator weights from checkpoint")

pipeline.generator.model.eval().requires_grad_(False)
if bool(getattr(config, "fp8_quant", False)):
    from utils.fp8 import quantize_model_fp8

    quantize_model_fp8(pipeline.generator.model, verbose=(local_rank == 0))
configure_generator_torch_compile(pipeline, config)

vae_device_str = getattr(config, "vae_device", None)
use_dedicated_vae_device = bool(getattr(config, "streaming_vae", False)) and bool(vae_device_str)
if use_dedicated_vae_device:
    vae_device = torch.device(vae_device_str)
    pipeline.vae.to(device="cpu")
    pipeline.vae.to(device=vae_device)
    if hasattr(pipeline.vae, "mean"):
        pipeline.vae.mean = pipeline.vae.mean.to(device=vae_device)
        pipeline.vae.std = pipeline.vae.std.to(device=vae_device)
    if local_rank == 0:
        print(f"[inference] VAE on {vae_device}, diffusion on {device}")
else:
    pipeline.vae.to(device=device)
    if vae_device_str and local_rank == 0:
        print(f"[inference] Ignoring vae_device={vae_device_str} because streaming_vae is false")

# Create dataset
nfpb = getattr(config, 'num_frame_per_block', 8)
data_path = config.data_path
chunks_per_shot = getattr(config, 'chunks_per_shot', 0)
scene_cut_prefix = getattr(config, 'scene_cut_prefix', "The scene transitions. ")
if getattr(config, "i2v", False):
    model_name = config.model_kwargs.model_name
    frame_raw_height = list(config.image_or_video_shape)[3] * wan_default_config[model_name]["spatial_compression_ratio"]
    frame_raw_width = list(config.image_or_video_shape)[4] * wan_default_config[model_name]["spatial_compression_ratio"]
    temporal_compression_ratio = wan_default_config[model_name]["temporal_compression_ratio"]
    total_frames = (config.num_output_frames - 1) * temporal_compression_ratio + 1
    dataset = MultiVideoConcatDataset(
        data_dir=data_path,
        video_size=(frame_raw_height, frame_raw_width),
        total_frames=total_frames,
        deterministic=True,
        num_frame_per_block=nfpb,
        temporal_compression_ratio=temporal_compression_ratio,
        target_fps=24 if "5B" in model_name else 16,
        allow_padding=getattr(config, "allow_padding", False),
        min_latent_frames=getattr(config, "min_latent_frames", 0),
        single_video_only=getattr(config, "uniform_prompt", False),
        independent_first_frame=getattr(config, "independent_first_frame", False),
        return_image=True,
        max_chunks_per_shot=getattr(config, "max_chunks_per_shot", 0),
        scene_cut_prefix=scene_cut_prefix,
    )
    collate_fn = multi_video_collate_fn
    num_blocks = config.num_output_frames // nfpb
else:
    num_blocks = config.num_output_frames // nfpb
    dataset = MultiTextConcatDataset(
        data_path=data_path,
        num_blocks=num_blocks,
        chunks_per_shot=chunks_per_shot,
        scene_cut_prefix=scene_cut_prefix,
        deterministic=True,
    )
    collate_fn = eval_collate_fn
if local_rank == 0:
    print(f"[data] data_path={data_path}, mode={getattr(dataset, '_mode', dataset.__class__.__name__)}, num_blocks={num_blocks}")
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, collate_fn=collate_fn)

# Create output directory (only on main process to avoid race conditions)
if local_rank == 0:
    os.makedirs(config.output_folder, exist_ok=True)

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


def encode(self, videos: torch.Tensor) -> torch.Tensor:
    device, dtype = videos[0].device, videos[0].dtype
    scale = [self.mean.to(device=device, dtype=dtype),
             1.0 / self.std.to(device=device, dtype=dtype)]
    output = [
        self.model.encode(u.unsqueeze(0), scale).float().squeeze(0)
        for u in videos
    ]

    output = torch.stack(output, dim=0)
    return output


for i, batch_data in tqdm(enumerate(dataloader), disable=(local_rank != 0)):
    idx = batch_data['idx'].item()

    # For DataLoader batch_size=1, the batch_data is already a single item, but in a batch container
    # Unpack the batch data for convenience
    if isinstance(batch_data, dict):
        batch = batch_data
    elif isinstance(batch_data, list):
        batch = batch_data[0]  # First (and only) item in the batch

    all_video = []

    # MultiTextConcatDataset + eval_collate_fn: prompts[0] is List[str].
    block_prompts = list(batch['prompts'][0])
    prompt = block_prompts[0]  # for filename
    prompts = [block_prompts] * config.num_samples

    shape = config.image_or_video_shape
    sampled_noise = torch.randn(
        [config.num_samples, config.num_output_frames, shape[2], shape[3], shape[4]], device=device, dtype=torch.bfloat16
    )
    initial_latent = None
    if getattr(config, "i2v", False):
        image = batch["image"].to(device=device, dtype=torch.bfloat16)
        if image.ndim == 4:
            image = image.unsqueeze(2)
        elif image.ndim != 5:
            raise ValueError(f"Expected i2v image with shape [B,C,H,W] or [B,C,T,H,W], got {tuple(image.shape)}")
        initial_latent = pipeline.vae.encode_to_latent(image).to(device=device, dtype=torch.bfloat16)
        if initial_latent.shape[0] != config.num_samples:
            initial_latent = initial_latent.repeat(config.num_samples, 1, 1, 1, 1)
        if config.num_output_frames <= initial_latent.shape[1]:
            raise ValueError(
                f"num_output_frames must exceed the i2v conditioning frames; "
                f"got {config.num_output_frames} and {initial_latent.shape[1]}"
            )
    print("sampled_noise.device", sampled_noise.device)
    print("prompts", prompts)
    print('sampled_noise.shape', sampled_noise.shape, 'prompts', prompts)
    save_latents_only = section_get(
        config,
        "inference",
        "save_latents_only",
        getattr(config, "save_latents_only", getattr(config, "save_latent_only", False)),
        aliases=("save_latent_only", "return_latents"),
    )
    inference_kwargs = dict(
        noise=sampled_noise,
        text_prompts=prompts,
        return_latents=save_latents_only,
    )
    if initial_latent is not None:
        inference_kwargs["initial_latent"] = initial_latent
    with torch.inference_mode():
        generated = pipeline.inference(**inference_kwargs)

    if not save_latents_only:
        current_video = rearrange(generated, 'b t c h w -> b t h w c').cpu()
        all_video.append(current_video)

        # Final output video
        video = 255.0 * torch.cat(all_video, dim=1)

        # Clear VAE cache
        pipeline.vae.model.clear_cache()
    else:
        latents = generated

    if dist.is_initialized():
        rank = dist.get_rank()
    else:
        rank = 0

    # Save the video if the current prompt is not a dummy prompt
    if idx < num_prompts:
        # 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:
                base_name = f'rank{rank}-{idx}-{seed_idx}_{model_type}'
            else:
                base_name = f'rank{rank}-{prompt[:100]}-{seed_idx}_{model_type}'

            if save_latents_only:
                latent_path = os.path.join(config.output_folder, f'{base_name}.pt')
                torch.save(latents[seed_idx].cpu(), latent_path)
            else:
                output_path = os.path.join(config.output_folder, f'{base_name}.mp4')
                fps = 24 if '5B' in config.model_kwargs.model_name else 16
                write_video(output_path, video[seed_idx], fps=fps)

            prompt_txt_path = os.path.join(config.output_folder, f'{base_name}_prompts.txt')
            save_prompts_to_txt(
                prompts[seed_idx] if isinstance(prompts[seed_idx], list) else [prompts[seed_idx]],
                prompt_txt_path,
                is_main_process=(rank == 0),
            )

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