Upload folder using huggingface_hub

inference.py

@@ -1,10 +1,11 @@
 import os
 from typing import List, Optional
 
+import traceback
 import torch
 
 import internals.util.prompt as prompt_util
-from internals.data.dataAccessor import update_db
+from internals.data.dataAccessor import update_db, update_db_source_failed
 from internals.data.task import Task, TaskType
 from internals.pipelines.commons import Img2Img, Text2Img
 from internals.pipelines.controlnets import ControlNet
@@ -18,11 +19,15 @@ from internals.pipelines.replace_background import ReplaceBackground
 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
+from internals.util.cache import auto_clear_cuda_and_gc, clear_cuda, clear_cuda_and_gc
 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_config,
-                                   set_root_dir)
+from internals.util.config import (
+    get_model_dir,
+    num_return_sequences,
+    set_configs_from_task,
+    set_model_config,
+    set_root_dir,
+)
 from internals.util.failure_hander import FailureHandler
 from internals.util.lora_style import LoraStyle
 from internals.util.model_loader import load_model_from_config
@@ -80,7 +85,7 @@ def canny(task: Task):
 
     width, height = get_intermediate_dimension(task)
 
-    controlnet.load_canny()
+    controlnet.load_model("canny")
 
     # pipe2 is used for canny and pose
     lora_patcher = lora_style.get_patcher(
@@ -88,7 +93,7 @@ def canny(task: Task):
     )
     lora_patcher.patch()
 
