inference_lance.py

# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# coding: utf-8

import warnings
warnings.filterwarnings("ignore", message=".*pkg_resources is deprecated.*", category=UserWarning)
warnings.filterwarnings("ignore", category=FutureWarning, module="diffusers.models.transformers.transformer_2d")
import os
import time
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"

import os.path as osp
from copy import deepcopy
import json
from typing import Tuple, cast, Optional
import torch
import torch.distributed as dist
from torch.utils.data import DataLoader
from transformers import HfArgumentParser, set_seed
from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLVisionConfig
from safetensors.torch import load_file

from data.dataset_base import DataConfig, simple_custom_collate
from data.data_utils import add_special_tokens
from modeling.vae.wan.model import WanVideoVAE
from modeling.lance import LanceConfig, Lance, Qwen2ForCausalLM
from modeling.qwen2 import Qwen2Tokenizer
from modeling.qwen2.modeling_qwen2 import Qwen2Config
from modeling.vit.qwen2_5_vl_vit import Qwen2_5_VisionTransformerPretrainedModel
from common.utils.misc import tuple_mul, AutoEncoderParams
from common.utils.distributed import get_global_rank
from common.utils.logging import get_logger
from common.val.utils import make_padded_latent, decode_video_tensor
from data.datasets_custom import ValidationDataset
from config.config_factory import (
    ModelArguments,
    DataArguments,
    InferenceArguments,
    get_model_path,
)

from tqdm import trange


# Constants
MAX_GENERATION_LENGTH = 256
PROMPT_JSON_FILENAME = "prompt.json"
RESULT_JSON_FILENAME = "result.json"
INTERNAL_VALIDATION_MAX_SAMPLES = 100000
TASK_T2V = "t2v"
TASK_T2I = "t2i"
TASK_X2T_IMAGE = "x2t_image"
TASK_X2T_VIDEO = "x2t_video"
TASK_IMAGE_EDIT = "image_edit"
TASK_VIDEO_EDIT = "video_edit"
GENERATION_TASKS = {
    TASK_T2V,
    TASK_T2I,
    TASK_IMAGE_EDIT,
    TASK_VIDEO_EDIT,
}
UNDERSTANDING_TASKS = {
    TASK_X2T_IMAGE,
    TASK_X2T_VIDEO,
}
TASK_DEFAULT_CONFIGS = {
    TASK_T2I: {
        "model_family": "image",
        "example_json": "config/examples/t2i_example.json",
        "save_path_prefix": "results/t2i_sample",
    },
    TASK_T2V: {
        "model_family": "video",
        "example_json": "config/examples/t2v_example.json",
        "save_path_prefix": "results/t2v_sample",
    },
    TASK_IMAGE_EDIT: {
        "model_family": "image",
        "example_json": "config/examples/image_edit_example.json",
        "save_path_prefix": "results/image_edit_sample",
    },
    TASK_VIDEO_EDIT: {
        "model_family": "video",
        "example_json": "config/examples/video_edit_example.json",
        "save_path_prefix": "results/video_edit_sample",
    },
    TASK_X2T_IMAGE: {
        "model_family": "image",
        "example_json": "config/examples/x2t_image_example.json",
        "save_path_prefix": "results/x2t_image_sample",
    },
    TASK_X2T_VIDEO: {
        "model_family": "video",
        "example_json": "config/examples/x2t_video_example.json",
        "save_path_prefix": "results/x2t_video_sample",
    },
}

