Update app.py

app.py

@@ -1,574 +1,525 @@
 # app.py
-# -*- coding: utf-8 -*-
-"""
-LongCat-Video-Avatar | Hugging Face Spaces (ZeroGPU) | 单文件 Gradio 应用
-- 自动：clone 推理代码仓库 + 下载权重到 ./weights
-- ZeroGPU：用 @spaces.GPU 在 fork 子进程里执行 CUDA 推理
-- 输入：单人/双人（多音频）模板 JSON 自动填充
-"""
+# 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}
 
 import os
+import re
 import sys
 import json
 import time
 import shutil
-import copy
-import re
+import zipfile
+import hashlib
 import subprocess
 from pathlib import Path
-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)
+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))
     subprocess.check_call(cmd)
 
-def _ensure_bootstrap() -> None:
+def _ensure_imports():
     """
-    为了“单文件”，这里做最小自举：
-    1) 确保 gradio/spaces/huggingface_hub 等可用
-    2) clone repo 后，安装官方 requirements（若首次启动）
+    只安装本 app 直接需要的包。
+    LongCat-Video 自身依赖很多（官方 requirements），这里不强制全量预装，
+    而是交给官方脚本在运行时 import；若缺包会在日志里体现，再按需加到下面列表。
     """
-    if BOOTSTRAP_MARK.exists():
-        return
-
-    # 先装运行必须的基础包
-    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:
-        _pip_install(base_pkgs)
-    except Exception as e:
-        # 如果某些包已存在或网络抖动，仍继续尝试后续步骤
-        print("[bootstrap] base pip install warning:", repr(e), flush=True)
+        import gradio as gr  # noqa
+    except Exception:
+        _pip_install(["gradio>=4.0.0"])
 
-    # clone 代码仓库（若未 clone）
-    _ensure_repo_cloned()
+    try:
+        import requests  # noqa
+    except Exception:
+        _pip_install(["requests>=2.31.0"])
 
-    # 安装官方 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)
+    try:
+        from huggingface_hub import snapshot_download  # noqa
+    except Exception:
+        _pip_install(["huggingface_hub[cli]>=0.24.0"])
 
-    # librosa/ffmpeg 在官方说明里是 conda 安装。Space 没有 conda，这里用 pip + imageio-ffmpeg 兜底。:contentReference[oaicite:6]{index=6}
+    # ZeroGPU 推荐的 spaces 包：多数 ZeroGPU 环境自带；没有就装
     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)
+        import spaces  # noqa
+    except Exception:
+        _pip_install(["spaces>=0.27.0"])
 
-    BOOTSTRAP_MARK.write_text(f"ok {time.time()}\n", encoding="utf-8")
+_ensure_imports()
 
+import gradio as gr
+import requests
+from huggingface_hub import snapshot_download
 
-# -------------------- Repo/权重准备 --------------------
-def _ensure_repo_cloned() -> None:
-    if REPO_DIR.exists() and (REPO_DIR / ".git").exists():
-        return
+# spaces 在非 ZeroGPU 环境也应可安全使用；若导入失败已在上面安装
+import spaces
 
-    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)
 
-    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)
+# ----------------------------
+# 配置区（可按需改）
+# ----------------------------
+GITHUB_ZIP_URL = "https://github.com/meituan-longcat/LongCat-Video/archive/refs/heads/main.zip"
 
-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
+# HF 权重（模型卡说明的目录） :contentReference[oaicite:7]{index=7}
+HF_MODEL_LONGCAT_VIDEO = "meituan-longcat/LongCat-Video"
+HF_MODEL_LONGCAT_AVATAR = "meituan-longcat/LongCat-Video-Avatar"
 
-    local_dir.mkdir(parents=True, exist_ok=True)
-    print(f"[hf] downloading {repo_id} -> {local_dir}", flush=True)
+# 本地缓存目录：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"
 
-    # 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,
-    )
-    print(f"[hf] done: {repo_id}", flush=True)
+# 为了减少 torch CUDA 内存碎片（有时有用）
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
 
-def _ensure_weights_downloaded() -> None:
-    WEIGHTS_DIR.mkdir(parents=True, exist_ok=True)
 
-    # LongCat-Video
-    if not (WEIGHTS_LONGCAT_VIDEO.exists() and any(WEIGHTS_LONGCAT_VIDEO.iterdir())):
-        _hf_snapshot_download(HF_REPO_LONGCAT_VIDEO, WEIGHTS_LONGCAT_VIDEO)
+# ----------------------------
+# 工具函数
+# ----------------------------
+def _sha1(s: str) -> str:
+    return hashlib.sha1(s.encode("utf-8")).hexdigest()[:10]
 
-    # LongCat-Video-Avatar
-    if not (WEIGHTS_AVATAR.exists() and any(WEIGHTS_AVATAR.iterdir())):
-        _hf_snapshot_download(HF_REPO_AVATAR, WEIGHTS_AVATAR)
+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,
+    )
+    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
 
