Update gradio_app.py

gradio_app.py

@@ -41,6 +41,14 @@ TRAINED_MODEL_REPO = os.getenv("TRAINED_MODEL_REPO", "")
 # 优先读取官方变量名，其次兼容 HF_TOKEN
 HF_AUTH_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN") or os.getenv("HF_TOKEN")
 
+# 需要的权重文件清单
+REQUIRED_WEIGHT_FILENAMES = [
+    "pytorch_model.bin",
+    "motion_module-4140000.pth",
+    "pytorch_model_1.bin",
+    "pytorch_model_2.bin",
+]
+
 
 # -----------------------------------------------------------------------------
 # Utilities
@@ -165,7 +173,11 @@ def _download_models() -> Tuple[Optional[str], Optional[str], Optional[str]]:
             token=HF_AUTH_TOKEN,
         )
         # Symlink to ./trained_model so downstream code can load from there
-        _ = _ensure_symlink(tm_snap, os.path.abspath("trained_model"))
+        tm_linked = _ensure_symlink(tm_snap, os.path.abspath("trained_model"))
+        # If the repo contains a nested pretrain/ folder, also expose it at ./pretrain
+        nested_pretrain = os.path.join(tm_linked, "pretrain")
+        if os.path.isdir(nested_pretrain):
+            _ensure_symlink(nested_pretrain, os.path.abspath("pretrain"))
 
     return sd15_path, hairmapper_dir, ffhq_dir
 
@@ -201,13 +213,42 @@ SD15_PATH, _, _ = _download_models()
 # -----------------------------------------------------------------------------
 # Global model loading (CPU) so GPU task only does inference
 # -----------------------------------------------------------------------------
+def _has_all_weights(dir_path: str) -> bool:
+    return all(os.path.isfile(os.path.join(dir_path, name)) for name in REQUIRED_WEIGHT_FILENAMES)
+
+
 def _resolve_trained_model_dir() -> str:
-    tm_dir = os.path.abspath("trained_model") if os.path.isdir("trained_model") else None
-    if tm_dir is None and os.path.isdir("pretrain"):
-        tm_dir = os.path.abspath("pretrain")
-    if tm_dir is None:
-        raise RuntimeError("Missing trained model weights. Provide TRAINED_MODEL_REPO or include ./pretrain.")
-    return tm_dir
+    pretrain_dir = os.path.abspath("pretrain") if os.path.isdir("pretrain") else None
+    trained_dir = os.path.abspath("trained_model") if os.path.isdir("trained_model") else None
+    trained_dir_nested = os.path.join(trained_dir, "pretrain") if trained_dir else None
+
+    # 优先使用 pretrain（你已说明文件在此），并校验文件齐全
+    if pretrain_dir and _has_all_weights(pretrain_dir):
+        return pretrain_dir
+
+    # 其次尝试 trained_model，并校验文件齐全
+    if trained_dir and _has_all_weights(trained_dir):
+        return trained_dir
+    # 再尝试 trained_model/pretrain 子目录
+    if trained_dir_nested and os.path.isdir(trained_dir_nested) and _has_all_weights(trained_dir_nested):
+        return trained_dir_nested
+
+    # 构造更友好的报错信息
+    def _missing_list(dir_path: str) -> str:
+        if not dir_path:
+            return "目录不存在"
+        missing = [n for n in REQUIRED_WEIGHT_FILENAMES if not os.path.isfile(os.path.join(dir_path, n))]
+        if not missing:
+            return "文件齐全"
+        return "缺少: " + ", ".join(missing)
+
+    msg = (
+        "Missing trained model weights. Provide TRAINED_MODEL_REPO or include ./pretrain.\n"
+        f"pretrain 状态: {_missing_list(pretrain_dir)}\n"
+        f"trained_model 状态: {_missing_list(trained_dir)}\n"
+        f"trained_model/pretrain 状态: {_missing_list(trained_dir_nested)}"
+    )
+    raise RuntimeError(msg)
 
 
 # Lazy globals