Update core/image_generator.py

core/image_generator.py

@@ -247,6 +247,7 @@
 
 #     print(f"[ImageGen] Generated {len(images)} image(s) successfully.")
 #     return images
+
 import os
 from pathlib import Path
 import gc
@@ -341,12 +342,22 @@ def unload_pipelines():
 
 
 def safe_load_pipeline(pipeline_cls, model_path: Path):
-    """Load model from disk safely."""
-    return pipeline_cls.from_single_file(
-        str(model_path),
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        local_files_only=True,
-    )
+    try:
+        # Try local-only load first
+        return pipeline_cls.from_single_file(
+            str(model_path),
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            local_files_only=True,
+            cache_dir=str(HF_CACHE_DIR)  # <-- ensure all caches go to /tmp/hf_cache
+        )
+    except Exception as e:
+        print(f"[WARN] Local-only load failed ({e}). Retrying with network access...")
+        return pipeline_cls.from_single_file(
+            str(model_path),
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            local_files_only=False,  # allow network download
+            cache_dir=str(HF_CACHE_DIR)  # <-- must specify this
+        )
 
 
 def load_pipeline():