multi_view/test.py

import torch
from PIL import Image
from einops import rearrange
import numpy as np
from typing import Optional, List, Tuple, Callable
import json
import math

from tqdm import tqdm
import os
import argparse
import torch.distributed as dist
import torch.nn.functional as F
from diffsynth.models import ModelManager
from diffsynth.models.utils import load_state_dict
import torch
from PIL import Image
from diffsynth import save_video
from diffsynth.pipelines.wan_video_new import WanVideoPipeline, ModelConfig
from modelscope import dataset_snapshot_download
import yaml
import torch, os, json
from diffsynth.pipelines.wan_video_new import WanVideoPipeline, ModelConfig
from diffsynth.trainers.utils import DiffusionTrainingModule, ModelLogger, launch_training_task, wan_parser
from datasets.videodataset import MulltiShot_MultiView_Dataset
from PIL import Image, ImageOps


def test_video(args):
    checkpoint_path = os.path.join(args.output_path, args.visual_log_project_name, f"checkpoint-step-{args.infer_step}-epoch-{args.epoch_id}", "weights.safetensors")
    output_path =  os.path.join("./output", args.visual_log_project_name)
    os.makedirs(f"{output_path}/ref_images", exist_ok=True)
    os.makedirs(f"{output_path}/video", exist_ok=True)
    print(checkpoint_path)
    pipe = WanVideoPipeline.from_pretrained(
        torch_dtype=torch.bfloat16,
        device="cuda",
        model_configs=[
            ModelConfig(path=os.path.join(args.local_model_path, "Wan2.2-TI2V-5B/models_t5_umt5-xxl-enc-bf16.pth"),  offload_device="cuda"),
            ModelConfig(path=os.path.join(args.local_model_path, "Wan2.2-TI2V-5B/Wan2.2_VAE.pth"), offload_device="cuda"),
            ModelConfig(path=checkpoint_path, offload_device="cuda"),
        ],
        redirect_common_files = False
    )
    pipe.enable_vram_management()

    with open(args.train_yaml, "r", encoding="utf-8") as f:
        conf_info = yaml.safe_load(f)   # 用 safe_load 更安全
    dataset = MulltiShot_MultiView_Dataset(
        dataset_base_path=args.dataset_base_path,
        resolution=(args.height, args.width),
        ref_num=args.ref_num,
        training=False
    )

    log_file_name = "output_log.txt"
    import pdb; pdb.set_trace()
    # v_indexs = [0, 10, 30, 50, 70, 100, 130, 150, 180, 200]
    v_indexs = [0, 5, 15, 20]
    with open(os.path.join(output_path, log_file_name), "w") as f:
        for v_index in v_indexs: 
            metadata = dataset[v_index]
            video, _ = pipe(
                args = args,
                prompt = [metadata["single_caption"]], #prompt, #"两只狗在擂台上打拳击", ### 手动变成batch  = 1  的list
                ref_images = [metadata["ref_images"]],
                negative_prompt=["色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"],
                seed=42, tiled=True,
                height=args.height, width=args.width,
                num_frames=args.num_frames,
                cfg_scale_face = 5.,
                num_ref_images = metadata["ref_num"]
            )
            for r_index, img in enumerate(metadata["ref_images"]):
                img.save(f"{output_path}/ref_images/{v_index}-{r_index}.png")
            
            save_video(video, f"{output_path}/video/{v_index}.mp4", fps=15, quality=10)
               
            f.write(f"{metadata['single_caption']}\n")
    
