benchmarks/image_gen/DPG/sample_DPG.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
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"

import os.path as osp
from copy import deepcopy
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 PIL import Image
from torchvision.utils import make_grid
import numpy as np
from tqdm import trange

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
from data.datasets_custom import ValidationDataset
from config.config_factory import ModelArguments, DataArguments, TrainingArguments, EvaluationArguments, get_model_path


def init_from_vlm_if_needed(model: Qwen2ForCausalLM, model_args: ModelArguments, log_rank0):
    # NOTE: 初始化加载VLM模型走这里
    def load_safetensors_state_dict(folder_path):
        # 只选取safetensors文件，按文件名排序保证顺序
        safetensor_files = sorted(
            f for f in os.listdir(folder_path) if f.endswith(".safetensors")
        )
        state_dict = {}
        for filename in safetensor_files:
            file_path = osp.join(folder_path, filename)
            state_dict.update(load_file(file_path))
        return state_dict

    state_dict = load_safetensors_state_dict(model_args.llm_path)

    # 参数名的更改以适配Lance的参数名
    for k in list(state_dict.keys()):
        if "visual" in k:  # ViT and connector
            state_dict[k.replace("visual", "vit_model")] = state_dict.pop(k)
        else:
            # 添加language_model前缀
            state_dict["language_model." + k] = state_dict.pop(k)

    result = model.load_state_dict(state_dict, strict=False)

    clean_memory(state_dict)


def init_from_model_path_if_needed(model: Qwen2ForCausalLM, model_args: ModelArguments):
    # 统一从 model_path 加载训练好的 Lance checkpoint。
    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)

    clean_memory(model_state_dict)

    return msg


def clean_memory(*objects):
    """清理内存并释放 GPU 缓存"""
    for obj in objects:
        del obj
    import gc
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def decode_video_tensor_for_dpg(v_list):
    """
    专门为 DPG 解码视频张量，保持原有的保存格式
    """
    N_target = len(v_list)
    if N_target != 1:
        from einops import rearrange
        padded_videos_latent = [v.permute(1, 0, 2, 3) for v in v_list]
        v_tc_hw = rearrange(padded_videos_latent, "n t c h w -> t c h (n w)")
    else:
        v_tc_hw = v_list[0].permute(1, 0, 2, 3)

    v_tc_hw = v_tc_hw.float().clip(-1, 1).mul_(0.5).add_(0.5).mul_(255).round().clamp(0, 255).to(torch.uint8)
    return v_tc_hw


def resolve_dpg_paths(
    model_args: ModelArguments,
    data_args: DataArguments,
) -> None:
    if not model_args.model_path:
        raise ValueError("DPG requires --model_path to be provided explicitly.")

    if not 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:
        data_args.val_dataset_config_file = get_model_path("dpg.data")


