Update core/image_generator.py

core/image_generator.py

@@ -118,24 +118,20 @@
 
 # core/image_generator.py
 # ---------------- CACHE & MODEL DIRECTORIES (FIXED) ----------------
+# core/image_generator.py
 import os
 from pathlib import Path
-import torch
-from diffusers import StableDiffusionXLPipeline
-from huggingface_hub import hf_hub_download
-from typing import List
-from io import BytesIO
-import base64
-from PIL import Image
 
-# Force all Hugging Face caches to /tmp/hf_cache
+# --------------------------------------------------------------
+# ✅ FORCE ALL HUGGINGFACE + DIFFUSERS CACHES TO /tmp/hf_cache
+# --------------------------------------------------------------
 HF_CACHE_DIR = Path("/tmp/hf_cache")
-HF_CACHE_DIR.mkdir(parents=True, exist_ok=True)
-
 MODEL_DIR = Path("/tmp/models/realvisxl_v4")
-MODEL_DIR.mkdir(parents=True, exist_ok=True)
 
-# MUST be set before importing diffusers/transformers
+for d in [HF_CACHE_DIR, MODEL_DIR]:
+    d.mkdir(parents=True, exist_ok=True)
+
+# Set all relevant environment variables *before* any imports
 os.environ["HF_HOME"] = str(HF_CACHE_DIR)
 os.environ["HF_HUB_CACHE"] = str(HF_CACHE_DIR)
 os.environ["DIFFUSERS_CACHE"] = str(HF_CACHE_DIR)
@@ -143,22 +139,33 @@ os.environ["TRANSFORMERS_CACHE"] = str(HF_CACHE_DIR)
 os.environ["XDG_CACHE_HOME"] = str(HF_CACHE_DIR)
 os.environ["HF_DATASETS_CACHE"] = str(HF_CACHE_DIR)
 os.environ["HF_MODULES_CACHE"] = str(HF_CACHE_DIR)
+os.environ["TMPDIR"] = str(HF_CACHE_DIR)
+os.environ["CACHE_DIR"] = str(HF_CACHE_DIR)
 
 print("[DEBUG] Hugging Face cache directory set to:", HF_CACHE_DIR)
 print("[DEBUG] Model directory set to:", MODEL_DIR)
-# ---------------- MODEL CONFIG ---------------- 
-MODEL_REPO = "SG161222/RealVisXL_V4.0" 
-MODEL_FILENAME = "RealVisXL_V4.0.safetensors"
-
 
+# --------------------------------------------------------------
+# IMPORTS (must be done after env vars are set)
+# --------------------------------------------------------------
+import torch
+from diffusers import StableDiffusionXLPipeline
+from huggingface_hub import hf_hub_download
+from typing import List
+from io import BytesIO
+import base64
+from PIL import Image
 
+# --------------------------------------------------------------
+# MODEL CONFIG
+# --------------------------------------------------------------
+MODEL_REPO = "SG161222/RealVisXL_V4.0"
+MODEL_FILENAME = "RealVisXL_V4.0.safetensors"
 
-# ---------------- MODEL DOWNLOAD ----------------
+# --------------------------------------------------------------
+# MODEL DOWNLOAD
+# --------------------------------------------------------------
 def download_model() -> Path:
-    """
-    Downloads RealVisXL V4.0 model if not present.
-    Returns local path.
-    """
     model_path = MODEL_DIR / MODEL_FILENAME
     if not model_path.exists():
         print("[ImageGen] Downloading RealVisXL V4.0 model...")
@@ -168,7 +175,7 @@ def download_model() -> Path:
                 filename=MODEL_FILENAME,
                 cache_dir=str(HF_CACHE_DIR),
                 force_download=False,
-                resume_download=True,  # safer if download interrupted
+                resume_download=True,
             )
         )
         print(f"[ImageGen] Model downloaded to: {model_path}")
@@ -176,17 +183,17 @@ def download_model() -> Path:
         print("[ImageGen] Model already exists. Skipping download.")
     return model_path
 
-# ---------------- PIPELINE LOAD ----------------
+# --------------------------------------------------------------
+# PIPELINE LOAD
+# --------------------------------------------------------------
 def load_pipeline() -> StableDiffusionXLPipeline:
-    """
-    Loads the RealVisXL V4.0 model for image generation.
-    """
     model_path = download_model()
     print("[ImageGen] Loading model into pipeline...")
 
     pipe = StableDiffusionXLPipeline.from_single_file(
         str(model_path),
         torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        local_files_only=True,  # ✅ ensures it never writes to /.cache
     )
 
     if torch.cuda.is_available():
@@ -194,10 +201,8 @@ def load_pipeline() -> StableDiffusionXLPipeline:
     else:
         pipe.to("cpu")
 
-    # Optional: skip safety checker to save memory/performance
-    pipe.safety_checker = None  
-    # Enable attention slicing for memory-efficient CPU usage
-    pipe.enable_attention_slicing()  
+    pipe.safety_checker = None
+    pipe.enable_attention_slicing()
 
     print("[ImageGen] Model ready.")
     return pipe