-    images, has_nsfw = controlnet.process_canny(
+    images, has_nsfw = controlnet.process(
         prompt=prompt,
         imageUrl=task.get_imageUrl(),
         seed=task.get_seed(),
@@ -132,12 +137,12 @@ def tile_upscale(task: Task):
 
     prompt = get_patched_prompt_tile_upscale(task)
 
-    controlnet.load_tile_upscaler()
+    controlnet.load_model("tile_upscaler")
 
     lora_patcher = lora_style.get_patcher(controlnet.pipe, task.get_style())
     lora_patcher.patch()
 
-    images, has_nsfw = controlnet.process_tile_upscaler(
+    images, has_nsfw = controlnet.process(
         imageUrl=task.get_imageUrl(),
         seed=task.get_seed(),
         steps=task.get_steps(),
@@ -169,14 +174,14 @@ def scribble(task: Task):
 
     width, height = get_intermediate_dimension(task)
 
-    controlnet.load_scribble()
+    controlnet.load_model("scribble")
 
     lora_patcher = lora_style.get_patcher(
         [controlnet.pipe2, high_res.pipe], task.get_style()
     )
     lora_patcher.patch()
 
-    images, has_nsfw = controlnet.process_scribble(
+    images, has_nsfw = controlnet.process(
         imageUrl=task.get_imageUrl(),
         seed=task.get_seed(),
         steps=task.get_steps(),
@@ -215,14 +220,14 @@ def linearart(task: Task):
 
     width, height = get_intermediate_dimension(task)
 
-    controlnet.load_linearart()
+    controlnet.load_model("linearart")
 
     lora_patcher = lora_style.get_patcher(
         [controlnet.pipe2, high_res.pipe], task.get_style()
     )
     lora_patcher.patch()
 
-    images, has_nsfw = controlnet.process_linearart(
+    images, has_nsfw = controlnet.process(
         imageUrl=task.get_imageUrl(),
         seed=task.get_seed(),
         steps=task.get_steps(),
@@ -261,7 +266,7 @@ def pose(task: Task, s3_outkey: str = "_pose", poses: Optional[list] = None):
 
     width, height = get_intermediate_dimension(task)
 
-    controlnet.load_pose()
+    controlnet.load_model("pose")
 
     # pipe2 is used for canny and pose
     lora_patcher = lora_style.get_patcher(
@@ -291,7 +296,7 @@ def pose(task: Task, s3_outkey: str = "_pose", poses: Optional[list] = None):
     )
     condition_image = ControlNet.linearart_condition_image(src_image)
 
-    images, has_nsfw = controlnet.process_pose(
+    images, has_nsfw = controlnet.process(
         prompt=prompt,
         image=poses,
         condition_image=[condition_image] * num_return_sequences,
@@ -440,7 +445,7 @@ def inpaint(task: Task):
 
     generated_image_urls = upload_images(images, "_inpaint", task.get_taskId())
 
-    clear_cuda()
+    clear_cuda_and_gc()
 
     return {"modified_prompts": prompt, "generated_image_urls": generated_image_urls}
 
@@ -469,12 +474,13 @@ def replace_bg(task: Task):
         product_scale_width=task.get_image_scale(),
         apply_high_res=task.get_high_res_fix(),
         conditioning_scale=task.rbg_controlnet_conditioning_scale(),
+        model_type=task.get_modelType(),
     )
 
     generated_image_urls = upload_images(images, "_replace_bg", task.get_taskId())
 
     lora_patcher.cleanup()
-    clear_cuda()
+    clear_cuda_and_gc()
 
     return {
         "modified_prompts": prompt,
@@ -484,38 +490,33 @@ def replace_bg(task: Task):
 
 
 def load_model_by_task(task: Task):
-    high_res.load()
-
-    if (
-        task.get_type()
-        in [
-            TaskType.TEXT_TO_IMAGE,
-            TaskType.IMAGE_TO_IMAGE,
-            TaskType.INPAINT,
-        ]
-        and not text2img_pipe.is_loaded()
-    ):
+    if not text2img_pipe.is_loaded():
         text2img_pipe.load(get_model_dir())
         img2img_pipe.create(text2img_pipe)
-        inpainter.load()
         high_res.load(img2img_pipe)
 
+        inpainter.init(text2img_pipe)
+        controlnet.init(text2img_pipe)
+
         safety_checker.apply(text2img_pipe)
         safety_checker.apply(img2img_pipe)
+
+    if task.get_type() == TaskType.INPAINT:
+        inpainter.load()
         safety_checker.apply(inpainter)
     elif task.get_type() == TaskType.REPLACE_BG:
         replace_background.load(inpainter=inpainter, high_res=high_res)
     else:
         if task.get_type() == TaskType.TILE_UPSCALE:
-            controlnet.load_tile_upscaler()
+            controlnet.load_model("tile_upscaler")
         elif task.get_type() == TaskType.CANNY:
-            controlnet.load_canny()
+            controlnet.load_model("canny")
         elif task.get_type() == TaskType.SCRIBBLE:
-            controlnet.load_scribble()
+            controlnet.load_model("scribble")
         elif task.get_type() == TaskType.LINEARART:
-            controlnet.load_linearart()
+            controlnet.load_model("linearart")
         elif task.get_type() == TaskType.POSE:
-            controlnet.load_pose()
+            controlnet.load_model("pose")
 
         safety_checker.apply(controlnet)
 
@@ -589,7 +590,8 @@ def predict_fn(data, pipe):
         else:
             raise Exception("Invalid task type")
     except Exception as e:
-        print(f"Error: {e}")
         slack.error_alert(task, e)
         controlnet.cleanup()
+        traceback.print_exc()
+        update_db_source_failed(task.get_sourceId(), task.get_userId())
         return None

inference2.py

@@ -13,17 +13,19 @@ 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_config, 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_config,
+    set_root_dir,
+)
 from internals.util.failure_hander import FailureHandler
 from internals.util.lora_style import LoraStyle
 from internals.util.model_loader import load_model_from_config
@@ -65,7 +67,7 @@ def tile_upscale(task: Task):
 
     prompt = get_patched_prompt_tile_upscale(task)
 
-    controlnet.load_tile_upscaler()
+    controlnet.load_model("tile_upscaler")
 
     lora_patcher = lora_style.get_patcher(controlnet.pipe, task.get_style())
     lora_patcher.patch()
@@ -98,7 +100,9 @@ def tile_upscale(task: Task):
 @slack.auto_send_alert
 def remove_bg(task: Task):
     # remove_background = RemoveBackground()
-    output_image = remove_background_v2.remove(task.get_imageUrl())
+    output_image = remove_background_v2.remove(
+        task.get_imageUrl(), model_type=task.get_modelType()
+    )
 
     output_key = "crecoAI/{}_rmbg.png".format(task.get_taskId())
     upload_image(output_image, output_key)
@@ -173,6 +177,7 @@ def replace_bg(task: Task):
         extend_object=task.rbg_extend_object(),
         product_scale_width=task.get_image_scale(),
         conditioning_scale=task.rbg_controlnet_conditioning_scale(),
+        model_type=task.get_modelType(),
     )
 
     generated_image_urls = upload_images(images, "_replace_bg", task.get_taskId())
@@ -231,6 +236,7 @@ def model_fn(model_dir):
     upscaler.load()
     inpainter.load()
     high_res.load()
+    controlnet.init(high_res)
 
     replace_background.load(
         upscaler=upscaler, remove_background=remove_background_v2, high_res=high_res
@@ -242,7 +248,7 @@ def model_fn(model_dir):
 
 def load_model_by_task(task: Task):
     if task.get_type() == TaskType.TILE_UPSCALE:
-        controlnet.load_tile_upscaler()
+        controlnet.load_model("tile_upscaler")
 
     safety_checker.apply(controlnet)
 

internals/data/dataAccessor.py

@@ -1,6 +1,7 @@
 import traceback
 from typing import Dict, List, Optional
 
+from requests.adapters import Retry, HTTPAdapter
 import requests
 from pydash import includes
 
@@ -9,6 +10,14 @@ from internals.util.config import api_endpoint, api_headers
 from internals.util.slack import Slack
 
 
+class RetryRequest:
+    def __new__(cls):
+        obj = Retry(total=5, backoff_factor=2, status_forcelist=[500, 502, 503, 504])
+        session = requests.Session()
+        session.mount("https://", HTTPAdapter(max_retries=obj))
+        return session
+
+
 def updateSource(sourceId, userId, state):
     print("update source is called")
     url = api_endpoint() + f"/autodraft-crecoai/source/{sourceId}"
@@ -21,7 +30,8 @@ def updateSource(sourceId, userId, state):
     data = {"state": state}
 
     try:
-        response = requests.patch(url, headers=headers, json=data, timeout=10)
+        with RetryRequest() as session:
+            response = session.patch(url, headers=headers, json=data, timeout=10)
         print("update source response", response)
     except requests.exceptions.Timeout:
         print("Request timed out while updating source")
@@ -47,7 +57,8 @@ def saveGeneratedImages(sourceId, userId, has_nsfw: bool):
     data = {"state": "ACTIVE", "has_nsfw": has_nsfw}
 
     try:
-        requests.patch(url, headers=headers, json=data)
+        with RetryRequest() as session:
+            session.patch(url, headers=headers, json=data)
         # print("save generation response", response)
     except requests.exceptions.Timeout:
         print("Request timed out while saving image")
@@ -61,11 +72,12 @@ def getStyles() -> Optional[Dict]:
     url = api_endpoint() + "/autodraft-crecoai/style"
     print(url)
     try:
-        response = requests.get(
-            url,
-            timeout=10,
-            headers={"x-api-key": "kGyEMp)oHB(zf^E5>-{o]I%go", **api_headers()},
-        )
+        with RetryRequest() as session:
+            response = session.get(
+                url,
+                timeout=10,
+                headers={"x-api-key": "kGyEMp)oHB(zf^E5>-{o]I%go", **api_headers()},
+            )
         return response.json()
     except requests.exceptions.Timeout:
         print("Request timed out while fetching styles")
@@ -78,9 +90,10 @@ def getStyles() -> Optional[Dict]:
 def getCharacters(model_id: str) -> Optional[List]:
     url = api_endpoint() + "/autodraft-crecoai/model/{}".format(model_id)
     try:
-        response = requests.get(url, timeout=10, headers=api_headers())
-        response = response.json()
-        response = response["data"]["characters"]
+        with RetryRequest() as session:
+            response = session.get(url, timeout=10, headers=api_headers())
+            response = response.json()
+            response = response["data"]["characters"]
         return response
     except requests.exceptions.Timeout:
         print("Request timed out while fetching characters")
@@ -89,6 +102,10 @@ def getCharacters(model_id: str) -> Optional[List]:
     return None
 
 
+def update_db_source_failed(sourceId, userId):
+    updateSource(sourceId, userId, "FAILED")
+
+
 def update_db(func):
     def caller(*args, **kwargs):
         if type(args[0]) is not Task:

internals/pipelines/commons.py

@@ -2,12 +2,16 @@ from dataclasses import dataclass
 from typing import Any, Callable, Dict, List, Optional, Union
 
 import torch
-from diffusers import StableDiffusionImg2ImgPipeline
+from diffusers import (
+    StableDiffusionImg2ImgPipeline,
+    StableDiffusionXLPipeline,
+    StableDiffusionXLImg2ImgPipeline,
+)
 
 from internals.data.result import Result
 from internals.pipelines.twoStepPipeline import two_step_pipeline
 from internals.util.commons import disable_safety_checker, download_image
-from internals.util.config import get_hf_token, num_return_sequences
+from internals.util.config import get_hf_token, num_return_sequences, get_is_sdxl
 
 
 class AbstractPipeline:
@@ -27,9 +31,17 @@ class Text2Img(AbstractPipeline):
         prompt_right: List[str] = None
 
     def load(self, model_dir: str):
-        self.pipe = two_step_pipeline.from_pretrained(
-            model_dir, torch_dtype=torch.float16, use_auth_token=get_hf_token()
-        ).to("cuda")
+        if get_is_sdxl():
+            self.pipe = StableDiffusionXLPipeline.from_pretrained(
+                model_dir,
+                torch_dtype=torch.float16,
+                use_auth_token=get_hf_token(),
+                use_safetensors=True,
+            ).to("cuda")
+        else:
+            self.pipe = two_step_pipeline.from_pretrained(
+                model_dir, torch_dtype=torch.float16, use_auth_token=get_hf_token()
+            ).to("cuda")
         self.__patch()
 
     def is_loaded(self):
@@ -38,10 +50,16 @@ class Text2Img(AbstractPipeline):
         return False
 
     def create(self, pipeline: AbstractPipeline):
-        self.pipe = two_step_pipeline(**pipeline.pipe.components).to("cuda")
+        if get_is_sdxl():
+            self.pipe = StableDiffusionXLPipeline(**pipeline.pipe.components).to("cuda")
+        else:
+            self.pipe = two_step_pipeline(**pipeline.pipe.components).to("cuda")
         self.__patch()
 
     def __patch(self):
+        if get_is_sdxl():
+            self.pipe.enable_vae_tiling()
+            self.pipe.enable_vae_slicing()
         self.pipe.enable_xformers_memory_efficient_attention()
 
     @torch.inference_mode()
@@ -92,9 +110,19 @@ class Text2Img(AbstractPipeline):
             # two step pipeline
             modified_prompt = params.modified_prompt
 
-            result = self.pipe.two_step_pipeline(
-                prompt=prompt,
-                modified_prompts=modified_prompt,
+            if get_is_sdxl():
+                print("Warning: Two step pipeline is not supported on SDXL")
+                kwargs = {
+                    "prompt": modified_prompt,
+                }
+            else:
+                kwargs = {
+                    "prompt": prompt,
+                    "modified_prompts": modified_prompt,
+                    "iteration": iteration,
+                }
+
+            result = self.pipe.__call__(
                 height=height,
                 width=width,
                 num_inference_steps=num_inference_steps,
@@ -111,7 +139,7 @@ class Text2Img(AbstractPipeline):
                 callback=callback,
                 callback_steps=callback_steps,
                 cross_attention_kwargs=cross_attention_kwargs,
-                iteration=iteration,
+                **kwargs
             )
 
         return Result.from_result(result)
@@ -124,22 +152,38 @@ class Img2Img(AbstractPipeline):
         if self.__loaded:
             return
 
-        self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
-            model_dir, torch_dtype=torch.float16, use_auth_token=get_hf_token()
-        ).to("cuda")
+        if get_is_sdxl():
+            self.pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+                model_dir,
+                torch_dtype=torch.float16,
+                use_auth_token=get_hf_token(),
+                use_safetensors=True,
+            ).to("cuda")
+        else:
+            self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+                model_dir, torch_dtype=torch.float16, use_auth_token=get_hf_token()
+            ).to("cuda")
         self.__patch()
 
         self.__loaded = True
 
     def create(self, pipeline: AbstractPipeline):
-        self.pipe = StableDiffusionImg2ImgPipeline(**pipeline.pipe.components).to(
-            "cuda"
-        )
+        if get_is_sdxl():
+            self.pipe = StableDiffusionXLImg2ImgPipeline(**pipeline.pipe.components).to(
+                "cuda"
+            )
+        else:
+            self.pipe = StableDiffusionImg2ImgPipeline(**pipeline.pipe.components).to(
+                "cuda"
+            )
         self.__patch()
 
         self.__loaded = True
 
     def __patch(self):
+        if get_is_sdxl():
+            self.pipe.enable_vae_tiling()
+            self.pipe.enable_vae_slicing()
         self.pipe.enable_xformers_memory_efficient_attention()
 
     @torch.inference_mode()

internals/pipelines/controlnets.py

@@ -1,14 +1,20 @@
-from typing import List, Union
+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)
-from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import \
-    MultiControlNetModel
+from diffusers import (
+    ControlNetModel,
+    DiffusionPipeline,
+    StableDiffusionControlNetPipeline,
+    UniPCMultistepScheduler,
+    StableDiffusionXLControlNetPipeline,
+)
+from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import (
+    MultiControlNetModel,
+)
 from PIL import Image
 from torch.nn import Linear
 from tqdm import gui
@@ -18,156 +24,127 @@ import internals.util.image as ImageUtil
 from external.midas import apply_midas
 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
-from internals.util.config import get_hf_cache_dir, get_hf_token, get_model_dir
+from internals.util.config import (
+    get_hf_cache_dir,
+    get_hf_token,
+    get_model_dir,
+    get_is_sdxl,
+)
+
+
+CONTROLNET_TYPES = Literal["pose", "canny", "scribble", "linearart", "tile_upscaler"]
 
 
 class ControlNet(AbstractPipeline):
     __current_task_name = ""
     __loaded = False
 
-    def load(self):
-        "Should not be called externally"
-        if self.__loaded:
-            return
+    __pipeline: AbstractPipeline
 
-        if not hasattr(self, "controlnet"):
-            self.load_pose()
+    def init(self, pipeline: AbstractPipeline):
+        self.__pipeline = pipeline
 
-        # controlnet pipeline for tile upscaler
-        pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-            get_model_dir(),
-            controlnet=self.controlnet,
-            torch_dtype=torch.float16,
-            use_auth_token=get_hf_token(),
-            cache_dir=get_hf_cache_dir(),
-        ).to("cuda")
-        # pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
-        pipe.enable_model_cpu_offload()
-        pipe.enable_xformers_memory_efficient_attention()
-        self.pipe = pipe
-
-        # controlnet pipeline for canny and pose
-        pipe2 = StableDiffusionControlNetPipeline(**pipe.components).to("cuda")
-        pipe2.scheduler = UniPCMultistepScheduler.from_config(pipe2.scheduler.config)
-        pipe2.enable_xformers_memory_efficient_attention()
-        self.pipe2 = pipe2
-
-        self.__loaded = True
-
-    def load_canny(self):
-        if self.__current_task_name == "canny":
+    def load_model(self, task_name: CONTROLNET_TYPES):
+        config = self.__model_sdxl if get_is_sdxl() else self.__model_normal
+        if self.__current_task_name == task_name:
             return
-        canny = ControlNetModel.from_pretrained(
-            "lllyasviel/control_v11p_sd15_canny",
+        model = config[task_name]
+        if not model:
+            raise Exception(f"ControlNet is not supported for {task_name}")
+        while model in list(config.keys()):
+            task_name = config[model]  # pyright: ignore
+            model = config[task_name]
+
+        controlnet = ControlNetModel.from_pretrained(
+            model,
             torch_dtype=torch.float16,
             cache_dir=get_hf_cache_dir(),
         ).to("cuda")
-        self.__current_task_name = "canny"
-        self.controlnet = canny
+        self.__current_task_name = task_name
+        self.controlnet = controlnet
 
-        self.load()
+        self.__load()
 
         if hasattr(self, "pipe"):
-            self.pipe.controlnet = canny
+            self.pipe.controlnet = controlnet
         if hasattr(self, "pipe2"):
-            self.pipe2.controlnet = canny
+            self.pipe2.controlnet = controlnet
         clear_cuda_and_gc()
 
-    def load_pose(self):
-        if self.__current_task_name == "pose":
+    def __load(self):
+        "Should not be called externally"
+        if self.__loaded:
             return
-        pose = ControlNetModel.from_pretrained(
-            "lllyasviel/control_v11p_sd15_openpose",
-            torch_dtype=torch.float16,
-            cache_dir=get_hf_cache_dir(),
-        ).to("cuda")
-        # lineart = ControlNetModel.from_pretrained(
-        #     "ControlNet-1-1-preview/control_v11p_sd15_lineart",
-        #     torch_dtype=torch.float16,
-        #     cache_dir=get_hf_cache_dir(),
-        # ).to("cuda")
-        self.__current_task_name = "pose"
-        self.controlnet = MultiControlNetModel([pose]).to("cuda")
-
-        self.load()
 
-        if hasattr(self, "pipe"):
-            self.pipe.controlnet = self.controlnet
-        if hasattr(self, "pipe2"):
-            self.pipe2.controlnet = self.controlnet
-        clear_cuda_and_gc()
-
-    def load_tile_upscaler(self):
-        if self.__current_task_name == "tile_upscaler":
-            return
-        tile_upscaler = ControlNetModel.from_pretrained(
-            "lllyasviel/control_v11f1e_sd15_tile",
-            torch_dtype=torch.float16,
-            cache_dir=get_hf_cache_dir(),
-        ).to("cuda")
-        self.__current_task_name = "tile_upscaler"
-        self.controlnet = tile_upscaler
+        if not hasattr(self, "controlnet"):
+            self.load_model("pose")
 
-        self.load()
+        # controlnet pipeline for tile upscaler
+        if get_is_sdxl():
+            print("Warning: Tile upscale is not supported on SDXL")
+
+            if self.__pipeline:
+                pipe = StableDiffusionXLControlNetPipeline(
+                    controlnet=self.controlnet, **self.__pipeline.pipe.components
+                ).to("cuda")
+            else:
+                pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
+                    get_model_dir(),
+                    controlnet=self.controlnet,
+                    torch_dtype=torch.float16,
+                    use_auth_token=get_hf_token(),
+                    cache_dir=get_hf_cache_dir(),
+                    use_safetensors=True,
+                ).to("cuda")
+            pipe.enable_vae_tiling()
+            pipe.enable_vae_slicing()
+            pipe.enable_xformers_memory_efficient_attention()
+            self.pipe2 = pipe
+        else:
+            if hasattr(self, "__pipeline"):
+                pipe = StableDiffusionControlNetImg2ImgPipeline(
+                    controlnet=self.controlnet, **self.__pipeline.pipe.components
+                ).to("cuda")
+            else:
+                pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+                    get_model_dir(),
+                    controlnet=self.controlnet,
+                    torch_dtype=torch.float16,
+                    use_auth_token=get_hf_token(),
+                    cache_dir=get_hf_cache_dir(),
+                ).to("cuda")
+            # pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+            pipe.enable_model_cpu_offload()
+            pipe.enable_xformers_memory_efficient_attention()
+            self.pipe = pipe
+
+            # controlnet pipeline for canny and pose
+            pipe2 = StableDiffusionControlNetPipeline(**pipe.components).to("cuda")
+            pipe2.scheduler = UniPCMultistepScheduler.from_config(
+                pipe2.scheduler.config
+            )
+            pipe2.enable_xformers_memory_efficient_attention()
+            self.pipe2 = pipe2
 
-        if hasattr(self, "pipe"):
-            self.pipe.controlnet = tile_upscaler
-        if hasattr(self, "pipe2"):
-            self.pipe2.controlnet = tile_upscaler
-        clear_cuda_and_gc()
+        self.__loaded = True
 
-    def load_scribble(self):
+    def process(self, **kwargs):
+        if self.__current_task_name == "pose":
+            return self.process_pose(**kwargs)
+        if self.__current_task_name == "canny":
+            return self.process_canny(**kwargs)
         if self.__current_task_name == "scribble":
-            return
-        scribble = ControlNetModel.from_pretrained(
-            "lllyasviel/control_v11p_sd15_scribble",
-            torch_dtype=torch.float16,
-            cache_dir=get_hf_cache_dir(),
-        ).to("cuda")
-        self.__current_task_name = "scribble"
-        self.controlnet = scribble
-
-        self.load()
-
-        if hasattr(self, "pipe"):
-            self.pipe.controlnet = scribble
-        if hasattr(self, "pipe2"):
-            self.pipe2.controlnet = scribble
-        clear_cuda_and_gc()
-
-    def load_linearart(self):
+            return self.process_scribble(**kwargs)
         if self.__current_task_name == "linearart":
-            return
-        linearart = ControlNetModel.from_pretrained(
-            "ControlNet-1-1-preview/control_v11p_sd15_lineart",
-            torch_dtype=torch.float16,
-            cache_dir=get_hf_cache_dir(),
-        ).to("cuda")
-        self.__current_task_name = "linearart"
-        self.controlnet = linearart
-
-        self.load()
-
-        if hasattr(self, "pipe"):
-            self.pipe.controlnet = linearart
-        if hasattr(self, "pipe2"):
-            self.pipe2.controlnet = linearart
-        clear_cuda_and_gc()
-
-    def cleanup(self):
-        if hasattr(self, "pipe"):
-            self.pipe.controlnet = None
-        if hasattr(self, "pipe2"):
-            self.pipe2.controlnet = None
-        self.controlnet = None
-        del self.controlnet
-        self.__current_task_name = ""
-
-        clear_cuda_and_gc()
+            return self.process_linearart(**kwargs)
+        if self.__current_task_name == "tile_upscaler":
+            return self.process_tile_upscaler(**kwargs)
+        raise Exception("ControlNet is not loaded with any model")
 
     @torch.inference_mode()
     def process_canny(
@@ -228,7 +205,6 @@ class ControlNet(AbstractPipeline):
             guidance_scale=guidance_scale,
             height=height,
             width=width,
-            controlnet_conditioning_scale=[1.0],
         )
         return Result.from_result(result)
 
@@ -333,6 +309,17 @@ class ControlNet(AbstractPipeline):
         )
         return Result.from_result(result)
 
+    def cleanup(self):
+        if hasattr(self, "pipe") and hasattr(self.pipe, "controlnet"):
+            del self.pipe.controlnet
+        if hasattr(self, "pipe2") and hasattr(self.pipe2, "controlnet"):
+            del self.pipe2.controlnet
+        if hasattr(self, "controlnet"):
+            del self.controlnet
+        self.__current_task_name = ""
+
+        clear_cuda_and_gc()
+
     def detect_pose(self, imageUrl: str) -> Image.Image:
         detector = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
         image = download_image(imageUrl)
@@ -381,3 +368,18 @@ class ControlNet(AbstractPipeline):
         W = int(round(W / 64.0)) * 64
         img = input_image.resize((W, H), resample=Image.LANCZOS)
         return img
+
+    __model_normal = {
+        "pose": "lllyasviel/control_v11p_sd15_openpose",
+        "canny": "lllyasviel/control_v11p_sd15_canny",
+        "linearart": "lllyasviel/control_v11p_sd15_lineart",
+        "scribble": "lllyasviel/control_v11p_sd15_scribble",
+        "tile_upscaler": "lllyasviel/control_v11f1e_sd15_tile",
+    }
+    __model_sdxl = {
+        "pose": "thibaud/controlnet-openpose-sdxl-1.0",
+        "canny": "diffusers/controlnet-canny-sdxl-1.0",
+        "linearart": "canny",
+        "scribble": "canny",
+        "tile_upscaler": None,
+    }

internals/pipelines/high_res.py

@@ -42,7 +42,7 @@ class HighRes(AbstractPipeline):
 
     @staticmethod
     def get_intermediate_dimension(target_width: int, target_height: int):
-        def_size = 512
+        def_size = 1024
 
         desired_pixel_count = def_size * def_size
         actual_pixel_count = target_width * target_height

internals/pipelines/inpainter.py

@@ -1,38 +1,74 @@
 from typing import List, Union
 
 import torch
-from diffusers import StableDiffusionInpaintPipeline
+from diffusers import StableDiffusionInpaintPipeline, StableDiffusionXLInpaintPipeline
 
 from internals.pipelines.commons import AbstractPipeline
 from internals.util.commons import disable_safety_checker, download_image
-from internals.util.config import (get_hf_cache_dir, get_hf_token,
-                                   get_inpaint_model_path)
+from internals.util.config import (
+    get_hf_cache_dir,
+    get_hf_token,
+    get_is_sdxl,
+    get_inpaint_model_path,
+    get_model_dir,
+)
 
 
 class InPainter(AbstractPipeline):
     __loaded = False
 
+    def init(self, pipeline: AbstractPipeline):
+        self.__base = pipeline
+
     def load(self):
         if self.__loaded:
             return
 
-        self.pipe = StableDiffusionInpaintPipeline.from_pretrained(
-            get_inpaint_model_path(),
-            torch_dtype=torch.float16,
-            cache_dir=get_hf_cache_dir(),
-            use_auth_token=get_hf_token(),
-        ).to("cuda")
+        if hasattr(self, "__base") and get_inpaint_model_path() == get_model_dir():
+            self.create(self.__base)
+            self.__loaded = True
+            return
+
+        if get_is_sdxl():
+            self.pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+                get_inpaint_model_path(),
+                torch_dtype=torch.float16,
+                cache_dir=get_hf_cache_dir(),
+                use_auth_token=get_hf_token(),
+            ).to("cuda")
+        else:
+            self.pipe = StableDiffusionInpaintPipeline.from_pretrained(
+                get_inpaint_model_path(),
+                torch_dtype=torch.float16,
+                cache_dir=get_hf_cache_dir(),
+                use_auth_token=get_hf_token(),
+            ).to("cuda")
 
         disable_safety_checker(self.pipe)
 
+        self.__patch()
+
         self.__loaded = True
 
     def create(self, pipeline: AbstractPipeline):
-        self.pipe = StableDiffusionInpaintPipeline(**pipeline.pipe.components).to(
-            "cuda"
-        )
+        if get_is_sdxl():
+            self.pipe = StableDiffusionXLInpaintPipeline(**pipeline.pipe.components).to(
+                "cuda"
+            )
+        else:
+            self.pipe = StableDiffusionInpaintPipeline(**pipeline.pipe.components).to(
+                "cuda"
+            )
         disable_safety_checker(self.pipe)
 
+        self.__patch()
+
+    def __patch(self):
+        if get_is_sdxl():
+            self.pipe.enable_vae_tiling()
+            self.pipe.enable_vae_slicing()
+        self.pipe.enable_xformers_memory_efficient_attention()
+
     @torch.inference_mode()
     def process(
         self,

internals/pipelines/remove_background.py

@@ -1,15 +1,20 @@
 import io
 from pathlib import Path
 from typing import Union
+import numpy as np
+import cv2
 
 import torch
 import torch.nn.functional as F
 from PIL import Image
 from rembg import remove
+from internals.data.task import ModelType
 
 import internals.util.image as ImageUtil
 from carvekit.api.high import HiInterface
 from internals.util.commons import download_image, read_url
+import onnxruntime as rt
+import huggingface_hub
 
 
 class RemoveBackground:
@@ -23,6 +28,11 @@ class RemoveBackground:
 
 class RemoveBackgroundV2:
     def __init__(self):
+        model_path = huggingface_hub.hf_hub_download("skytnt/anime-seg", "isnetis.onnx")
+        self.anime_rembg = rt.InferenceSession(
+            model_path, providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
+        )
+
         self.interface = HiInterface(
             object_type="object",  # Can be "object" or "hairs-like".
             batch_size_seg=5,
@@ -36,16 +46,51 @@ class RemoveBackgroundV2:
             fp16=False,
         )
 
-    def remove(self, image: Union[str, Image.Image]) -> Image.Image:
-        img_path = Path.home() / ".cache" / "rm_bg.png"
+    def remove(
+        self, image: Union[str, Image.Image], model_type: ModelType = ModelType.REAL
+    ) -> Image.Image:
         if type(image) is str:
             image = download_image(image)
 
-        w, h = image.size
-        if max(w, h) > 1536:
-            image = ImageUtil.resize_image(image, dimension=1024)
+        if model_type == ModelType.ANIME or model_type == ModelType.COMIC:
+            print("Using Anime Background remover")
+            _, img = self.__rmbg_fn(np.array(image))
+
+            return Image.fromarray(img)
+        else:
+            print("Using Real Background remover")
+            img_path = Path.home() / ".cache" / "rm_bg.png"
+
+            w, h = image.size
+            if max(w, h) > 1536:
+                image = ImageUtil.resize_image(image, dimension=1024)
+
+            image.save(img_path)
+            images_without_background = self.interface([img_path])
+            out = images_without_background[0]
+            return out
+
+    def __get_mask(self, img, s=1024):
+        img = (img / 255).astype(np.float32)
+        h, w = h0, w0 = img.shape[:-1]
+        h, w = (s, int(s * w / h)) if h > w else (int(s * h / w), s)
+        ph, pw = s - h, s - w
+        img_input = np.zeros([s, s, 3], dtype=np.float32)
+        img_input[ph // 2 : ph // 2 + h, pw // 2 : pw // 2 + w] = cv2.resize(
+            img, (w, h)
+        )
+        img_input = np.transpose(img_input, (2, 0, 1))
+        img_input = img_input[np.newaxis, :]
+        mask = self.anime_rembg.run(None, {"img": img_input})[0][0]
+        mask = np.transpose(mask, (1, 2, 0))
+        mask = mask[ph // 2 : ph // 2 + h, pw // 2 : pw // 2 + w]
+        mask = cv2.resize(mask, (w0, h0))[:, :, np.newaxis]
+        return mask
 
-        image.save(img_path)
-        images_without_background = self.interface([img_path])
-        out = images_without_background[0]
-        return out
+    def __rmbg_fn(self, img):
+        mask = self.__get_mask(img)
+        img = (mask * img + 255 * (1 - mask)).astype(np.uint8)
+        mask = (mask * 255).astype(np.uint8)
+        img = np.concatenate([img, mask], axis=2, dtype=np.uint8)
+        mask = mask.repeat(3, axis=2)
+        return mask, img

internals/pipelines/replace_background.py

@@ -3,10 +3,14 @@ from typing import List, Optional, Union
 
 import torch
 from cv2 import inpaint
-from diffusers import (ControlNetModel,
-                       StableDiffusionControlNetInpaintPipeline,
-                       StableDiffusionInpaintPipeline, UniPCMultistepScheduler)
+from diffusers import (
+    ControlNetModel,
+    StableDiffusionControlNetInpaintPipeline,
+    StableDiffusionInpaintPipeline,
+    UniPCMultistepScheduler,
+)
 from PIL import Image, ImageFilter, ImageOps
+from internals.data.task import ModelType
 
 import internals.util.image as ImageUtil
 from internals.data.result import Result
@@ -17,8 +21,12 @@ from internals.pipelines.inpainter import InPainter
 from internals.pipelines.remove_background import RemoveBackgroundV2
 from internals.pipelines.upscaler import Upscaler
 from internals.util.commons import download_image
-from internals.util.config import (get_hf_cache_dir, get_hf_token,
-                                   get_inpaint_model_path, get_model_dir)
+from internals.util.config import (
+    get_hf_cache_dir,
+    get_hf_token,
+    get_inpaint_model_path,
+    get_model_dir,
+)
 
 
 class ReplaceBackground(AbstractPipeline):
@@ -52,7 +60,8 @@ class ReplaceBackground(AbstractPipeline):
                 cache_dir=get_hf_cache_dir(),
                 use_auth_token=get_hf_token(),
             )
-        pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+        pipe.enable_xformers_memory_efficient_attention()
+        pipe.enable_vae_slicing()
         pipe.to("cuda")
 
         self.pipe = pipe
@@ -87,6 +96,7 @@ class ReplaceBackground(AbstractPipeline):
         seed: int,
         steps: int,
         apply_high_res: bool = False,
+        model_type: ModelType = ModelType.REAL,
     ):
         # image = Image.open("original.png")
         if type(image) is str:
@@ -98,7 +108,7 @@ class ReplaceBackground(AbstractPipeline):
         image = image.convert("RGB")
         if max(image.size) > 1024:
             image = ImageUtil.resize_image(image, dimension=1024)
-        image = self.remove_background.remove(image)
+        image = self.remove_background.remove(image, model_type=model_type)
 
         width = int(width)
         height = int(height)

internals/pipelines/twoStepPipeline.py

@@ -12,7 +12,7 @@ from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
 
 class two_step_pipeline(StableDiffusionPipeline):
     @torch.no_grad()
-    def two_step_pipeline(
+    def __call__(
         self,
         prompt: Union[str, List[str]] = None,
         modified_prompts: Union[str, List[str]] = None,

internals/util/cache.py

@@ -1,15 +1,25 @@
 import gc
-
+import os
+import psutil
 import torch
 
 
+def print_memory_usage():
+    process = psutil.Process(os.getpid())
+    print(f"Memory usage: {process.memory_info().rss / 1024 ** 2:2f} MB")
+
+
 def clear_cuda_and_gc():
-    clear_cuda()
+    print_memory_usage()
+    print("Clearing cuda and gc")
     clear_gc()
+    clear_cuda()
+    print_memory_usage()
 
 
 def clear_cuda():
-    torch.cuda.empty_cache()
+    with torch.no_grad():
+        torch.cuda.empty_cache()
 
 
 def clear_gc():

internals/util/commons.py

@@ -150,9 +150,9 @@ def upload_image(image: Union[Image.Image, BytesIO], out_path):
     return image_url
 
 
-def download_image(url) -> Image.Image:
+def download_image(url, mode="RGB") -> Image.Image:
     response = requests.get(url)
-    return Image.open(BytesIO(response.content)).convert("RGB")
+    return Image.open(BytesIO(response.content)).convert(mode)
 
 
 def download_file(url, out_path: Path):

internals/util/config.py

@@ -61,6 +61,11 @@ def get_inpaint_model_path():
     return model_config.base_inpaint_model_path  # pyright: ignore
 
 
+def get_is_sdxl():
+    global model_config
+    return model_config.is_sdxl  # pyright: ignore
+
+
 def get_root_dir():
     global root_dir
     return root_dir

internals/util/lora_style.py

@@ -10,6 +10,7 @@ from lora_diffusion import patch_pipe, tune_lora_scale
 from pydash import chain
 
 from internals.data.dataAccessor import getStyles
+from internals.util.config import get_is_sdxl
 from internals.util.commons import download_file
 
 
@@ -112,6 +113,10 @@ class LoraStyle:
     ) -> Union[LoraPatcher, LoraDiffuserPatcher, EmptyLoraPatcher]:
         "Returns a lora patcher for the given `key` and `pipe`. `pipe` can also be a list of pipes"
         pipe = [pipe] if not isinstance(pipe, list) else pipe
+        if get_is_sdxl():
+            print("Warning: Lora is not supported on SDXL")
+            return self.EmptyLoraPatcher(pipe)
+
         if key in self.__styles:
             style = self.__styles[key]
             if style["type"] == "diffuser":

internals/util/model_loader.py

@@ -14,6 +14,7 @@ from tqdm import tqdm
 class ModelConfig:
     base_model_path: str
     base_inpaint_model_path: str
+    is_sdxl: bool = False
 
 
 def load_model_from_config(path):
@@ -23,9 +24,11 @@ def load_model_from_config(path):
             config = json.loads(f.read())
             model_path = config.get("model_path", path)
             inpaint_model_path = config.get("inpaint_model_path", path)
+            is_sdxl = config.get("is_sdxl", False)
 
             m_config.base_model_path = model_path
             m_config.base_inpaint_model_path = inpaint_model_path
+            m_config.is_sdxl = is_sdxl
 
             #
             # if config.get("model_type") == "huggingface":

pyproject.toml

@@ -1,4 +1,4 @@
 [tool.pyright]
-venvPath = "/Users/devel/Documents/WebProjects/creco-inference"
+venvPath = "."
 venv = "env"
 exclude = ["env"]

requirements.txt

@@ -15,6 +15,7 @@ realesrgan==0.3.0
 compel==1.0.4
 scikit-image>=0.19.3
 six==1.16.0
+psutil
 tifffile==2021.8.30
 easydict==1.9.0
 albumentations
@@ -32,10 +33,13 @@ xformers==0.0.21
 scikit-image==0.19.3
 omegaconf==2.3.0
 webdataset==0.2.48
+invisible-watermark 
 https://comic-assets.s3.ap-south-1.amazonaws.com/packages/mmcv_full-1.7.0-cp39-cp39-linux_x86_64.whl
 python-dateutil==2.8.2
 PyYAML
 invisible-watermark
 torchvision==0.15.2
+onnx
+onnxruntime-gpu
 imgaug==0.4.0
 tqdm==4.64.1