Update core/image_generator.py

core/image_generator.py

@@ -9,20 +9,27 @@ from typing import Dict, Any
 from PIL import Image
 from io import BytesIO
 import base64
+import tempfile
 
 # --------------------------------------------------------------
-# ✅ PERSISTENT STORAGE SETUP (for Hugging Face Spaces)
+# 🚨 ABSOLUTE FIX FOR PermissionError('/.cache')
 # --------------------------------------------------------------
 HF_CACHE_DIR = Path("/tmp/hf_cache")
-MODEL_DIR = Path("/tmp/models/realvisxl_v4")
-SEED_DIR = Path("/tmp/seed_images")
-TMP_DIR = Path("/tmp/generated_images")
+HF_CACHE_DIR.mkdir(parents=True, exist_ok=True)
 
-for d in [HF_CACHE_DIR, MODEL_DIR, SEED_DIR, TMP_DIR]:
-    d.mkdir(parents=True, exist_ok=True)
+# Patch expanduser BEFORE any library imports that might touch ~/.cache
+import os.path
+if not hasattr(os.path, "expanduser_original"):
+    os.path.expanduser_original = os.path.expanduser
 
+def safe_expanduser(path):
+    if path.startswith("~") or path.startswith("/.cache"):
+        return str(HF_CACHE_DIR)
+    return os.path.expanduser_original(path)
 
-# Ensure all relevant environment variables are correctly set
+os.path.expanduser = safe_expanduser
+
+# Set environment variables AFTER patching
 os.environ.update({
     "HF_HOME": str(HF_CACHE_DIR),
     "HF_HUB_CACHE": str(HF_CACHE_DIR),
@@ -35,23 +42,19 @@ os.environ.update({
     "CACHE_DIR": str(HF_CACHE_DIR),
 })
 
-# Force Python's tempfile to use /tmp/hf_cache too
-import tempfile
 tempfile.tempdir = str(HF_CACHE_DIR)
 
-# 🔒 Extra layer: patch os.path.expanduser to block “/.cache”
-import os.path
-def safe_expanduser(path):
-    if path.startswith("~") or path.startswith("/.cache"):
-        return str(HF_CACHE_DIR)
-    return os.path.expanduser_original(path)
-
-os.path.expanduser_original = os.path.expanduser
-os.path.expanduser = safe_expanduser
-
 print("[DEBUG] ✅ Hugging Face and Diffusers cache fully redirected to:", HF_CACHE_DIR)
 
+# --------------------------------------------------------------
+# ✅ PERSISTENT STORAGE SETUP (for Hugging Face Spaces)
+# --------------------------------------------------------------
+MODEL_DIR = Path("/tmp/models/realvisxl_v4")
+SEED_DIR = Path("/tmp/seed_images")
+TMP_DIR = Path("/tmp/generated_images")
 
+for d in [MODEL_DIR, SEED_DIR, TMP_DIR]:
+    d.mkdir(parents=True, exist_ok=True)
 
 print("[DEBUG] ✅ Using persistent Hugging Face cache at:", HF_CACHE_DIR)
 print("[DEBUG] ✅ Model directory:", MODEL_DIR)
@@ -88,7 +91,6 @@ def download_model() -> Path:
         print("[ImageGen] ✅ Model already exists at:", model_path)
     return model_path
 
-
 # --------------------------------------------------------------
 # MEMORY-SAFE PIPELINE MANAGER
 # --------------------------------------------------------------
@@ -116,10 +118,27 @@ def unload_pipelines(target="all"):
         torch.cuda.empty_cache()
     print("[ImageGen] ✅ Memory cleared.")
 
+def safe_load_pipeline(pipeline_class, model_path):
+    """Safely load a pipeline with retry logic and memory handling."""
+    try:
+        print(f"[ImageGen] 🔄 Loading {pipeline_class.__name__} from {model_path} ...")
+        pipe = pipeline_class.from_single_file(
+            model_path,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        )
+        print(f"[ImageGen] ✅ Successfully loaded {pipeline_class.__name__}.")
+        return pipe
+    except Exception as e:
+        print(f"[ImageGen] ❌ Failed to load {pipeline_class.__name__}: {e}")
+        unload_pipelines()
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        raise e
 
 def load_pipeline():
     global pipe
-    unload_pipelines(target="pipe")  # only clear old txt2img
+    unload_pipelines(target="pipe")
     model_path = download_model()
     print("[ImageGen] Loading main (txt2img) pipeline...")
     pipe = safe_load_pipeline(StableDiffusionXLPipeline, model_path)
@@ -130,10 +149,9 @@ def load_pipeline():
     print("[ImageGen] ✅ Text-to-image pipeline ready.")
     return pipe
 
-
 def load_img2img_pipeline():
     global img2img_pipe
-    unload_pipelines(target="img2img_pipe")  # only clear old img2img
+    unload_pipelines(target="img2img_pipe")
     model_path = download_model()
     print("[ImageGen] Loading img2img pipeline...")
     img2img_pipe = safe_load_pipeline(StableDiffusionXLImg2ImgPipeline, model_path)
@@ -144,41 +162,6 @@ def load_img2img_pipeline():
     print("[ImageGen] ✅ Img2Img pipeline ready.")
     return img2img_pipe
 
-
-
-def load_img2img_pipeline():
-    """Load img2img pipeline into RAM."""
-    global img2img_pipe
-    unload_pipelines()  # Ensure txt2img is removed first
-    model_path = download_model()
-    print("[ImageGen] Loading img2img pipeline...")
-
-    img2img_pipe = safe_load_pipeline(StableDiffusionXLImg2ImgPipeline, model_path)
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    img2img_pipe.to(device)
-    img2img_pipe.safety_checker = None
-    img2img_pipe.enable_attention_slicing()
-    print("[ImageGen] ✅ Img2Img pipeline ready.")
-    return img2img_pipe
-
-def safe_load_pipeline(pipeline_class, model_path):
-    """Safely load a pipeline with retry logic and memory handling."""
-    try:
-        print(f"[ImageGen] 🔄 Loading {pipeline_class.__name__} from {model_path} ...")
-        pipe = pipeline_class.from_single_file(
-            model_path,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        )
-        print(f"[ImageGen] ✅ Successfully loaded {pipeline_class.__name__}.")
-        return pipe
-    except Exception as e:
-        print(f"[ImageGen] ❌ Failed to load {pipeline_class.__name__}: {e}")
-        unload_pipelines()
-        gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        raise e
-
 # --------------------------------------------------------------
 # UTILITY: PIL → BASE64
 # --------------------------------------------------------------
@@ -187,7 +170,6 @@ def pil_to_base64(img: Image.Image) -> str:
     img.save(buffered, format="PNG")
     return f"data:image/png;base64,{base64.b64encode(buffered.getvalue()).decode()}"
 
-
 # --------------------------------------------------------------
 # UNIFIED IMAGE GENERATION FUNCTION
 # --------------------------------------------------------------