Upload folder using huggingface_hub

inference.py

@@ -1,7 +1,7 @@
 import os
+import traceback
 from typing import List, Optional
 
-import traceback
 import torch
 
 import internals.util.prompt as prompt_util
@@ -305,9 +305,14 @@ def pose(task: Task, s3_outkey: str = "_pose", poses: Optional[list] = None):
     else:
         poses = [controlnet.detect_pose(task.get_imageUrl())] * num_return_sequences
 
+    depth = download_image(task.get_auxilary_imageUrl()).resize(
+        (task.get_width(), task.get_height())
+    )
+    depth = ControlNet.depth_image(depth)
+
     kwargs = {
         "prompt": prompt,
-        "image": poses,
+        "image": [depth, poses[0]],
         "seed": task.get_seed(),
         "num_inference_steps": task.get_steps(),
         "negative_prompt": [task.get_negative_prompt()] * num_return_sequences,
@@ -331,6 +336,7 @@ def pose(task: Task, s3_outkey: str = "_pose", poses: Optional[list] = None):
         images, _ = high_res.apply(**kwargs)
 
     upload_image(poses[0], "crecoAI/{}_pose.png".format(task.get_taskId()))
+    upload_image(depth, "crecoAI/{}_depth.png".format(task.get_taskId()))
 
     generated_image_urls = upload_images(images, s3_outkey, task.get_taskId())
 

internals/pipelines/controlnets.py

@@ -2,20 +2,20 @@ from typing import List, Literal, Union
 
 import cv2
 import numpy as np
-from pydash import has
 import torch
 from controlnet_aux import HEDdetector, LineartDetector, OpenposeDetector
 from diffusers import (
     ControlNetModel,
     DiffusionPipeline,
     StableDiffusionControlNetPipeline,
-    UniPCMultistepScheduler,
     StableDiffusionXLControlNetPipeline,
+    UniPCMultistepScheduler,
 )
 from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import (
     MultiControlNetModel,
 )
 from PIL import Image
+from pydash import has
 from torch.nn import Linear
 from tqdm import gui
 from transformers import pipeline
@@ -32,11 +32,10 @@ from internals.util.commons import download_image
 from internals.util.config import (
     get_hf_cache_dir,
     get_hf_token,
-    get_model_dir,
     get_is_sdxl,
+    get_model_dir,
 )
 
-
 CONTROLNET_TYPES = Literal["pose", "canny", "scribble", "linearart", "tile_upscaler"]
 
 
@@ -60,11 +59,25 @@ class ControlNet(AbstractPipeline):
             task_name = model  # pyright: ignore
             model = config[task_name]
 
-        controlnet = ControlNetModel.from_pretrained(
-            model,
-            torch_dtype=torch.float16,
-            cache_dir=get_hf_cache_dir(),
-        ).to("cuda")
+        # Multi controlnet
+        if "," in model:
+            model_names = [m.strip() for m in model.split(",")]
+            controlnets = []
+            for name in model_names:
+                cn = ControlNetModel.from_pretrained(
+                    name,
+                    torch_dtype=torch.float16,
+                    cache_dir=get_hf_cache_dir(),
+                ).to("cuda")
+                controlnets.append(cn)
+            controlnet = MultiControlNetModel(controlnets).to("cuda")
+        # Single controlnet
+        else:
+            controlnet = ControlNetModel.from_pretrained(
+                model,
+                torch_dtype=torch.float16,
+                cache_dir=get_hf_cache_dir(),
+            ).to("cuda")
         self.__current_task_name = task_name
         self.controlnet = controlnet
 
@@ -202,15 +215,17 @@ class ControlNet(AbstractPipeline):
 
         kwargs = {
             "prompt": prompt[0],
-            "image": [image[0]],
+            "image": image,
             "num_images_per_prompt": 4,
             "num_inference_steps": num_inference_steps,
             "negative_prompt": negative_prompt[0],
             "guidance_scale": guidance_scale,
+            "control_guidance_end": [0.5, 1.0],
             "height": height,
             "width": width,
             **kwargs,
         }
+        print(kwargs)
         result = self.pipe2.__call__(**kwargs)
         return Result.from_result(result)
 
@@ -386,7 +401,7 @@ class ControlNet(AbstractPipeline):
         return img
 
     __model_normal = {
-        "pose": "lllyasviel/control_v11p_sd15_openpose",
+        "pose": "lllyasviel/control_v11f1p_sd15_depth, lllyasviel/control_v11p_sd15_openpose",
         "canny": "lllyasviel/control_v11p_sd15_canny",
         "linearart": "lllyasviel/control_v11p_sd15_lineart",
         "scribble": "lllyasviel/control_v11p_sd15_scribble",

internals/pipelines/replace_background.py

@@ -102,18 +102,19 @@ class ReplaceBackground(AbstractPipeline):
         torch.manual_seed(seed)
         torch.cuda.manual_seed(seed)
 
+        image = image.convert("RGB")
+        if max(image.size) > 1024:
+            image = ImageUtil.resize_image(image, dimension=1024)
+        image = self.remove_background.remove(image, model_type=model_type)
+
         width = int(width)
         height = int(height)
 
         resolution = max(width, height)
 
-        image = image.convert("RGB")
-
         image = ImageUtil.resize_image(image, resolution)
         image = ImageUtil.padd_image(image, width, height)
 
-        image = self.remove_background.remove(image, model_type=model_type)
-
         mask = image.copy()
         pixdata = mask.load()
 

internals/util/slack.py

@@ -31,6 +31,8 @@ class Slack:
         for key, value in raw.items():
             if value:
                 if type(value) == list:
+                    if type(value[0]) == float or type(value[0]) == int:
+                        value = str(value)
                     message += f"*{key}*: {', '.join(value)}\n"
                 else:
                     message += f"*{key}*: {value}\n"