def validate_on_fixed_batch(
    fsdp_model: Lance,
    vae_model: Optional[WanVideoVAE],
    tokenizer: Qwen2Tokenizer,
    val_data_cpu: dict,
    training_args: TrainingArguments,
    model_args: ModelArguments,
    data_args: DataArguments,
    inference_args: EvaluationArguments,
    curr_step: int,
    logger,
    new_token_ids,
    image_token_id: int,
    device: int,
    save_source_video: bool = False,
    save_path_gen: str = "",
    save_path_gt: str = "",
    sample_num_per_prompt: int = 1,
):
    """
    验证逻辑，保持与原文件相同的保存格式
    """
    # 检查是否初始化了分布式环境
    if dist.is_initialized():
        is_rank0 = (dist.get_rank() == 0)
    else:
        is_rank0 = True
    
    log_rank0 = logger.info if is_rank0 else (lambda *_: None)
    val_data = val_data_cpu.cuda(device).to_dict()

    with torch.no_grad(), torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16):
        # 计算 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)

        # -------------------- GEN 分支 --------------------
        tensor_list_for_grid = []
        loop_iterator = trange(sample_num_per_prompt) if is_rank0 else range(sample_num_per_prompt)

        # 支持断点重新生成
        save_name = f"{save_path_gen}/{val_data['index']}.png"
        if os.path.exists(save_name):
            return None

        for sample_num_per_prompt_index in loop_iterator:
            # 采样生成（保持原参数）
            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 + sample_num_per_prompt_index,
                "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),
                "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 training_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)

            # 解码 + 保存
            for i_val, latent in enumerate(denoise_latent):
                v_list = [vae_model.vae_decode([latent_])[0] for latent_ in latent]

                # 保持与原文件相同的保存格式
                v_thwc = decode_video_tensor_for_dpg(v_list)

                # 直接取第0帧
                if v_thwc.shape[0] == 1:
                    tensor_list_for_grid.append(v_thwc.squeeze(0).cpu())
                else:
                    raise NotImplementedError("需要保存图像")

    # 保持原有的保存格式
    grid_tensor = make_grid(tensor_list_for_grid, nrow=int(np.sqrt(sample_num_per_prompt)), padding=0, pad_value=255)
    grid_numpy = grid_tensor.permute(1, 2, 0).numpy()
    Image.fromarray(grid_numpy).save(save_name)


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, EvaluationArguments))
    model_args, data_args, inference_args = cast(
        Tuple[ModelArguments, DataArguments, EvaluationArguments],
        parser.parse_args_into_dataclasses(),
    )
    training_args = inference_args

    # ========================= DPG 路径解析 ==============================
    resolve_dpg_paths(model_args, data_args)

    # NOTE validation_noise_seed 与 validation_data_seed 相同
    training_args.validation_noise_seed = inference_args.evaluation_seed
    training_args.validation_data_seed = inference_args.evaluation_seed
    logger = get_logger()
    log_rank0 = print if GLOBAL_RANK == 0 else (lambda *_: None)

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

    # ========================= LLM model setup ==============================
    llm_config: Qwen2Config = Qwen2Config.from_json_file(osp.join(model_args.model_path, "llm_config.json"))

    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

    language_model: Qwen2ForCausalLM = Qwen2ForCausalLM(llm_config)

    if training_args.visual_und:
        if model_args.vit_type in ("qwen2_5_vl", "qwen_2_5_vl_original"):
            vit_config = Qwen2_5_VLVisionConfig.from_pretrained(model_args.vit_path)
            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)
        else:
            raise ValueError(f"Unsupported vit_type: {model_args.vit_type}")

        clean_memory(vit_weights)

    if training_args.visual_gen:
        vae_model = WanVideoVAE()
        vae_config: AutoEncoderParams = deepcopy(vae_model.vae_config)
    else:
        vae_model = None
        vae_config = None

    # Lance的配置
    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,
    )
    model = model.to(DEVICE)

    # Setup tokenizer for model:
    tokenizer: Qwen2Tokenizer = Qwen2Tokenizer.from_pretrained(model_args.model_path)

    tokenizer, new_token_ids, num_new_tokens = add_special_tokens(tokenizer)

    # 在加载ckpt前，初始化moe
    if training_args.copy_init_moe:
        language_model.init_moe()

    init_from_model_path_if_needed(model, model_args)

    # 现在再 resize
    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
    new_token_ids.update({"image_token_id": image_token_id})
    model.update_tokenizer(tokenizer=tokenizer)

    if model_args.tie_word_embeddings:
        model.language_model.untie_lm_head()
        model.language_model.copy_new_token_rows_to_lm_head(num_new_tokens)

        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_world_embeddings 冲突'

    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 - 直接初始化简单的 DataConfig 对象
    dataset_config = DataConfig(grouped_datasets={})

    # 配置基本参数
    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

    # 配置 VIT 相关参数
    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
        dataset_config.vit_max_num_patch_per_side = model_args.vit_max_num_patch_per_side

    # 配置 VAE 相关参数
    if training_args.visual_gen and vae_config:
        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)
        )
        dataset_config.latent_patch_size = model_args.latent_patch_size
        dataset_config.vae_downsample = vae_downsample
        dataset_config.max_latent_size = model_args.max_latent_size
        dataset_config.max_num_frames = model_args.max_num_frames

    # fix: 共享dropout
    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

    # 创建数据集
    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,
        world_size=WORLD_SIZE,
    )

    val_loader = DataLoader(
            val_dataset,
            batch_size=1,
            num_workers=0,
            pin_memory=True,
            collate_fn=simple_custom_collate,
            drop_last=True,
            prefetch_factor=None,
            persistent_workers=False,
            multiprocessing_context=None,
        )

    val_loader_iter = iter(val_loader)

    if not os.path.exists(inference_args.save_path_gen):
        os.makedirs(inference_args.save_path_gen, exist_ok=True)

    # 主循环
    from tqdm import tqdm
    import time
    from datetime import datetime, timedelta

    total_batches = len(val_loader)
    pbar = tqdm(total=total_batches, desc="Validating", unit="batch", leave=True, ncols=120, disable=(GLOBAL_RANK != 0))
    start_time = time.time()

    for i in range(total_batches):
        val_data_cpu = next(val_loader_iter)

        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,
            data_args=data_args,
            inference_args=inference_args,
            curr_step=0,
            logger=logger,
            new_token_ids=new_token_ids,
            image_token_id=image_token_id,
            device=DEVICE,
            save_source_video=False,
            save_path_gen=inference_args.save_path_gen,
            save_path_gt="",
            sample_num_per_prompt=inference_args.sample_num_per_prompt,
        )

        if GLOBAL_RANK == 0:
            elapsed = time.time() - start_time
            avg_time = elapsed / (i + 1)
            eta_seconds = avg_time * (total_batches - i - 1)
            expected_finish = datetime.now() + timedelta(seconds=eta_seconds)
            finish_str = expected_finish.strftime('%Y-%m-%d %H:%M:%S')

            pbar.set_postfix_str(f"ETA: {timedelta(seconds=int(eta_seconds))} | Finish: {finish_str}")
            pbar.update(1)

    if GLOBAL_RANK == 0:
        pbar.close()

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


if __name__ == "__main__":
    main()