Update core/image_generator.py

core/image_generator.py

@@ -116,7 +116,7 @@
 #     return images
 
 
-
+# core/image_generator.py
 import os
 import torch
 from diffusers import StableDiffusionXLPipeline
@@ -127,15 +127,15 @@ from io import BytesIO
 import base64
 from PIL import Image
 
-# ---------------- CACHE & MODEL DIRS ----------------
+# ---------------- CACHE & MODEL DIRECTORIES ----------------
 HF_CACHE_DIR = Path("/tmp/hf_cache")
 MODEL_DIR = Path("/tmp/models/realvisxl_v4")
 
+# Create directories safely (no chmod)
 for d in [HF_CACHE_DIR, MODEL_DIR]:
     d.mkdir(parents=True, exist_ok=True)
-    os.chmod(d, 0o777)
 
-# make sure env vars are applied *before* any huggingface import usage
+# Apply environment variables BEFORE any Hugging Face usage
 os.environ.update({
     "HF_HOME": str(HF_CACHE_DIR),
     "TRANSFORMERS_CACHE": str(HF_CACHE_DIR),
@@ -150,7 +150,10 @@ MODEL_FILENAME = "realvisxlV40_v40LightningBakedvae.safetensors"
 
 # ---------------- MODEL DOWNLOAD ----------------
 def download_model() -> Path:
-    """Downloads RealVisXL V4.0 model if not present, returns local 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...")
@@ -161,7 +164,7 @@ def download_model() -> Path:
                 local_dir=str(MODEL_DIR),
                 cache_dir=str(HF_CACHE_DIR),
                 force_download=False,
-                resume_download=True,   # <— safer for interrupted downloads
+                resume_download=True,  # safer if download interrupted
             )
         )
         print(f"[ImageGen] Model downloaded to: {model_path}")
@@ -171,7 +174,9 @@ def download_model() -> Path:
 
 # ---------------- PIPELINE LOAD ----------------
 def load_pipeline() -> StableDiffusionXLPipeline:
-    """Loads the RealVisXL V4.0 model for image generation."""
+    """
+    Loads the RealVisXL V4.0 model for image generation.
+    """
     model_path = download_model()
     print("[ImageGen] Loading model into pipeline...")
 
@@ -185,8 +190,10 @@ def load_pipeline() -> StableDiffusionXLPipeline:
     else:
         pipe.to("cpu")
 
-    pipe.safety_checker = None  # optional: skip safety for performance
-    pipe.enable_attention_slicing()  # <— reduces memory use on 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()  
 
     print("[ImageGen] Model ready.")
     return pipe
@@ -194,8 +201,11 @@ def load_pipeline() -> StableDiffusionXLPipeline:
 # ---------------- GLOBAL PIPELINE CACHE ----------------
 pipe: StableDiffusionXLPipeline | None = None
 
-# ---------------- UTILITY: PIL → Base64 ----------------
+# ---------------- UTILITY: PIL → BASE64 ----------------
 def pil_to_base64(img: Image.Image) -> str:
+    """
+    Converts PIL image to base64 string for frontend.
+    """
     buffered = BytesIO()
     img.save(buffered, format="PNG")
     return f"data:image/png;base64,{base64.b64encode(buffered.getvalue()).decode()}"
@@ -204,13 +214,13 @@ def pil_to_base64(img: Image.Image) -> str:
 def generate_images(prompt: str, seed: int | None = None, num_images: int = 3) -> List[str]:
     """
     Generates high-quality images using RealVisXL V4.0.
-    Returns list of base64-encoded PNGs.
+    Returns a list of base64-encoded PNGs.
     """
     global pipe
     if pipe is None:
         pipe = load_pipeline()
 
-    print(f"[ImageGen] Generating {num_images} image(s) for prompt: '{prompt}'  seed={seed}")
+    print(f"[ImageGen] Generating {num_images} image(s) for prompt: '{prompt}' seed={seed}")
     images: List[str] = []
 
     for i in range(num_images):