wehub/bytedance--lance
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
9af80cc3e43f1d2456a67ccfc0d608dd6b2cd87c
bytedance--lance/inference_lance.py
T
wehub-resource-sync e710618da0 chore: import upstream snapshot with attribution
2026-07-13 13:16:54 +08:00

689 lines
29 KiB
Python
Raw Blame History

# 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.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_I2V = "i2v"
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_I2V,
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_I2V: {
"model_family": "video",
"example_json": "config/examples/i2v_example.json",
"save_path_prefix": "results/i2v_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()
torch.cuda.ipc_collect()
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:
save_fps = int(val_data.get("save_fps", 12))
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_I2V, 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, save_fps=save_fps)
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, save_fps=save_fps)
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:
params = {
"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"],
}
if inference_args.use_KVcache:
generated_sequence_all, captions, index = fsdp_model.validation_und_KVcache(**params)
else:
generated_sequence_all, captions, index = fsdp_model.validation_video_to_text(**params)
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, params
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.load_from_lance_checkpoint 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:
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
dataset_config.enhance_prompt = inference_args.enhance_prompt
if inference_args.enhance_prompt:
if inference_args.task not in {TASK_T2V, TASK_I2V}:
log_rank0("[startup] enhance_prompt is enabled but only applies to t2v and i2v; skipping prompt rewrite for this task.")
else:
log_rank0(f"[startup] enhance_prompt is enabled for {inference_args.task} prompts. Configure API_KEY in common/utils/caption_rewrite.py.")
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()
Reference in New Issue View Git Blame Copy Permalink