Update core/image_generator.py

core/image_generator.py

@@ -272,12 +272,11 @@
 
 
 
-
 import os
 from pathlib import Path
 import gc
 import torch
-from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
+from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline, AutoPipelineForText2Image
 from huggingface_hub import hf_hub_download
 from typing import Dict, Any
 from PIL import Image
@@ -322,7 +321,6 @@ def safe_expanduser(path):
     return os.path.expanduser_original(path)
 
 os.path.expanduser = safe_expanduser
-
 tempfile.tempdir = str(HF_CACHE_DIR)
 
 print("[DEBUG] ✅ Hugging Face, Diffusers, Datasets and Torch cache fully redirected to:", HF_CACHE_DIR)
@@ -344,34 +342,11 @@ print("[DEBUG] ✅ Seed directory:", SEED_DIR)
 # --------------------------------------------------------------
 # MODEL CONFIG
 # --------------------------------------------------------------
-MODEL_REPO = "Lykon/dreamshaper-8"
-MODEL_FILENAME = "dreamshaper_8.safetensors"
-
+MODEL_REPO = "lykon/dreamshaper-8"  # Use Hugging Face repo
 # ---------------- GLOBAL PIPELINE CACHE ----------------
-pipe: StableDiffusionXLPipeline | None = None
+pipe: StableDiffusionXLPipeline | AutoPipelineForText2Image | None = None
 img2img_pipe: StableDiffusionXLImg2ImgPipeline | None = None
 
-# --------------------------------------------------------------
-# MODEL DOWNLOAD
-# --------------------------------------------------------------
-def download_model() -> Path:
-    model_path = MODEL_DIR / MODEL_FILENAME
-    if not model_path.exists():
-        print("[ImageGen] Downloading DreamShaper SD1.5 model...")
-        model_path = Path(
-            hf_hub_download(
-                repo_id=MODEL_REPO,
-                filename=MODEL_FILENAME,
-                cache_dir=str(HF_CACHE_DIR),
-                force_download=False,
-                resume_download=True,
-            )
-        )
-        print(f"[ImageGen] ✅ Model downloaded to: {model_path}")
-    else:
-        print("[ImageGen] ✅ Model already exists at:", model_path)
-    return model_path
-
 # --------------------------------------------------------------
 # MEMORY-SAFE PIPELINE MANAGER
 # --------------------------------------------------------------
@@ -399,18 +374,22 @@ 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."""
+def safe_load_pipeline(pretrained_model_name):
+    """Load DreamShaper SD1.5 safely via from_pretrained."""
     try:
-        print(f"[ImageGen] 🔄 Loading {pipeline_class.__name__} from {model_path} ...")
-        pipe = pipeline_class.from_single_file(
-            model_path,
+        print(f"[ImageGen] 🔄 Loading model {pretrained_model_name} ...")
+        pipe = AutoPipelineForText2Image.from_pretrained(
+            pretrained_model_name,
             torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            variant="fp16"  # use fp16 if possible
         )
-        print(f"[ImageGen] ✅ Successfully loaded {pipeline_class.__name__}.")
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        pipe = pipe.to(device)
+        pipe.enable_attention_slicing()
+        print(f"[ImageGen] ✅ Successfully loaded {pretrained_model_name}.")
         return pipe
     except Exception as e:
-        print(f"[ImageGen] ❌ Failed to load {pipeline_class.__name__}: {e}")
+        print(f"[ImageGen] ❌ Failed to load {pretrained_model_name}: {e}")
         unload_pipelines()
         gc.collect()
         if torch.cuda.is_available():
@@ -420,26 +399,17 @@ def safe_load_pipeline(pipeline_class, model_path):
 def load_pipeline():
     global pipe
     unload_pipelines(target="pipe")
-    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)
-    pipe.safety_checker = None
-    pipe.enable_attention_slicing()
+    pipe = safe_load_pipeline(MODEL_REPO)
     print("[ImageGen] ✅ Text-to-image pipeline ready.")
     return pipe
 
 def load_img2img_pipeline():
     global img2img_pipe
     unload_pipelines(target="img2img_pipe")
-    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()
+    # For DreamShaper, img2img uses the same pipeline
+    img2img_pipe = safe_load_pipeline(MODEL_REPO)
     print("[ImageGen] ✅ Img2Img pipeline ready.")
     return img2img_pipe
 
@@ -529,13 +499,15 @@ async def generate_images(prompt_or_json, seed: int | None = None, num_images: i
     pipe = load_pipeline()
     images = []
     for i in range(num_images):
-      gen = torch.Generator(device).manual_seed(seed + i) if seed is not None else None
-      try:
-        img = pipe(prompt_or_json, num_inference_steps=30, generator=gen).images[0]
-        img_path = TMP_DIR / f"prompt_{i}.png"
-        img.save(img_path)
-        images.append(pil_to_base64(img))
-      except Exception as e:
-        print(f"[ImageGen] ⚠️ Failed on image {i}: {e}")
-
- 
+        gen = torch.Generator(device).manual_seed(seed + i) if seed is not None else None
+        try:
+            img = pipe(prompt_or_json, num_inference_steps=30, generator=gen).images[0]
+            img_path = TMP_DIR / f"prompt_{i}.png"
+            img.save(img_path)
+            images.append(pil_to_base64(img))
+        except Exception as e:
+            print(f"[ImageGen] ⚠️ Failed on image {i}: {e}")
+
+    unload_pipelines(target="pipe")
+    print(f"[ImageGen] ✅ Generated {len(images)} image(s) successfully.")
+    return images