Update app.py

app.py

@@ -1,525 +1,574 @@
 # app.py
-# HuggingFace Spaces (Gradio + ZeroGPU) 单文件示例：
-# - 自动下载 LongCat-Video GitHub 代码（zip）
-# - 自动下载 LongCat-Video / LongCat-Video-Avatar 权重（HF Hub）
-# - 通过 spaces.GPU 在 ZeroGPU 环境下按需申请 GPU 执行推理
-# - 支持单人：AT2V / AI2V
-#
-# 说明：
-# 1) 官方示例使用 torchrun nproc=2（多进程/可能更快）：
-#    这里默认改为 nproc=1 + context_parallel_size=1，更适合 Spaces。
-# 2) FlashAttention 默认在 config 开启，但在 Spaces 上未必能顺利安装；
-#    本示例会尝试把 config 里所有包含 "flash" 的 attention backend 字段递归替换为 "sdpa"。
-#
-# 参考：
-# - ZeroGPU 官方用法：@spaces.GPU(duration=...)  :contentReference[oaicite:5]{index=5}（用户侧不需要引用，代码内不写引用）
-# - LongCat-Video-Avatar 模型卡：推理命令/参数/权重目录结构 :contentReference[oaicite:6]{index=6}
+# -*- coding: utf-8 -*-
+"""
+LongCat-Video-Avatar | Hugging Face Spaces (ZeroGPU) | 单文件 Gradio 应用
+- 自动：clone 推理代码仓库 + 下载权重到 ./weights
+- ZeroGPU：用 @spaces.GPU 在 fork 子进程里执行 CUDA 推理
+- 输入：单人/双人（多音频）模板 JSON 自动填充
+"""
 
 import os
-import re
 import sys
 import json
 import time
 import shutil
-import zipfile
-import hashlib
+import copy
+import re
 import subprocess
 from pathlib import Path
-from datetime import datetime
-from typing import Any, Dict, Tuple, Optional
-
-# ----------------------------
-# 运行时“尽量单文件”的依赖安装
-# ----------------------------
-def _pip_install(pkgs):
-    """在 Spaces 里尽量避免反复安装：用一个标记文件 + 简单 import 探测。"""
-    cmd = [sys.executable, "-m", "pip", "install", "-U"] + pkgs
-    print("[pip]", " ".join(cmd))
+from typing import Any, Dict, List, Tuple, Optional
+
+# -------------------- 基础路径 --------------------
+ROOT = Path(__file__).resolve().parent
+REPO_DIR = ROOT / "LongCat-Video"
+WEIGHTS_DIR = ROOT / "weights"
+WEIGHTS_LONGCAT_VIDEO = WEIGHTS_DIR / "LongCat-Video"
+WEIGHTS_AVATAR = WEIGHTS_DIR / "LongCat-Video-Avatar"
+
+# Hugging Face 仓库（权重）
+HF_REPO_LONGCAT_VIDEO = "meituan-longcat/LongCat-Video"
+HF_REPO_AVATAR = "meituan-longcat/LongCat-Video-Avatar"
+
+# GitHub 代码仓库（推理脚本/实现）
+GIT_REPO_URL = "https://github.com/meituan-longcat/LongCat-Video.git"
+GIT_BRANCH = "main"
+
+# 自举标记：避免每次启动都 pip install
+BOOTSTRAP_MARK = ROOT / ".bootstrap_done"
+
+# -------------------- 依赖自举（单文件策略） --------------------
+def _pip_install(args: List[str]) -> None:
+    cmd = [sys.executable, "-m", "pip", "install", "--no-cache-dir"] + args
+    print("[pip]", " ".join(cmd), flush=True)
     subprocess.check_call(cmd)
 
-def _ensure_imports():
+def _ensure_bootstrap() -> None:
     """
-    只安装本 app 直接需要的包。
-    LongCat-Video 自身依赖很多（官方 requirements），这里不强制全量预装，
-    而是交给官方脚本在运行时 import；若缺包会在日志里体现，再按需加到下面列表。
+    为了“单文件”，这里做最小自举：
+    1) 确保 gradio/spaces/huggingface_hub 等可用
+    2) clone repo 后，安装官方 requirements（若首次启动）
     """
-    try:
-        import gradio as gr  # noqa
-    except Exception:
-        _pip_install(["gradio>=4.0.0"])
-
-    try:
-        import requests  # noqa
-    except Exception:
-        _pip_install(["requests>=2.31.0"])
-
-    try:
-        from huggingface_hub import snapshot_download  # noqa
-    except Exception:
-        _pip_install(["huggingface_hub[cli]>=0.24.0"])
+    if BOOTSTRAP_MARK.exists():
+        return
 
-    # ZeroGPU 推荐的 spaces 包：多数 ZeroGPU 环境自带；没有就装
+    # 先装运行必须的基础包
+    base_pkgs = [
+        "gradio>=5.0.0",
+        "huggingface_hub[cli]>=0.24.0",
+        "gitpython>=3.1.0",
+        "spaces>=0.33.0",
+        "imageio-ffmpeg>=0.5.0",  # 提供 ffmpeg 可执行文件，避免系统缺 ffmpeg
+    ]
     try:
-        import spaces  # noqa
-    except Exception:
-        _pip_install(["spaces>=0.27.0"])
+        _pip_install(base_pkgs)
+    except Exception as e:
+        # 如果某些包已存在或网络抖动，仍继续尝试后续步骤
+        print("[bootstrap] base pip install warning:", repr(e), flush=True)
 
