app.py

# -*- coding: utf-8 -*-
import csv
import gc
import os
import random
import shutil
import subprocess
import sys
import uuid
from pathlib import Path

import gradio as gr
import spaces
import torch
from huggingface_hub import snapshot_download


ROOT = Path(__file__).resolve().parent
SOURCE_DIR = ROOT / "HunyuanVideo-Avatar"
WEIGHTS_DIR = ROOT / "weights"
OUTPUT_DIR = ROOT / "outputs"
SOURCE_REPO = "https://github.com/Tencent-Hunyuan/HunyuanVideo-Avatar.git"
MODEL_REPO = "tencent/HunyuanVideo-Avatar"
FPS = 25
FRAME_OPTIONS = {
    "约 2 秒（49 帧，推荐首测）": 49,
    "约 3 秒（73 帧）": 73,
    "约 4 秒（97 帧）": 97,
    "约 5 秒（129 帧，官方配置）": 129,
}

os.environ.setdefault("GRADIO_SSR_MODE", "0")
os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")


def ensure_source():
    if not (SOURCE_DIR / "hymm_sp" / "sample_gpu_poor.py").is_file():
        print("[SETUP] Cloning official HunyuanVideo-Avatar source...", flush=True)
        subprocess.run(
            ["git", "clone", "--depth", "1", SOURCE_REPO, str(SOURCE_DIR)],
            check=True,
        )

    # The official `--infer-min` path hard-codes 129 frames. Make it honor the
    # requested evaluation length so a Space can run short, lower-cost tests.
    sample_file = SOURCE_DIR / "hymm_sp" / "sample_gpu_poor.py"
    source = sample_file.read_text(encoding="utf-8")
    patched = source.replace(
        'batch["audio_len"][0] = 129',
        'batch["audio_len"][0] = args.sample_n_frames',
    )
    if patched != source:
        sample_file.write_text(patched, encoding="utf-8")
        print("[SETUP] Patched infer-min to honor --sample-n-frames", flush=True)


ensure_source()
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)


MODEL_PATTERNS = [
    "ckpts/config.json",
    "ckpts/det_align/detface.pt",
    "ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8.pt",
    "ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8_map.pt",
    "ckpts/hunyuan-video-t2v-720p/vae/config.json",
    "ckpts/hunyuan-video-t2v-720p/vae/pytorch_model.pt",
    "ckpts/llava_llama_image/*.json",
    "ckpts/llava_llama_image/*.safetensors",
    "ckpts/llava_llama_image/*.model",
    "ckpts/text_encoder_2/config.json",
    "ckpts/text_encoder_2/model.safetensors",
    "ckpts/text_encoder_2/*.json",
    "ckpts/text_encoder_2/*.txt",
    "ckpts/whisper-tiny/config.json",
    "ckpts/whisper-tiny/model.safetensors",
    "ckpts/whisper-tiny/preprocessor_config.json",
]


def ensure_weights():
    checkpoint = (
        WEIGHTS_DIR
        / "ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8.pt"
    )
    if checkpoint.is_file():
        print(f"[MODEL CACHE] FP8 checkpoint ready: {checkpoint}", flush=True)
        return checkpoint

    WEIGHTS_DIR.mkdir(parents=True, exist_ok=True)
    print(
        "[MODEL DOWNLOAD START] repo=tencent/HunyuanVideo-Avatar, "
        "profile=FP8-minimal, expected_size≈45GB",
        flush=True,
    )
    snapshot_download(
        repo_id=MODEL_REPO,
        local_dir=WEIGHTS_DIR,
        allow_patterns=MODEL_PATTERNS,
    )
    if not checkpoint.is_file():
        raise RuntimeError("FP8 checkpoint download did not complete")
    print(f"[MODEL DOWNLOAD DONE] checkpoint={checkpoint}", flush=True)
    return checkpoint


def estimate_gpu_duration(_image, _audio, _prompt, frame_profile, *_args):
    frames = FRAME_OPTIONS.get(frame_profile, 49)
    # Includes first-run model download and CPU-offloaded inference.
    return max(900, min(3600, 1200 + frames * 16))


def normalize_media(image_path, audio_path, work_dir):
    image_target = work_dir / "character.png"
    audio_target = work_dir / "speech.wav"
    shutil.copy2(image_path, image_target)
    subprocess.run(
        [
            "ffmpeg", "-y", "-i", str(audio_path), "-ac", "1", "-ar", "16000",
            "-c:a", "pcm_s16le", str(audio_target),
        ],
        check=True,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.PIPE,
    )
    return image_target, audio_target


