Upload folder using huggingface_hub

inference.py

@@ -20,13 +20,9 @@ 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
 from internals.util.commons import download_image, upload_image, upload_images
-from internals.util.config import (
-    get_model_dir,
-    num_return_sequences,
-    set_configs_from_task,
-    set_model_dir,
-    set_root_dir,
-)
+from internals.util.config import (get_model_dir, num_return_sequences,
+                                   set_configs_from_task, set_model_dir,
+                                   set_root_dir)
 from internals.util.failure_hander import FailureHandler
 from internals.util.lora_style import LoraStyle
 from internals.util.slack import Slack
@@ -468,7 +464,7 @@ def replace_bg(task: Task):
         width=task.get_width(),
         height=task.get_height(),
         steps=task.get_steps(),
-        resize_dimension=task.get_resize_dimension(),
+        extend_object=task.rbg_extend_object(),
         product_scale_width=task.get_image_scale(),
         apply_high_res=task.get_high_res_fix(),
         conditioning_scale=task.rbg_controlnet_conditioning_scale(),
@@ -477,6 +473,7 @@ def replace_bg(task: Task):
     generated_image_urls = upload_images(images, "_replace_bg", task.get_taskId())
 
     lora_patcher.cleanup()
+    clear_cuda()
 
     return {
         "modified_prompts": prompt,
@@ -486,6 +483,8 @@ def replace_bg(task: Task):
 
 
 def load_model_by_task(task: Task):
+    high_res.load()
+
     if (
         task.get_type()
         in [
@@ -516,8 +515,6 @@ def load_model_by_task(task: Task):
         elif task.get_type() == TaskType.POSE:
             controlnet.load_pose()
 
-        high_res.load()
-
         safety_checker.apply(controlnet)
 
 

inference2.py

@@ -13,19 +13,17 @@ from internals.pipelines.img_to_text import Image2Text
 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_model_dir,
-    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_model_dir, set_root_dir)
 from internals.util.failure_hander import FailureHandler
 from internals.util.lora_style import LoraStyle
 from internals.util.slack import Slack
@@ -171,7 +169,7 @@ def replace_bg(task: Task):
         width=task.get_width(),
         height=task.get_height(),
         steps=task.get_steps(),
-        resize_dimension=task.get_resize_dimension(),
+        extend_object=task.rbg_extend_object(),
         product_scale_width=task.get_image_scale(),
         conditioning_scale=task.rbg_controlnet_conditioning_scale(),
     )
@@ -232,7 +230,9 @@ def model_fn(model_dir):
     inpainter.load()
     high_res.load()
 
-    replace_background.load(upscaler, remove_background_v2)
+    replace_background.load(
+        upscaler=upscaler, remove_background=remove_background_v2, high_res=high_res
+    )
 
     print("Logs: model loaded ....")
     return

internals/data/task.py

@@ -141,6 +141,9 @@ class Task:
     def rbg_controlnet_conditioning_scale(self) -> float:
         return self.__data.get("rbg_conditioning_scale", 0.5)
 
+    def rbg_extend_object(self) -> bool:
+        return self.__data.get("rbg_extend_object", False)
+
     def get_nsfw_threshold(self) -> float:
         return self.__data.get("nsfw_threshold", 0.03)
 

internals/pipelines/replace_background.py

@@ -2,12 +2,9 @@ from io import BytesIO
 from typing import List, Optional, Union
 
 import torch
-from diffusers import (
-    ControlNetModel,
-    StableDiffusionControlNetInpaintPipeline,
-    StableDiffusionInpaintPipeline,
-    UniPCMultistepScheduler,
-)
+from diffusers import (ControlNetModel,
+                       StableDiffusionControlNetInpaintPipeline,
+                       StableDiffusionInpaintPipeline, UniPCMultistepScheduler)
 from PIL import Image, ImageFilter, ImageOps
 
 import internals.util.image as ImageUtil
@@ -46,7 +43,7 @@ class ReplaceBackground(AbstractPipeline):
             pipe.controlnet = controlnet_model
         else:
             pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
-                get_model_dir(),
+                "runwayml/stable-diffusion-inpainting",
                 controlnet=controlnet_model,
                 torch_dtype=torch.float16,
                 cache_dir=get_hf_cache_dir(),
@@ -81,12 +78,13 @@ class ReplaceBackground(AbstractPipeline):
         product_scale_width: float,
         prompt: List[str],
         negative_prompt: List[str],
-        resize_dimension: int,
+        extend_object: bool,
         conditioning_scale: float,
         seed: int,
         steps: int,
         apply_high_res: bool = False,
     ):
