Stage preloaded models locally

app.py

@@ -227,11 +227,6 @@ BUILTIN_ASSETS = {
         "image_encoder/model.safetensors"
     ),
 }
-PERSISTENT_MODEL_DIR = (
-    BUCKET_MOUNT
-    if BUCKET_MOUNT.is_dir()
-    else DATA_DIR / "models" if HAS_DATA_MOUNT else None
-)
 MATRIX_WIDTH = 1152
 MATRIX_HEIGHT = 896
 MATRIX_SAMPLER = "dpmpp_2m"
@@ -348,11 +343,6 @@ def download(url, target):
 
 
 def model_path(kind, name):
-    mounted = BUCKET_MOUNT / kind / name
-    if mounted.is_file():
-        return mounted
-    if PERSISTENT_MODEL_DIR is not None and (kind, name) in BUILTIN_ASSETS:
-        return PERSISTENT_MODEL_DIR / kind / name
     return LOCAL_MODEL_DIR / kind / name
 
 
@@ -704,40 +694,15 @@ def init_comfy():
     for kind in MODEL_KINDS:
         local = LOCAL_MODEL_DIR / kind
         local.mkdir(parents=True, exist_ok=True)
-        comfy_kind = COMFY_KINDS.get(kind, kind)
         folder_paths.add_model_folder_path(
-            comfy_kind,
+            COMFY_KINDS.get(kind, kind),
             str(local),
             is_default=True,
         )
-        mounted = BUCKET_MOUNT / kind
-        if mounted.is_dir():
-            folder_paths.add_model_folder_path(
-                comfy_kind,
-                str(mounted),
-                is_default=True,
-            )
-    if PERSISTENT_MODEL_DIR is not None and PERSISTENT_MODEL_DIR != BUCKET_MOUNT:
-        for kind in {kind for kind, _ in BUILTIN_ASSETS}:
-            persistent = PERSISTENT_MODEL_DIR / kind
-            if persistent == LOCAL_MODEL_DIR / kind:
-                continue
-            persistent.mkdir(parents=True, exist_ok=True)
-            folder_paths.add_model_folder_path(
-                COMFY_KINDS.get(kind, kind),
-                str(persistent),
-                is_default=True,
-            )
     folder_paths.add_model_folder_path(
         "ultralytics_bbox",
         str(LOCAL_MODEL_DIR / "ultralytics"),
     )
-    mounted_ultralytics = BUCKET_MOUNT / "ultralytics"
-    if mounted_ultralytics.is_dir():
-        folder_paths.add_model_folder_path(
-            "ultralytics_bbox",
-            str(mounted_ultralytics),
-        )
 
     folder_paths.add_model_folder_path("custom_nodes", str(CUSTOM_NODES_DIR))
     import_custom_nodes()
@@ -3135,6 +3100,18 @@ def login(p):
     )
 
 
+def copy_mounted_asset(kind, name):
+    target = model_path(kind, name)
+    if target.is_file():
+        return
+    log(f"Copying {kind}/{name}")
+    target.parent.mkdir(parents=True, exist_ok=True)
+    temp = target.with_suffix(target.suffix + ".part")
+    shutil.copy2(BUCKET_MOUNT / kind / name, temp)
+    temp.replace(target)
+    log(f"Copied {kind}/{name}: {target.stat().st_size // MIB} MiB")
+
+
 def preload_assets():
     assets = []
     for kind, asset_ids in STARTUP_ASSET_IDS.items():
@@ -3158,7 +3135,10 @@ def preload_assets():
     }
     for kind, name in assets:
         log(f"Preloading {kind}/{name}")
-        stage_model(kind, name)
+        if (BUCKET_MOUNT / kind / name).is_file():
+            copy_mounted_asset(kind, name)
+        else:
+            stage_model(kind, name)
         if kind not in STYLE_MODEL_KINDS:
             loaders[kind](name)
         log(f"Preloaded {kind}/{name}")