Update handler.py

handler.py

@@ -3,7 +3,8 @@ import torch
 import numpy as np
 import PIL
 from PIL import Image
-from typing import Tuple
+from typing import Tuple, List
+from pydantic import BaseModel
 import diffusers
 from diffusers.utils import load_image
 from diffusers.models import ControlNetModel
@@ -26,6 +27,22 @@ 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
@@ -146,50 +163,33 @@ class EndpointHandler:
             p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
             return p.replace("{prompt}", positive), n + " " + negative
 
-        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()
+        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
 
-        scheduler_class_name = scheduler.split("-")[0]
+        self.pipe.disable_lora() if not enable_LCM else self.pipe.enable_lora()
 
-        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)
+        scheduler_class_name = "EulerDiscreteScheduler"
+
+        self.pipe.scheduler = diffusers.EulerDiscreteScheduler.from_config(self.pipe.scheduler.config)
 
         # 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,7 +246,6 @@ 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(
@@ -261,6 +260,7 @@ class EndpointHandler:
             height=height,
             width=width,
             generator=generator,
+            enhance_face_region=enhance_face_region
         ).images
 
         # Convert the output image to base64