-_ensure_imports()
+    # clone 代码仓库（若未 clone）
+    _ensure_repo_cloned()
 
-import gradio as gr
-import requests
-from huggingface_hub import snapshot_download
+    # 安装官方 requirements（可能很大；只在首次启动做）
+    # 注意：官方 README/Model Card 提到 requirements.txt + requirements_avatar.txt。:contentReference[oaicite:5]{index=5}
+    req_main = REPO_DIR / "requirements.txt"
+    req_avatar = REPO_DIR / "requirements_avatar.txt"
+    if req_main.exists():
+        try:
+            _pip_install(["-r", str(req_main)])
+        except Exception as e:
+            print("[bootstrap] install requirements.txt warning:", repr(e), flush=True)
+    if req_avatar.exists():
+        try:
+            _pip_install(["-r", str(req_avatar)])
+        except Exception as e:
+            print("[bootstrap] install requirements_avatar.txt warning:", repr(e), flush=True)
 
-# spaces 在非 ZeroGPU 环境也应可安全使用；若导入失败已在上面安装
-import spaces
+    # librosa/ffmpeg 在官方说明里是 conda 安装。Space 没有 conda，这里用 pip + imageio-ffmpeg 兜底。:contentReference[oaicite:6]{index=6}
+    try:
+        _pip_install(["librosa>=0.10.0", "soundfile>=0.12.0"])
+    except Exception as e:
+        print("[bootstrap] install librosa/soundfile warning:", repr(e), flush=True)
 
+    BOOTSTRAP_MARK.write_text(f"ok {time.time()}\n", encoding="utf-8")
 
-# ----------------------------
-# 配置区（可按需改）
-# ----------------------------
-GITHUB_ZIP_URL = "https://github.com/meituan-longcat/LongCat-Video/archive/refs/heads/main.zip"
 
-# HF 权重（模型卡说明的目录） :contentReference[oaicite:7]{index=7}
-HF_MODEL_LONGCAT_VIDEO = "meituan-longcat/LongCat-Video"
-HF_MODEL_LONGCAT_AVATAR = "meituan-longcat/LongCat-Video-Avatar"
+# -------------------- Repo/权重准备 --------------------
+def _ensure_repo_cloned() -> None:
+    if REPO_DIR.exists() and (REPO_DIR / ".git").exists():
+        return
 
-# 本地缓存目录：Spaces 上建议放到 /home/user 或当前目录
-BASE_DIR = Path(__file__).parent.resolve()
-CACHE_DIR = BASE_DIR / "_cache"
-REPO_DIR = CACHE_DIR / "LongCat-Video-main"  # zip 解压后的目录名
-WEIGHTS_DIR = CACHE_DIR / "weights"
-WEIGHTS_LONGCAT_VIDEO = WEIGHTS_DIR / "LongCat-Video"
-WEIGHTS_LONGCAT_AVATAR = WEIGHTS_DIR / "LongCat-Video-Avatar"
-OUTPUT_DIR = CACHE_DIR / "outputs"
-TMP_DIR = CACHE_DIR / "tmp"
+    REPO_DIR.mkdir(parents=True, exist_ok=True)
+    # 如果目录非空，先清理，避免 git clone 失败
+    if any(REPO_DIR.iterdir()):
+        shutil.rmtree(REPO_DIR)
+        REPO_DIR.mkdir(parents=True, exist_ok=True)
 
-# 为了减少 torch CUDA 内存碎片（有时有用）
-os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+    print("[git] cloning repo...", flush=True)
+    subprocess.check_call([
+        "git", "clone", "--single-branch", "--branch", GIT_BRANCH, GIT_REPO_URL, str(REPO_DIR)
+    ])
+    print("[git] cloned:", REPO_DIR, flush=True)
 
+def _hf_snapshot_download(repo_id: str, local_dir: Path) -> None:
+    """
+    使用 huggingface_hub 下载权重到本地目录（会自动缓存并增量更新）。
+    官方模型卡建议用 huggingface-cli download 到 ./weights/... :contentReference[oaicite:7]{index=7}
+    """
+    from huggingface_hub import snapshot_download
 
-# ----------------------------
-# 工具函数
-# ----------------------------
-def _sha1(s: str) -> str:
-    return hashlib.sha1(s.encode("utf-8")).hexdigest()[:10]
+    local_dir.mkdir(parents=True, exist_ok=True)
+    print(f"[hf] downloading {repo_id} -> {local_dir}", flush=True)
 
