Update handler.py

handler.py

@@ -3,7 +3,7 @@ import torch
 import numpy as np
 import PIL
 from PIL import Image
-from typing import Tuple, List
+from typing import Tuple, List, Optional
 from pydantic import BaseModel
 import diffusers
 from diffusers.utils import load_image
@@ -47,6 +47,8 @@ class GenerateImageRequest(BaseModel):
     seed: int
     enable_LCM: bool
     enhance_face_region: bool
+    face_image_path: Optional[str] = None
+    pose_image_path: Optional[str] = None
 
 class EndpointHandler:
     def __init__(self, model_dir):
@@ -195,38 +197,38 @@ class EndpointHandler:
         # apply the style template
         inputs, negative_prompt = apply_style(style_name, inputs, 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")
+        # Load face image
+        face_image_path = data.get("face_image_path", "raw/input.jpg")
+        if os.path.exists(face_image_path):
+            face_image = load_image(face_image_path)
+        else:
+            raise FileNotFoundError(f"Face image not found at path: {face_image_path}")
 
-        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)
         height, width, _ = face_image_cv2.shape
 
-        # Extract face features
-        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_path is not None:
+        # Load pose image if provided
+        pose_image_path = data.get("pose_image_path")
+        if pose_image_path and os.path.exists(pose_image_path):
             pose_image = load_image(pose_image_path)
             pose_image = resize_img(pose_image, max_side=1024)
             img_controlnet = 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"])
 
             width, height = face_kps.size
+        else:
+            img_controlnet = face_image
+
+        # Extract face features
+        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"])
 
         control_mask = np.zeros([height, width, 3])
         x1, y1, x2, y2 = face_info["bbox"]