Update handler.py

handler.py

@@ -3,8 +3,7 @@ import torch
 import numpy as np
 import PIL
 from PIL import Image
-from typing import Tuple, List
-from pydantic import BaseModel
+from typing import Tuple
 import diffusers
 from diffusers.utils import load_image
 from diffusers.models import ControlNetModel
@@ -27,22 +26,6 @@ dtype = torch.float16 if str(device).__contains__("cuda") else torch.float32
 STYLE_NAMES = list(styles.keys())
 DEFAULT_STYLE_NAME = "Spring Festival"
 
-class GenerateImageRequest(BaseModel):
-    prompt: str
-    negative_prompt: str
-    style: str
-    num_steps: int
-    identitynet_strength_ratio: float
-    adapter_strength_ratio: float
-    pose_strength: float
-    canny_strength: float
-    depth_strength: float
-    controlnet_selection: List[str]
-    guidance_scale: float
-    seed: int
-    enable_LCM: bool
-    enhance_face_region: bool
-
 class EndpointHandler:
     def __init__(self, model_dir):
         # Ensure the necessary files are downloaded
@@ -163,33 +146,50 @@ class EndpointHandler:
             p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
             return p.replace("{prompt}", positive), n + " " + negative
 
-        request = GenerateImageRequest(**data)
-        prompt = request.prompt
-        negative_prompt = request.negative_prompt
-        style_name = request.style
-        identitynet_strength_ratio = request.identitynet_strength_ratio
-        adapter_strength_ratio = request.adapter_strength_ratio
-        pose_strength = request.pose_strength
-        canny_strength = request.canny_strength
-        num_steps = request.num_steps
-        guidance_scale = request.guidance_scale
-        controlnet_selection = request.controlnet_selection
-        seed = request.seed
-        enhance_face_region = request.enhance_face_region
-        enable_LCM = request.enable_LCM
-
-        self.pipe.disable_lora() if not enable_LCM else self.pipe.enable_lora()
+        face_image_path = data.pop("face_image_path", "https://i.ibb.co/GQzm527/examples-musk-resize.jpg")
+        pose_image_path = data.pop("pose_image_path", "https://i.ibb.co/TRCK4MS/examples-poses-pose2.jpg")
+        style_name = data.pop("style_name", DEFAULT_STYLE_NAME)
+        prompt = data.pop("inputs", "a man flying in the sky in Mars")
+        negative_prompt = data.pop("negative_prompt", "(lowres, low quality, worst quality:1.2), (text:1.2), watermark, (frame:1.2), deformed, ugly, deformed eyes, blur, out of focus, blurry, deformed cat, deformed, photo, anthropomorphic cat, monochrome, pet collar, gun, weapon, blue, 3d, drones, drone, buildings in background, green")
+
+        identitynet_strength_ratio = data.pop("identitynet_strength_ratio", 0.8)
+        adapter_strength_ratio = data.pop("adapter_strength_ratio", 0.8)
+        pose_strength = data.pop("pose_strength", 0.5)
+        canny_strength = data.pop("canny_strength", 0.3)
+        num_steps = data.pop("num_steps", 20)
+        guidance_scale = data.pop("guidance_scale", 5.0)
+        controlnet_selection = data.pop("controlnet_selection", ["pose", "canny"])
+        scheduler = data.pop("scheduler", "EulerDiscreteScheduler")
+        enable_fast_inference = data.pop("enable_fast_inference", False)
+        enhance_non_face_region = data.pop("enhance_non_face_region", False)
+        seed = data.pop("seed", 42)
+
+        # Ensure required fields are present
+        data.setdefault("prompt", prompt)
+        data.setdefault("style", style_name)
+        data.setdefault("num_steps", num_steps)
+        data.setdefault("enable_LCM", enable_fast_inference)
+        data.setdefault("enhance_face_region", enhance_non_face_region)
+
+        # Enable LCM if fast inference is enabled
+        if enable_fast_inference:
+            self.pipe.enable_lora()
+        else:
+            self.pipe.disable_lora()
 
-        scheduler_class_name = "EulerDiscreteScheduler"
+        scheduler_class_name = scheduler.split("-")[0]
 
-        self.pipe.scheduler = diffusers.EulerDiscreteScheduler.from_config(self.pipe.scheduler.config)
+        add_kwargs = {}
+        if len(scheduler.split("-")) > 1:
+            add_kwargs["use_karras_sigmas"] = True
+        if len(scheduler.split("-")) > 2:
+            add_kwargs["algorithm_type"] = "sde-dpmsolver++"
+        scheduler = getattr(diffusers, scheduler_class_name)
+        self.pipe.scheduler = scheduler.from_config(self.pipe.scheduler.config, **add_kwargs)
 
         # apply the style template
         prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)
 
-        face_image_path = data.get("face_image_path", "https://i.ibb.co/GQzm527/examples-musk-resize.jpg")
-        pose_image_path = data.get("pose_image_path", "https://i.ibb.co/TRCK4MS/examples-poses-pose2.jpg")
-
         face_image = load_image(face_image_path)
         face_image = resize_img(face_image, max_side=1024)
         face_image_cv2 = convert_from_image_to_cv2(face_image)
@@ -246,6 +246,7 @@ class EndpointHandler:
 
         print("Start inference...")
         print(f"[Debug] Prompt: {prompt}, \n[Debug] Neg Prompt: {negative_prompt}")
+        print(f"[Debug] Number of Inference Steps: {num_steps}")
 
         self.pipe.set_ip_adapter_scale(adapter_strength_ratio)
         images = self.pipe(
@@ -260,7 +261,6 @@ class EndpointHandler:
             height=height,
             width=width,
             generator=generator,
-            enhance_face_region=enhance_face_region
         ).images
 
         # Convert the output image to base64