def init_from_model_path_if_needed(model: Qwen2ForCausalLM, model_args: ModelArguments):
    # Always load the trained Lance checkpoint from model_path.
    path_dir = model_args.model_path
    ema_path = osp.join(path_dir, "ema.safetensors")
    model_path = osp.join(path_dir, "model.safetensors")


    model_path_ft = None
    if osp.exists(model_path):
        model_path_ft = model_path
    elif osp.exists(ema_path):
        model_path_ft = ema_path

    if model_path_ft:
        model_state_dict = load_file(model_path_ft, device="cpu")
    else:
        raise FileNotFoundError(
            f"Fine-tuning failed: No valid checkpoint ('ema.safetensors' or 'model.safetensors') found in {path_dir}"
        )

    # NOTE: position embeds are fixed sinusoidal embeddings, so we can just pop it off,
    # which makes it easier to adapt to different resolutions.
    if 'latent_pos_embed.pos_embed' in model_state_dict:
        model_state_dict.pop('latent_pos_embed.pos_embed')

    msg = model.load_state_dict(model_state_dict, strict=False)  # strict = True | False
    clean_memory(model_state_dict)

    return msg


def clean_memory(*objects):
    """Clear temporary container references and release unused GPU allocator cache."""
    for obj in objects:
        if isinstance(obj, dict):
            obj.clear()
        elif isinstance(obj, (list, set)):
            obj.clear()
    import gc
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def apply_inference_defaults(
    model_args: ModelArguments,
    data_args: DataArguments,
    inference_args: InferenceArguments,
) -> None:
    if inference_args.task not in TASK_DEFAULT_CONFIGS:
        raise ValueError(f"Unsupported inference task: {inference_args.task}")

    task_config = TASK_DEFAULT_CONFIGS[inference_args.task]
    default_inference_args = InferenceArguments()

    model_family = task_config.get("model_family", "")
    if not model_args.model_path and model_family:
        model_args.model_path = get_model_path(f"lance.{model_family}")
    if not getattr(model_args, "llm_path", ""):
        model_args.llm_path = model_args.model_path
    if not model_args.vit_path:
        model_args.vit_path = get_model_path("vit.qwen2_5_vl")

    if not data_args.val_dataset_config_file and task_config.get("example_json"):
        data_args.val_dataset_config_file = task_config["example_json"]

    if inference_args.save_path_gen == default_inference_args.save_path_gen and task_config.get("save_path_prefix"):
        inference_args.save_path_gen = task_config["save_path_prefix"]
    if inference_args.validation_max_samples == default_inference_args.validation_max_samples:
        inference_args.validation_max_samples = INTERNAL_VALIDATION_MAX_SAMPLES
    if inference_args.video_height == default_inference_args.video_height:
        inference_args.video_height = int(task_config.get("video_height", default_inference_args.video_height))
    if inference_args.video_width == default_inference_args.video_width:
        inference_args.video_width = int(task_config.get("video_width", default_inference_args.video_width))
    if inference_args.resolution == default_inference_args.resolution:
        inference_args.resolution = task_config.get("resolution", default_inference_args.resolution)
    if inference_args.text_template == default_inference_args.text_template:
        inference_args.text_template = bool(task_config.get("text_template", default_inference_args.text_template))


def save_prompt_results(prompt_data_dict, save_path_gen, logger):
    """Save validation results to a JSON file."""
    prompt_json_path = os.path.join(save_path_gen, PROMPT_JSON_FILENAME)
    with open(prompt_json_path, 'w', encoding='utf-8') as f:
        json.dump(prompt_data_dict, f, ensure_ascii=False, indent=2)


def normalize_understanding_answer(text: Optional[str]) -> str:
    """Normalize generated understanding text before exporting it."""
    if text is None:
        return ""
    return text.replace("<|im_end|>", "").strip()


def save_understanding_results(
    prompt_data_dict: dict,
    dataset_config_file: str,
    save_path_gen: str,
) -> None:
    """Save x2t results as a structured result.json file."""
    with open(dataset_config_file, "r", encoding="utf-8") as f:
        dataset_samples = json.load(f)

    result_entries = []
    for sample_key, sample in dataset_samples.items():
        interleave_array = sample.get("interleave_array", [])
        element_dtype_array = sample.get("element_dtype_array", [])
        if len(interleave_array) < 2 or not element_dtype_array:
            continue

        visual_path = interleave_array[0]
        text_payload = interleave_array[1]
        question = text_payload[1] if isinstance(text_payload, list) and len(text_payload) > 1 else ""
        modality = element_dtype_array[0]

        lookup_keys = [os.path.basename(visual_path), sample_key]
        generated_answer = ""
        for lookup_key in lookup_keys:
            if lookup_key in prompt_data_dict:
                generated_answer = prompt_data_dict[lookup_key]
                break

        result_entries.append(
            {
                modality: visual_path,
                "question": question,
                "answer": normalize_understanding_answer(generated_answer),
            }
        )

    result_json_path = os.path.join(save_path_gen, RESULT_JSON_FILENAME)
    with open(result_json_path, "w", encoding="utf-8") as f:
        json.dump(result_entries, f, ensure_ascii=False, indent=2)


