grmchn commited on
Commit
9b07278
·
1 Parent(s): 50b4cf7

chore(setup): add model pre-download utility and CLI flag; update README usage

Browse files
Files changed (3) hide show
  1. README.md +14 -0
  2. app.py +13 -2
  3. utils/model_setup.py +65 -0
README.md CHANGED
@@ -40,6 +40,20 @@ python app.py
40
 
41
  ブラウザで表示されるGradio UIからポーズ編集を行えます。
42
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
43
  ## 使い方(Hugging Face Spaces)
44
 
45
  スペース版の基本フローは「入力 → 編集 → 出力」です。
 
40
 
41
  ブラウザで表示されるGradio UIからポーズ編集を行えます。
42
 
43
+ ### 初回セットアップ(モデルの事前ダウンロード)
44
+
45
+ 初回起動時に DWPose の ONNX モデルを自動ダウンロードしますが、起動前に取得しておきたい場合は次のどちらかを実行してください。
46
+
47
+ ```bash
48
+ # 1) app.pyのオプションでダウンロードだけ実施
49
+ python app.py --setup-models
50
+
51
+ # 2) ユーティリティを直接実行(同等の動作)
52
+ python -m utils.model_setup
53
+ ```
54
+
55
+ デフォルトでは `./models` にキャッシュします。企業ネットワーク等でトークンが必要な場合は、先に `huggingface-cli login` を実施してください。
56
+
57
  ## 使い方(Hugging Face Spaces)
58
 
59
  スペース版の基本フローは「入力 → 編集 → 出力」です。
app.py CHANGED
@@ -6,6 +6,7 @@ from utils.notifications import notify_success, notify_error, NotificationMessag
6
  from utils.coordinate_system import update_coordinate_system
7
  from utils.image_processing import process_uploaded_image
8
  from utils.export_utils import export_pose_as_image, export_pose_as_json, get_timestamp_filename
 
9
  import json
10
  import tempfile
11
  import base64
@@ -959,5 +960,15 @@ def main():
959
  return demo
960
 
961
  if __name__ == "__main__":
962
- demo = main()
963
- demo.launch()
 
 
 
 
 
 
 
 
 
 
 
6
  from utils.coordinate_system import update_coordinate_system
7
  from utils.image_processing import process_uploaded_image
8
  from utils.export_utils import export_pose_as_image, export_pose_as_json, get_timestamp_filename
9
+ from utils.model_setup import download_models as _download_models_setup
10
  import json
11
  import tempfile
12
  import base64
 
960
  return demo
961
 
962
  if __name__ == "__main__":
963
+ import argparse, sys
964
+ parser = argparse.ArgumentParser()
965
+ parser.add_argument("--setup-models", action="store_true", help="Download required models into ./models and exit")
966
+ args, _ = parser.parse_known_args()
967
+
968
+ if args.setup_models:
969
+ code = _download_models_setup()
970
+ # exit with code so CI/automation can check
971
+ sys.exit(code)
972
+ else:
973
+ demo = main()
974
+ demo.launch()
utils/model_setup.py ADDED
@@ -0,0 +1,65 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Lightweight model setup utility for dwpose-editor.
3
+
4
+ Provides a CLI to pre-download required ONNX models into the local cache
5
+ directory used by the app (./models), without initializing ONNX sessions.
6
+ """
7
+
8
+ from typing import List
9
+ import os
10
+ import sys
11
+
12
+ try:
13
+ from huggingface_hub import hf_hub_download
14
+ except Exception as e: # pragma: no cover
15
+ hf_hub_download = None
16
+
17
+
18
+ DEFAULT_REPO_ID = "yzd-v/DWPose"
19
+ DEFAULT_CACHE_DIR = "./models"
20
+ REQUIRED_FILES: List[str] = [
21
+ "yolox_l.onnx",
22
+ "dw-ll_ucoco_384.onnx",
23
+ ]
24
+
25
+
26
+ def ensure_dir(path: str) -> None:
27
+ os.makedirs(path, exist_ok=True)
28
+
29
+
30
+ def download_models(repo_id: str = DEFAULT_REPO_ID, cache_dir: str = DEFAULT_CACHE_DIR) -> int:
31
+ """Download required model files into cache_dir. Returns 0 on success, non-zero on failure."""
32
+ if hf_hub_download is None:
33
+ print("[ERROR] huggingface-hub is not installed. Run: pip install huggingface-hub")
34
+ return 2
35
+
36
+ ensure_dir(cache_dir)
37
+
38
+ ok = True
39
+ for filename in REQUIRED_FILES:
40
+ try:
41
+ print(f"[SETUP] Downloading {filename} from {repo_id} → {cache_dir} ...")
42
+ local_path = hf_hub_download(repo_id=repo_id, filename=filename, cache_dir=cache_dir)
43
+ print(f"[SETUP] OK: {local_path}")
44
+ except Exception as e:
45
+ ok = False
46
+ print(f"[SETUP] ERROR: failed to download {filename}: {e}")
47
+
48
+ return 0 if ok else 1
49
+
50
+
51
+ def main(argv: List[str] | None = None) -> int:
52
+ """CLI entry point: python -m utils.model_setup [--repo REPO] [--cache-dir DIR]"""
53
+ import argparse
54
+
55
+ parser = argparse.ArgumentParser(description="Pre-download DWPose models into ./models")
56
+ parser.add_argument("--repo", default=DEFAULT_REPO_ID, help="Hugging Face repo id")
57
+ parser.add_argument("--cache-dir", default=DEFAULT_CACHE_DIR, help="Local cache directory")
58
+ args = parser.parse_args(argv)
59
+
60
+ return download_models(repo_id=args.repo, cache_dir=args.cache_dir)
61
+
62
+
63
+ if __name__ == "__main__": # pragma: no cover
64
+ sys.exit(main())
65
+