@spaces.GPU(size="xlarge", duration=estimate_gpu_duration)
def generate(image, audio, prompt, frame_profile, seed, progress=gr.Progress()):
    if not image:
        raise gr.Error("请上传一张角色图片")
    if not audio:
        raise gr.Error("请上传驱动语音")

    frames = FRAME_OPTIONS.get(frame_profile, 49)
    actual_seed = random.randint(0, 2**31 - 1) if int(seed) < 0 else int(seed)
    job_id = uuid.uuid4().hex
    job_dir = OUTPUT_DIR / job_id
    result_dir = job_dir / "result"
    job_dir.mkdir(parents=True, exist_ok=True)
    result_dir.mkdir(parents=True, exist_ok=True)

    try:
        progress(0.02, desc="检查并下载官方 FP8 模型（首次约 45GB）...")
        checkpoint = ensure_weights()
        progress(0.12, desc="预处理图片和音频...")
        image_path, audio_path = normalize_media(image, audio, job_dir)

        meta_path = job_dir / "input.csv"
        safe_prompt = (prompt or "A cartoon character speaks to the camera.").strip()
        with meta_path.open("w", encoding="utf-8-sig", newline="") as handle:
            writer = csv.DictWriter(
                handle,
                fieldnames=["videoid", "image", "audio", "prompt", "fps"],
            )
            writer.writeheader()
            writer.writerow(
                {
                    "videoid": job_id,
                    "image": str(image_path),
                    "audio": str(audio_path),
                    "prompt": safe_prompt,
                    "fps": FPS,
                }
            )

        command = [
            sys.executable,
            str(SOURCE_DIR / "hymm_sp" / "sample_gpu_poor.py"),
            "--input", str(meta_path),
            "--ckpt", str(checkpoint),
            "--sample-n-frames", str(frames),
            "--seed", str(actual_seed),
            "--image-size", "704",
            "--cfg-scale", "7.5",
            "--infer-steps", "50",
            "--use-deepcache", "1",
            "--flow-shift-eval-video", "5.0",
            "--save-path", str(result_dir),
            "--use-fp8",
            "--cpu-offload",
            "--infer-min",
        ]
        env = os.environ.copy()
        env.update(
            {
                "MODEL_BASE": str(WEIGHTS_DIR),
                "CPU_OFFLOAD": "1",
                "DISABLE_SP": "1",
                "PYTHONPATH": str(SOURCE_DIR),
            }
        )
        print("[INFERENCE START] " + " ".join(command), flush=True)
        progress(0.18, desc="加载 FP8 模型并生成视频，可能需要较长时间...")
        process = subprocess.run(
            command,
            cwd=SOURCE_DIR,
            env=env,
            text=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            timeout=3500,
        )
        print(process.stdout, flush=True)
        if process.returncode != 0:
            tail = "\n".join(process.stdout.splitlines()[-25:])
            raise RuntimeError(f"官方推理进程退出码 {process.returncode}\n{tail}")

        output = result_dir / f"{job_id}_audio.mp4"
        if not output.is_file():
            candidates = sorted(result_dir.glob("*_audio.mp4"))
            if not candidates:
                raise RuntimeError("推理完成，但没有找到带音频的 MP4 输出")
            output = candidates[-1]

        progress(1.0, desc="生成完成")
        info = (
            f"完成：{frames} 帧 / {FPS}fps（约 {frames / FPS:.1f} 秒），"
            f"704px，50 steps，FP8 + CPU offload。"
        )
        return str(output), actual_seed, info
    except subprocess.TimeoutExpired as exc:
        raise gr.Error("生成超过 ZeroGPU 最长执行时间，请改用 49 帧重试") from exc
    except gr.Error:
        raise
    except Exception as exc:
        print(f"[ERROR] {exc}", flush=True)
        raise gr.Error(f"生成失败：{exc}") from exc
    finally:
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()


with gr.Blocks(title="HunyuanVideo-Avatar 卡通数字人测试") as demo:
    gr.Markdown(
        """
        # HunyuanVideo-Avatar 卡通数字人测试
        上传一张卡通/3D/拟人角色图片和一段语音，评估角色一致性、口型、表情和身体动作。

        **首次运行需要下载约 45GB 官方权重。建议先选 49 帧短片。**
        """
    )
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="角色图片", type="filepath")
            audio_input = gr.Audio(label="驱动语音", type="filepath")
            prompt_input = gr.Textbox(
                label="英文画面描述",
                value="A cute cartoon character speaks naturally to the camera with subtle gestures.",
                lines=3,
                info="描述角色、构图、背景和期望动作；不要写与原图冲突的外观。",
            )
            frame_input = gr.Dropdown(
                choices=list(FRAME_OPTIONS),
                value=list(FRAME_OPTIONS)[0],
                label="测试时长",
            )
            seed_input = gr.Number(label="随机种子（-1 为随机）", value=-1, precision=0)
            generate_button = gr.Button("生成测试视频", variant="primary")
        with gr.Column():
            video_output = gr.Video(label="生成结果")
            seed_output = gr.Number(label="实际种子", precision=0)
            info_output = gr.Textbox(label="生成信息")

    gr.Markdown(
        """
        ### 素材建议
        - 单个角色、脸部清晰；正面或轻微侧脸更稳定。
        - 半身或全身卡通图都可测试，避免文字、水印和多人画面。
        - 语音尽量清晰、无背景音乐；当前测试输出最多约 5 秒。
        - 官方模型很大，ZeroGPU 冷启动与 CPU offload 都会显著增加等待时间。
        """
    )

    generate_button.click(
        fn=generate,
        inputs=[image_input, audio_input, prompt_input, frame_input, seed_input],
        outputs=[video_output, seed_output, info_output],
        api_name="generate",
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1, max_size=8).launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
    )