Update app.py

app.py

@@ -10,7 +10,13 @@ from pathlib import Path
 from typing import Optional, Dict, Any
 import uuid, shutil, cv2, json, time, urllib.parse, sys
 import threading
+# <<<< HUGINFACE PATCH: WARMUP IMPORTS START >>>>
+from typing import List
+from huggingface_hub import snapshot_download
+import traceback
+# <<<< HUGINFACE PATCH: WARMUP IMPORTS END >>>>
 import subprocess
+
 import shutil as _shutil
 # --- POINTEUR DE BACKEND (lit l'URL actuelle depuis une source externe) ------
 import os
@@ -72,8 +78,103 @@ async def proxy_all(full_path: str, request: Request):
 # -------------------------------------------------------------------------------
 # Global progress dict (vid_stem -> {percent, logs, done})
 progress_data: Dict[str, Dict[str, Any]] = {}
+# <<<< HUGINFACE PATCH: WARMUP STATE+HELPERS START >>>>
+# État global du warm-up (progression, logs, etc.)
+warmup_state: Dict[str, Any] = {
+    "running": False,
+    "percent": 0,
+    "logs": [],
+    "done": False,
+    "current": None,
+    "total": 0,
+    "idx": 0,
+    "job_id": "default",
+}
+warmup_lock = threading.Lock()
+warmup_stop = threading.Event()
+
+def _default_model_list() -> List[str]:
+    """
+    Liste par défaut lue depuis l'env WARMUP_MODELS (JSON array) sinon vide.
+    Exemple env:
+      WARMUP_MODELS=["runwayml/stable-diffusion-v1-5","facebook/sam2-hiera-base"]
+    """
+    env = (os.getenv("WARMUP_MODELS") or "").strip()
+    if env:
+        try:
+            lst = json.loads(env)
+            if isinstance(lst, list):
+                return [str(x).strip() for x in lst if str(x).strip()]
+        except Exception:
+            pass
+    return []
 
+def _log_warmup(msg: str):
+    print(f"[WARMUP] {msg}", file=sys.stdout)
+    with warmup_lock:
+        warmup_state["logs"].append(msg)
+        if len(warmup_state["logs"]) > 400:
+            warmup_state["logs"] = warmup_state["logs"][-400:]
+
+def _download_one(repo_id: str, tries: int = 3) -> bool:
+    """
+    Télécharge un repo HF en réessayant si besoin.
+    """
+    cache_home = os.path.expanduser(os.getenv("HF_HOME", "/home/user/.cache/huggingface"))
+    local_dir = os.path.join(cache_home, "models", repo_id.replace("/", "__"))
+    for attempt in range(1, tries + 1):
+        if warmup_stop.is_set():
+            return False
+        try:
+            snapshot_download(
+                repo_id,
+                local_dir=local_dir,
+                local_dir_use_symlinks=False,
+                resume_download=True,
+            )
+            return True
+        except Exception as e:
+            _log_warmup(f"{repo_id} -> tentative {attempt}/{tries} échouée: {e}")
+            time.sleep(min(10, 2 * attempt))
+    return False
+
+def _warmup_thread(models: List[str]):
+    with warmup_lock:
+        warmup_state.update({
+            "running": True,
+            "percent": 0,
+            "done": False,
+            "logs": [],
+            "current": None,
+            "total": len(models),
+            "idx": 0,
+        })
+    warmup_stop.clear()
+    ok_count = 0
+    for i, repo in enumerate(models):
+        if warmup_stop.is_set():
+            _log_warmup("Arrêt demandé par l’utilisateur.")
+            break
+        with warmup_lock:
+            warmup_state["idx"] = i
+            warmup_state["current"] = repo
+            warmup_state["percent"] = int((i / max(1, len(models))) * 100)
+        _log_warmup(f"Téléchargement: {repo}")
+        ok = _download_one(repo)
+        if ok:
+            ok_count += 1
+            _log_warmup(f"OK: {repo}")
+        else:
+            _log_warmup(f"ÉCHEC: {repo}")
+
+    with warmup_lock:
+        warmup_state["percent"] = 100
+        warmup_state["done"] = True
+        warmup_state["running"] = False
+    _log_warmup(f"Terminé: {ok_count}/{len(models)} modèles.")
+# <<<< HUGINFACE PATCH: WARMUP STATE+HELPERS END >>>>
 # ---------- Helpers ----------
+
 def _is_video(p: Path) -> bool:
     return p.suffix.lower() in {".mp4", ".mov", ".mkv", ".webm"}
 
@@ -361,22 +462,36 @@ def window(vid: str, center: int = 0, count: int = 21):
     center = max(0, min(int(center), max(0, frames-1)))
     if frames <= 0:
         print(f"[WINDOW] frames=0 for {vid}", file=sys.stdout)
-        return {"vid": vid, "start": 0, "count": 0, "selected": 0, "items": [], "frames": 0}
-    if frames <= count:
-        start = 0; sel = center; n = frames
-    else:
-        start = max(0, min(center - (count//2), frames - count))
-        n = count; sel = center - start
-    items = []
-    bust = int(time.time()*1000)
-    for i in range(n):
-        idx = start + i
-        url = f"/thumbs/f_{v.stem}_{idx}.jpg?b={bust}"
-        items.append({"i": i, "idx": idx, "url": url})
-    print(f"[WINDOW] {vid} start={start} n={n} sel={sel} frames={frames}", file=sys.stdout)
-    return {"vid": vid, "start": start, "count": n, "selected": sel, "items": items, "frames": frames}
+            return {"vid": vid, "start": start, "count": n, "selected": sel, "items": items, "frames": frames}
+# <<<< HUGINFACE PATCH: WARMUP ROUTES START >>>>
+@app.post("/warmup/start", tags=["warmup"])
+async def warmup_start(payload: Optional[Dict[str, Any]] = Body(None)):
+    """
+    Lance un téléchargement séquentiel d'une liste de modèles HF.
+    Body JSON: {"models": ["org/modelA","org/modelB", ...]}  (optionnel)
+    À défaut, lit WARMUP_MODELS (env).
+    """
+    models = (payload or {}).get("models") or _default_model_list()
+    if not isinstance(models, list) or not models:
+        raise HTTPException(400, "Liste 'models' vide. Fournir JSON {models:[\"org/model\"]} ou variable d'environnement WARMUP_MODELS.")
+    if warmup_state.get("running"):
+        return {"started": False, "already_running": True, "state": warmup_state}
+    t = threading.Thread(target=_warmup_thread, args=(models,), daemon=True)
+    t.start()
+    return {"started": True, "count": len(models)}
+
+@app.get("/warmup/status", tags=["warmup"])
+def warmup_status():
+    with warmup_lock:
+        return dict(warmup_state)
 
+@app.post("/warmup/stop", tags=["warmup"])
+def warmup_stop_api():
+    warmup_stop.set()
+    return {"stopping": True}
+# <<<< HUGINFACE PATCH: WARMUP ROUTES END >>>>
 # ----- Masques -----
+
 @app.post("/mask", tags=["mask"])
 async def save_mask(payload: Dict[str, Any] = Body(...)):
     vid = payload.get("vid")