Update download_assets.py

download_assets.py

@@ -1,119 +1,112 @@
-# download_assets.py — fetch models from a Google Drive FOLDER (no zips)
-# Expects a folder with subfolders: Text/, Image/, Audio/
-# Env: GDRIVE_FOLDER_ID=<your folder id>
-# Coder: Amir
-
-import os, sys, shutil, pathlib
-from typing import Optional
-from filelock import FileLock
-import gdown
-
-BASE = pathlib.Path(__file__).resolve().parent
-CACHE_DIR = BASE / "cache"
-DL_DIR = CACHE_DIR / "gdrive_folder"   # where we download the Drive folder
-
-TARGETS = {
-    "text": {
-        "dst": BASE / "Text",
-        "must": ["config.json", "tokenizer.json", "model.safetensors", "vocab.json"],
-        "folder_name": "Text",
-    },
-    "image": {
-        "dst": BASE / "Image",
-        "must": ["resnet_safety_classifier.pth", "clip_safety_classifier.pth"],
-        "folder_name": "Image",
-    },
-    "audio": {
-        "dst": BASE / "Audio",
-        "must": ["text_pipeline_balanced.joblib"],
-        "folder_name": "Audio",
-    },
-}
-
-def _ready(kind: str) -> bool:
-    info = TARGETS[kind]
-    d = info["dst"]
-    return d.exists() and all((d / m).exists() for m in info["must"])
-
-def _find_subdir(root: pathlib.Path, name: str) -> Optional[pathlib.Path]:
-    lname = name.lower()
-    for p in root.rglob("*"):
-        if p.is_dir() and p.name.lower() == lname:
-            return p
-    return None
-
-def _merge_copy(src: pathlib.Path, dst: pathlib.Path):
-    dst.mkdir(parents=True, exist_ok=True)
-    for root, dirs, files in os.walk(src):
-        rroot = pathlib.Path(root)
-        rel = rroot.relative_to(src)
-        # create subdirs
-        for d in dirs:
-            (dst / rel / d).mkdir(parents=True, exist_ok=True)
-        # copy files
-        for f in files:
-            s = rroot / f
-            t = dst / rel / f
-            t.parent.mkdir(parents=True, exist_ok=True)
-            shutil.copy2(s, t)
-
-def _download_folder(folder_id: str):
-    DL_DIR.mkdir(parents=True, exist_ok=True)
-    lock = FileLock(str(DL_DIR) + ".lock")
-    with lock:
-        # If already downloaded once, skip re-download
-        if any(DL_DIR.iterdir()):
-            print("[assets] Drive folder already present; skipping fresh download.")
-            return
-        print(f"[assets] downloading Drive folder {folder_id} …")
-        gdown.download_folder(
-            id=folder_id,
-            output=str(DL_DIR),
-            quiet=False,
-            use_cookies=False,
-            remaining_ok=True,
-        )
-
-def ensure_all_assets():
-    # If everything already exists, return fast
-    if all(_ready(k) for k in TARGETS):
-        print("[assets] all bundles already present.")
-        return
-
-    folder_id = os.getenv("GDRIVE_FOLDER_ID", "").strip()
-    if not folder_id:
-        raise RuntimeError("[assets] Please set GDRIVE_FOLDER_ID (top-level Drive folder with Text/Image/Audio)")
-
-    _download_folder(folder_id)
-
-    # Drive may nest one extra directory level; detect the top
-    # Pick the largest immediate subdir if needed, else use DL_DIR itself
-    root = DL_DIR
-    # If there is exactly one child dir and it contains our subfolders, use it
-    subdirs = [p for p in DL_DIR.iterdir() if p.is_dir()]
-    if len(subdirs) == 1:
-        maybe = subdirs[0]
-        if all(_find_subdir(maybe, TARGETS[k]["folder_name"]) for k in TARGETS):
-            root = maybe
-
-    # For each target, locate its folder and merge-copy into app dirs
-    for kind, info in TARGETS.items():
-        if _ready(kind):
-            print(f"[assets] {kind}: already ready.")
-            continue
-
-        src_dir = _find_subdir(root, info["folder_name"])
-        if not src_dir:
-            raise RuntimeError(f"[assets] {kind}: could not find '{info['folder_name']}' folder in Drive download")
-
-        print(f"[assets] {kind}: copying from {src_dir} → {info['dst']}")
-        _merge_copy(src_dir, info["dst"])
-
-        # sanity check
-        missing = [m for m in info["must"] if not (info["dst"] / m).exists()]
-        if missing:
-            raise RuntimeError(f"[assets] {kind}: missing files after copy → {missing}")
-
-        print(f"[assets] {kind}: ready.")
-
-    print("[assets] all bundles ready.")
+# download_assets.py — fetch models from a Google Drive FOLDER (no zips)
+# Expects a folder with subfolders: Text/, Image/, Audio/
+# Env: GDRIVE_FOLDER_ID=<your folder id>
+# Author: Amir
+
+import os, sys, shutil, pathlib
+from typing import Optional
+from filelock import FileLock
+import gdown
+
+BASE = pathlib.Path(__file__).resolve().parent
+CACHE_DIR = BASE / "cache"
+DL_DIR = CACHE_DIR / "gdrive_folder"   # where we download the Drive folder
+
+TARGETS = {
+    "text": {
+        "dst": BASE / "Text",
+        "must": ["config.json", "tokenizer.json", "model.safetensors", "vocab.json"],
+        "folder_name": "Text",
+    },
+    "image": {
+        "dst": BASE / "Image",
+        "must": ["resnet_safety_classifier.pth", "clip_safety_classifier.pth"],
+        "folder_name": "Image",
+    },
+    "audio": {
+        "dst": BASE / "Audio",
+        "must": ["text_pipeline_balanced.joblib"],
+        "folder_name": "Audio",
+    },
+}
+
+def _ready(kind: str) -> bool:
+    info = TARGETS[kind]
+    d = info["dst"]
+    return d.exists() and all((d / m).exists() for m in info["must"])
+
+def _find_subdir(root: pathlib.Path, name: str) -> Optional[pathlib.Path]:
+    lname = name.lower()
+    for p in root.rglob("*"):
+        if p.is_dir() and p.name.lower() == lname:
+            return p
+    return None
+
+def _merge_copy(src: pathlib.Path, dst: pathlib.Path):
+    dst.mkdir(parents=True, exist_ok=True)
+    for root, dirs, files in os.walk(src):
+        rroot = pathlib.Path(root)
+        rel = rroot.relative_to(src)
+        for d in dirs:
+            (dst / rel / d).mkdir(parents=True, exist_ok=True)
+        for f in files:
+            s = rroot / f
+            t = dst / rel / f
+            t.parent.mkdir(parents=True, exist_ok=True)
+            shutil.copy2(s, t)
+
+def _download_folder(folder_id: str):
+    DL_DIR.mkdir(parents=True, exist_ok=True)
+    lock = FileLock(str(DL_DIR) + ".lock")
+    with lock:
+        if any(DL_DIR.iterdir()):
+            print("[assets] Drive folder already present; skipping fresh download.")
+            return
+        print(f"[assets] downloading Drive folder {folder_id} …")
+        gdown.download_folder(
+            id=folder_id,
+            output=str(DL_DIR),
+            quiet=False,
+            use_cookies=False,
+            remaining_ok=True,
+        )
+
+def ensure_all_assets():
+    # If everything already exists, return fast
+    if all(_ready(k) for k in TARGETS):
+        print("[assets] all bundles already present.")
+        return
+
+    folder_id = os.getenv("GDRIVE_FOLDER_ID", "").strip()
+    if not folder_id:
+        raise RuntimeError("[assets] Please set GDRIVE_FOLDER_ID (top-level Drive folder with Text/Image/Audio)")
+
+    _download_folder(folder_id)
+
+    # Detect an extra nesting level that gdown sometimes creates
+    root = DL_DIR
+    subdirs = [p for p in DL_DIR.iterdir() if p.is_dir()]
+    if len(subdirs) == 1:
+        maybe = subdirs[0]
+        if all(_find_subdir(maybe, TARGETS[k]["folder_name"]) for k in TARGETS):
+            root = maybe
+
+    for kind, info in TARGETS.items():
+        if _ready(kind):
+            print(f"[assets] {kind}: already ready.")
+            continue
+
+        src_dir = _find_subdir(root, info["folder_name"])
+        if not src_dir:
+            raise RuntimeError(f"[assets] {kind}: could not find '{info['folder_name']}' folder in Drive download")
+
+        print(f"[assets] {kind}: copying from {src_dir} → {info['dst']}")
+        _merge_copy(src_dir, info["dst"])
+
+        missing = [m for m in info["must"] if not (info["dst"] / m).exists()]
+        if missing:
+            raise RuntimeError(f"[assets] {kind}: missing files after copy → {missing}")
+
+        print(f"[assets] {kind}: ready.")
+
+    print("[assets] all bundles ready.")