Update core/image_generator.py

core/image_generator.py

@@ -321,52 +321,36 @@ def download_model() -> Path:
 # --------------------------------------------------------------
 # MEMORY-SAFE PIPELINE MANAGER
 # --------------------------------------------------------------
-def unload_pipelines():
-    """Unload any existing pipelines from RAM and GPU."""
+def unload_pipelines(target="all"):
+    """Unload specific or all pipelines."""
     global pipe, img2img_pipe
-    print("[ImageGen] 🧹 Clearing old pipelines from memory...")
-    try:
-        del pipe
+    print("[ImageGen] 🧹 Clearing pipelines from memory...")
+
+    if target in ("pipe", "all"):
+        try:
+            del pipe
+        except:
+            pass
         pipe = None
-    except:
-        pass
-    try:
-        del img2img_pipe
+
+    if target in ("img2img_pipe", "all"):
+        try:
+            del img2img_pipe
+        except:
+            pass
         img2img_pipe = None
-    except:
-        pass
+
     gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
     print("[ImageGen] ✅ Memory cleared.")
 
 
-def safe_load_pipeline(pipeline_cls, model_path: Path):
-    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():
-    """Load text-to-image pipeline into RAM."""
     global pipe
-    unload_pipelines()  # Ensure no duplicate pipelines in memory
+    unload_pipelines(target="pipe")  # only clear old txt2img
     model_path = download_model()
     print("[ImageGen] Loading main (txt2img) pipeline...")
-
     pipe = safe_load_pipeline(StableDiffusionXLPipeline, model_path)
     device = "cuda" if torch.cuda.is_available() else "cpu"
     pipe.to(device)
@@ -376,6 +360,21 @@ def load_pipeline():
     return pipe
 
 
+def load_img2img_pipeline():
+    global img2img_pipe
+    unload_pipelines(target="img2img_pipe")  # only clear old img2img
+    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 load_img2img_pipeline():
     """Load img2img pipeline into RAM."""
     global img2img_pipe