Update handler.py

handler.py

@@ -1,13 +1,15 @@
 import cv2
 import torch
 import numpy as np
+import PIL
 from PIL import Image
-from typing import Tuple, List, Optional, Dict, Any
+from typing import Tuple, List, Optional
 from pydantic import BaseModel
 import diffusers
 from diffusers.utils import load_image
 from diffusers.models import ControlNetModel
 from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
+from insightface.app import FaceAnalysis
 from style_template import styles
 from pipeline_stable_diffusion_xl_instantid_full import StableDiffusionXLInstantIDPipeline, draw_kps
 from controlnet_aux import OpenposeDetector
@@ -17,10 +19,10 @@ import os
 from huggingface_hub import hf_hub_download
 import base64
 import io
+import json
 from transformers import CLIPProcessor, CLIPModel
-import onnxruntime as ort
 
-# Global variables
+# global variable
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16 if str(device).__contains__("cuda") else torch.float32
 STYLE_NAMES = list(styles.keys())
@@ -50,42 +52,58 @@ class GenerateImageRequest(BaseModel):
     pose_image_base64: Optional[str] = None
 
 class EndpointHandler:
-    def __init__(self, model_dir=""):
+    def __init__(self, model_dir):
         # Ensure the necessary files are downloaded
         controlnet_config = hf_hub_download(repo_id="InstantX/InstantID", filename="ControlNetModel/config.json", local_dir=os.path.join(model_dir, "checkpoints"))
         controlnet_model = hf_hub_download(repo_id="InstantX/InstantID", filename="ControlNetModel/diffusion_pytorch_model.safetensors", local_dir=os.path.join(model_dir, "checkpoints"))
         face_adapter = hf_hub_download(repo_id="InstantX/InstantID", filename="ip-adapter.bin", local_dir=os.path.join(model_dir, "checkpoints"))
 
