Upload folder using huggingface_hub

inference.py

@@ -14,13 +14,11 @@ from internals.pipelines.prompt_modifier import PromptModifier
 from internals.pipelines.safety_checker import SafetyChecker
 from internals.util.args import apply_style_args
 from internals.util.avatar import Avatar
-from internals.util.cache import auto_clear_cuda_and_gc, clear_cuda, clear_cuda_and_gc
+from internals.util.cache import (auto_clear_cuda_and_gc, clear_cuda,
+                                  clear_cuda_and_gc)
 from internals.util.commons import pickPoses, upload_image, upload_images
-from internals.util.config import (
-    num_return_sequences,
-    set_configs_from_task,
-    set_root_dir,
-)
+from internals.util.config import (num_return_sequences, set_configs_from_task,
+                                   set_root_dir)
 from internals.util.failure_hander import FailureHandler
 from internals.util.lora_style import LoraStyle
 from internals.util.slack import Slack
@@ -295,17 +293,17 @@ def pose(task: Task, s3_outkey: str = "_pose", poses: Optional[list] = None):
     lora_patcher = lora_style.get_patcher(controlnet.pipe2, task.get_style())
     lora_patcher.patch()
 
-    if poses is None:
-        if task.get_pose_coordinates():
-            infered_pose = pose_detector.transform(
-                image=task.get_imageUrl(),
-                client_coordinates=task.get_pose_coordinates(),
-                width=task.get_width(),
-                height=task.get_height(),
-            )
-            poses = [infered_pose] * num_return_sequences
-        else:
-            poses = [controlnet.detect_pose(task.get_imageUrl())] * num_return_sequences
+    try:
+        infered_pose = pose_detector.transform(
+            image=task.get_imageUrl(),
+            client_coordinates=task.get_pose_coordinates(),
+            width=task.get_width(),
+            height=task.get_height(),
+        )
+        poses = [infered_pose] * num_return_sequences
+    except Exception as e:
+        print("Failed to detect pose, using Open Pose detector", e)
+        poses = [controlnet.detect_pose(task.get_imageUrl())] * num_return_sequences
 
     images, has_nsfw = controlnet.process_pose(
         prompt=prompt,

internals/pipelines/pose_detector.py

@@ -4,7 +4,7 @@ from typing import Optional, Union
 from PIL import Image, ImageDraw
 from torch import ge
 
-from internals.util.commons import download_file, download_image
+from internals.util.commons import download_file, download_image, safe_index
 from internals.util.config import get_root_dir
 from models.pose.body import Body
 
@@ -77,16 +77,18 @@ class PoseDetector:
         image = Image.new("RGB", (width, height), "black")
         draw = ImageDraw.Draw(image)
 
-        points = data["candidate"]
+        points: list = data["candidate"]
         for pair in self.__pose_logical_map:
-            xy = points[pair[0] - 1]
-            x1y1 = points[pair[1] - 1]
+            xy = safe_index(points, pair[0] - 1)
+            x1y1 = safe_index(points, pair[1] - 1)
+
+            if xy and x1y1:
+                draw.line(
+                    (xy[0], xy[1], x1y1[0], x1y1[1]),
+                    fill=pair[2],
+                    width=4,
+                )
 
-            draw.line(
-                (xy[0], xy[1], x1y1[0], x1y1[1]),
-                fill=pair[2],
-                width=4,
-            )
         for i, point in enumerate(points):
             x = point[0]
             y = point[1]
@@ -99,7 +101,7 @@ class PoseDetector:
         subset = []
 
         if type(image) == str:
-            image = download_image(imageUrl)
+            image = download_image(image)
 
         image = image.resize((width, height))
 

internals/util/commons.py

@@ -5,7 +5,7 @@ import random
 import re
 from io import BytesIO
 from pathlib import Path
-from typing import Union
+from typing import Optional, Union
 
 import boto3
 import requests
@@ -191,6 +191,14 @@ def construct_default_s3_url(key):
     return "https://comic-assets.s3.ap-south-1.amazonaws.com/" + key
 
 
+def safe_index(array, index) -> Optional:
+    if index < 0:
+        return None
+    if index >= len(array):
+        return None
+    return array[index]
+
+
 def read_url(url: str):
     with urllib.request.urlopen(url) as u:
         return u.read()