Update handler.py

handler.py

@@ -2,26 +2,32 @@ import torch
 import base64
 import io
 from PIL import Image
-from diffusers import StableDiffusionXLImg2ImgPipeline
+from diffusers import StableDiffusionXLImg2ImgPipeline, DPMSolverMultistepScheduler
 import os
 
 class EndpointHandler():
     def __init__(self, path=""):
-        # 'path' is the folder where HF automatically loaded your repo files
-        print("Loading ARX Pipeline...")
+        print("Loading ARX Elite Pipeline...")
         
-        # 1. Point directly to the model you uploaded to the repo
         model_path = os.path.join(path, "biglust.safetensors")
         
-        # 2. Load the pipeline
+        # 1. Load the core pipeline
         self.pipe = StableDiffusionXLImg2ImgPipeline.from_single_file(
             model_path,
             torch_dtype=torch.float16,
             use_safetensors=True,
             safety_checker=None
         )
+
+        # 2. THE SECRET SAUCE: Upgrade the Scheduler to DPM++ 2M Karras
+        # This makes the images much sharper and helps maintain facial features
+        self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(
+            self.pipe.scheduler.config, 
+            use_karras_sigmas=True,
+            algorithm_type="sde-dpmsolver++"
+        )
         
-        # 3. Load IP-Adapter (HF will download this from the public hub automatically)
+        # 3. Load IP-Adapter
         self.pipe.load_ip_adapter(
             "h94/IP-Adapter", 
             subfolder="sdxl_models", 
@@ -29,7 +35,7 @@ class EndpointHandler():
         )
         
         self.pipe.to("cuda")
-        print("ARX Pipeline Ready.")
+        print("ARX Elite Ready.")
 
     def decode_base64_image(self, image_string):
         if "," in image_string:
@@ -43,31 +49,29 @@ class EndpointHandler():
         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))
+        prompt = inputs.get("prompt", "masterpiece, best quality, highly detailed")
+        negative_prompt = inputs.get("negative_prompt", "blurry, lowres, bad anatomy, worst quality, ugly, deformed eyes")
+        
+        # We increase steps slightly to 30 for better quality with this scheduler
+        strength = float(inputs.get("strength", 0.55))
+        guidance_scale = float(inputs.get("guidance_scale", 7.5))
+        num_inference_steps = int(inputs.get("steps", 30)) 
+        ip_adapter_scale = float(inputs.get("ip_adapter_scale", 0.75))
         
         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."}
+            return {"error": "Missing image inputs."}
 
         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!
+        # Generate with the new high-precision scheduler
         result = self.pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,