-        # Load the ONNX model
-        onnx_model_path = os.path.join(model_dir, "models", "version-RFB-320.onnx")
-        if not os.path.exists(onnx_model_path):
-            print(f"Model path {onnx_model_path} does not exist. Please ensure the model is available.")
-        self.ort_session = ort.InferenceSession(onnx_model_path)
+        dir_path = os.path.join(model_dir, "models", "face_detection_yunet_2023mar_int8.onnx")
+        if not os.path.exists(dir_path):
+            print(f"Model path {dir_path} does not exist. Attempting to download.")
+            self.app = FaceAnalysis(name='antelopev2', root=model_dir, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
+        else:
+            print(f"Model path {dir_path} exists. Skipping download.")
+            self.app = FaceAnalysis(name='antelopev2', root=model_dir, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
 
-        self.openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
+        self.app.prepare(ctx_id=0, det_size=(640, 640))
+        openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
 
         # Path to InstantID models
         controlnet_path = os.path.join(model_dir, "checkpoints", "ControlNetModel")
 
         # Load pipeline face ControlNetModel
-        self.controlnet_identitynet = ControlNetModel.from_pretrained(controlnet_path, torch_dtype=dtype)
+        self.controlnet_identitynet = ControlNetModel.from_pretrained(
+            controlnet_path, torch_dtype=dtype
+        )
+
+        # controlnet-pose
+        controlnet_pose_model = "thibaud/controlnet-openpose-sdxl-1.0"
+        controlnet_canny_model = "diffusers/controlnet-canny-sdxl-1.0"
+
+        controlnet_pose = ControlNetModel.from_pretrained(
+            controlnet_pose_model, torch_dtype=dtype
+        ).to(device)
+        controlnet_canny = ControlNetModel.from_pretrained(
+            controlnet_canny_model, torch_dtype=dtype
+        ).to(device)
 
-        # Load custom ControlNet models
-        self.controlnet_pose = ControlNetModel.from_pretrained("thibaud/controlnet-openpose-sdxl-1.0", torch_dtype=dtype).to(device)
-        self.controlnet_canny = ControlNetModel.from_pretrained("diffusers/controlnet-canny-sdxl-1.0", torch_dtype=dtype).to(device)
+        def get_canny_image(image, t1=100, t2=200):
+            image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+            edges = cv2.Canny(image, t1, t2)
+            return Image.fromarray(edges, "L")
 
-        # ControlNet map
         self.controlnet_map = {
-            "pose": self.controlnet_pose,
-            "canny": self.controlnet_canny
+            "pose": controlnet_pose,
+            "canny": controlnet_canny
         }
 
         self.controlnet_map_fn = {
-            "pose": self.openpose,
-            "canny": self.get_canny_image
+            "pose": openpose,
+            "canny": get_canny_image
         }
 
-        pretrained_model_name_or_path = "stablediffusionapi/protovision-xl-high-fidel"
+        pretrained_model_name_or_path = "wangqixun/YamerMIX_v8"
 
         self.pipe = StableDiffusionXLInstantIDPipeline.from_pretrained(
             pretrained_model_name_or_path,
@@ -99,7 +117,7 @@ class EndpointHandler:
             self.pipe.scheduler.config
         )
 
-        # Load and disable LCM
+        # load and disable LCM
         self.pipe.load_lora_weights(lcm_lora_path)
         self.pipe.fuse_lora()
         self.pipe.disable_lora()
@@ -113,161 +131,186 @@ class EndpointHandler:
         self.clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
         self.clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
 
-    def get_canny_image(self, image, t1=100, t2=200):
-        image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-        edges = cv2.Canny(image, t1, t2)
-        return Image.fromarray(edges, "L")
-
     def is_nsfw(self, image: Image.Image) -> bool:
+        """
+        Check if an image contains NSFW content using CLIP model.
+
+        Args:
+            image (Image.Image): PIL image to check.
+
+        Returns:
+            bool: True if the image is NSFW, False otherwise.
+        """
         inputs = self.clip_processor(text=["NSFW", "SFW"], images=image, return_tensors="pt", padding=True)
         inputs = {k: v.to(device) for k, v in inputs.items()}
         outputs = self.clip_model(**inputs)
-        logits_per_image = outputs.logits_per_image  # image-text similarity score
-        probs = logits_per_image.softmax(dim=1)  # probabilities
+        logits_per_image = outputs.logits_per_image  # this is the image-text similarity score
+        probs = logits_per_image.softmax(dim=1)  # we take the softmax to get the probabilities
         nsfw_prob = probs[0, 0].item()  # probability of "NSFW" label
-        return nsfw_prob > 0.9  # threshold for NSFW detection
-
-    def preprocess(self, image):
-        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-        image = cv2.resize(image, (320, 240))
-        image_mean = np.array([127, 127, 127])
-        image = (image - image_mean) / 128
-        image = np.transpose(image, [2, 0, 1])
-        image = np.expand_dims(image, axis=0)
-        image = image.astype(np.float32)
-        return image
-
-    def get_face_info(self, image):
-        preprocessed_image = self.preprocess(image)
-        input_name = self.ort_session.get_inputs()[0].name
-        confidences, boxes = self.ort_session.run(None, {input_name: preprocessed_image})
-
-        print(f"Confidences shape: {confidences.shape}, Boxes shape: {boxes.shape}")
-
-        face_info_list = []
-        for i in range(len(boxes)):
-            box = boxes[i]
-            conf = confidences[i]
-            if conf[0] > 0.7:  # Fixing the out-of-bounds issue
-                x1, y1, x2, y2 = box[0] * 320, box[1] * 240, box[2] * 320, box[3] * 240
-                face_info_list.append({"bbox": [x1, y1, x2, y2]})
-        return face_info_list
-
-    def __call__(self, data: Any) -> Dict[str, Any]:
-        request = GenerateImageRequest(**data)
+        return nsfw_prob > 0.8  # Adjusted threshold for NSFW detection
+
+    def __call__(self, data):
+
+        def convert_from_cv2_to_image(img: np.ndarray) -> Image:
+            return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+
+        def convert_from_image_to_cv2(img: Image) -> np.ndarray:
+            return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+
+        def resize_img(
+            input_image,
+            max_side=1280,
+            min_side=1024,
+            size=None,
+            pad_to_max_side=False,
+            mode=PIL.Image.BILINEAR,
+            base_pixel_number=64,
+        ):
+            w, h = input_image.size
+            if size is not None:
+                w_resize_new, h_resize_new = size
+            else:
+                ratio = min_side / min(h, w)
+                w, h = round(ratio * w), round(ratio * h)
+                ratio = max_side / max(h, w)
+                input_image = input_image.resize([round(ratio * w), round(ratio * h)], mode)
+                w_resize_new = (round(ratio * w) // base_pixel_number) * base_pixel_number
+                h_resize_new = (round(ratio * h) // base_pixel_number) * base_pixel_number
+            input_image = input_image.resize([w_resize_new, h_resize_new], mode)
+
+            if pad_to_max_side:
+                res = np.ones([max_side, max_side, 3], dtype=np.uint8) * 255
+                offset_x = (max_side - w_resize_new) // 2
+                offset_y = (max_side - h_resize_new) // 2
+                res[
+                    offset_y : offset_y + h_resize_new, offset_x : offset_x + w_resize_new
+                ] = np.array(input_image)
+                input_image = Image.fromarray(res)
+            return input_image
+
+        def apply_style(
+            style_name: str, positive: str, negative: str = ""
+        ) -> Tuple[str, str]:
+            p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
+            return p.replace("{prompt}", positive), n + " " + negative
 
-        if request.enable_LCM:
+        request = GenerateImageRequest(**data)
+        inputs = request.inputs
+        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
+
+        if enable_LCM:
             self.pipe.enable_lora()
             self.pipe.scheduler = diffusers.LCMScheduler.from_config(self.pipe.scheduler.config)
-            guidance_scale = min(max(request.guidance_scale, 0), 1)
+            guidance_scale = min(max(guidance_scale, 0), 1)
         else:
             self.pipe.disable_lora()
             self.pipe.scheduler = diffusers.EulerDiscreteScheduler.from_config(self.pipe.scheduler.config)
 
-        # Apply style
-        inputs, negative_prompt = self.apply_style(request.style, request.inputs, request.negative_prompt)
+        # apply the style template
+        inputs, negative_prompt = apply_style(style_name, inputs, negative_prompt)
+
+        # Decode base64 image
+        face_image_base64 = data.get("face_image_base64")
+        face_image_data = base64.b64decode(face_image_base64)
+        face_image = Image.open(io.BytesIO(face_image_data))
 
-        # Decode base64 images
-        face_image = self.decode_base64_image(request.face_image_base64)
-        pose_image = self.decode_base64_image(request.pose_image_base64) if request.pose_image_base64 else None
+        pose_image_base64 = data.get("pose_image_base64")
+        pose_image = None
+        if pose_image_base64:
+            pose_image_data = base64.b64decode(pose_image_base64)
+            pose_image = Image.open(io.BytesIO(pose_image_data))
 
-        face_image = self.resize_img(face_image, max_side=1024)
-        face_image_cv2 = self.convert_from_image_to_cv2(face_image)
+        face_image = resize_img(face_image, max_side=1024)
+        face_image_cv2 = convert_from_image_to_cv2(face_image)
         height, width, _ = face_image_cv2.shape
 
         # Extract face features
-        face_info_list = self.get_face_info(face_image_cv2)
-        if len(face_info_list) == 0:
-            return {"error": "No faces detected."}
-
-        face_info = max(face_info_list, key=lambda x: (x["bbox"][2] - x["bbox"][0]) * (x["bbox"][3] - x["bbox"][1]))
-        face_kps = draw_kps(self.convert_from_cv2_to_image(face_image_cv2), face_info["bbox"])
+        face_info = self.app.get(face_image_cv2)
+
+        face_info = sorted(
+            face_info,
+            key=lambda x: (x["bbox"][2] - x["bbox"][0]) * x["bbox"][3] - x["bbox"][1],
+        )[
+            -1
+        ]  # only use the maximum face
+        face_emb = face_info["embedding"]
+        face_kps = draw_kps(convert_from_cv2_to_image(face_image_cv2), face_info["kps"])
         img_controlnet = face_image
-
         if pose_image:
-            pose_image = self.resize_img(pose_image, max_side=1024)
+            pose_image = resize_img(pose_image, max_side=1024)
             img_controlnet = pose_image
-            pose_image_cv2 = self.convert_from_image_to_cv2(pose_image)
-            face_info_list = self.get_face_info(pose_image_cv2)
-            if len(face_info_list) == 0:
-                return {"error": "No faces detected in pose image."}
+            pose_image_cv2 = convert_from_image_to_cv2(pose_image)
+
+            face_info = self.app.get(pose_image_cv2)
+
+            face_info = face_info[-1]
+            face_kps = draw_kps(pose_image, face_info["kps"])
 
-            face_info = max(face_info_list, key=lambda x: (x["bbox"][2] - x["bbox"][0]) * (x["bbox"][3] - x["bbox"][1]))
-            face_kps = draw_kps(pose_image, face_info["bbox"])
             width, height = face_kps.size
 
-        control_mask = np.zeros([height, width, 3], dtype=np.uint8)
-        x1, y1, x2, y2 = map(int, face_info["bbox"])
+        control_mask = np.zeros([height, width, 3])
+        x1, y1, x2, y2 = face_info["bbox"]
+        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
         control_mask[y1:y2, x1:x2] = 255
-        control_mask = Image.fromarray(control_mask)
+        control_mask = Image.fromarray(control_mask.astype(np.uint8))
 
-        controlnet_scales = {"pose": request.pose_strength, "canny": request.canny_strength}
+        controlnet_scales = {
+            "pose": pose_strength,
+            "canny": canny_strength
+        }
         self.pipe.controlnet = MultiControlNetModel(
-            [self.controlnet_identitynet] + [self.controlnet_map[s] for s in request.controlnet_selection]
+            [self.controlnet_identitynet]
+            + [self.controlnet_map[s] for s in controlnet_selection]
         )
-        control_scales = [float(request.identitynet_strength_ratio)] + [controlnet_scales[s] for s in request.controlnet_selection]
-        control_images = [face_kps] + [self.controlnet_map_fn[s](img_controlnet).resize((width, height)) for s in request.controlnet_selection]
+        control_scales = [float(identitynet_strength_ratio)] + [
+            controlnet_scales[s] for s in controlnet_selection
+        ]
+        control_images = [face_kps] + [
+            self.controlnet_map_fn[s](img_controlnet).resize((width, height))
+            for s in controlnet_selection
+        ]
+
+        generator = torch.Generator(device=device).manual_seed(seed)
 
-        generator = torch.Generator(device=device).manual_seed(request.seed)
+        print("Start inference...")
+        print(f"[Debug] Prompt: {inputs}, \n[Debug] Neg Prompt: {negative_prompt}")
 
+        self.pipe.set_ip_adapter_scale(adapter_strength_ratio)
         outputs = self.pipe(
             prompt=inputs,
             negative_prompt=negative_prompt,
+            image_embeds=face_emb,
             image=control_images,
             control_mask=control_mask,
             controlnet_conditioning_scale=control_scales,
-            num_inference_steps=request.num_steps,
-            guidance_scale=request.guidance_scale,
+            num_inference_steps=num_steps,
+            guidance_scale=guidance_scale,
             height=height,
             width=width,
             generator=generator,
-            enhance_face_region=request.enhance_face_region,
+            enhance_face_region=enhance_face_region
         )
-
+       
         images = outputs.images
 
+        # Check for NSFW content
         if self.is_nsfw(images[0]):
             return {"error": "Generated image contains NSFW content and was discarded."}
 
-        # Convert the image to base64
+        # Convert the output image to base64
         buffered = io.BytesIO()
         images[0].save(buffered, format="JPEG")
         img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
 
         return {"generated_image_base64": img_str}
-
-    def decode_base64_image(self, image_string):
-        base64_image = base64.b64decode(image_string)
-        buffer = io.BytesIO(base64_image)
-        return Image.open(buffer)
-
-    def convert_from_cv2_to_image(self, img: np.ndarray) -> Image:
-        return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
-
-    def convert_from_image_to_cv2(self, img: Image) -> np.ndarray:
-        return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
-
-    def resize_img(self, input_image, max_side=1280, min_side=1024, size=None, pad_to_max_side=False, mode=Image.BILINEAR, base_pixel_number=64):
-        w, h = input_image.size
-        if size is not None:
-            w_resize_new, h_resize_new = size
-        else:
-            ratio = min_side / min(h, w)
-            w, h = round(ratio * w), round(ratio * h)
-            ratio = max_side / max(h, w)
-            input_image = input_image.resize([round(ratio * w), round(ratio * h)], mode)
-            w_resize_new = (round(ratio * w) // base_pixel_number) * base_pixel_number
-            h_resize_new = (round(ratio * h) // base_pixel_number) * base_pixel_number
-        input_image = input_image.resize([w_resize_new, h_resize_new], mode)
-
-        if pad_to_max_side:
-            res = np.ones([max_side, max_side, 3], dtype=np.uint8) * 255
-            offset_x = (max_side - w_resize_new) // 2
-            offset_y = (max_side - h_resize_new) // 2
-            res[offset_y: offset_y + h_resize_new, offset_x: offset_x + w_resize_new] = np.array(input_image)
-            input_image = Image.fromarray(res)
-        return input_image
-
-    def apply_style(self, style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
-        p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
-        return p.replace("{prompt}", positive), n + " " + negative