+        # image = Image.open("original.png")
         if type(image) is str:
             image = download_image(image)
 
@@ -94,8 +92,8 @@ class ReplaceBackground(AbstractPipeline):
         torch.cuda.manual_seed(seed)
 
         image = image.convert("RGB")
-        if max(image.size) > 1536:
-            image = ImageUtil.resize_image(image, dimension=1536)
+        if max(image.size) > 1024:
+            image = ImageUtil.resize_image(image, dimension=1024)
         image = self.remove_background.remove(image)
 
         width = int(width)
@@ -106,11 +104,15 @@ class ReplaceBackground(AbstractPipeline):
 
         print(width, height, n_width, n_height)
 
-        image = ImageUtil.padd_image(image, n_width, n_height)
+        if extend_object:
+            condition_image = ControlNet.linearart_condition_image(image).resize(
+                (n_width, n_height)
+            )
+            condition_image = ImageUtil.padd_image(condition_image, width, height)
+            condition_image = condition_image.convert("RGB")
 
-        f_image = Image.new("RGBA", (width, height), (0, 0, 0, 0))
-        f_image.paste(image, ((width - n_width) // 2, (height - n_height) // 2))
-        image = f_image
+        image = image.resize((n_width, n_height))
+        image = ImageUtil.padd_image(image, width, height)
 
         mask = image.copy()
         pixdata = mask.load()
@@ -124,13 +126,13 @@ class ReplaceBackground(AbstractPipeline):
                 else:
                     pixdata[x, y] = (0, 0, 0, 255)
 
+        if not extend_object:
+            condition_image = ControlNet.linearart_condition_image(image)
         mask = mask.convert("RGB")
 
-        condition_image = ControlNet.linearart_condition_image(image)
-
         if apply_high_res and hasattr(self, "high_res"):
-            (w, h) = self.high_res.get_intermediate_dimension(width, height)
-            images = self.pipe.__call__(
+            w, h = HighRes.get_intermediate_dimension(width, height)
+            result = self.pipe.__call__(
                 prompt=prompt,
                 negative_prompt=negative_prompt,
                 image=image,
@@ -142,15 +144,17 @@ class ReplaceBackground(AbstractPipeline):
                 num_inference_steps=steps,
                 height=w,
                 width=h,
-            ).images
-            result = self.high_res.apply(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                images=images,
-                width=width,
-                height=width,
-                steps=steps,
             )
+            for i, _ in enumerate(result.images):
+                out_bytes = self.upscaler.upscale(
+                    image=result.images[i],
+                    width=w,
+                    height=h,
+                    face_enhance=False,
+                    resize_dimension=max(width, height),
+                )
+                result.images[i] = Image.open(BytesIO(out_bytes)).convert("RGB")
+            result = Result.from_result(result)
         else:
             result = self.pipe.__call__(
                 prompt=prompt,

internals/util/image.py

@@ -50,7 +50,7 @@ def padd_image(image: Image.Image, to_width: int, to_height: int) -> Image.Image
     # resize Image
     if iw > ih:
         image = image.resize((value, int(value * ih / iw)))
-    else:
+    elif ih > iw:
         image = image.resize((int(value * iw / ih), value))
 
     # padd Image