Create handler.py

handler.py

@@ -0,0 +1,81 @@
+import torch
+import base64
+import io
+from PIL import Image
+from diffusers import StableDiffusionXLImg2ImgPipeline
+import os
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        # 'path' is the folder where HF automatically loaded your repo files
+        print("Loading ARX Pipeline...")
+        
+        # 1. Point directly to the model you uploaded to the repo
+        model_path = os.path.join(path, "biglust.safetensors")
+        
+        # 2. Load the pipeline
+        self.pipe = StableDiffusionXLImg2ImgPipeline.from_single_file(
+            model_path,
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+            safety_checker=None
+        )
+        
+        # 3. Load IP-Adapter (HF will download this from the public hub automatically)
+        self.pipe.load_ip_adapter(
+            "h94/IP-Adapter", 
+            subfolder="sdxl_models", 
+            weight_name="ip-adapter_sdxl.bin"
+        )
+        
+        self.pipe.to("cuda")
+        print("ARX Pipeline Ready.")
+
+    def decode_base64_image(self, image_string):
+        if "," in image_string:
+            image_string = image_string.split(",")[1]
+        image_bytes = base64.b64decode(image_string)
+        return Image.open(io.BytesIO(image_bytes)).convert("RGB")
+
+    def encode_image_base64(self, image):
+        buffered = io.BytesIO()
+        image.save(buffered, format="PNG")
+        return base64.b64encode(buffered.getvalue()).decode('utf-8')
+
+    def __call__(self, data):
+        """
+        data param format: {"inputs": { "prompt": "...", "init_image": "..." }}
+        """
+        # HF wraps payloads in an "inputs" key automatically
+        inputs = data.pop("inputs", data)
+
+        prompt = inputs.get("prompt", "masterpiece, best quality")
+        negative_prompt = inputs.get("negative_prompt", "lowres, bad anatomy, worst quality, ugly")
+        strength = float(inputs.get("strength", 0.65))
+        guidance_scale = float(inputs.get("guidance_scale", 7.0))
+        num_inference_steps = int(inputs.get("steps", 25))
+        ip_adapter_scale = float(inputs.get("ip_adapter_scale", 0.50))
+        
+        init_image_b64 = inputs.get("init_image")
+        ip_adapter_image_b64 = inputs.get("ip_adapter_image")
+
+        if not init_image_b64 or not ip_adapter_image_b64:
+            return {"error": "Both init_image and ip_adapter_image must be provided."}
+
+        init_image = self.decode_base64_image(init_image_b64).resize((1024, 1024))
+        ip_image = self.decode_base64_image(ip_adapter_image_b64).resize((1024, 1024))
+
+        self.pipe.set_ip_adapter_scale(ip_adapter_scale)
+
+        # Generate!
+        result = self.pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=init_image,
+            ip_adapter_image=ip_image,
+            strength=strength,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps
+        ).images[0]
+
+        return {"image": self.encode_image_base64(result)}