pose

handler.py

@@ -101,13 +101,13 @@ class EndpointHandler():
 
 
         # controlnet-pose/canny/depth
-        # controlnet_pose_model = "thibaud/controlnet-openpose-sdxl-1.0"
+        controlnet_pose_model = "thibaud/controlnet-openpose-sdxl-1.0"
         # controlnet_canny_model = "diffusers/controlnet-canny-sdxl-1.0"
         # controlnet_depth_model = "diffusers/controlnet-depth-sdxl-1.0-small"
 
-        # controlnet_pose = ControlNetModel.from_pretrained(
-        #     controlnet_pose_model, torch_dtype=dtype
-        # ).to(device)
+        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)
@@ -141,21 +141,21 @@ class EndpointHandler():
 
         #     return depth_image
         
-        # self.controlnet_map = {
-        #     "pose": controlnet_pose,
+        self.controlnet_map = {
+            "pose": controlnet_pose,
         #     "canny": controlnet_canny,
         #     "depth": controlnet_depth,
-        # }
+        }
 
-        # openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
+        openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
         # depth_anything = DepthAnything.from_pretrained('LiheYoung/depth_anything_vitl14').to(device).eval()
 
 
-        # self.controlnet_map_fn = {
-        #     "pose": openpose,
+        self.controlnet_map_fn = {
+            "pose": openpose,
         #     "canny": get_canny_image,
         #     "depth": get_depth_map,
-        # }
+        }
         
         self.app = FaceAnalysis(name="buffalo_l", root="./", providers=["CPUExecutionProvider"])
         self.app.prepare(ctx_id=0, det_size=(640, 640))
@@ -166,9 +166,10 @@ class EndpointHandler():
         
         adapter_strength_ratio = 0.8
         identitynet_strength_ratio = 0.8
-        # pose_strength = 0.4
+        pose_strength = 0.4
         # canny_strength = 0.3
         # depth_strength = 0.5
+        controlnet_selection = ["pose"]
         # controlnet_selection = ["pose", "canny", "depth"]
         
         face_image_path = "https://i.ibb.co/SKg69dD/kaifu-resize.png"
@@ -276,27 +277,27 @@ class EndpointHandler():
         control_mask[y1:y2, x1:x2] = 255
         control_mask = Image.fromarray(control_mask.astype(np.uint8))
 
-        # if len(controlnet_selection) > 0:
-        #     controlnet_scales = {
-        #         "pose": pose_strength,
-        #         "canny": canny_strength,
-        #         "depth": depth_strength,
-        #     }
-        #     self.pipe.controlnet = MultiControlNetModel(
-        #         [self.controlnet_identitynet]
-        #         + [self.controlnet_map[s] for s in 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
-        #     ]
-        # else:
-        self.pipe.controlnet = self.controlnet_identitynet
-        control_scales = float(identitynet_strength_ratio)
-        control_images = face_kps
+        if len(controlnet_selection) > 0:
+            controlnet_scales = {
+                "pose": pose_strength,
+                # "canny": canny_strength,
+                # "depth": depth_strength,
+            }
+            self.pipe.controlnet = MultiControlNetModel(
+                [self.controlnet_identitynet]
+                + [self.controlnet_map[s] for s in 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
+            ]
+        else:
+            self.pipe.controlnet = self.controlnet_identitynet
+            control_scales = float(identitynet_strength_ratio)
+            control_images = face_kps
 
         generator = torch.Generator(device=device.type).manual_seed(3)