Create handler.py

handler.py

@@ -0,0 +1,57 @@
+import torch
+from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel
+from diffusers.utils import load_image
+import base64
+import io
+from PIL import Image
+
+class EndpointHandler:
+    def __init__(self, model_dir=None):
+        # Load ControlNet (Depth)
+        self.controlnet = ControlNetModel.from_pretrained(
+            "diffusers/controlnet-depth-sdxl-1.0",
+            torch_dtype=torch.float16
+        )
+
+        # Load Juggernaut XL as base model (from HuggingFace)
+        self.pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
+            "RunDiffusion/Juggernaut-XL-v9",
+            controlnet=self.controlnet,
+            torch_dtype=torch.float16
+        ).to("cuda")
+
+        self.pipe.enable_xformers_memory_efficient_attention()
+
+    def __call__(self, data):
+        # Parse prompt
+        prompt = data.get("prompt", "")
+
+        # Optional negative prompt
+        negative_prompt = data.get("negative_prompt", "bad quality, distortions")
+
+        # Handle uploaded image (base64 or raw bytes)
+        image_bytes = data.get("image")
+
+        if isinstance(image_bytes, str):
+            image_bytes = base64.b64decode(image_bytes)
+
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+
+        # Generate 10 images
+        results = self.pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=image,
+            num_inference_steps=35,
+            guidance_scale=7.5,
+            controlnet_conditioning_scale=1.0,
+            num_images_per_prompt=10
+        )
+
+        output_images = []
+        for img in results.images:
+            buffer = io.BytesIO()
+            img.save(buffer, format="JPEG")
+            output_images.append(base64.b64encode(buffer.getvalue()).decode())
+
+        return { "images": output_images }