Upload demo/app.py with huggingface_hub

demo/app.py

@@ -42,13 +42,13 @@ if str(PROJECT_ROOT) not in sys.path:
 
 from demo.real_world_pipeline import (  # noqa: E402
     DEFAULT_BBOX_MODEL,
-    DEFAULT_MODEL_REPO_ID,
+    DEFAULT_STAGE2_MODEL_REPO_ID,
     DEFAULT_REAL_CONFIG_PATH,
     DEFAULT_RUN_NAME,
     DEFAULT_WORK_DIR,
     run_real_world_pipeline,
 )
-from demo.hf_repo_assets import ensure_repo_assets  # noqa: E402
+from demo.hf_repo_assets import ensure_repo_assets, get_stage2_model_repo_id  # noqa: E402
 
 
 DEFAULT_EXAMPLE_DIR = Path(
@@ -91,10 +91,10 @@ ZERO_GPU_DURATION = clamp(
 
 MODEL_PREFETCH_STATUS = {
     "enabled": os.environ.get("SYNLAYERS_DISABLE_PREFETCH", "0") != "1",
-    "bbox_model": str(DEFAULT_BBOX_MODEL),
-    "main_model": str(os.environ.get("SYNLAYERS_MODEL_REPO") or DEFAULT_MODEL_REPO_ID),
+    "stage1_bbox_model": str(DEFAULT_BBOX_MODEL),
+    "stage2_model_repo": get_stage2_model_repo_id(),
     "bbox_done": False,
-    "main_done": False,
+    "stage2_done": False,
     "error": "",
 }
 
@@ -181,20 +181,28 @@ def prefetch_model_assets() -> None:
     if not MODEL_PREFETCH_STATUS["enabled"]:
         return
 
-    main_model = os.environ.get("SYNLAYERS_MODEL_REPO") or DEFAULT_MODEL_REPO_ID
-    bbox_ok = prefetch_one_model(DEFAULT_BBOX_MODEL, "bbox model")
-    main_ok = prefetch_one_model(main_model, "main model")
+    stage1_bbox_model = (
+        os.environ.get("SYNLAYERS_BBOX_MODEL")
+        or os.environ.get("SYNLAYERS_BBOX_MODEL_REPO")
+        or DEFAULT_BBOX_MODEL
+    )
+    stage2_model_repo = get_stage2_model_repo_id()
+    MODEL_PREFETCH_STATUS["stage1_bbox_model"] = str(stage1_bbox_model)
+    MODEL_PREFETCH_STATUS["stage2_model_repo"] = str(stage2_model_repo)
+
+    bbox_ok = prefetch_one_model(stage1_bbox_model, "Stage 1 bbox model")
+    stage2_ok = prefetch_one_model(stage2_model_repo, "Stage 2 model repo")
 
     try:
-        ensure_repo_assets(main_model)
+        ensure_repo_assets(stage2_model_repo)
     except Exception as exc:
         MODEL_PREFETCH_STATUS["error"] += (
-            f"\n- Failed to prefetch runtime assets from `{main_model}`: {exc}"
+            f"\n- Failed to prefetch Stage 2 runtime assets from `{stage2_model_repo}`: {exc}"
         )
-        main_ok = False
+        stage2_ok = False
 
     MODEL_PREFETCH_STATUS["bbox_done"] = bool(bbox_ok)
-    MODEL_PREFETCH_STATUS["main_done"] = bool(main_ok)
+    MODEL_PREFETCH_STATUS["stage2_done"] = bool(stage2_ok)
 
 
 # Run prefetch during Space startup, outside the ZeroGPU-decorated function.
@@ -252,7 +260,12 @@ def is_zero_gpu_space() -> bool:
 def get_runtime_status_markdown() -> str:
     accelerator = os.environ.get("ACCELERATOR", "unknown")
     space_id = os.environ.get("SPACE_ID", "local")
-    model_repo = os.environ.get("SYNLAYERS_MODEL_REPO") or DEFAULT_MODEL_REPO_ID
+    stage1_bbox_model = (
+        os.environ.get("SYNLAYERS_BBOX_MODEL")
+        or os.environ.get("SYNLAYERS_BBOX_MODEL_REPO")
+        or DEFAULT_BBOX_MODEL
+    )
+    stage2_model_repo = get_stage2_model_repo_id()
     zero_gpu_enabled = is_zero_gpu_space()
 
     lines = [
@@ -260,14 +273,15 @@ def get_runtime_status_markdown() -> str:
         f"- `SPACE_ID`: `{space_id}`",
         f"- `ACCELERATOR`: `{accelerator}`",
         f"- `HF_HOME`: `{os.environ.get('HF_HOME', '')}`",
-        f"- `SYNLAYERS_MODEL_REPO`: `{model_repo}`",
+        f"- `Stage 1 bbox repo/path`: `{stage1_bbox_model}`",
+        f"- `Stage 2 repo`: `{stage2_model_repo}`",
         "",
         "## Model Asset Prefetch",
         f"- `Prefetch enabled`: `{MODEL_PREFETCH_STATUS['enabled']}`",
-        f"- `BBox model`: `{MODEL_PREFETCH_STATUS['bbox_model']}`",
-        f"- `Main model`: `{MODEL_PREFETCH_STATUS['main_model']}`",
-        f"- `BBox model files prefetched`: `{MODEL_PREFETCH_STATUS['bbox_done']}`",
-        f"- `Main model files prefetched`: `{MODEL_PREFETCH_STATUS['main_done']}`",
+        f"- `Stage 1 bbox model`: `{MODEL_PREFETCH_STATUS['stage1_bbox_model']}`",
+        f"- `Stage 2 model repo`: `{MODEL_PREFETCH_STATUS['stage2_model_repo']}`",
+        f"- `Stage 1 files prefetched`: `{MODEL_PREFETCH_STATUS['bbox_done']}`",
+        f"- `Stage 2 assets prefetched`: `{MODEL_PREFETCH_STATUS['stage2_done']}`",
     ]
 
     if MODEL_PREFETCH_STATUS["error"]: