app.py

# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import argparse
import logging
import os
import sys
import json
import warnings
from datetime import datetime
import gradio as gr
import spaces
warnings.filterwarnings('ignore')

import random
import math
import torch
import torch.distributed as dist
from PIL import Image
from huggingface_hub import snapshot_download

# 检查 GPU 可用性（参考 Meigen-MultiTalk）
# 参考: https://huggingface.co/spaces/fffiloni/Meigen-MultiTalk/blob/main/app.py
is_shared_ui = True if os.environ.get('SPACE_ID', '').startswith('C4G-HKUST/AnyTalker') else False
is_gpu_associated = torch.cuda.is_available()

if is_gpu_associated:
    try:
        num_gpus = torch.cuda.device_count()
        print(f"GPU AVAILABLE: {num_gpus} GPU(s)")
        if num_gpus > 0:
            gpu_name = torch.cuda.get_device_name(0)
            print(f"GPU Name: {gpu_name}")
    except Exception as e:
        print(f"GPU detection error: {e}")
        is_gpu_associated = False
else:
    print("No CUDA-compatible GPU found. Will use CPU (slower).")

# 导入 AnyTalker 相关的模块
import wan
from wan.configs import SIZE_CONFIGS, SUPPORTED_SIZES, WAN_CONFIGS, MAX_AREA_CONFIGS
from wan.utils.utils import cache_video, str2bool
from wan.utils.infer_utils import calculate_frame_num_from_audio
from utils.get_face_bbox import FaceInference

# 获取 HuggingFace 预加载模型的路径
def get_model_path(repo_id):
    """获取 HF 模型的本地缓存路径"""
    try:
        return snapshot_download(repo_id=repo_id, local_files_only=True)
    except:
        # 如果缓存不存在，则下载
        return snapshot_download(repo_id=repo_id)

# 模型路径（优先使用 HF cache，否则使用本地 checkpoints）
def get_ckpt_dir():
    try:
        return get_model_path("alibaba-pai/Wan2.1-Fun-V1.1-1.3B-InP")
    except:
        return "./checkpoints/Wan2.1-Fun-1.3B-Inp"

def get_anytalker_path():
    try:
        path = get_model_path("zzz66/AnyTalker-1.3B")
        return os.path.join(path, "1_3B-single-v1.pth")
    except:
        return "./checkpoints/AnyTalker/1_3B-single-v1.pth"

def get_dit_config_path():
    try:
        path = get_model_path("zzz66/AnyTalker-1.3B")
        return os.path.join(path, "config_af2v_1_3B.json")
    except:
        return "./checkpoints/AnyTalker/config_af2v_1_3B.json"


def str2bool(v):
    """字符串转布尔值工具函数"""
    if isinstance(v, bool):
        return v
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')

def _validate_args(args):
    # 设置模型路径（从 HF cache 或本地 checkpoints）
    if args.ckpt_dir is None:
        args.ckpt_dir = get_ckpt_dir()
    if args.post_trained_checkpoint_path is None:
        args.post_trained_checkpoint_path = get_anytalker_path()
    if args.dit_config is None:
        args.dit_config = get_dit_config_path()
    
    # Basic check
    assert args.ckpt_dir is not None, "Please specify the checkpoint directory."
    assert args.task in WAN_CONFIGS, f"Unsupport task: {args.task}"
    
    # The default sampling steps are 40 for image-to-video tasks and 50 for text-to-video tasks.
    if args.sample_steps is None:
        if any(key in args.task for key in ["i2v", "a2v"]):
            args.sample_steps = 40
        else:
            args.sample_steps = 50

    if args.sample_shift is None:
        args.sample_shift = 5.0
        if any(key in args.task for key in ["i2v", "a2v"]) and args.size in ["832*480", "480*832"]:
            args.sample_shift = 3.0

    # For a2v tasks, frame_num will be determined by audio length if not specified
    if args.frame_num is None:
        args.frame_num = None

    args.base_seed = args.base_seed if args.base_seed >= 0 else random.randint(0, sys.maxsize)
    
    # Size check
    assert args.size in SUPPORTED_SIZES[args.task], f"Unsupport size {args.size} for task {args.task}, supported sizes are: {', '.join(SUPPORTED_SIZES[args.task])}"

