Update handler.py

handler.py

@@ -1,87 +1,47 @@
 import base64
-import io
-from PIL import Image
+from io import BytesIO
+from diffusers import StableDiffusionXLPipeline
 import torch
-from diffusers import (
-    StableDiffusionXLControlNetPipeline,
-    ControlNetModel,
-    AutoencoderKL,
-    EulerAncestralDiscreteScheduler
-)
+from PIL import Image
 
 class EndpointHandler:
-    def __init__(self, model_dir):
-
-        # -------------------------------
-        # Load actual model from OTHER repo
-        # -------------------------------
-        BASE_MODEL = "Gjm1234/juggernaut-sfw"   # MAIN MODEL REPO
-        CONTROLNET_REPO = "Gjm1234/juggernaut-sfw/controlnet"  # CONTROLNET INSIDE MODEL REPO
-
-        print("🔧 Loading ControlNet from:", CONTROLNET_REPO)
-        self.controlnet = ControlNetModel.from_pretrained(
-            CONTROLNET_REPO,
-            torch_dtype=torch.float16
-        )
-
-        print("🔧 Loading base model pipeline from:", BASE_MODEL)
-        self.pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
-            BASE_MODEL,
-            controlnet=self.controlnet,
-            torch_dtype=torch.float16
-        )
-
-        # Memory-friendly settings
-        self.pipe.to("cuda")
+    def __init__(self, path=""):
+        print("🔧 Loading Juggernaut-SFW model...")
+        self.pipe = StableDiffusionXLPipeline.from_pretrained(
+            "Gjm1234/juggernaut-sfw",
+            torch_dtype=torch.float16,
+        ).to("cuda")
+
+        # Disable controlnet entirely (you are not loading any ControlNet weights)
+        print("⚠️ ControlNet disabled — no weights provided.")
+
+        # Memory optimisation
+        self.pipe.enable_attention_slicing()
         self.pipe.enable_model_cpu_offload()
-        self.pipe.enable_xformers_memory_efficient_attention()
-
-        print("✅ Pipeline loaded successfully.")
 
     def __call__(self, data):
-
-        prompt = data.get("prompt", "")
-        num_images = int(data.get("num_images", 1))
-
-        # Optional image input
-        image_b64 = data.get("image", None)
-
-        if image_b64:
-            try:
-                image_bytes = base64.b64decode(image_b64)
-                init_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-            except:
-                init_image = None
-        else:
-            init_image = None
-
-        print("📥 Prompt:", prompt)
-        print("📥 Image provided:", init_image is not None)
-
-        # ---- Handle optional image ----
-        if init_image is None:
-            # No image → use text-only diffusion
-            print("🎨 Running TEXT-ONLY generation...")
-            out = self.pipe(
-                prompt=prompt,
-                num_inference_steps=25,
-                num_images_per_prompt=num_images
-            )
-        else:
-            # ControlNet uses the image
-            print("🎛 Running ControlNet IMAGE + TEXT...")
-            out = self.pipe(
-                prompt=prompt,
-                image=init_image,
-                num_inference_steps=25,
-                num_images_per_prompt=num_images
-            )
-
-        images = []
-
-        for img in out.images:
-            buffer = io.BytesIO()
-            img.save(buffer, format="JPEG")
-            images.append(base64.b64encode(buffer.getvalue()).decode("utf-8"))
-
-        return {"images": images}
+        prompt = data.get("prompt", None)
+        if prompt is None:
+            return {"error": "prompt is required"}
+
+        num_images = int(data.get("num_images", 4))
+        if num_images < 1:
+            num_images = 1
+
+        print("🎨 Generating images…")
+
+        images = self.pipe(
+            prompt=prompt,
+            num_inference_steps=25,
+            guidance_scale=6,
+            num_images_per_prompt=num_images
+        ).images
+
+        # Encode all images to Base64
+        encoded = []
+        for img in images:
+            buf = BytesIO()
+            img.save(buf, format="PNG")
+            encoded.append(base64.b64encode(buf.getvalue()).decode("utf-8"))
+
+        return {"images": encoded}