+    # 清理旧目录
+    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。"""
+    WEIGHTS_DIR.mkdir(parents=True, exist_ok=True)
 
-# -------------------- JSON 模板读取与“灌参” --------------------
-def _load_json(p: Path) -> Any:
-    return json.loads(p.read_text(encoding="utf-8"))
+    # 使用 token（若你在 Space Secrets 里配置了 HF_TOKEN）
+    token = os.environ.get("HF_TOKEN", None)
 
-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。"
-            "请确认仓库结构与官方一致。"
+    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,
+        )
+
+    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,
         )
-    return single, multi
 
-def _collect_string_nodes(obj: Any, path: str = "") -> List[Tuple[str, str]]:
+def _recursive_patch_attention_backend(obj: Any) -> Any:
     """
-    收集所有字符串叶子节点：返回 (json_path, value)
+    递归把 config 里疑似 flash-attn backend 的字段替换为 sdpa。
+    不依赖具体 key 名，尽量“宽松匹配”：
+      - key 或 value 里出现 flash / flashattn / flash_attn => 改成 "sdpa"
     """
-    out = []
     if isinstance(obj, dict):
+        new = {}
         for k, v in obj.items():
-            out.extend(_collect_string_nodes(v, f"{path}.{k}" if path else str(k)))
+            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
     elif isinstance(obj, list):
-        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 _set_by_path(obj: Any, path: str, value: Any) -> None:
-    """
-    按类似 a.b[0].c 的路径写入值
-    """
-    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
+        return [_recursive_patch_attention_backend(x) for x in obj]
     else:
-        cur[last_val] = value
-
-def _patch_template_with_inputs(
-    template: Any,
-    prompt: str,
-    image_paths: List[str],
-    audio_paths: List[str],
-) -> Any:
+        # 普通标量
+        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():
     """
-    不依赖 schema 的通用替换：
-    - 按出现顺序替换模板里第 N 个“像音频路径”的字符串为 audio_paths[N]
-    - 按出现顺序替换模板里第 N 个“像图片路径”的字符串为 image_paths[N]
-    - 尝试替换常见 prompt 字段
+    官方说明：avatar_single/config.json 和 avatar_multi/config.json 默认��用 FlashAttention-2 :contentReference[oaicite:8]{index=8}
+    这里尽量替换为 sdpa，避免必须安装 flash-attn。
     """
-    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:
+    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
+
+def _recursive_replace_first_match(data: Any, key_pred, value_pred, new_value) -> Tuple[Any, bool]:
     """
-    使用 imageio-ffmpeg 提供的 ffmpeg，把它加入 PATH。
+    在任意 JSON 结构中，找到第一个满足条件的 (key, value) 并替换 value。
+    返回 (new_data, replaced?)
     """
-    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)
+    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
+    else:
+        return data, False
 
-def _run_subprocess(cmd: List[str], cwd: Path) -> Tuple[int, str]:
+def _build_input_json_single(
+    mode: str,
+    audio_path: Path,
+    prompt: str,
+    ref_image_path: Optional[Path],
+    seed: int,
+    resolution: str
+) -> Path:
     """
-    运行命令并收集输出（stdout+stderr）
+    基于 assets/avatar/single_example_1.json 模板生成 input_json。
+    官方脚本以 --input_json 读取参数 :contentReference[oaicite:9]{index=9}
     """
-    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,
+    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."
     )
-    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])
-
-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]
 
-def _prepare_runtime() -> None:
-    """
-    启动阶段准备：
-    - 自举依赖
-    - clone repo
-    - 下载权重
-    - ffmpeg 兜底
-    - sys.path 加入 repo（以便脚本 import）
-    """
-    _ensure_bootstrap()
-    _ensure_repo_cloned()
-    _ensure_weights_downloaded()
-    _ensure_ffmpeg_in_path()
+    # 替换 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)
+    )
 
-    if str(REPO_DIR) not in sys.path:
-        sys.path.insert(0, str(REPO_DIR))
+    # 替换 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)
+        )
 
+    # 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)
+    )
 
-# -------------------- Gradio / ZeroGPU --------------------
-_prepare_runtime()
+    # 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)
+    )
 
-import gradio as gr
-import spaces
+    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]
 
-def _save_upload_to_dir(upload_path: str, dst_dir: Path, prefix: str) -> str:
-    """
-    将 gradio 上传的临时文件复制到工作目录，返回新路径
-    """
-    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:
+def _ensure_ready() -> str:
     """