def validate_on_fixed_batch(
    fsdp_model: Lance,
    vae_model: Optional[WanVideoVAE],
    tokenizer: Qwen2Tokenizer,
    val_data_cpu: dict,
    training_args: InferenceArguments,
    model_args: ModelArguments,
    inference_args: InferenceArguments,
    new_token_ids,
    image_token_id: int,
    device: int,
    save_source_video: bool = False,
    save_path_gen: str = "",
    save_path_gt: str = "",
):
    val_data = val_data_cpu.cuda(device).to_dict()
    fsdp_model = fsdp_model.to(device=device, dtype=torch.bfloat16)

    with torch.no_grad(), torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16):
        # Compute padded_latent.
        if "padded_videos" in val_data.keys():
            val_data["padded_latent"] = make_padded_latent(val_data["padded_videos"], val_data["vae_data_mode"], vae_model)

        # -------------------- Generation branch --------------------
        if inference_args.task in GENERATION_TASKS:
            params = {
                "val_packed_text_ids": val_data["packed_text_ids"],
                "val_packed_text_indexes": val_data["packed_text_indexes"],
                "val_sample_lens": val_data["sample_lens"],
                "val_packed_position_ids": val_data["packed_position_ids"],
                "val_split_lens": val_data["split_lens"],
                "val_attn_modes": val_data["attn_modes"],
                "val_sample_N_target": val_data["sample_N_target"],
                "val_packed_vae_token_indexes": val_data["packed_vae_token_indexes"],
                "timestep_shift": training_args.validation_timestep_shift,
                "num_timesteps": training_args.validation_num_timesteps,
                "val_mse_loss_indexes": val_data.get("mse_loss_indexes", None),
                "val_padded_latent": val_data["padded_latent"],
                "video_sizes": val_data["video_sizes"],
                "cfg_text_scale": model_args.cfg_text_scale,
                "cfg_interval": training_args.cfg_interval,
                "cfg_renorm_min": training_args.cfg_renorm_min,
                "cfg_renorm_type": training_args.cfg_renorm_type,
                "device": device,
                "dtype": torch.bfloat16,
                "new_token_ids": new_token_ids,
                "max_samples": training_args.validation_max_samples,
                "validation_noise_seed": training_args.validation_noise_seed,
                "apply_chat_template": training_args.apply_chat_template,
                "apply_qwen_2_5_vl_pos_emb": training_args.apply_qwen_2_5_vl_pos_emb,
                "image_token_id": image_token_id,
                "val_packed_vit_token_indexes": val_data.get("packed_vit_token_indexes", None),
                "val_packed_vit_tokens": val_data.get("packed_vit_tokens", None),
                "vit_video_grid_thw": val_data.get("vit_video_grid_thw", None),
                "vae_video_grid_thw": val_data["vae_video_grid_thw"],
                "video_grid_thw": val_data.get("video_grid_thw", None),
                "caption": val_data.get("caption", None),  # The dataset uses "caption" as the default caption field.
                "sample_task": val_data["sample_task"],
                "sample_modality": val_data["sample_modality"],
                "cfg_type": training_args.cfg_type,
                "cfg_uncond_token_id": training_args.cfg_uncond_token_id,
                "index": val_data["index"],
                "val_padded_videos": val_data["padded_videos"] if save_source_video else None,
            }
            if inference_args.use_KVcache:
                denoise_latent, captions, padded_videos, index = fsdp_model.validation_gen_KVcache(**params)
            else:
                denoise_latent, captions, padded_videos, index = fsdp_model.validation_gen(**params)

            # Decode.
            for i_val, latent in enumerate(denoise_latent):
                if inference_args.task in {TASK_IMAGE_EDIT, TASK_VIDEO_EDIT}:
                    target_latents = [latent[-1]]
                else:
                    target_latents = latent

                v_list = []
                for latent_ in target_latents:
                    v_list.append(vae_model.vae_decode([latent_])[0])

                save_item_name = f"{index:06d}" if isinstance(index, int) else index
                v_thwc = decode_video_tensor(v_list, save_path=save_path_gen, save_half=False, save_item_name=save_item_name)

                if v_thwc.shape[0] > 1:
                    prompt_data_path = f"{save_item_name}.mp4"
                else:
                    prompt_data_path = f"{save_item_name}.png"
                inference_args.prompt_data_dict[prompt_data_path] = captions[i_val]

                if save_source_video:
                    curr_padded_videos = padded_videos[i_val * 2 : (i_val + 1) * 2]
                    v_thwc_gt = decode_video_tensor(curr_padded_videos[-1:], save_path=save_path_gt, save_item_name=save_item_name)
                    del curr_padded_videos, v_thwc_gt

                del v_list, v_thwc, latent, target_latents
                clean_memory()

            del denoise_latent, captions, padded_videos, params
            clean_memory()

        elif inference_args.task in UNDERSTANDING_TASKS:
            generated_sequence_all, captions, index = fsdp_model.validation_video_to_text(
                val_packed_text_ids=val_data["packed_text_ids"],
                val_packed_text_indexes=val_data["packed_text_indexes"],
                val_packed_position_ids=val_data["packed_position_ids"],
                val_sample_N_target=val_data["sample_N_target"],
                val_split_lens=val_data["split_lens"],
                val_attn_modes=val_data["attn_modes"],
                val_sample_lens=val_data["sample_lens"],
                val_sample_type=val_data["sample_type"],
                val_packed_vit_tokens=val_data["packed_vit_tokens"],
                val_vit_video_grid_thw=val_data["vit_video_grid_thw"],
                val_ce_loss_indexes=val_data["ce_loss_indexes"],
                max_samples=training_args.validation_max_samples,
                max_length=MAX_GENERATION_LENGTH,
                device=device,
                dtype=torch.bfloat16,
                new_token_ids=new_token_ids,
                pad_token_id=tokenizer.pad_token_id,
                vocab_size=len(tokenizer),
                caption=val_data.get("caption_cn", None),
                tokenizer=tokenizer,
                apply_chat_template=training_args.apply_chat_template,
                apply_qwen_2_5_vl_pos_emb=training_args.apply_qwen_2_5_vl_pos_emb,
                do_sample=False,
                image_token_id=image_token_id,
                index=val_data["index"],
            )

            for i_val, generated_sequence in enumerate(generated_sequence_all):
                cap = tokenizer.decode(generated_sequence[:, 0])
                # inference_args.prompt_data_dict[index] = f"target_caption: {captions} /// generated_caption: {cap} "
                inference_args.prompt_data_dict[index] = f"{cap}"
                del generated_sequence

            del generated_sequence_all, captions
            clean_memory()

    del val_data
    clean_memory()


