Upload demo/hf_repo_assets.py with huggingface_hub

demo/hf_repo_assets.py

@@ -15,6 +15,14 @@ def get_cache_dir() -> str | None:
     return os.environ.get("SYNLAYERS_HF_CACHE")
 
 
+def should_download_repo_flux() -> bool:
+    """Download FLUX from the model repo only when no external base model is configured."""
+    base_model = os.environ.get("SYNLAYERS_BASE_MODEL", "").strip()
+    if not base_model:
+        return True
+    return base_model.startswith(get_model_repo_id() or "")
+
+
 @lru_cache(maxsize=4)
 def ensure_repo_assets(repo_id: str | None = None) -> Path | None:
     """Download required runtime assets from the uploaded model repo when configured."""
@@ -23,13 +31,14 @@ def ensure_repo_assets(repo_id: str | None = None) -> Path | None:
         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/**",
     ]
+    if should_download_repo_flux():
+        allow_patterns.insert(0, "SynLayers_checkpoints/FLUX.1-dev/**")
 
     local_root = snapshot_download(
         repo_id=resolved_repo_id,
@@ -46,12 +55,9 @@ def build_repo_asset_overrides(repo_id: str | None = None) -> dict[str, str]:
     if local_root is None:
         return {}
 
-    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"
@@ -61,3 +67,8 @@ def build_repo_asset_overrides(repo_id: str | None = None) -> dict[str, str]:
         "pretrained_lora_dir": str(local_root / "ckpt" / "pre_trained_LoRA"),
         "artplus_lora_dir": str(local_root / "ckpt" / "prism_ft_LoRA"),
     }
+    if should_download_repo_flux():
+        overrides["pretrained_model_name_or_path"] = str(
+            local_root / "SynLayers_checkpoints" / "FLUX.1-dev"
+        )
+    return overrides