Upload folder using huggingface_hub

inference.py

@@ -14,11 +14,18 @@ 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.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.cache import auto_clear_cuda_and_gc, clear_cuda, clear_cuda_and_gc
+from internals.util.commons import (
+    download_image,
+    pickPoses,
+    upload_image,
+    upload_images,
+)
+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
@@ -293,7 +300,12 @@ 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 task.get_pose_coordinates():
+    if not task.get_pose_estimation():
+        pose = download_image(task.get_imageUrl()).resize(
+            (task.get_width(), task.get_height())
+        )
+        poses = [pose] * num_return_sequences
+    elif task.get_pose_coordinates():
         infered_pose = pose_detector.transform(
             image=task.get_imageUrl(),
             client_coordinates=task.get_pose_coordinates(),

inference2.py

@@ -7,17 +7,18 @@ from internals.data.task import ModelType, Task, TaskType
 from internals.pipelines.inpainter import InPainter
 from internals.pipelines.object_remove import ObjectRemoval
 from internals.pipelines.prompt_modifier import PromptModifier
-from internals.pipelines.remove_background import (RemoveBackground,
-                                                   RemoveBackgroundV2)
+from internals.pipelines.remove_background import RemoveBackground, RemoveBackgroundV2
 from internals.pipelines.replace_background import ReplaceBackground
 from internals.pipelines.safety_checker import SafetyChecker
 from internals.pipelines.upscaler import Upscaler
 from internals.util.avatar import Avatar
 from internals.util.cache import auto_clear_cuda_and_gc, clear_cuda
-from internals.util.commons import (construct_default_s3_url, upload_image,
-                                    upload_images)
-from internals.util.config import (num_return_sequences, set_configs_from_task,
-                                   set_root_dir)
+from internals.util.commons import construct_default_s3_url, upload_image, upload_images
+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.slack import Slack
 
@@ -41,8 +42,8 @@ replace_background = ReplaceBackground()
 @update_db
 @slack.auto_send_alert
 def remove_bg(task: Task):
-    remove_background = RemoveBackground()
-    output_image = remove_background.remove(task.get_imageUrl())
+    # remove_background = RemoveBackground()
+    output_image = remove_background_v2.remove(task.get_imageUrl())
 
     output_key = "crecoAI/{}_rmbg.png".format(task.get_taskId())
     upload_image(output_image, output_key)

internals/data/task.py

@@ -97,6 +97,9 @@ class Task:
     def get_pose_coordinates(self) -> dict:
         return self.__data.get("pose_coordinates", None)
 
+    def get_pose_estimation(self) -> bool:
+        return self.__data.get("pose_estimation", True)
+
     def get_negative_prompt(self) -> str:
         return self.__data.get("negative_prompt", "")
 

internals/pipelines/controlnets.py

@@ -4,21 +4,17 @@ import cv2
 import numpy as np
 import torch
 from controlnet_aux import HEDdetector, LineartDetector, OpenposeDetector
-from diffusers import (
-    ControlNetModel,
-    DiffusionPipeline,
-    StableDiffusionControlNetPipeline,
-    UniPCMultistepScheduler,
-)
+from diffusers import (ControlNetModel, DiffusionPipeline,
+                       StableDiffusionControlNetPipeline,
+                       UniPCMultistepScheduler)
 from PIL import Image
 from torch.nn import Linear
 from tqdm import gui
 
 from internals.data.result import Result
 from internals.pipelines.commons import AbstractPipeline
-from internals.pipelines.tileUpscalePipeline import (
-    StableDiffusionControlNetImg2ImgPipeline,
-)
+from internals.pipelines.tileUpscalePipeline import \
+    StableDiffusionControlNetImg2ImgPipeline
 from internals.util.cache import clear_cuda_and_gc
 from internals.util.commons import download_image
 
@@ -291,7 +287,7 @@ class ControlNet(AbstractPipeline):
     def detect_pose(self, imageUrl: str) -> Image.Image:
         detector = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
         image = download_image(imageUrl)
-        image = detector.__call__(image, hand_and_face=True)
+        image = detector.__call__(image)
         return image
 
     def __scribble_condition_image(self, image: Image.Image) -> Image.Image:

internals/pipelines/pose_detector.py

@@ -79,7 +79,7 @@ class PoseDetector:
         return {"candidate": new_points, "subset": data["subset"]}
 
     def create_pose(self, data: dict, width: int, height: int) -> Image.Image:
-        image = Image.new("RGB", (width, height), "black")
+        image = Image.new("RGBA", (width, height), "black")
         draw = ImageDraw.Draw(image)
 
         points: list = data["candidate"]
@@ -164,24 +164,43 @@ class PoseDetector:
 
     __kim = [0, 17, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3]
     __pose_logical_map = [
-        [1, 2, "#000099"],
-        [1, 16, "#330099"],
-        [1, 15, "#660099"],
-        [16, 18, "#990099"],
-        [15, 17, "#990066"],
-        [2, 3, "#990001"],
-        [2, 6, "#993301"],
-        [3, 4, "#996502"],
-        [4, 5, "#999900"],
-        [6, 7, "#669900"],
-        [7, 8, "#349900"],
-        [2, 9, "#009900"],
-        [2, 12, "#009999"],
-        [9, 10, "#009966"],
-        [10, 11, "#009966"],
-        [12, 13, "#006699"],
-        [13, 14, "#013399"],
+        [2, 3, (255, 0, 0, 153)],
+        [3, 4, (255, 85, 0, 153)],
+        [4, 5, (255, 170, 0, 153)],
+        [2, 6, (255, 255, 0, 153)],
+        [6, 7, (170, 255, 0, 153)],
+        [7, 8, (85, 255, 0, 153)],
+        [2, 9, (0, 255, 0, 153)],
+        [9, 10, (0, 255, 85, 153)],
+        [10, 11, (0, 255, 170, 153)],
+        [2, 12, (0, 255, 255, 153)],
+        [12, 13, (0, 170, 255, 153)],
+        [13, 14, (0, 85, 255, 153)],
+        [2, 1, (0, 0, 255, 153)],
+        [1, 15, (85, 0, 255, 153)],
+        [15, 17, (170, 0, 255, 153)],
+        [1, 16, (255, 0, 255, 153)],
+        [16, 18, (255, 0, 170, 153)],
     ]
+    # __pose_logical_map = [
+    #     [1, 2, "#000099"],
+    #     [1, 16, "#330099"],
+    #     [1, 15, "#660099"],
+    #     [16, 18, "#990099"],
+    #     [15, 17, "#990066"],
+    #     [2, 3, "#990001"],
+    #     [2, 6, "#993301"],
+    #     [3, 4, "#996502"],
+    #     [4, 5, "#999900"],
+    #     [6, 7, "#669900"],
+    #     [7, 8, "#349900"],
+    #     [2, 9, "#009900"],
+    #     [2, 12, "#009999"],
+    #     [9, 10, "#009966"],
+    #     [10, 11, "#009966"],
+    #     [12, 13, "#006699"],
+    #     [13, 14, "#013399"],
+    # ]
     __points_color = [
         "#ff0000",
         "#ff5600",