def specify_video(args):
    def process_ref_images(ref_images, height, width):
        ref_images_new = []
        for ref_image in ref_images:
            h = height
            w = width
            ref_image = ref_image.convert("RGB")
            # Calculate the required size to keep aspect ratio and fill the rest with padding.
            img_ratio = ref_image.width / ref_image.height
            target_ratio = w / h
            
            if img_ratio > target_ratio:  # Image is wider than target
                new_width = w
                new_height = int(new_width / img_ratio)
            else:  # Image is taller than target
                new_height = h
                new_width = int(new_height * img_ratio)
            
            # img = img.resize((new_width, new_height), Image.ANTIALIAS)
            ref_image = ref_image.resize((new_width, new_height), Image.Resampling.LANCZOS)

            # Create a new image with the target size and place the resized image in the center
            delta_w = w - ref_image.size[0]
            delta_h = h - ref_image.size[1]
            padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2))
            new_img = ImageOps.expand(ref_image, padding, fill=(255, 255, 255))
            ref_images_new.append(new_img)
        return ref_images_new

    checkpoint_path = os.path.join(args.output_path, args.visual_log_project_name, f"checkpoint-step-{args.infer_step}-epoch-{args.epoch_id}", "weights.safetensors")
    output_path =  os.path.join("./output", args.visual_log_project_name)
    os.makedirs(f"{output_path}/ref_images", exist_ok=True)
    os.makedirs(f"{output_path}/video", exist_ok=True)
    print(checkpoint_path)
    pipe = WanVideoPipeline.from_pretrained(
        torch_dtype=torch.bfloat16,
        device="cuda",
        model_configs=[
            ModelConfig(path=os.path.join(args.local_model_path, "Wan2.2-TI2V-5B/models_t5_umt5-xxl-enc-bf16.pth"),  offload_device="cuda"),
            ModelConfig(path=os.path.join(args.local_model_path, "Wan2.2-TI2V-5B/Wan2.2_VAE.pth"), offload_device="cuda"),
            ModelConfig(path=checkpoint_path, offload_device="cuda"),
        ],
        redirect_common_files = False
    )
    pipe.enable_vram_management()
    ref_images=[
                    Image.open("/root/paddlejob/workspace/qizipeng/baidu/personal-code/Multi-view/multi_view/datasets/cl_0.png"),
                    # Image.open("/root/paddlejob/workspace/qizipeng/baidu/personal-code/Multi-view/multi_view/datasets/cl_2.png"),
                    # Image.open("/root/paddlejob/workspace/qizipeng/baidu/personal-code/Multi-view/multi_view/datasets/cl_3.png"),
                    ]
    
    ref_images = process_ref_images(ref_images, args.height, args.width)


    video, _ = pipe(
            args = args,
            prompt = ["An elderly man with short gray hair and glasses stands in a softly lit indoor hallway. The shot begins with a frontal view of his face, his expression calm and attentive as he looks straight ahead. Then, he turns his head to his right, responding to someone standing beside him. His gaze shifts fully toward the other person as his expression becomes more engaged. The movement continues until he reaches a complete side profile, fully turning his face toward the person he is interacting with. Smooth and natural head rotation, warm indoor lighting."], #prompt, #"两只狗在擂台上打拳击", ### 手动变成batch  = 1  的list
            ref_images = [ref_images],
            negative_prompt=["色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"],
            seed=42, tiled=True,
            height=args.height, width=args.width,
            num_frames=args.num_frames,
            cfg_scale_face = 5.,
            num_ref_images = len(ref_images)
        )
    save_video(video, f"{output_path}/video/cl.mp4", fps=15, quality=10)
            

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="长视频分镜头连续生成脚本")
    
    # --- 核心路径参数 ---
    parser = wan_parser()
    args = parser.parse_args()

    args, unknown = parser.parse_known_args()
    print("❗ Unknown arguments:", unknown)
    ### 执行过pip install -e . 的话diffsynth 里的东西修改后要重新安装
    # import pdb; pdb.set_trace()
    ###下面是解析train.yaml里的内容
    with open(args.train_yaml, "r", encoding="utf-8") as f:
        conf_info = yaml.safe_load(f)   # 用 safe_load 更安全
    print(conf_info)
    args.dataset_base_path  = conf_info["dataset_args"]["base_path"]
    args.max_checkpoints_to_keep = conf_info["train_args"]["max_checkpoints_to_keep"]
    args.resume_from_checkpoint =  conf_info["train_args"]["resume_from_checkpoint"]
    args.visual_log_project_name = conf_info["train_args"]["visual_log_project_name"]
    args.seed = conf_info["train_args"]["seed"]
    args.output_path = conf_info["train_args"]["output_path"]
    args.save_steps = conf_info["train_args"]["save_steps"]
    args.save_epoches = conf_info["train_args"]["save_epoches"]
    args.batch_size = conf_info["train_args"]["batch_size"]
    args.local_model_path = conf_info["train_args"]["local_model_path"]
    args.height = conf_info["dataset_args"]["height"] 
    args.width = conf_info["dataset_args"]["width"]
    args.num_frames = conf_info["dataset_args"]["num_frames"]
    args.ref_num = conf_info["dataset_args"]["ref_num"]
    args.infer_step = conf_info["infer_args"]["infer_step"]
    args.epoch_id = conf_info["infer_args"]["epoch_id"]
    args.split_rope = conf_info["train_args"]["split_rope"]
    args.split1 = conf_info["train_args"]["split1"]
    args.split2 = conf_info["train_args"]["split2"]
    args.split3 = conf_info["train_args"]["split3"]
        
    test_video(args)
    # specify_video(args)