Update handler.py

handler.py

@@ -42,7 +42,7 @@ class EndpointHandler:
             self.app = FaceAnalysis(name='antelopev2', root=model_dir, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
         
         self.app.prepare(ctx_id=0, det_size=(640, 640))
-        openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
+        self.openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
 
         # Path to InstantID models
         controlnet_path = os.path.join(model_dir, "checkpoints", "ControlNetModel")
@@ -63,19 +63,14 @@ class EndpointHandler:
             controlnet_canny_model, 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")
-            
         self.controlnet_map = {
             "pose": controlnet_pose,
             "canny": controlnet_canny
         }
 
         self.controlnet_map_fn = {
-            "pose": openpose,
-            "canny": get_canny_image
+            "pose": self.openpose,
+            "canny": self.get_canny_image
         }
 
         pretrained_model_name_or_path = "wangqixun/YamerMIX_v8"
@@ -92,15 +87,14 @@ class EndpointHandler:
             self.pipe.scheduler.config
         )
 
-        # load and disable LCM
+        # Load LoRA weights for LCM
         self.pipe.load_lora_weights("latent-consistency/lcm-lora-sdxl")
-        self.pipe.disable_lora()
-
-        self.pipe.cuda()
-        self.pipe.load_ip_adapter_instantid(face_adapter)
-        self.pipe.image_proj_model.to("cuda")
-        self.pipe.unet.to("cuda")
         
+    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 __call__(self, data):
 
         def convert_from_cv2_to_image(img: np.ndarray) -> Image:
@@ -152,7 +146,7 @@ class EndpointHandler:
             return image
 
         # Debug: Print incoming data
-        print("Incoming data:", json.dumps(data, indent=4))
+        print("Incoming data:", data)
 
         face_image_base64 = data.pop("face_image_base64", None)
         if face_image_base64 is not None:
@@ -182,22 +176,24 @@ class EndpointHandler:
         enhance_non_face_region = data.pop("enhance_non_face_region", False)
         seed = data.pop("seed", 42)
 
-        print(f"[Debug] Style Name: {style_name}")
-        print(f"[Debug] Prompt: {prompt}")
-        print(f"[Debug] Negative Prompt: {negative_prompt}")
+        # Debug: Print additional settings
         print(f"[Debug] IdentityNet Strength Ratio: {identitynet_strength_ratio}")
         print(f"[Debug] Adapter Strength Ratio: {adapter_strength_ratio}")
         print(f"[Debug] Pose Strength: {pose_strength}")
         print(f"[Debug] Canny Strength: {canny_strength}")
         print(f"[Debug] Num Steps: {num_steps}")
         print(f"[Debug] Guidance Scale: {guidance_scale}")
-        print(f"[Debug] ControlNet Selection: {controlnet_selection}")
         print(f"[Debug] Scheduler: {scheduler}")
         print(f"[Debug] Enable Fast Inference: {enable_fast_inference}")
         print(f"[Debug] Enhance Non-Face Region: {enhance_non_face_region}")
         print(f"[Debug] Seed: {seed}")
 
-        self.pipe.disable_lora()
+        # Enable LCM if fast inference is enabled
+        if enable_fast_inference:
+            self.pipe.enable_lora()
+        else:
+            self.pipe.disable_lora()
+
         scheduler_class_name = scheduler.split("-")[0]
 
         add_kwargs = {}