def main():
    # ========================= Env setup ==============================
    assert torch.cuda.is_available()
    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
        dist.init_process_group("nccl")
        GLOBAL_RANK = dist.get_rank()
        WORLD_SIZE = dist.get_world_size()
    else:
        GLOBAL_RANK = 0
        WORLD_SIZE = 1

    LOCAL_RANK = GLOBAL_RANK % torch.cuda.device_count()
    DEVICE = LOCAL_RANK
    torch.cuda.set_device(DEVICE)

    # ========================= Args and logger setup ==============================
    parser = HfArgumentParser((ModelArguments, DataArguments, InferenceArguments))
    model_args, data_args, inference_args = cast(
        Tuple[ModelArguments, DataArguments, InferenceArguments],
        parser.parse_args_into_dataclasses(),
    )
    training_args = inference_args

    # ========================= Load task paths and example JSONs from defaults ==============================
    apply_inference_defaults(model_args, data_args, inference_args)
    training_args.validation_noise_seed = training_args.validation_data_seed

    logger = get_logger()
    log_rank0 = print if GLOBAL_RANK == 0 else (lambda *_: None)  # Only print on rank 0.

    def log_stage(stage_name: str, start_time: float, extra: str = ""):
        elapsed = time.perf_counter() - start_time
        suffix = f" | {extra}" if extra else ""
        log_rank0(f"[startup] {stage_name} done in {elapsed:.2f}s{suffix}")

    # Set seed:
    seed = training_args.global_seed * WORLD_SIZE + GLOBAL_RANK
    set_seed(seed)

    # ========================= LLM model setup ==============================
    stage_start = time.perf_counter()
    log_rank0(f"[startup] Loading LLM config: {osp.join(model_args.model_path, 'llm_config.json')}")
    llm_config: Qwen2Config = Qwen2Config.from_json_file(osp.join(model_args.model_path, "llm_config.json"))
    log_stage("LLM config load", stage_start)

    llm_config.layer_module = model_args.layer_module
    llm_config.qk_norm = model_args.llm_qk_norm
    llm_config.qk_norm_und = model_args.llm_qk_norm_und
    llm_config.qk_norm_gen = model_args.llm_qk_norm_gen

    llm_config.tie_word_embeddings = model_args.tie_word_embeddings
    llm_config.freeze_und = training_args.freeze_und
    llm_config.apply_qwen_2_5_vl_pos_emb = training_args.apply_qwen_2_5_vl_pos_emb

    stage_start = time.perf_counter()
    log_rank0(f"[startup] Initializing LLM weights: {model_args.model_path}")
    language_model: Qwen2ForCausalLM = Qwen2ForCausalLM(llm_config)
    log_stage("LLM weight init", stage_start)

    if training_args.visual_und:
        if model_args.vit_type in ("qwen2_5_vl", "qwen_2_5_vl_original"):
            stage_start = time.perf_counter()
            log_rank0(f"[startup] Loading VIT config: {model_args.vit_path}")
            vit_config = Qwen2_5_VLVisionConfig.from_pretrained(model_args.vit_path)
            log_stage("VIT config load", stage_start)

            stage_start = time.perf_counter()
            log_rank0(f"[startup] Loading VIT weights: {osp.join(model_args.vit_path, 'vit.safetensors')}")
            vit_model = Qwen2_5_VisionTransformerPretrainedModel(vit_config)
            vit_weights = load_file(osp.join(model_args.vit_path, "vit.safetensors"))
            vit_model.load_state_dict(vit_weights, strict=True)
            log_stage("VIT weight load", stage_start)
        else:
            raise ValueError(f"Unsupported vit_type: {model_args.vit_type}")

        clean_memory(vit_weights)

    if training_args.visual_gen:
        stage_start = time.perf_counter()
        log_rank0("[startup] Initializing VAE")
        vae_model = WanVideoVAE()
        vae_config: AutoEncoderParams = deepcopy(vae_model.vae_config)
        log_stage("VAE init", stage_start)
    else:
        vae_model = None
        vae_config = None

    # Lance configuration
    config = LanceConfig(
        visual_gen=training_args.visual_gen,
        visual_und=training_args.visual_und,
        llm_config=llm_config,
        vit_config=vit_config if training_args.visual_und else None,
        vae_config=vae_config if training_args.visual_gen else None,
        latent_patch_size=model_args.latent_patch_size,
        max_num_frames=model_args.max_num_frames,
        max_latent_size=model_args.max_latent_size,
        vit_max_num_patch_per_side=model_args.vit_max_num_patch_per_side,
        connector_act=model_args.connector_act,
        interpolate_pos=model_args.interpolate_pos,
        timestep_shift=training_args.timestep_shift,
    )
    model: Lance = Lance(
        language_model=language_model,
        vit_model=vit_model if training_args.visual_und else None,
        vit_type=model_args.vit_type,
        config=config,
        training_args=training_args,
    )
    stage_start = time.perf_counter()
    log_rank0(f"[startup] Moving Lance model to GPU {DEVICE}")
    model = model.to(DEVICE)
    log_stage("Lance model move to GPU", stage_start)

    # Setup tokenizer for model:
    stage_start = time.perf_counter()
    log_rank0(f"[startup] Loading tokenizer: {model_args.model_path}")
    tokenizer: Qwen2Tokenizer = Qwen2Tokenizer.from_pretrained(model_args.model_path)

    tokenizer, new_token_ids, num_new_tokens = add_special_tokens(tokenizer)
    log_stage("tokenizer load and special token init", stage_start, extra=f"num_new_tokens={num_new_tokens}")

    # Initialize MoE before loading the checkpoint.
    if training_args.copy_init_moe:
        language_model.init_moe()

    init_from_model_path_if_needed(model, model_args)

    # Resize afterward to avoid checkpoint shape mismatches or overwritten weights.
    if num_new_tokens > 0:
        model.language_model.resize_token_embeddings(len(tokenizer))
        model.config.llm_config.vocab_size = len(tokenizer)
        model.language_model.config.vocab_size = len(tokenizer)

    if model_args.vit_type.lower() == "qwen2_5_vl":
        from common.model.hacks import hack_qwen2_5_vl_config
        language_model = hack_qwen2_5_vl_config(language_model)

    image_token_id = language_model.config.video_token_id # image_token_id # <|image_pad|>
    new_token_ids.update({"image_token_id": image_token_id})
    model.update_tokenizer(tokenizer=tokenizer)

    if model_args.tie_word_embeddings: # and training_args.finetune_from_hf is False:
        # HACK: Handle the tying logic manually.
        model.language_model.untie_lm_head() # NOTE: untied lm head weights
        model.language_model.copy_new_token_rows_to_lm_head(num_new_tokens) # NOTE: copy the new token rows into lm_head

        # Make sure this stays False.
        model_args.tie_word_embeddings = False
        llm_config.tie_word_embeddings = False
    else: # HACK!!!
        assert model.language_model.get_input_embeddings().weight.data.data_ptr() != model.language_model.get_output_embeddings().weight.data.data_ptr(), 'tie_word_embeddings conflict'

    model = model.to(device=DEVICE, dtype=torch.bfloat16)
    model.eval()
    if vae_model is not None and hasattr(vae_model, "eval"):
        vae_model.eval()

    # Setup packed dataloader
    stage_start = time.perf_counter()
    log_rank0(f"[startup] Loading dataset config and validation set: {data_args.val_dataset_config_file}")
    dataset_config = DataConfig.from_yaml(data_args.val_dataset_config_file)

    # NOTE: This block performs in-place assignments. ⚠️
    if training_args.visual_und:
        dataset_config.vit_patch_size = model_args.vit_patch_size
        dataset_config.vit_patch_size_temporal = model_args.vit_patch_size_temporal # TODO: fix
        dataset_config.vit_max_num_patch_per_side = model_args.vit_max_num_patch_per_side
        # dataset_config.vit_downsample = vit_downsample # NOTE: need to update !
    if training_args.visual_gen:
        assert len(model_args.latent_patch_size) == 3, "len(latent_patch_size) must be 3"
        vae_downsample = tuple_mul(
            model_args.latent_patch_size, (vae_config.downsample_temporal, vae_config.downsample_spatial, vae_config.downsample_spatial)
        )  # NOTE: This already includes patch_size.
        dataset_config.latent_patch_size = model_args.latent_patch_size
        dataset_config.vae_downsample = vae_downsample  # NOTE: update !
        dataset_config.max_latent_size = model_args.max_latent_size  # NOTE: update!
        dataset_config.max_num_frames = model_args.max_num_frames  # NOTE: update!

    # Fix: share dropout settings.
    dataset_config.text_cond_dropout_prob = model_args.text_cond_dropout_prob
    dataset_config.vae_cond_dropout_prob = model_args.vae_cond_dropout_prob
    dataset_config.vit_cond_dropout_prob = model_args.vit_cond_dropout_prob

    # Load inference parameters.
    dataset_config.num_frames = inference_args.num_frames
    dataset_config.H = inference_args.video_height
    dataset_config.W = inference_args.video_width
    dataset_config.task = inference_args.task
    dataset_config.resolution = inference_args.resolution
    dataset_config.text_template = inference_args.text_template
    val_dataset = ValidationDataset(
        jsonl_path= data_args.val_dataset_config_file,
        tokenizer=tokenizer,
        data_args=data_args,
        model_args=model_args,
        training_args=training_args,
        new_token_ids=new_token_ids,
        dataset_config=dataset_config,
        local_rank=GLOBAL_RANK,  # global rank, not local rank
        world_size=WORLD_SIZE,
    )

    val_loader = DataLoader(
        val_dataset,
        batch_size=1,
        num_workers=0,
        pin_memory=True,
        collate_fn=simple_custom_collate,     # Top-level function
        drop_last=True,
        prefetch_factor=None,
        persistent_workers=False,
        multiprocessing_context=None,
    )
    log_stage("validation set and DataLoader init", stage_start, extra=f"dataset_size={len(val_dataset)}")

    # Prepare the validation data loader iterator.
    val_loader_iter = iter(val_loader)

    # Initialize a local dictionary to avoid accumulating stale data.
    if not hasattr(inference_args, "prompt_data_dict"):
        inference_args.prompt_data_dict = {}

    if not os.path.exists(inference_args.save_path_gen):
        os.makedirs(inference_args.save_path_gen)

    for epoch in trange(len(val_loader), desc="Validating", unit="batch", leave=True, ncols=80, disable=(GLOBAL_RANK != 0)):
        try:
            val_data_cpu = next(val_loader_iter)
        except StopIteration:
            break

        validate_on_fixed_batch(
            fsdp_model=model,
            vae_model=vae_model,
            tokenizer=tokenizer,
            val_data_cpu=val_data_cpu,
            training_args=training_args,
            model_args=model_args,
            inference_args=inference_args,
            new_token_ids=new_token_ids,
            image_token_id=image_token_id,
            device=DEVICE,
            save_source_video=False, # Whether to save the GT video
            save_path_gen=inference_args.save_path_gen, # Generated video path
            save_path_gt="", # GT video path
        )
        del val_data_cpu
        clean_memory()

    # Final gather after all generation loops
    if dist.is_initialized():
        dist.barrier()
        gathered = [None for _ in range(dist.get_world_size())]
        dist.all_gather_object(gathered, inference_args.prompt_data_dict)

        if GLOBAL_RANK == 0:
            merged = {}
            for d in gathered:
                merged.update(d)
            inference_args.prompt_data_dict = merged
            save_prompt_results(inference_args.prompt_data_dict, inference_args.save_path_gen, logger)
            if inference_args.task in UNDERSTANDING_TASKS:
                save_understanding_results(
                    prompt_data_dict=inference_args.prompt_data_dict,
                    dataset_config_file=data_args.val_dataset_config_file,
                    save_path_gen=inference_args.save_path_gen,
                )

    elif GLOBAL_RANK == 0:
        save_prompt_results(inference_args.prompt_data_dict, inference_args.save_path_gen, logger)
        if inference_args.task in UNDERSTANDING_TASKS:
            save_understanding_results(
                prompt_data_dict=inference_args.prompt_data_dict,
                dataset_config_file=data_args.val_dataset_config_file,
                save_path_gen=inference_args.save_path_gen,
            )

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


if __name__ == "__main__":
    main()