-def _run(cmd, cwd: Optional[Path] = None, env: Optional[Dict[str, str]] = None) -> Tuple[int, str]:
-    """运行命令并返回 (code, stdout+stderr)。"""
-    print("[run]", " ".join(cmd))
-    p = subprocess.Popen(
-        cmd,
-        cwd=str(cwd) if cwd else None,
-        env=env,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.STDOUT,
-        text=True,
-        bufsize=1,
-        universal_newlines=True,
+    # local_dir_use_symlinks=False：在 Spaces 环境里更稳
+    snapshot_download(
+        repo_id=repo_id,
+        local_dir=str(local_dir),
+        local_dir_use_symlinks=False,
+        resume_download=True,
     )
-    out_lines = []
-    while True:
-        line = p.stdout.readline()
-        if not line and p.poll() is not None:
-            break
-        if line:
-            out_lines.append(line)
-    code = p.wait()
-    return code, "".join(out_lines)
-
-def _download_and_extract_repo():
-    """下载并解压 GitHub zip 到 CACHE_DIR。"""
-    CACHE_DIR.mkdir(parents=True, exist_ok=True)
-    zip_path = CACHE_DIR / "LongCat-Video-main.zip"
-
-    if REPO_DIR.exists() and (REPO_DIR / "run_demo_avatar_single_audio_to_video.py").exists():
-        return
+    print(f"[hf] done: {repo_id}", flush=True)
 
-    # 清理旧目录
-    if REPO_DIR.exists():
-        shutil.rmtree(REPO_DIR, ignore_errors=True)
-
-    # 下载 zip
-    if not zip_path.exists():
-        r = requests.get(GITHUB_ZIP_URL, stream=True, timeout=120)
-        r.raise_for_status()
-        with open(zip_path, "wb") as f:
-            for chunk in r.iter_content(chunk_size=1024 * 1024):
-                if chunk:
-                    f.write(chunk)
-
-    # 解压
-    with zipfile.ZipFile(zip_path, "r") as zf:
-        zf.extractall(CACHE_DIR)
-
-    # 基本校验
-    if not (REPO_DIR / "run_demo_avatar_single_audio_to_video.py").exists():
-        raise RuntimeError("仓库解压后未找到 run_demo_avatar_single_audio_to_video.py，可能 GitHub 结构变化。")
-
-def _download_weights():
-    """下载 HF 权重到 WEIGHTS_DIR。"""
+def _ensure_weights_downloaded() -> None:
     WEIGHTS_DIR.mkdir(parents=True, exist_ok=True)
 
-    # 使用 token（若你在 Space Secrets 里配置了 HF_TOKEN）
-    token = os.environ.get("HF_TOKEN", None)
+    # LongCat-Video
+    if not (WEIGHTS_LONGCAT_VIDEO.exists() and any(WEIGHTS_LONGCAT_VIDEO.iterdir())):
+        _hf_snapshot_download(HF_REPO_LONGCAT_VIDEO, WEIGHTS_LONGCAT_VIDEO)
 
-    if not WEIGHTS_LONGCAT_VIDEO.exists():
-        snapshot_download(
-            repo_id=HF_MODEL_LONGCAT_VIDEO,
-            local_dir=str(WEIGHTS_LONGCAT_VIDEO),
-            token=token,
-            local_dir_use_symlinks=False,
-        )
+    # LongCat-Video-Avatar
+    if not (WEIGHTS_AVATAR.exists() and any(WEIGHTS_AVATAR.iterdir())):
+        _hf_snapshot_download(HF_REPO_AVATAR, WEIGHTS_AVATAR)
 
-    if not WEIGHTS_LONGCAT_AVATAR.exists():
-        snapshot_download(
-            repo_id=HF_MODEL_LONGCAT_AVATAR,
-            local_dir=str(WEIGHTS_LONGCAT_AVATAR),
-            token=token,
-            local_dir_use_symlinks=False,
+
+# -------------------- JSON 模板读取与“灌参” --------------------
+def _load_json(p: Path) -> Any:
+    return json.loads(p.read_text(encoding="utf-8"))
+
+def _find_avatar_templates() -> Tuple[Path, Path]:
+    """
+    读取官方 repo 里自带的模板 JSON：
+      - assets/avatar/single_example_1.json
+      - assets/avatar/multi_example_1.json
+    """
+    single = REPO_DIR / "assets" / "avatar" / "single_example_1.json"
+    multi = REPO_DIR / "assets" / "avatar" / "multi_example_1.json"
+    if not single.exists() or not multi.exists():
+        raise FileNotFoundError(
+            "未找到 assets/avatar/single_example_1.json 或 multi_example_1.json。"
+            "请确认仓库结构与官方一致。"
         )
+    return single, multi
 
-def _recursive_patch_attention_backend(obj: Any) -> Any:
+def _collect_string_nodes(obj: Any, path: str = "") -> List[Tuple[str, str]]:
     """
-    递归把 config 里疑似 flash-attn backend 的字段替换为 sdpa。
-    不依赖具体 key 名，尽量“宽松匹配”：
-      - key 或 value 里出现 flash / flashattn / flash_attn => 改成 "sdpa"
+    收集所有字符串叶子节点：返回 (json_path, value)
     """
+    out = []
     if isinstance(obj, dict):
-        new = {}
         for k, v in obj.items():
-            lk = str(k).lower()
-            if any(x in lk for x in ["attn", "attention", "backend"]):
-                # 先递归处理 value
-                vv = _recursive_patch_attention_backend(v)
-                # 再判断是否需要替换
-                if isinstance(vv, str) and ("flash" in vv.lower() or "flash_attn" in vv.lower() or "flashattn" in vv.lower()):
-                    new[k] = "sdpa"
-                else:
-                    new[k] = vv
-            else:
-                new[k] = _recursive_patch_attention_backend(v)
-        return new
+            out.extend(_collect_string_nodes(v, f"{path}.{k}" if path else str(k)))
     elif isinstance(obj, list):
-        return [_recursive_patch_attention_backend(x) for x in obj]
-    else:
-        # 普通标量
-        if isinstance(obj, str):
-            lo = obj.lower()
-            if "flash_attn" in lo or "flashattn" in lo or lo.strip() == "flash" or "flash" == lo.strip():
-                return "sdpa"
-        return obj
-
-def _try_patch_avatar_configs():
-    """
-    官方说明：avatar_single/config.json 和 avatar_multi/config.json 默认启用 FlashAttention-2 :contentReference[oaicite:8]{index=8}
-    这里尽量替换为 sdpa，避免必须安装 flash-attn。
-    """
-    cfgs = [
-        WEIGHTS_LONGCAT_AVATAR / "avatar_single" / "config.json",
-        WEIGHTS_LONGCAT_AVATAR / "avatar_multi" / "config.json",
-    ]
-    for cfg in cfgs:
-        if not cfg.exists():
-            continue
-        try:
-            raw = json.loads(cfg.read_text(encoding="utf-8"))
-            patched = _recursive_patch_attention_backend(raw)
-            if patched != raw:
-                cfg.write_text(json.dumps(patched, ensure_ascii=False, indent=2), encoding="utf-8")
-        except Exception as e:
-            print(f"[warn] patch config failed: {cfg} -> {e}")
-
-def _load_template_json(template_path: Path) -> Dict[str, Any]:
-    data = json.loads(template_path.read_text(encoding="utf-8"))
-    if not isinstance(data, dict):
-        raise ValueError("模板 JSON 不是 dict 结构，无法安全修改。")
-    return data
+        for i, v in enumerate(obj):
+            out.extend(_collect_string_nodes(v, f"{path}[{i}]"))
+    elif isinstance(obj, str):
+        out.append((path, obj))
+    return out
 
-def _recursive_replace_first_match(data: Any, key_pred, value_pred, new_value) -> Tuple[Any, bool]:
+def _set_by_path(obj: Any, path: str, value: Any) -> None:
     """
-    在任意 JSON 结构中，找到第一个满足条件的 (key, value) 并替换 value。
-    返回 (new_data, replaced?)
+    按类似 a.b[0].c 的路径写入值
     """
-    if isinstance(data, dict):
-        out = {}
-        replaced = False
-        for k, v in data.items():
-            if (not replaced) and key_pred(k) and value_pred(v):
-                out[k] = new_value
-                replaced = True
-            else:
-                nv, r = _recursive_replace_first_match(v, key_pred, value_pred, new_value)
-                out[k] = nv
-                replaced = replaced or r
-        return out, replaced
-    elif isinstance(data, list):
-        out_list = []
-        replaced = False
-        for item in data:
-            if replaced:
-                out_list.append(item)
-                continue
-            nv, r = _recursive_replace_first_match(item, key_pred, value_pred, new_value)
-            out_list.append(nv)
-            replaced = replaced or r
-        return out_list, replaced
+    cur = obj
+    # 分割 tokens：key / [idx]
+    tokens = []
+    i = 0
+    while i < len(path):
+        if path[i] == "[":
+            j = path.index("]", i)
+            tokens.append(("idx", int(path[i+1:j])))
+            i = j + 1
+        elif path[i] == ".":
+            i += 1
+        else:
+            j = i
+            while j < len(path) and path[j] not in ".[":
+                j += 1
+            tokens.append(("key", path[i:j]))
+            i = j
+
+    for ttype, tval in tokens[:-1]:
+        if ttype == "key":
+            cur = cur[tval]
+        else:
+            cur = cur[tval]
+
+    last_type, last_val = tokens[-1]
+    if last_type == "key":
+        cur[last_val] = value
     else:
-        return data, False
+        cur[last_val] = value
 
-def _build_input_json_single(
-    mode: str,
-    audio_path: Path,
+def _patch_template_with_inputs(
+    template: Any,
     prompt: str,
-    ref_image_path: Optional[Path],
-    seed: int,
-    resolution: str
-) -> Path:
+    image_paths: List[str],
+    audio_paths: List[str],
+) -> Any:
     """
-    基于 assets/avatar/single_example_1.json 模板生成 input_json。
-    官方脚本以 --input_json 读取参数 :contentReference[oaicite:9]{index=9}
+    不依赖 schema 的通用替换：
+    - 按出现顺序替换模板里第 N 个“像音频路径”的字符串为 audio_paths[N]
+    - 按出现顺序替换模板里第 N 个“像图片路径”的字符串为 image_paths[N]
+    - 尝试替换常见 prompt 字段
     """
-    template = REPO_DIR / "assets" / "avatar" / "single_example_1.json"
-    if not template.exists():
-        raise RuntimeError("未找到模板 assets/avatar/single_example_1.json（仓库结构可能变化）。")
-
-    data = _load_template_json(template)
-
-    # 替换 prompt：优先找 key 包含 prompt/text 之类
-    data, _ = _recursive_replace_first_match(
-        data,
-        key_pred=lambda k: "prompt" in str(k).lower() or "text" in str(k).lower(),
-        value_pred=lambda v: isinstance(v, str),
-        new_value=prompt.strip() if prompt else "A person is talking."
-    )
+    patched = copy.deepcopy(template)
+    string_nodes = _collect_string_nodes(patched)
+
+    # 识别“可能是音频/图片路径”的节点（按出现顺序）
+    audio_like = []
+    image_like = []
+    prompt_like = []
+
+    for pth, val in string_nodes:
+        low = val.lower()
+        # 音频后缀或路径特征
+        if any(low.endswith(ext) for ext in [".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg"]):
+            audio_like.append((pth, val))
+        # 图片后缀
+        if any(low.endswith(ext) for ext in [".png", ".jpg", ".jpeg", ".webp", ".bmp"]):
+            image_like.append((pth, val))
+        # prompt 字段（按路径名判断更靠谱）
+        if re.search(r"(prompt|caption|text|instruction)$", pth, re.IGNORECASE):
+            prompt_like.append((pth, val))
+
+    # 替换音频（按顺序）
+    for idx, (pth, _) in enumerate(audio_like):
+        if idx < len(audio_paths):
+            _set_by_path(patched, pth, audio_paths[idx])
+
+    # 替换图片（按顺序）
+    for idx, (pth, _) in enumerate(image_like):
+        if idx < len(image_paths):
+            _set_by_path(patched, pth, image_paths[idx])
+
+    # 替换 prompt（如果模板里有多个 prompt 字段，就全写同一个）
+    if prompt.strip():
+        for pth, _ in prompt_like:
+            _set_by_path(patched, pth, prompt.strip())
+
+    return patched
+
+
+# -------------------- 推理执行（调用官方脚本） --------------------
+def _ensure_ffmpeg_in_path() -> None:
+    """
+    使用 imageio-ffmpeg 提供的 ffmpeg，把它加入 PATH。
+    """
+    try:
+        import imageio_ffmpeg
+        ffmpeg_exe = imageio_ffmpeg.get_ffmpeg_exe()
+        ffmpeg_dir = str(Path(ffmpeg_exe).parent)
+        os.environ["PATH"] = ffmpeg_dir + os.pathsep + os.environ.get("PATH", "")
+        os.environ["IMAGEIO_FFMPEG_EXE"] = ffmpeg_exe
+        print("[ffmpeg] using:", ffmpeg_exe, flush=True)
+    except Exception as e:
+        print("[ffmpeg] warning:", repr(e), flush=True)
 
-    # 替换 audio path：找 key 包含 audio 且 value 是字符串
-    data, _ = _recursive_replace_first_match(
-        data,
-        key_pred=lambda k: "audio" in str(k).lower(),
-        value_pred=lambda v: isinstance(v, str),
-        new_value=str(audio_path)
+def _run_subprocess(cmd: List[str], cwd: Path) -> Tuple[int, str]:
+    """
+    运行命令并收集输出（stdout+stderr）
+    """
+    print("[cmd]", " ".join(cmd), flush=True)
+    p = subprocess.Popen(
+        cmd,
+        cwd=str(cwd),
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        bufsize=1,
+        universal_newlines=True,
     )
+    lines = []
+    assert p.stdout is not None
+    for line in p.stdout:
+        lines.append(line)
+    p.wait()
+    out = "".join(lines)
+    return p.returncode, out
+
+def _find_latest_video(search_dir: Path) -> Optional[str]:
+    mp4s = list(search_dir.rglob("*.mp4"))
+    if not mp4s:
+        return None
+    mp4s.sort(key=lambda p: p.stat().st_mtime, reverse=True)
+    return str(mp4s[0])
 
-    # 替换 image path（仅 AI2V）
-    if mode == "ai2v" and ref_image_path is not None:
-        data, _ = _recursive_replace_first_match(
-            data,
-            key_pred=lambda k: ("image" in str(k).lower()) or ("ref" in str(k).lower()),
-            value_pred=lambda v: isinstance(v, str),
-            new_value=str(ref_image_path)
-        )
+def _extract_video_path_from_log(log: str) -> Optional[str]:
+    # 从日志里提取类似 xxx.mp4 的路径
+    cand = re.findall(r"([^\s\"']+\.mp4)", log)
+    if not cand:
+        return None
+    # 取最后一个更可能是输出
+    return cand[-1]
 
-    # seed（若模板里有）
-    data, _ = _recursive_replace_first_match(
-        data,
-        key_pred=lambda k: "seed" in str(k).lower(),
-        value_pred=lambda v: isinstance(v, (int, float, str)),
-        new_value=int(seed)
-    )
+def _prepare_runtime() -> None:
+    """
+    启动阶段准备：
+    - 自举依赖
+    - clone repo
+    - 下载权重
+    - ffmpeg 兜底
+    - sys.path 加入 repo（以便脚本 import）
+    """
+    _ensure_bootstrap()
+    _ensure_repo_cloned()
+    _ensure_weights_downloaded()
+    _ensure_ffmpeg_in_path()
 
-    # resolution（若模板里有）
-    data, _ = _recursive_replace_first_match(
-        data,
-        key_pred=lambda k: "resolution" in str(k).lower(),
-        value_pred=lambda v: isinstance(v, str),
-        new_value=str(resolution)
-    )
+    if str(REPO_DIR) not in sys.path:
+        sys.path.insert(0, str(REPO_DIR))
 
-    TMP_DIR.mkdir(parents=True, exist_ok=True)
-    out_path = TMP_DIR / f"single_{mode}_{_sha1(str(audio_path) + prompt + str(seed) + str(time.time()))}.json"
-    out_path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
-    return out_path
 
-def _find_latest_mp4(since_ts: float) -> Optional[Path]:
-    if not OUTPUT_DIR.exists():
-        return None
-    candidates = []
-    for p in OUTPUT_DIR.rglob("*.mp4"):
-        try:
-            if p.stat().st_mtime >= since_ts - 2:
-                candidates.append(p)
-        except Exception:
-            pass
-    if not candidates:
-        return None
-    candidates.sort(key=lambda x: x.stat().st_mtime, reverse=True)
-    return candidates[0]
+# -------------------- Gradio / ZeroGPU --------------------
+_prepare_runtime()
 
-def _ensure_ready() -> str:
+import gradio as gr
+import spaces
+
+
+def _save_upload_to_dir(upload_path: str, dst_dir: Path, prefix: str) -> str:
     """
-    准备：
-    - 下载 repo
-    - 下载权重
-    - 尝试 patch attention backend
+    将 gradio 上传的临时文件复制到工作目录，返回新路径
     """
-    _download_and_extract_repo()
-    _download_weights()
-    _try_patch_avatar_configs()
-
-    # 输出目录
-    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    dst_dir.mkdir(parents=True, exist_ok=True)
+    src = Path(upload_path)
+    ext = src.suffix
+    dst = dst_dir / f"{prefix}_{int(time.time()*1000)}{ext}"
+    shutil.copy2(src, dst)
+    return str(dst)
+
+def _build_input_json_file(
+    mode: str,
+    prompt: str,
+    images: List[Optional[str]],
+    audios: List[Optional[str]],
+) -> str:
+    """
+    mode: "single" or "multi"
+    """
+    single_tpl_path, multi_tpl_path = _find_avatar_templates()
+    tpl = _load_json(single_tpl_path if mode == "single" else multi_tpl_path)
+
+    work_dir = REPO_DIR / "assets" / "avatar" / "custom_inputs"
+    img_paths = []
+    aud_paths = []
+
+    # 复制输入文件到 repo 内（避免脚本用相对路径时找不到）
+    for i, p in enumerate(images):
+        if p:
+            img_paths.append(_save_upload_to_dir(p, work_dir, f"img{i+1}"))
+    for i, p in enumerate(audios):
+        if p:
+            aud_paths.append(_save_upload_to_dir(p, work_dir, f"aud{i+1}"))
+
+    patched = _patch_template_with_inputs(
+        template=tpl,
+        prompt=prompt or "",
+        image_paths=img_paths,
+        audio_paths=aud_paths,
+    )
 
-    return "准备完成：代码与权重已就绪。"
+    out_json = work_dir / f"input_{mode}_{int(time.time()*1000)}.json"
+    out_json.write_text(json.dumps(patched, ensure_ascii=False, indent=2), encoding="utf-8")
+    return str(out_json)
 
 
-# ----------------------------
-# GPU 推理函数（ZeroGPU 核心）
-# ----------------------------
-@spaces.GPU(duration=900)  # 生成视频通常 >60s，给足时间；你可视情况调小/调大 :contentReference[oaicite:10]{index=10}
+@spaces.GPU
 def generate_single(
-    mode: str,
-    audio_file: str,
+    stage_1: str,
     prompt: str,
-    ref_image_file: Optional[str],
-    seed: int,
+    image_path: Optional[str],
+    audio_path: str,
     resolution: str,
     num_segments: int,
     ref_img_index: int,
     mask_frame_range: int,
+    nproc: int,
+    context_parallel_size: int,
 ) -> Tuple[Optional[str], str]:
     """
-    返回：(mp4路径 or None, 日志文本)
+    单人：Audio-Text-to-Video (at2v) 或 Audio-Image-to-Video (ai2v)
+    注意：官方示例使用 torchrun nproc=2 / context_parallel_size=2。:contentReference[oaicite:8]{index=8}
+    ZeroGPU 下默认先尝试 nproc=1。
     """
-    t0 = time.time()
-
-    # 文件落盘路径（Gradio 传入的是本地临时文件路径字符串）
-    audio_path = Path(audio_file).resolve()
-    ref_image_path = Path(ref_image_file).resolve() if ref_image_file else None
-
-    # 构造 input_json
-    input_json = _build_input_json_single(
-        mode=mode,
-        audio_path=audio_path,
+    input_json = _build_input_json_file(
+        mode="single",
         prompt=prompt,
-        ref_image_path=ref_image_path,
-        seed=seed,
-        resolution=resolution,
+        images=[image_path] if image_path else [],
+        audios=[audio_path],
     )
 
-    # 运行官方脚本（单进程 torchrun）
-    # 官方示例：torchrun --nproc_per_node=2 ... --context_parallel_size=2 ... :contentReference[oaicite:11]{index=11}
-    # 这里适配 Space：nproc=1, context_parallel_size=1
+    # 分辨率参数（官方说明可 480P/720P）:contentReference[oaicite:9]{index=9}
+    # 这里不假设脚本参数名，直接透传 --resolution
     cmd = [
-        sys.executable, "-m", "torch.distributed.run",
-        "--nproc_per_node=1",
+        "torchrun",
+        f"--nproc_per_node={nproc}",
         "run_demo_avatar_single_audio_to_video.py",
-        "--context_parallel_size=1",
-        f"--checkpoint_dir={WEIGHTS_LONGCAT_AVATAR}",
-        f"--stage_1={mode}",
+        f"--context_parallel_size={context_parallel_size}",
+        f"--checkpoint_dir={str(WEIGHTS_AVATAR)}",
+        f"--stage_1={stage_1}",
         f"--input_json={input_json}",
         f"--resolution={resolution}",
+        f"--num_segments={num_segments}",
+        f"--ref_img_index={ref_img_index}",
+        f"--mask_frame_range={mask_frame_range}",
     ]
 
-    # 续写参数（用户设置 >1 才启用）
-    if num_segments and int(num_segments) > 1:
-        cmd += [
-            f"--num_segments={int(num_segments)}",
-            f"--ref_img_index={int(ref_img_index)}",
-            f"--mask_frame_range={int(mask_frame_range)}",
-        ]
+    rc, log = _run_subprocess(cmd, cwd=REPO_DIR)
+    if rc != 0:
+        return None, f"运行失败（exit={rc}）。日志如下：\n{log}"
 
-    # 环境变量：让脚本能找到模块
-    env = dict(os.environ)
-    env["PYTHONPATH"] = str(REPO_DIR) + (os.pathsep + env["PYTHONPATH"] if env.get("PYTHONPATH") else "")
-    env["HF_HOME"] = str(CACHE_DIR / "hf_home")
-    env["TORCH_HOME"] = str(CACHE_DIR / "torch_home")
+    vid = _extract_video_path_from_log(log)
+    if vid:
+        # 相对路径转绝对
+        p = Path(vid)
+        if not p.is_absolute():
+            p = (REPO_DIR / p).resolve()
+        if p.exists():
+            return str(p), log
 
-    # 约定输出目录（若脚本支持/或脚本默认输出在当前目录下的 outputs）
-    # 我们用 cwd + 输出扫描兜底
-    env["OUTPUT_DIR"] = str(OUTPUT_DIR)
+    # 兜底：找最近生成的 mp4
+    fallback = _find_latest_video(REPO_DIR)
+    return fallback, log
 
-    code, log = _run(cmd, cwd=REPO_DIR, env=env)
 
-    # 尝试找到最新 mp4
-    mp4 = _find_latest_mp4(t0)
-    if mp4 is None:
-        # 兜底：在 repo 内也扫一下
-        repo_candidates = list(REPO_DIR.rglob("*.mp4"))
-        repo_candidates.sort(key=lambda x: x.stat().st_mtime, reverse=True)
-        if repo_candidates and repo_candidates[0].stat().st_mtime >= t0 - 2:
-            mp4 = repo_candidates[0]
+@spaces.GPU
+def generate_multi(
+    prompt: str,
+    image1: str,
+    image2: str,
+    audio1: str,
+    audio2: str,
+    audio_type: str,  # para/add
+    resolution: str,
+    num_segments: int,
+    ref_img_index: int,
+    mask_frame_range: int,
+    nproc: int,
+    context_parallel_size: int,
+) -> Tuple[Optional[str], str]:
+    """
+    多人：Audio-Image-to-Video（官方示例）:contentReference[oaicite:10]{index=10}
+    """
+    input_json = _build_input_json_file(
+        mode="multi",
+        prompt=prompt,
+        images=[image1, image2],
+        audios=[audio1, audio2],
+    )
 
-    if code != 0:
-        return None, f"执行失败（exit={code}）。日志如下：\n\n{log}"
+    cmd = [
+        "torchrun",
+        f"--nproc_per_node={nproc}",
+        "run_demo_avatar_multi_audio_to_video.py",
+        f"--context_parallel_size={context_parallel_size}",
+        f"--checkpoint_dir={str(WEIGHTS_AVATAR)}",
+        f"--input_json={input_json}",
+        f"--audio_type={audio_type}",
+        f"--resolution={resolution}",
+        f"--num_segments={num_segments}",
+        f"--ref_img_index={ref_img_index}",
+        f"--mask_frame_range={mask_frame_range}",
+    ]
 
-    if mp4 is None or not mp4.exists():
-        return None, f"执行完成，但未找到 mp4 输出文件。日志如下：\n\n{log}"
+    rc, log = _run_subprocess(cmd, cwd=REPO_DIR)
+    if rc != 0:
+        return None, f"运行失败（exit={rc}）。日志如下：\n{log}"
 
-    return str(mp4), f"执行成功：{mp4}\n\n日志如下：\n\n{log}"
+    vid = _extract_video_path_from_log(log)
+    if vid:
+        p = Path(vid)
+        if not p.is_absolute():
+            p = (REPO_DIR / p).resolve()
+        if p.exists():
+            return str(p), log
 
+    fallback = _find_latest_video(REPO_DIR)
+    return fallback, log
 
-# ----------------------------
-# Gradio UI
-# ----------------------------
-def ui_prepare() -> str:
-    try:
-        return _ensure_ready()
-    except Exception as e:
-        return f"准备失败：{e}"
 
-with gr.Blocks(title="LongCat-Video-Avatar (ZeroGPU) - Single File Space") as demo:
+# -------------------- UI --------------------
+with gr.Blocks(title="LongCat-Video-Avatar (ZeroGPU)", fill_height=True) as demo:
     gr.Markdown(
-        """
-# LongCat-Video-Avatar（ZeroGPU / 单文件 Space）
-
-- 单人模式：**AT2V（音频+文本）** / **AI2V（音频+图片）**
-- 续写（Video Continuation）：把 **num_segments** 设为 > 1 即可（官方参数：ref_img_index / mask_frame_range）  
-- 提示：为了更自然的口型，prompt 里建议包含 talking/speaking 等动作词（模型卡建议）  
-"""
+        "## LongCat-Video-Avatar (ZeroGPU)\n"
+        "- 启动后会自动下载权重到 `./weights` 并准备环境。\n"
+        "- ZeroGPU 会在点击生成时，按需分配 GPU 执行（@spaces.GPU）。\n"
+        "- 如果你发现必须 2 卡才能跑通，可把 **nproc/context_parallel_size** 改为 2。"
     )
 
-    with gr.Row():
-        btn_prepare = gr.Button("一键准备（下载代码+权重）", variant="primary")
-        prep_status = gr.Textbox(label="准备状态", value="尚未准备。首次准备会下载较大权重。", lines=2)
-
-    btn_prepare.click(fn=ui_prepare, outputs=prep_status)
-
-    with gr.Row():
-        mode = gr.Radio(
-            choices=[("Audio-Text-to-Video (AT2V)", "at2v"), ("Audio-Image-to-Video (AI2V)", "ai2v")],
-            value="at2v",
-            label="模式"
-        )
-
-    with gr.Row():
-        audio_in = gr.Audio(label="输入音频（wav/mp3等）", type="filepath")
-        ref_img = gr.Image(label="参考图（仅 AI2V 需要）", type="filepath")
-
-    prompt = gr.Textbox(
-        label="Prompt（建议包含 talking/speaking 等动作词）",
-        value="A young person is talking naturally, realistic style.",
-        lines=2
-    )
+    with gr.Accordion("高级参数（默认先按 ZeroGPU 更稳的 1 卡尝试）", open=False):
+        nproc = gr.Slider(1, 2, value=1, step=1, label="torchrun --nproc_per_node")
+        cps = gr.Slider(1, 2, value=1, step=1, label="--context_parallel_size")
+        resolution = gr.Radio(["480p", "720p"], value="480p", label="resolution")
+        num_segments = gr.Slider(1, 8, value=1, step=1, label="num_segments（>1 启用续写/长视频段）")
+        ref_img_index = gr.Slider(-24, 48, value=10, step=1, label="ref_img_index（减少重复动作/增强一致性）")
+        mask_frame_range = gr.Slider(0, 12, value=3, step=1, label="mask_frame_range（过大可能出伪影）")
+
+    with gr.Tabs():
+        with gr.Tab("单人（AT2V / AI2V）"):
+            stage_1 = gr.Radio(["at2v", "ai2v"], value="ai2v", label="stage_1")
+            prompt = gr.Textbox(
+                label="文本提示（建议包含 talking/speaking 等动词提示）",
+                value="A realistic person is speaking naturally, talking to the camera.",
+                lines=2,
+            )
+            img = gr.Image(type="filepath", label="参考图（ai2v 必填，at2v 可不传）")
+            aud = gr.Audio(type="filepath", label="音频（必填）")
+            btn = gr.Button("生成", variant="primary")
+            out_v = gr.Video(label="输出视频")
+            out_log = gr.Textbox(label="日志", lines=12)
+
+            btn.click(
+                fn=generate_single,
+                inputs=[stage_1, prompt, img, aud, resolution, num_segments, ref_img_index, mask_frame_range, nproc, cps],
+                outputs=[out_v, out_log],
+                api_name="generate_single",
+            )
+
+        with gr.Tab("双人（Multi）"):
+            prompt_m = gr.Textbox(
+                label="文本提示（可选）",
+                value="Two people are talking in turns naturally, facing the camera.",
+                lines=2,
+            )
+            c1, c2 = gr.Row(), gr.Row()
+            with c1:
+                img1 = gr.Image(type="filepath", label="人物1参考图（必填）")
+                aud1 = gr.Audio(type="filepath", label="人物1音频（必填）")
+            with c2:
+                img2 = gr.Image(type="filepath", label="人物2参考图（必填）")
+                aud2 = gr.Audio(type="filepath", label="人物2音频（必填）")
+
+            audio_type = gr.Radio(
+                ["para", "add"],
+                value="add",
+                label="双音频模式：para=混合(等长) / add=拼接(可不等长)",
+            )
+            btn2 = gr.Button("生成（双人）", variant="primary")
+            out_v2 = gr.Video(label="输出视频")
+            out_log2 = gr.Textbox(label="日志", lines=12)
+
+            btn2.click(
+                fn=generate_multi,
+                inputs=[prompt_m, img1, img2, aud1, aud2, audio_type, resolution, num_segments, ref_img_index, mask_frame_range, nproc, cps],
+                outputs=[out_v2, out_log2],
+                api_name="generate_multi",
+            )
 
-    with gr.Row():
-        seed = gr.Number(label="Seed", value=0, precision=0)
-        resolution = gr.Dropdown(label="分辨率", choices=["480P", "720P"], value="480P")
-
-    with gr.Accordion("高级参数（续写/一致性/防重复）", open=False):
-        num_segments = gr.Slider(label="num_segments（>1 启用续写）", minimum=1, maximum=8, step=1, value=1)
-        ref_img_index = gr.Slider(label="ref_img_index（默认 10）", minimum=-30, maximum=60, step=1, value=10)
-        mask_frame_range = gr.Slider(label="mask_frame_range（默认 3）", minimum=1, maximum=12, step=1, value=3)
-
-    btn = gr.Button("生成视频", variant="primary")
-
-    out_video = gr.Video(label="输出视频（mp4）")
-    out_log = gr.Textbox(label="运行日志", lines=18)
-
-    def _validate(mode_v, audio_fp, img_fp):
-        if not audio_fp:
-            raise gr.Error("请先上传音频。")
-        if mode_v == "ai2v" and not img_fp:
-            raise gr.Error("AI2V 模式必须上传参考图。")
-
-    def run(mode_v, audio_fp, prompt_v, img_fp, seed_v, res_v, seg_v, idx_v, mask_v):
-        _validate(mode_v, audio_fp, img_fp)
-        # seed=0 时也允许；如果想随机可自己改成 random
-        return generate_single(mode_v, audio_fp, prompt_v, img_fp, int(seed_v), res_v, int(seg_v), int(idx_v), int(mask_v))
-
-    btn.click(
-        fn=run,
-        inputs=[mode, audio_in, prompt, ref_img, seed, resolution, num_segments, ref_img_index, mask_frame_range],
-        outputs=[out_video, out_log],
+    gr.Markdown(
+        "### 重要说明\n"
+        "- 该模型当前 **没有 Inference Provider 托管**，因此 Space 必须本地跑推理代码与权重。:contentReference[oaicite:11]{index=11}\n"
+        "- ZeroGPU 的 CUDA 任务会在 `@spaces.GPU` 的函数调用时 fork 执行并释放。:contentReference[oaicite:12]{index=12}\n"
+        "- 官方示例的 Avatar 推理默认用 2 进程（nproc=2）。:contentReference[oaicite:13]{index=13}"
     )
 
-demo.queue(max_size=12).launch()
+demo.queue(concurrency_count=1).launch()