indentation

handler.py

@@ -148,64 +148,64 @@ class EndpointHandler():
 
         return os.path.join(os.getcwd(), output_path)
 
-def __call__(self, data: Any) -> Dict[str, str]:
-    inputs = data.get("inputs", {})
-    ref_image_base64 = inputs.get("ref_image", "")
-    pose_video_path = inputs.get("pose_video_path", "")
-    width = inputs.get("width", 512)
-    height = inputs.get("height", 768)
-    length = inputs.get("length", 24)
-    num_inference_steps = inputs.get("num_inference_steps", 25)
-    cfg = inputs.get("cfg", 3.5)
-    seed = inputs.get("seed", 123)
-
-    ref_image = Image.open(BytesIO(base64.b64decode(ref_image_base64)))
-
-    # Get the base directory of the current file
-    base_dir = os.path.dirname(os.path.abspath(__file__))
-    
-    # Update pose_video_path to use the base directory
-    pose_video_path = os.path.join(base_dir, pose_video_path)
-    
-    if not os.path.exists(pose_video_path):
-        raise FileNotFoundError(f"The pose video was not found at: {pose_video_path}")
-    
-    torch.manual_seed(seed)
-    pose_images = read_frames(pose_video_path)
-    src_fps = get_fps(pose_video_path)
-
-    pose_list = []
-    total_length = min(length, len(pose_images))
-    for pose_image_pil in pose_images[:total_length]:
-        pose_list.append(pose_image_pil)
-
-    video = self.pipeline(
-        ref_image,
-        pose_list,
-        width=width,
-        height=height,
-        video_length=total_length,
-        num_inference_steps=num_inference_steps,
-        guidance_scale=cfg
-    ).videos
-
-    save_dir = os.path.join(base_dir, "output", "gradio")
-    if not os.path.exists(save_dir):
-        os.makedirs(save_dir, exist_ok=True)
-    animation_path = os.path.join(save_dir, "animation_output.mp4")
-    save_videos_grid(video, animation_path, n_rows=1, fps=src_fps)
-
-    # Crop the face from the reference image and save it
-    cropped_face_path = os.path.join(save_dir, "cropped_face.jpg")
-    cropped_face = self._crop_face(ref_image, save_path=cropped_face_path)
-
-    # Perform face swapping
-    final_video_path = self._swap_face(cropped_face, animation_path)
-
-    # Encode the final video in base64
-    with open(final_video_path, "rb") as video_file:
-        video_base64 = base64.b64encode(video_file.read()).decode("utf-8")
-
-    torch.cuda.empty_cache()
-
-    return {"video": video_base64}
+    def __call__(self, data: Any) -> Dict[str, str]:
+        inputs = data.get("inputs", {})
+        ref_image_base64 = inputs.get("ref_image", "")
+        pose_video_path = inputs.get("pose_video_path", "")
+        width = inputs.get("width", 512)
+        height = inputs.get("height", 768)
+        length = inputs.get("length", 24)
+        num_inference_steps = inputs.get("num_inference_steps", 25)
+        cfg = inputs.get("cfg", 3.5)
+        seed = inputs.get("seed", 123)
+
+        ref_image = Image.open(BytesIO(base64.b64decode(ref_image_base64)))
+
+        # Get the base directory of the current file
+        base_dir = os.path.dirname(os.path.abspath(__file__))
+        
+        # Update pose_video_path to use the base directory
+        pose_video_path = os.path.join(base_dir, pose_video_path)
+        
+        if not os.path.exists(pose_video_path):
+            raise FileNotFoundError(f"The pose video was not found at: {pose_video_path}")
+        
+        torch.manual_seed(seed)
+        pose_images = read_frames(pose_video_path)
+        src_fps = get_fps(pose_video_path)
+
+        pose_list = []
+        total_length = min(length, len(pose_images))
+        for pose_image_pil in pose_images[:total_length]:
+            pose_list.append(pose_image_pil)
+
+        video = self.pipeline(
+            ref_image,
+            pose_list,
+            width=width,
+            height=height,
+            video_length=total_length,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=cfg
+        ).videos
+
+        save_dir = os.path.join(base_dir, "output", "gradio")
+        if not os.path.exists(save_dir):
+            os.makedirs(save_dir, exist_ok=True)
+        animation_path = os.path.join(save_dir, "animation_output.mp4")
+        save_videos_grid(video, animation_path, n_rows=1, fps=src_fps)
+
+        # Crop the face from the reference image and save it
+        cropped_face_path = os.path.join(save_dir, "cropped_face.jpg")
+        cropped_face = self._crop_face(ref_image, save_path=cropped_face_path)
+
+        # Perform face swapping
+        final_video_path = self._swap_face(cropped_face, animation_path)
+
+        # Encode the final video in base64
+        with open(final_video_path, "rb") as video_file:
+            video_base64 = base64.b64encode(video_file.read()).decode("utf-8")
+
+        torch.cuda.empty_cache()
+
+        return {"video": video_base64}