def _parse_args():
    parser = argparse.ArgumentParser(
        description="Generate a image or video from a text prompt or image using Wan"
    )
    parser.add_argument(
        "--task",
        type=str,
        default="a2v-1.3B",
        # choices=list(WAN_CONFIGS.keys()),
        help="The task to run.")
    parser.add_argument(
        "--size",
        type=str,
        default="832*480",
        # choices=list(SIZE_CONFIGS.keys()),
        help="The area (width*height) of the generated video. For the I2V task, the aspect ratio of the output video will follow that of the input image."
    )
    parser.add_argument(
        "--frame_num",
        type=int,
        default=None,
        help="How many frames to sample from a image or video. The number should be 4n+1. For a2v tasks, if not specified, frame number will be automatically determined based on audio length."
    )
    parser.add_argument(
        "--ckpt_dir",
        type=str,
        default=None,  # 将在运行时通过 get_ckpt_dir() 获取
        help="The path to the checkpoint directory.")
    parser.add_argument(
        "--post_trained_checkpoint_path",
        type=str,
        default=None,  # 将在运行时通过 get_anytalker_path() 获取
        help="The path to the posted-trained checkpoint file.")
    parser.add_argument(
        "--offload_model",
        type=str2bool,
        default=True,
        help="Whether to offload the model to CPU after each model forward, reducing GPU memory usage."
    )
    parser.add_argument(
        "--use_half",
        type=str2bool,
        default=True,
        help="Whether to use half precision for model inference, reducing GPU memory usage."
    )
    parser.add_argument(
        "--ulysses_size",
        type=int,
        default=1,
        help="The size of the ulysses parallelism in DiT.")
    parser.add_argument(
        "--ring_size",
        type=int,
        default=1,
        help="The size of the ring attention parallelism in DiT.")
    parser.add_argument(
        "--t5_fsdp",
        action="store_true",
        default=False,
        help="Whether to use FSDP for T5.")
    parser.add_argument(
        "--t5_cpu",
        action="store_true",
        default=False,
        help="Whether to place T5 model on CPU.")
    parser.add_argument(
        "--dit_fsdp",
        action="store_true",
        default=False,
        help="Whether to use FSDP for DiT.")
    parser.add_argument(
        "--save_file",
        type=str,
        default=None,
        help="The file to save the generated image or video to.")
    parser.add_argument(
        "--save_dir",
        type=str,
        default=None,
        help="The directory to save the generated image or video to.")  
    parser.add_argument(
        "--prompt",
        type=str,
        default=None,
        help="The prompt to generate the image or video from.")
    parser.add_argument(
        "--use_prompt_extend",
        action="store_true",
        default=False,
        help="Whether to use prompt extend.")
    parser.add_argument(
        "--prompt_extend_method",
        type=str,
        default="local_qwen",
        choices=["dashscope", "local_qwen"],
        help="The prompt extend method to use.")
    parser.add_argument(
        "--prompt_extend_model",
        type=str,
        default=None,
        help="The prompt extend model to use.")
    parser.add_argument(
        "--prompt_extend_target_lang",
        type=str,
        default="zh",
        choices=["zh", "en"],
        help="The target language of prompt extend.")
    parser.add_argument(
        "--base_seed",
        type=int,
        default=44,
        help="The seed to use for generating the image or video.")
    parser.add_argument(
        "--image",
        type=str,
        default=None,
        help="The image to generate the video from.")
    parser.add_argument(
        "--audio",
        type=str,
        default=None,
        help="The audio to generate the video from.")
    parser.add_argument(
        "--sample_solver",
        type=str,
        default='unipc',
        choices=['unipc', 'dpm++'],
        help="The solver used to sample.")
    parser.add_argument(
        "--sample_steps", type=int, default=None, help="The sampling steps.")
    parser.add_argument(
        "--sample_shift",
        type=float,
        default=None,
        help="Sampling shift factor for flow matching schedulers.")
    parser.add_argument(
        "--sample_guide_scale",
        type=float,
        default=4.5,
        help="Classifier free guidance scale.")
    parser.add_argument(
        "--cfg_zero",
        action="store_true",
        default=False,
        help="Whether to use adaptive CFG-Zero guidance instead of fixed guidance scale.")
    parser.add_argument(
        "--zero_init_steps",
        type=int,
        default=0,
        help="Number of initial steps to use zero guidance when using cfg_zero.")
    parser.add_argument(
        "--sample_fps",
        type=int,
        default=24,
        help="The frames per second (FPS) of the generated video. Overrides the default value from the config.")
    parser.add_argument(
        "--batch_gen_json",
        type=str,
        default=None,
        help="Path to prompts.json file for batch processing. Images and outputs are in the same directory.")
    parser.add_argument(
        "--batch_output",
        type=str,
        default=None,
        help="Directory to save generated videos when using batch processing. If not specified, defaults to the json filename (without extension) in the same directory.")
    parser.add_argument(
        "--dit_config",
        type=str,
        default=None,  # 将在运行时通过 get_dit_config_path() 获取
        help="The path to the dit config file.")
    parser.add_argument(
        "--det_thresh",
        type=float,
        default=0.10,
        help="Threshold for InsightFace face detection.")
    parser.add_argument(
        "--mode",
        type=str,
        default="pad",
        choices=["pad", "concat"],
        help="The mode to use for audio processing.")
    parser.add_argument(
        "--audio_save_dir",
        type=str,
        default='save_audio/gradio',
        help="The path to save the audio embedding.")
    args = parser.parse_args()
    _validate_args(args)
    return args


def _init_logging(rank):
    # logging
    if rank == 0:
        # set format
        logging.basicConfig(
            level=logging.INFO,
            format="[%(asctime)s] %(levelname)s: %(message)s",
            handlers=[logging.StreamHandler(stream=sys.stdout)])
    else:
        logging.basicConfig(level=logging.ERROR)