-    mode: "single" or "multi"
+    准备：
+    - 下载 repo
+    - 下载权重
+    - 尝试 patch attention backend
     """
-    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,
-    )
+    _download_and_extract_repo()
+    _download_weights()
+    _try_patch_avatar_configs()
 
-    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)
+    # 输出目录
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
 
+    return "准备完成：代码与权重已就绪。"
 
-@spaces.GPU
+
+# ----------------------------
+# GPU 推理函数（ZeroGPU 核心）
+# ----------------------------
+@spaces.GPU(duration=900)  # 生成视频通常 >60s，给足时间；你可视情况调小/调大 :contentReference[oaicite:10]{index=10}
 def generate_single(
-    stage_1: str,
+    mode: str,
+    audio_file: str,
     prompt: str,
-    image_path: Optional[str],
-    audio_path: str,
+    ref_image_file: Optional[str],
+    seed: int,
     resolution: str,
     num_segments: int,
     ref_img_index: int,
     mask_frame_range: int,
-    nproc: int,
-    context_parallel_size: int,
 ) -> Tuple[Optional[str], str]:
     """
-    单人：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。
+    返回：(mp4路径 or None, 日志文本)
     """
-    input_json = _build_input_json_file(
-        mode="single",
+    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,
         prompt=prompt,
-        images=[image_path] if image_path else [],
-        audios=[audio_path],
+        ref_image_path=ref_image_path,
+        seed=seed,
+        resolution=resolution,
     )
 
-    # 分辨率参数（官方说明可 480P/720P）:contentReference[oaicite:9]{index=9}
-    # 这里不假设脚本参数名，直接透传 --resolution
+    # 运行官方脚本（单进程 torchrun）
+    # 官方示例：torchrun --nproc_per_node=2 ... --context_parallel_size=2 ... :contentReference[oaicite:11]{index=11}
+    # 这里适配 Space：nproc=1, context_parallel_size=1
     cmd = [
-        "torchrun",
-        f"--nproc_per_node={nproc}",
+        sys.executable, "-m", "torch.distributed.run",
+        "--nproc_per_node=1",
         "run_demo_avatar_single_audio_to_video.py",
-        f"--context_parallel_size={context_parallel_size}",
-        f"--checkpoint_dir={str(WEIGHTS_AVATAR)}",
-        f"--stage_1={stage_1}",
+        "--context_parallel_size=1",
+        f"--checkpoint_dir={WEIGHTS_LONGCAT_AVATAR}",
+        f"--stage_1={mode}",
         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}",
     ]
 
-    rc, log = _run_subprocess(cmd, cwd=REPO_DIR)
-    if rc != 0:
-        return None, f"运行失败（exit={rc}）。日志如下：\n{log}"
+    # 续写参数（用户设置 >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)}",
+        ]
 
-    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
+    # 环境变量：让脚本能找到模块
+    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")
 
-    # 兜底：找最近生成的 mp4
-    fallback = _find_latest_video(REPO_DIR)
-    return fallback, log
+    # 约定输出目录（若脚本支持/或脚本默认输出在当前目录下的 outputs）
+    # 我们用 cwd + 输出扫描兜底
+    env["OUTPUT_DIR"] = str(OUTPUT_DIR)
 
+    code, log = _run(cmd, cwd=REPO_DIR, env=env)
 
-@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],
-    )
+    # 尝试找到最新 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]
 
-    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 code != 0:
+        return None, f"执行失败（exit={code}）。日志如下：\n\n{log}"
 
-    rc, log = _run_subprocess(cmd, cwd=REPO_DIR)
-    if rc != 0:
-        return None, f"运行失败（exit={rc}）。日志如下：\n{log}"
+    if mp4 is None or not mp4.exists():
+        return None, f"执行完成，但未找到 mp4 输出文件。日志如下：\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
+    return str(mp4), f"执行成功：{mp4}\n\n日志如下：\n\n{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}"
 
-# -------------------- UI --------------------
-with gr.Blocks(title="LongCat-Video-Avatar (ZeroGPU)", fill_height=True) as demo:
+with gr.Blocks(title="LongCat-Video-Avatar (ZeroGPU) - Single File Space") as demo:
     gr.Markdown(
-        "## LongCat-Video-Avatar (ZeroGPU)\n"
-        "- 启动后会自动下载权重到 `./weights` 并准备环境。\n"
-        "- ZeroGPU 会在点击生成时，按需分配 GPU 执行（@spaces.GPU）。\n"
-        "- 如果你发现必须 2 卡才能跑通，可把 **nproc/context_parallel_size** 改为 2。"
+        """
+# LongCat-Video-Avatar（ZeroGPU / 单文件 Space）
+
+- 单人模式：**AT2V（音频+文本）** / **AI2V（音频+图片）**
+- 续写（Video Continuation）：把 **num_segments** 设为 > 1 即可（官方参数：ref_img_index / mask_frame_range）  
+- 提示：为了更自然的口型，prompt 里建议包含 talking/speaking 等动作词（模型卡建议）  
+"""
     )
 
-    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():
+        btn_prepare = gr.Button("一键准备（下载代码+权重）", variant="primary")
+        prep_status = gr.Textbox(label="准备状态", value="尚未准备。首次准备会下载较大权重。", lines=2)
 
-    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}"
+    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.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],
     )
 
-demo.queue(concurrency_count=1).launch()
+demo.queue(max_size=12).launch()