Persister les features sur disque /tmp pour partage entre workers Gradio

backbone_utils.py

@@ -1,3 +1,5 @@
+import os
+
 import numpy as np
 import torch
 import torch.nn as nn
@@ -10,6 +12,9 @@ from config import HF_BACKBONE_REPO, HF_TOKEN
 _BACKBONE = None
 _FEATURES_CACHE = None
 
+# Partagé entre tous les workers Gradio (même process group)
+_DISK_CACHE_PATH = "/tmp/charcoal_features.npz"
+
 
 def load_backbone(device: torch.device) -> nn.Module:
     global _BACKBONE
@@ -17,12 +22,6 @@ def load_backbone(device: torch.device) -> nn.Module:
     if _BACKBONE is not None:
         return _BACKBONE.to(device)
 
-    if not HF_BACKBONE_REPO:
-        raise RuntimeError(
-            "HF_BACKBONE_REPO n'est pas configuré. "
-            "Ajoutez-le dans les Secrets du Space Hugging Face."
-        )
-
     pt_path = hf_hub_download(
         repo_id=HF_BACKBONE_REPO,
         filename="resnet18_charcoal_backbone.pt",
@@ -73,9 +72,32 @@ def extract_all_features(batch_size: int = 64):
         }
         counts[split_name] = len(cache[split_name]["y"])
 
+    # Sauvegarde sur disque pour que tous les workers Gradio y aient accès
+    np.savez(
+        _DISK_CACHE_PATH,
+        train_X=cache["train"]["X"],       train_y=cache["train"]["y"],
+        validation_X=cache["validation"]["X"], validation_y=cache["validation"]["y"],
+        test_X=cache["test"]["X"],         test_y=cache["test"]["y"],
+    )
+
     _FEATURES_CACHE = cache
     return cache, class_names, counts
 
 
 def get_cached_features():
-    return _FEATURES_CACHE
+    global _FEATURES_CACHE
+
+    if _FEATURES_CACHE is not None:
+        return _FEATURES_CACHE
+
+    # Essaye de charger depuis le disque (autre worker a peut-être déjà extrait)
+    if os.path.exists(_DISK_CACHE_PATH):
+        data = np.load(_DISK_CACHE_PATH)
+        _FEATURES_CACHE = {
+            "train":      {"X": data["train_X"],      "y": data["train_y"]},
+            "validation": {"X": data["validation_X"], "y": data["validation_y"]},
+            "test":       {"X": data["test_X"],       "y": data["test_y"]},
+        }
+        return _FEATURES_CACHE
+
+    return None