def run_graio_demo(args):
    # 设置 Gradio 临时文件目录
    gradio_temp_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'gradio_temp')
    os.makedirs(gradio_temp_dir, exist_ok=True)
    os.environ['GRADIO_TEMP_DIR'] = gradio_temp_dir
    
    rank = int(os.getenv("RANK", 0))
    world_size = int(os.getenv("WORLD_SIZE", 1))
    local_rank = int(os.getenv("LOCAL_RANK", 0))
    device = local_rank
    _init_logging(rank)

    if args.offload_model is None:
        args.offload_model = False if world_size > 1 else True
        logging.info(
            f"offload_model is not specified, set to {args.offload_model}.")
    if world_size > 1:
        torch.cuda.set_device(local_rank)
        dist.init_process_group(
            backend="nccl",
            init_method="env://",
            rank=rank,
            world_size=world_size)
    else:
        assert not (
            args.t5_fsdp or args.dit_fsdp
        ), f"t5_fsdp and dit_fsdp are not supported in non-distributed environments."
        assert not (
            args.ulysses_size > 1 or args.ring_size > 1
        ), f"context parallel are not supported in non-distributed environments."

    if args.ulysses_size > 1 or args.ring_size > 1:
        assert args.ulysses_size * args.ring_size == world_size, f"The number of ulysses_size and ring_size should be equal to the world size."
        from xfuser.core.distributed import (
            init_distributed_environment,
            initialize_model_parallel,
        )
        init_distributed_environment(
            rank=dist.get_rank(), world_size=dist.get_world_size())
        
        initialize_model_parallel(
            sequence_parallel_degree=dist.get_world_size(),
            ring_degree=args.ring_size,
            ulysses_degree=args.ulysses_size,
        )

    # 加载配置
    cfg = WAN_CONFIGS[args.task]
    if args.ulysses_size > 1:
        assert cfg.num_heads % args.ulysses_size == 0, f"`{cfg.num_heads=}` cannot be divided evenly by `{args.ulysses_size=}`."
    
    # 设置 fps
    cfg.fps = args.sample_fps if args.sample_fps is not None else cfg.fps
    
    logging.info(f"Generation job args: {args}")
    logging.info(f"Generation model config: {cfg}")

    if dist.is_initialized():
        base_seed = [args.base_seed] if rank == 0 else [None]
        dist.broadcast_object_list(base_seed, src=0)
        args.base_seed = base_seed[0]

    os.makedirs(args.audio_save_dir, exist_ok=True)

    # 参考 LivePortrait: 在 Stateless GPU 环境中，主进程不能初始化 CUDA
    # 参考: https://huggingface.co/spaces/KlingTeam/LivePortrait/blob/main/app.py
    # 在主进程中，我们使用 CPU 加载模型，GPU 将在 worker 进程中通过 @spaces.GPU 装饰器初始化
    is_stateless_gpu = os.environ.get("SPACE_ID") is not None
    
    if is_stateless_gpu:
        # Stateless GPU 环境：主进程不能初始化 CUDA，使用 CPU 加载模型
        logging.info("Stateless GPU environment detected. Loading models on CPU in main process.")
        logging.info("GPU will be initialized in worker process via @spaces.GPU decorator.")
        device = -1  # 使用 CPU 加载模型
    else:
        # 本地环境：可以正常检测和使用 GPU
        if torch.cuda.is_available():
            try:
                num_gpus = torch.cuda.device_count()
                if num_gpus > 0:
                    device = local_rank if world_size > 1 else 0
                    torch.cuda.set_device(device)
                    gpu_name = torch.cuda.get_device_name(device)
                    logging.info(f"GPU AVAILABLE: {num_gpus} GPU(s), Device ID: {device}, Name: {gpu_name}")
                else:
                    logging.warning("CUDA is available but no GPU devices found. Using CPU.")
                    device = -1  # 使用 CPU
            except Exception as e:
                logging.warning(f"GPU detection error: {e}. Using CPU.")
                device = -1  # 使用 CPU
        else:
            logging.warning("No CUDA-compatible GPU found. Using CPU (slower).")
            device = -1  # 使用 CPU

    logging.info("Creating AnyTalker pipeline.")
    # 加载模型
    wan_a2v = wan.WanAF2V(
        config=cfg,
        checkpoint_dir=args.ckpt_dir,
        device_id=device,
        rank=rank,
        use_half=args.use_half,
        t5_fsdp=args.t5_fsdp,
        dit_fsdp=args.dit_fsdp,
        t5_cpu=args.t5_cpu,
        post_trained_checkpoint_path=args.post_trained_checkpoint_path,
        dit_config=args.dit_config,
    )
    
    # 创建 InsightFace 人脸检测器
    # ctx_id=-1 表示使用 CPU，否则使用 GPU
    face_processor_ctx_id = -1 if device == -1 else local_rank
    face_processor = FaceInference(det_thresh=args.det_thresh, ctx_id=face_processor_ctx_id)
    logging.info("Model and face processor loaded successfully.")

    def generate_video(img2vid_image, img2vid_prompt, n_prompt, img2vid_audio_1, img2vid_audio_2, img2vid_audio_3,
                    sd_steps, seed, guide_scale, person_num_selector, audio_mode_selector, fixed_steps=None, trim_to_4s=False):
        # 参考 LivePortrait: 在 worker 进程中直接使用 cuda 设备
        # 参考: https://huggingface.co/spaces/KlingTeam/LivePortrait/blob/main/src/gradio_pipeline.py
        # @spaces.GPU 装饰器已经初始化了 GPU，这里直接使用即可
        if torch.cuda.is_available():
            try:
                current_device = torch.cuda.current_device()
                current_gpu_name = torch.cuda.get_device_name(current_device)
                logging.info(f"Using GPU device {current_device} ({current_gpu_name}) for inference")
            except Exception as e:
                logging.warning(f"GPU check error: {e}")
        
        input_data = {}
        input_data["prompt"] = img2vid_prompt
        input_data["cond_image"] = img2vid_image
        input_data["audio_mode"] = audio_mode_selector  # "pad" or "concat"
        
        # 根据人数收集音频路径
        audio_paths = []
        if person_num_selector == "1 Person":
            if img2vid_audio_1:
                audio_paths.append(img2vid_audio_1)
        elif person_num_selector == "2 Persons":
            if img2vid_audio_1:
                audio_paths.append(img2vid_audio_1)
            if img2vid_audio_2:
                audio_paths.append(img2vid_audio_2)
        elif person_num_selector == "3 Persons":
            if img2vid_audio_1:
                audio_paths.append(img2vid_audio_1)
            if img2vid_audio_2:
                audio_paths.append(img2vid_audio_2)
            if img2vid_audio_3:
                audio_paths.append(img2vid_audio_3)
        
        input_data["audio_paths"] = audio_paths

        logging.info(f"Generating video with {len(audio_paths)} audio(s), mode: {audio_mode_selector}")
        
        # 根据音频长度计算帧数
        current_frame_num = args.frame_num
        if current_frame_num is None:
            if audio_paths and len(audio_paths) > 0:
                # 使用 cfg 中的 fps，如果不可用则使用默认值 24
                fps = getattr(cfg, 'fps', 24)
                calculated_frame_num = calculate_frame_num_from_audio(audio_paths, fps, mode=audio_mode_selector)
                
                # Fast模式：如果trim_to_4s为True，强制限制为4秒对应的帧数
                if trim_to_4s:
                    # 4秒固定为97帧（4n+1格式：4秒*24fps=96帧，向上取整为97帧）
                    max_frames_4s = 97
                    current_frame_num = min(calculated_frame_num, max_frames_4s)
                    logging.warning(f"Fast mode: Audio duration exceeds 4 seconds. Trimming to 4 seconds ({max_frames_4s} frames). Original: {calculated_frame_num} frames")
                else:
                    current_frame_num = calculated_frame_num
                
                logging.info(f"Dynamically determined frame number: {current_frame_num} (mode: {audio_mode_selector})")
            else:
                # 没有音频时使用默认帧数
                current_frame_num = 81  # 默认帧数
                logging.info(f"No audio provided, using default frame number: {current_frame_num}")
        else:
            logging.info(f"Using specified frame number: {current_frame_num}")
        
        # 读取图片
        img = Image.open(input_data["cond_image"]).convert("RGB")
        
        # 如果提供了 fixed_steps，使用它；否则使用用户选择的 sd_steps
        actual_steps = fixed_steps if fixed_steps is not None else sd_steps
        if fixed_steps is not None:
            logging.info(f"Using fixed denoising steps: {fixed_steps}")
        else:
            logging.info(f"Using user-selected denoising steps: {sd_steps}")
        
        # 生成视频
        video = wan_a2v.generate(
            input_data["prompt"],
            img,
            audio=audio_paths[0] if audio_paths and len(audio_paths) > 0 else None,
            max_area=MAX_AREA_CONFIGS[args.size],
            frame_num=current_frame_num,
            shift=args.sample_shift,
            sample_solver=args.sample_solver,
            sampling_steps=actual_steps,
            guide_scale=guide_scale,
            seed=seed if seed >= 0 else args.base_seed,
            offload_model=args.offload_model,
            cfg_zero=args.cfg_zero,
            zero_init_steps=args.zero_init_steps,
            face_processor=face_processor,
            img_path=input_data["cond_image"],
            audio_paths=audio_paths,
            task_key="gradio_output",
            mode=audio_mode_selector,
            trim_to_4s=trim_to_4s,
        )
        
        if isinstance(video, dict):
            video = video['original']
        
        # 生成输出文件名（替换特殊字符避免 shell 解析问题）
        formatted_time = datetime.now().strftime("%Y%m%d_%H%M%S")
        formatted_prompt = input_data['prompt'].replace(" ", "_").replace("/", "_").replace(",", "").replace("*", "x")[:50]
        formatted_size = args.size.replace('*', 'x')
        save_file = f"outputs/{args.task}_{formatted_size}_{formatted_prompt}_{formatted_time}"
        
        # 确保输出目录存在
        os.makedirs("outputs", exist_ok=True)
        
        # 注意：cache_video 不会自动添加后缀，需要传入完整文件名
        output_file = save_file + '.mp4'
        
        logging.info(f"Saving generated video to {output_file}")
        cache_video(
            tensor=video[None],
            save_file=output_file,
            fps=args.sample_fps if args.sample_fps is not None else cfg.sample_fps,
            nrow=1,
            normalize=True,
            value_range=(-1, 1))
        
        # 如果有音频文件，进行音频合成
        if audio_paths:
            existing_audio_paths = [path for path in audio_paths if path and os.path.exists(path)]
            if existing_audio_paths:
                # 计算视频时长（用于Fast模式限制音频长度）
                fps = getattr(cfg, 'fps', 24)
                video_duration_seconds = current_frame_num / fps if current_frame_num and fps else 0
                
                # 构建输出文件名
                audio_names = [os.path.basename(path).split('.')[0] for path in existing_audio_paths]
                audio_suffix = "_".join([f"audio{i}_{name}" for i, name in enumerate(audio_names)])
                audio_video_path = save_file + f'_{audio_suffix}_cfg_{guide_scale}.mp4'
                
                # 构建 ffmpeg 命令
                if len(existing_audio_paths) == 1:
                    # 只有一个音频
                    if trim_to_4s and video_duration_seconds > 0:
                        # Fast模式：限制音频输入和输出时长为视频时长
                        ffmpeg_command = f'ffmpeg -i "{output_file}" -ss 0 -t {video_duration_seconds:.3f} -i "{existing_audio_paths[0]}" -t {video_duration_seconds:.3f} -vcodec libx264 -acodec aac -crf 18 -y "{audio_video_path}"'
                    else:
                        ffmpeg_command = f'ffmpeg -i "{output_file}" -i "{existing_audio_paths[0]}" -vcodec libx264 -acodec aac -crf 18 -shortest -y "{audio_video_path}"'
                else:
                    input_args = f'-i "{output_file}"'
                    if audio_mode_selector == "concat":
                        # concat 模式：串联音频
                        for audio_path in existing_audio_paths:
                            if trim_to_4s and video_duration_seconds > 0:
                                # Fast模式：限制每个音频输入的时长
                                input_args += f' -ss 0 -t {video_duration_seconds:.3f} -i "{audio_path}"'
                            else:
                                input_args += f' -i "{audio_path}"'
                        
                        num_audios = len(existing_audio_paths)
                        concat_inputs = ''.join([f'[{i+1}:a]' for i in range(num_audios)])
                        filter_complex = f'"{concat_inputs}concat=n={num_audios}:v=0:a=1[aout]"'
                        
                        if trim_to_4s and video_duration_seconds > 0:
                            # Fast模式：限制最终输出时长
                            ffmpeg_command = (
                                f'ffmpeg {input_args} -filter_complex {filter_complex} '
                                f'-map 0:v -map "[aout]" -t {video_duration_seconds:.3f} -vcodec libx264 -acodec aac -crf 18 -y "{audio_video_path}"'
                            )
                        else:
                            ffmpeg_command = (
                                f'ffmpeg {input_args} -filter_complex {filter_complex} '
                                f'-map 0:v -map "[aout]" -vcodec libx264 -acodec aac -crf 18 -y "{audio_video_path}"'
                            )
                    else:
                        # pad 模式：混合所有音频
                        filter_inputs = []
                        for i, audio_path in enumerate(existing_audio_paths):
                            if trim_to_4s and video_duration_seconds > 0:
                                # Fast模式：限制每个音频输入的时长
                                input_args += f' -ss 0 -t {video_duration_seconds:.3f} -i "{audio_path}"'
                            else:
                                input_args += f' -i "{audio_path}"'
                            filter_inputs.append(f'[{i+1}:a]')
                        
                        filter_complex = f'{"".join(filter_inputs)}amix=inputs={len(existing_audio_paths)}:duration=shortest[aout]'
                        if trim_to_4s and video_duration_seconds > 0:
                            # Fast模式：限制最终输出时长
                            ffmpeg_command = f'ffmpeg {input_args} -filter_complex "{filter_complex}" -map 0:v -map "[aout]" -t {video_duration_seconds:.3f} -vcodec libx264 -acodec aac -crf 18 -y "{audio_video_path}"'
                        else:
                            ffmpeg_command = f'ffmpeg {input_args} -filter_complex "{filter_complex}" -map 0:v -map "[aout]" -vcodec libx264 -acodec aac -crf 18 -y "{audio_video_path}"'
                
                logging.info(f"Adding audio: {ffmpeg_command}")
                os.system(ffmpeg_command)
                
                # 删除没有音频的原始视频文件
                if os.path.exists(audio_video_path):
                    os.remove(output_file)
                    output_file = audio_video_path
                    logging.info(f"Final video saved to: {output_file}")
                else:
                    logging.warning(f"Audio synthesis failed, keeping original video: {output_file}")
            else:
                logging.info(f"No valid audio files found, video saved to: {output_file}")
        else:
            logging.info(f"No audio files provided, video saved to: {output_file}")
        
        logging.info("Finished.")
        
        # 计算视频时长信息（用于quality模式的提示）
        fps = getattr(cfg, 'fps', 24)
        video_duration_seconds = current_frame_num / fps if current_frame_num and fps else 0
        return output_file, video_duration_seconds, actual_steps

    # 计算动态duration的函数
    def get_duration(video_seconds, steps):
        """
        计算quality模式所需的GPU duration
        duration = 60s (预处理时间) + 视频秒数 * 步数 * 3 秒
        """
        return int(60 + video_seconds * steps * 3)
    
    # 为quality模式创建动态duration计算函数
    def calculate_quality_duration(*args, **kwargs):
        """
        从函数参数中提取视频时长和步数，计算动态duration
        参数顺序: img2vid_image, img2vid_prompt, n_prompt, img2vid_audio_1, img2vid_audio_2, img2vid_audio_3, 
                 sd_steps, seed, guide_scale, person_num_selector, audio_mode_selector
        """
        if len(args) >= 11:
            img2vid_audio_1 = args[3]
            img2vid_audio_2 = args[4]
            img2vid_audio_3 = args[5]
            sd_steps = args[6]
            person_num_selector = args[9]
            audio_mode_selector = args[10]
            
            # 根据人数收集音频路径
            audio_paths = []
            if person_num_selector == "1 Person":
                if img2vid_audio_1:
                    audio_paths.append(img2vid_audio_1)
            elif person_num_selector == "2 Persons":
                if img2vid_audio_1:
                    audio_paths.append(img2vid_audio_1)
                if img2vid_audio_2:
                    audio_paths.append(img2vid_audio_2)
            elif person_num_selector == "3 Persons":
                if img2vid_audio_1:
                    audio_paths.append(img2vid_audio_1)
                if img2vid_audio_2:
                    audio_paths.append(img2vid_audio_2)
                if img2vid_audio_3:
                    audio_paths.append(img2vid_audio_3)
            
            # 计算预期的视频时长
            fps = getattr(cfg, 'fps', 24)
            expected_video_seconds = 8.0  # 默认值
            if audio_paths and len(audio_paths) > 0:
                try:
                    calculated_frame_num = calculate_frame_num_from_audio(audio_paths, fps, mode=audio_mode_selector)
                    expected_video_seconds = calculated_frame_num / fps
                except Exception as e:
                    logging.warning(f"Failed to calculate expected video duration for GPU allocation: {e}")
            
            # 计算并返回duration
            return get_duration(expected_video_seconds, sd_steps)
        else:
            # 如果参数不足，返回默认值
            return 720
    
    # 使用 @spaces.GPU 装饰器包装 generate_video 函数（参考 LivePortrait）
    # 参考: https://huggingface.co/spaces/KlingTeam/LivePortrait/blob/main/app.py
    # @spaces.GPU 装饰器会自动处理 GPU 初始化，不需要手动初始化
    
    # 快速生成模式：120秒，固定8步去噪
    @spaces.GPU(duration=120)
    def gpu_wrapped_generate_video_fast(*args, **kwargs):
        # 固定使用8步去噪，通过关键字参数传递
        kwargs['fixed_steps'] = 8
        
        # Fast模式音频长度检测：检查是否超过4秒
        # 参数顺序: img2vid_image, img2vid_prompt, n_prompt, img2vid_audio_1, img2vid_audio_2, img2vid_audio_3, 
        #          sd_steps, seed, guide_scale, person_num_selector, audio_mode_selector
        if len(args) >= 11:
            img2vid_image = args[0]
            img2vid_prompt = args[1]
            n_prompt = args[2]
            img2vid_audio_1 = args[3]
            img2vid_audio_2 = args[4]
            img2vid_audio_3 = args[5]
            sd_steps = args[6]
            seed = args[7]
            guide_scale = args[8]
            person_num_selector = args[9]
            audio_mode_selector = args[10]
            
            # 根据人数收集音频路径
            audio_paths = []
            if person_num_selector == "1 Person":
                if img2vid_audio_1:
                    audio_paths.append(img2vid_audio_1)
            elif person_num_selector == "2 Persons":
                if img2vid_audio_1:
                    audio_paths.append(img2vid_audio_1)
                if img2vid_audio_2:
                    audio_paths.append(img2vid_audio_2)
            elif person_num_selector == "3 Persons":
                if img2vid_audio_1:
                    audio_paths.append(img2vid_audio_1)
                if img2vid_audio_2:
                    audio_paths.append(img2vid_audio_2)
                if img2vid_audio_3:
                    audio_paths.append(img2vid_audio_3)
            
            # 检测音频长度是否超过4秒
            if audio_paths and len(audio_paths) > 0:
                fps = getattr(cfg, 'fps', 24)
                try:
                    calculated_frame_num = calculate_frame_num_from_audio(audio_paths, fps, mode=audio_mode_selector)
                    # 4秒固定为97帧（4n+1格式：4秒*24fps=96帧，向上取整为97帧）
                    max_frames_4s = 97
                    
                    if calculated_frame_num > max_frames_4s:
                        # 超过4秒，设置trim_to_4s标记
                        kwargs['trim_to_4s'] = True
                        calculated_duration = calculated_frame_num / fps
                        logging.warning(f"Fast mode: Audio duration ({calculated_duration:.2f}s) exceeds 4 seconds limit. Will trim to 4 seconds.")
                    else:
                        kwargs['trim_to_4s'] = False
                except Exception as e:
                    logging.warning(f"Failed to check audio duration: {e}")
                    kwargs['trim_to_4s'] = False
            else:
                kwargs['trim_to_4s'] = False
        
        return gpu_wrapped_generate_video_worker(*args, **kwargs)
    
    # 高质量生成模式：动态duration，根据视频时长和步数计算
    @spaces.GPU(duration=calculate_quality_duration)
    def gpu_wrapped_generate_video_quality(*args, **kwargs):
        return gpu_wrapped_generate_video_worker(*args, **kwargs)
    
    # 共享的 worker 函数，处理 GPU 移动逻辑
    def gpu_wrapped_generate_video_worker(*args, **kwargs):
        # 在 worker 进程中将模型移动到 GPU（如果模型在 CPU 上）
        # 参考 LivePortrait: 在 worker 进程中直接使用 .to("cuda")
        if torch.cuda.is_available() and device == -1:
            try:
                logging.info("Moving models from CPU to GPU in worker process...")
                cuda_device = torch.device("cuda")
                
                # 移动主模型到 GPU
                if hasattr(wan_a2v, 'model') and wan_a2v.model is not None:
                    wan_a2v.model = wan_a2v.model.to(cuda_device)
                
                # 移动 VAE 模型到 GPU
                if hasattr(wan_a2v, 'vae') and wan_a2v.vae is not None:
                    if hasattr(wan_a2v.vae, 'model'):
                        wan_a2v.vae.model = wan_a2v.vae.model.to(cuda_device)
                    # 移动 VAE 的 mean 和 std 张量
                    if hasattr(wan_a2v.vae, 'mean'):
                        wan_a2v.vae.mean = wan_a2v.vae.mean.to(cuda_device)
                    if hasattr(wan_a2v.vae, 'std'):
                        wan_a2v.vae.std = wan_a2v.vae.std.to(cuda_device)
                    # 重新创建 scale 列表，确保在 GPU 上
                    if hasattr(wan_a2v.vae, 'mean') and hasattr(wan_a2v.vae, 'std'):
                        wan_a2v.vae.scale = [wan_a2v.vae.mean, 1.0 / wan_a2v.vae.std]
                    # 更新 VAE 的设备属性
                    if hasattr(wan_a2v.vae, 'device'):
                        wan_a2v.vae.device = cuda_device
                
                # 移动 CLIP 模型到 GPU
                if hasattr(wan_a2v, 'clip') and wan_a2v.clip is not None:
                    if hasattr(wan_a2v.clip, 'model'):
                        wan_a2v.clip.model = wan_a2v.clip.model.to(cuda_device)
                
                # 移动 T5 encoder 到 GPU（如果不在 CPU 上）
                if hasattr(wan_a2v, 'text_encoder') and wan_a2v.text_encoder is not None:
                    if hasattr(wan_a2v.text_encoder, 'model'):
                        wan_a2v.text_encoder.model = wan_a2v.text_encoder.model.to(cuda_device)
                    # 更新 T5 encoder 的设备属性
                    if hasattr(wan_a2v.text_encoder, 'device'):
                        wan_a2v.text_encoder.device = cuda_device
                
                # 更新设备信息
                wan_a2v.device = cuda_device
                
                logging.info("Models moved to GPU successfully")
            except Exception as e:
                logging.warning(f"Failed to move models to GPU: {e}")
        
        result = generate_video(*args, **kwargs)
        # generate_video 现在返回 (output_file, video_duration_seconds, actual_steps)
        if isinstance(result, tuple) and len(result) == 3:
            return result
        # 兼容旧格式（如果返回的是单个值）
        return result, 0, 0

            
        
    with gr.Blocks() as demo:

        gr.Markdown("## 🎤 AnyTalker")
        
        gr.Markdown("Let your characters interact naturally")
        
        gr.HTML("""
        <div style="display:flex;column-gap:4px;flex-wrap:wrap;">
            <a href='https://hkust-c4g.github.io/AnyTalker-homepage/'><img src='https://img.shields.io/badge/Project-Page-blue'></a>
            <a href='https://huggingface.co/zzz66/AnyTalker-1.3B'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20HuggingFace-Model-yellow'></a>
            <a href='https://arxiv.org/abs/2511.23475/'><img src='https://img.shields.io/badge/Paper-Arxiv-red'></a>
        </div>
        """)
        
        gr.Markdown("""
        ⚠️ **Important Video Duration Limits**
        
        - **Fast Mode**: Maximum video duration should be less than 4 seconds. Audio inputs longer than 4 seconds will be automatically trimmed to 4 seconds.
        """)

        with gr.Row():
            with gr.Column(scale=1):
                img2vid_image = gr.Image(
                    type="filepath",
                    label="Upload Input Image",
                    elem_id="image_upload",
                )
                img2vid_prompt = gr.Textbox(
                    label="Prompt",
                    placeholder="Describe the video you want to generate",
                )
                
                
                with gr.Accordion("Audio Options", open=True):
                    person_num_selector = gr.Radio(
                        choices=["1 Person", "2 Persons", "3 Persons"],
                        label="Number of Persons (determined by audio inputs)",
                        value="1 Person"
                    )
                    audio_mode_selector = gr.Radio(
                        choices=["pad", "concat"],
                        label="Audio Processing Mode",
                        value="pad"
                    )
                    gr.Markdown("""
                    **Audio Mode Description:**
                    - **pad**: Select this if every audio input track has already been zero-padded to a common length.
                    - **concat**: Select this if you want the script to chain each speaker's clips together and then zero-pad the non-speaker segments to reach a uniform length.
                    """)
                    gr.Markdown("""
                    **Audio Binding Order:**
                    - Audio inputs are bound to persons based on their positions in the input image, from **left to right**.
                    - Person 1 corresponds to the leftmost person, Person 2 to the middle person (if any), and Person 3 to the rightmost person (if any).
                    """)
                    # 三个音频输入框始终可见，读取时根据 person_num_selector 只读取前 n 个
                    img2vid_audio_1 = gr.Audio(label="Audio for Person 1 (Leftmost)", type="filepath", visible=True)
                    img2vid_audio_2 = gr.Audio(label="Audio for Person 2 (Middle)", type="filepath", visible=True)
                    img2vid_audio_3 = gr.Audio(label="Audio for Person 3 (Rightmost)", type="filepath", visible=True)

                with gr.Accordion("Advanced Options", open=False):
                    with gr.Row():
                        sd_steps = gr.Slider(
                            label="Diffusion steps",
                            minimum=1,
                            maximum=1000,
                            value=25,
                            step=1)
                        seed = gr.Slider(
                            label="Seed",
                            minimum=-1,
                            maximum=2147483647,
                            step=1,
                            value=41)
                    with gr.Row():
                        guide_scale = gr.Slider(
                            label="Guide Scale",
                            minimum=0,
                            maximum=20,
                            value=4.5,
                            step=0.1)
                    # with gr.Row():
                    n_prompt = gr.Textbox(
                        label="Negative Prompt",
                        placeholder="Describe the negative prompt you want to add",
                        value="bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
                    )

                with gr.Row():
                    run_i2v_button_fast = gr.Button(
                        "Generate Video (Fast - 120s, 8 steps)", 
                        variant="secondary",
                        scale=1
                    )
                    run_i2v_button_quality = gr.Button(
                        "Generate Video (Quality - Dynamic, Custom steps)", 
                        variant="primary",
                        scale=1
                    )
                gr.Markdown("""
                **Generation Modes:**
                - **Fast Mode (120s GPU budget, suitable for any type of users)**: Fixed 8 denoising steps for quick generation. Maximum video duration: 4 seconds.
                - **Quality Mode (Dynamic GPU budget)**: Custom denoising steps (adjustable via "Diffusion steps" slider, default: 25 steps). GPU duration is dynamically calculated as: 60s (preprocessing) + video_seconds × steps × 3 s.
                
                *Note: Fast mode has a fixed 120s GPU budget. Quality mode dynamically allocates GPU time based on video length and denoising steps. Multi-person videos generally require longer duration and more Usage Quota for better quality.*
                """)

            with gr.Column(scale=2):
                result_gallery = gr.Video(
                    label='Generated Video', interactive=False, height=600, )
                
                # 创建一个函数来处理 examples 选择
                def handle_example_select(image, prompt, person_num, audio_mode, audio1, audio2, audio3):
                    # 三个音频输入框始终可见，只返回值，不改变可见性
                    # 读取时根据 person_num_selector 只读取前 n 个音频
                    return (
                        image, prompt, person_num, audio_mode,
                        audio1, audio2, audio3
                    )
                
                examples_component = gr.Examples(
                    examples = [
                        ["./input_example/images/1p-0.png", "The man stands in the dusty western street, backlit by the setting sun, and his determined gaze speaks of a rugged spirit.", "1 Person", "pad", "./input_example/audios/1p-0.wav", None, None],
                        ["./input_example/images/2p-1.png", "In a casual, intimate setting, a man and a woman are engaged in a heartfelt conversation inside a car. The man, sporting a denim jacket over a blue shirt, sits attentively with a seatbelt fastened, his gaze fixed on the woman beside him. The woman, wearing a black tank top and a denim jacket draped over her shoulders, smiles warmly, her eyes reflecting genuine interest and connection. The car's interior, with its beige seats and simple design, provides a backdrop that emphasizes their interaction. The scene captures a moment of shared understanding and connection, set against the soft, diffused light of an overcast day. A medium shot from a slightly angled perspective, focusing on their expressions and body language.", "2 Persons", "pad", "./input_example/audios/2p-1-left.wav", "./input_example/audios/2p-1-right.wav", None],
                        ["./input_example/images/2p-2.png", "In a cozy recording studio, a man and a woman are singing together. The man, with tousled brown hair, stands to the left, wearing a light green button-down shirt. His gaze is directed towards the woman, who is smiling warmly. She, with wavy dark hair, is dressed in a black floral dress and stands to the right, her eyes closed in enjoyment. Between them is a professional microphone, capturing their harmonious voices. The background features wooden panels and various audio equipment, creating an intimate and focused atmosphere. The lighting is soft and warm, highlighting their expressions and the intimate setting. A medium shot captures their interaction closely.", "2 Persons", "pad", "./input_example/audios/2p-2-left.wav", "./input_example/audios/2p-2-right.wav", None],
                    ],
                    inputs = [img2vid_image, img2vid_prompt, person_num_selector, audio_mode_selector, img2vid_audio_1, img2vid_audio_2, img2vid_audio_3],
                    outputs = [img2vid_image, img2vid_prompt, person_num_selector, audio_mode_selector, img2vid_audio_1, img2vid_audio_2, img2vid_audio_3],
                    fn=handle_example_select,
                    cache_examples=False,  # 禁用缓存以避免 JSON 解码错误
                )


        # 包装函数：处理警告信息显示
        def handle_fast_generation(img2vid_image, img2vid_prompt, n_prompt, img2vid_audio_1, img2vid_audio_2, img2vid_audio_3,
                                    sd_steps, seed, guide_scale, person_num_selector, audio_mode_selector):
            # 在开始生成前先检测音频长度，如果超过4秒立即显示警告
            # 根据人数收集音频路径
            audio_paths = []
            if person_num_selector == "1 Person":
                if img2vid_audio_1:
                    audio_paths.append(img2vid_audio_1)
            elif person_num_selector == "2 Persons":
                if img2vid_audio_1:
                    audio_paths.append(img2vid_audio_1)
                if img2vid_audio_2:
                    audio_paths.append(img2vid_audio_2)
            elif person_num_selector == "3 Persons":
                if img2vid_audio_1:
                    audio_paths.append(img2vid_audio_1)
                if img2vid_audio_2:
                    audio_paths.append(img2vid_audio_2)
                if img2vid_audio_3:
                    audio_paths.append(img2vid_audio_3)
            
            # 检测音频长度是否超过4秒
            if audio_paths and len(audio_paths) > 0:
                fps = getattr(cfg, 'fps', 24)
                try:
                    calculated_frame_num = calculate_frame_num_from_audio(audio_paths, fps, mode=audio_mode_selector)
                    # 4秒固定为97帧（4n+1格式：4秒*24fps=96帧，向上取整为97帧）
                    max_frames_4s = 97
                    
                    if calculated_frame_num > max_frames_4s:
                        # 超过4秒，立即显示警告
                        calculated_duration = calculated_frame_num / fps
                        warning_msg = f"⚠️ Warning: Your audio duration ({calculated_duration:.2f}s) exceeds the 4-second limit for Fast Mode. The audio will be automatically trimmed to 4 seconds to prevent timeout."
                        gr.Warning(warning_msg, duration=5)
                except Exception as e:
                    logging.warning(f"Failed to check audio duration: {e}")
            
            # 继续执行视频生成
            result = gpu_wrapped_generate_video_fast(
                img2vid_image, img2vid_prompt, n_prompt, img2vid_audio_1, img2vid_audio_2, img2vid_audio_3,
                sd_steps, seed, guide_scale, person_num_selector, audio_mode_selector
            )
            # 处理返回结果：可能是 (output_file, video_duration_seconds, actual_steps) 或 output_file
            if isinstance(result, tuple) and len(result) == 3:
                return result[0]  # 只返回视频文件
            return result
        
        def handle_quality_generation(img2vid_image, img2vid_prompt, n_prompt, img2vid_audio_1, img2vid_audio_2, img2vid_audio_3,
                                      sd_steps, seed, guide_scale, person_num_selector, audio_mode_selector):
            # 在生成前先计算预期的视频时长和duration
            # 根据人数收集音频路径
            audio_paths = []
            if person_num_selector == "1 Person":
                if img2vid_audio_1:
                    audio_paths.append(img2vid_audio_1)
            elif person_num_selector == "2 Persons":
                if img2vid_audio_1:
                    audio_paths.append(img2vid_audio_1)
                if img2vid_audio_2:
                    audio_paths.append(img2vid_audio_2)
            elif person_num_selector == "3 Persons":
                if img2vid_audio_1:
                    audio_paths.append(img2vid_audio_1)
                if img2vid_audio_2:
                    audio_paths.append(img2vid_audio_2)
                if img2vid_audio_3:
                    audio_paths.append(img2vid_audio_3)
            
            # 计算预期的视频时长
            fps = getattr(cfg, 'fps', 24)
            expected_video_seconds = 0
            if audio_paths and len(audio_paths) > 0:
                try:
                    calculated_frame_num = calculate_frame_num_from_audio(audio_paths, fps, mode=audio_mode_selector)
                    expected_video_seconds = calculated_frame_num / fps
                except Exception as e:
                    logging.warning(f"Failed to calculate expected video duration: {e}")
                    expected_video_seconds = 8.0  # 默认值
            
            # 计算动态duration
            expected_duration = get_duration(expected_video_seconds, sd_steps)
            
            # 执行生成
            result = gpu_wrapped_generate_video_quality(
                img2vid_image, img2vid_prompt, n_prompt, img2vid_audio_1, img2vid_audio_2, img2vid_audio_3,
                sd_steps, seed, guide_scale, person_num_selector, audio_mode_selector
            )
            
            # 处理返回结果并显示提示
            if isinstance(result, tuple) and len(result) == 3:
                output_file, actual_video_seconds, actual_steps = result
                # 计算实际使用的duration
                actual_duration = get_duration(actual_video_seconds, actual_steps)
                # 使用 gr.Info 提示用户
                info_msg = f"Video generation completed! Duration used: {actual_duration}s (60s preprocessing + {actual_video_seconds:.2f}s video × {actual_steps} steps × 3s)"
                gr.Info(info_msg)
                return output_file
            else:
                return result
        
        # 快速生成按钮：120秒，固定8步
        run_i2v_button_fast.click(
            fn=handle_fast_generation,
            inputs=[img2vid_image, img2vid_prompt, n_prompt, img2vid_audio_1, img2vid_audio_2, img2vid_audio_3, sd_steps, seed, guide_scale, person_num_selector, audio_mode_selector],
            outputs=[result_gallery],
        )
        
        # 高质量生成按钮：720秒，用户选择步数
        run_i2v_button_quality.click(
            fn=handle_quality_generation,
            inputs=[img2vid_image, img2vid_prompt, n_prompt, img2vid_audio_1, img2vid_audio_2, img2vid_audio_3, sd_steps, seed, guide_scale, person_num_selector, audio_mode_selector],
            outputs=[result_gallery],
        )
    # 参考 Meigen-MultiTalk 的成功配置
    # 在 Hugging Face Spaces 上，Gradio 会自动处理端口和服务器配置
    demo.queue(max_size=20).launch(show_error=True)

        


if __name__ == "__main__":
    args = _parse_args()
    run_graio_demo(args)