Upload demo/hf_repo_assets.py with huggingface_hub

demo/hf_repo_assets.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+import os
+from functools import lru_cache
+from pathlib import Path
+
+from huggingface_hub import snapshot_download
+
+
+def get_stage2_model_repo_id() -> str | None:
+    return os.environ.get("SYNLAYERS_STAGE2_MODEL_REPO") or os.environ.get("SYNLAYERS_MODEL_REPO")
+
+
+def get_cache_dir() -> str | None:
+    return os.environ.get("SYNLAYERS_HF_CACHE")
+
+
+@lru_cache(maxsize=4)
+def ensure_repo_assets(repo_id: str | None = None) -> Path | None:
+    """Download Stage 2 runtime assets from the configured model repo."""
+    resolved_repo_id = repo_id or get_stage2_model_repo_id()
+    if not resolved_repo_id:
+        return None
+
+    allow_patterns = [
+        "SynLayers_checkpoints/FLUX.1-dev/**",
+        "SynLayers_checkpoints/FLUX.1-dev-Controlnet-Inpainting-Alpha/**",
+        "SynLayers_ckpt/step_120000/**",
+        "ckpt/trans_vae/0008000.pt",
+        "ckpt/pre_trained_LoRA/**",
+        "ckpt/prism_ft_LoRA/**",
+    ]
+
+    local_root = snapshot_download(
+        repo_id=resolved_repo_id,
+        repo_type="model",
+        allow_patterns=allow_patterns,
+        cache_dir=get_cache_dir(),
+    )
+    return Path(local_root)
+
+
+def build_repo_asset_overrides(repo_id: str | None = None) -> dict[str, str]:
+    """Return repo-local Stage 2 asset paths after downloading the uploaded bundle."""
+    local_root = ensure_repo_assets(repo_id)
+    if local_root is None:
+        return {}
+
+    overrides = {
+        "repo_root": str(local_root),
+        "decomp_ckpt_root": str(local_root / "SynLayers_ckpt" / "step_120000"),
+        "pretrained_model_name_or_path": str(
+            local_root / "SynLayers_checkpoints" / "FLUX.1-dev"
+        ),
+        "pretrained_adapter_path": str(
+            local_root
+            / "SynLayers_checkpoints"
+            / "FLUX.1-dev-Controlnet-Inpainting-Alpha"
+        ),
+        "transp_vae_path": str(local_root / "ckpt" / "trans_vae" / "0008000.pt"),
+        "pretrained_lora_dir": str(local_root / "ckpt" / "pre_trained_LoRA"),
+        "artplus_lora_dir": str(local_root / "ckpt" / "prism_ft_LoRA"),
+    }
+    return overrides