Upload folder using huggingface_hub

.gitignore

@@ -1,6 +1,7 @@
 *.pyc
 .DS_Store
 .ipynb_checkpoints                            █
+.vscode
 env
 test.py
 *.jpeg

inference.py

@@ -4,26 +4,38 @@ from typing import List, Optional
 
 import pydash as _
 import torch
+from botocore.vendored.six import BytesIO
 from numpy import who
 
 import internals.util.prompt as prompt_util
 from internals.data.dataAccessor import update_db, update_db_source_failed
-from internals.data.task import Task, TaskType
+from internals.data.task import ModelType, Task, TaskType
 from internals.pipelines.commons import Img2Img, Text2Img
 from internals.pipelines.controlnets import ControlNet
 from internals.pipelines.high_res import HighRes
 from internals.pipelines.img_classifier import ImageClassifier
 from internals.pipelines.img_to_text import Image2Text
 from internals.pipelines.inpainter import InPainter
+from internals.pipelines.object_remove import ObjectRemoval
 from internals.pipelines.pose_detector import PoseDetector
 from internals.pipelines.prompt_modifier import PromptModifier
+from internals.pipelines.realtime_draw import RealtimeDraw
+from internals.pipelines.remove_background import RemoveBackgroundV2
 from internals.pipelines.replace_background import ReplaceBackground
 from internals.pipelines.safety_checker import SafetyChecker
 from internals.pipelines.sdxl_tile_upscale import SDXLTileUpscaler
+from internals.pipelines.upscaler import Upscaler
 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 download_image, upload_image, upload_images
+from internals.util.commons import (
+    base64_to_image,
+    construct_default_s3_url,
+    download_image,
+    image_to_base64,
+    upload_image,
+    upload_images,
+)
 from internals.util.config import (
     get_is_sdxl,
     get_model_dir,
@@ -43,11 +55,15 @@ torch.backends.cuda.matmul.allow_tf32 = True
 auto_mode = False
 
 prompt_modifier = PromptModifier(num_of_sequences=num_return_sequences)
+upscaler = Upscaler()
 pose_detector = PoseDetector()
 inpainter = InPainter()
 high_res = HighRes()
 img2text = Image2Text()
 img_classifier = ImageClassifier()
+object_removal = ObjectRemoval()
+replace_background = ReplaceBackground()
+remove_background_v2 = RemoveBackgroundV2()
 replace_background = ReplaceBackground()
 controlnet = ControlNet()
 lora_style = LoraStyle()
@@ -56,6 +72,7 @@ img2img_pipe = Img2Img()
 safety_checker = SafetyChecker()
 slack = Slack()
 avatar = Avatar()
+realtime_draw = RealtimeDraw()
 sdxl_tileupscaler = SDXLTileUpscaler()
 
 
@@ -149,7 +166,9 @@ def tile_upscale(task: Task):
     prompt = get_patched_prompt_tile_upscale(task)
 
     if get_is_sdxl():
-        lora_patcher = lora_style.get_patcher(sdxl_tileupscaler.pipe, task.get_style())
+        lora_patcher = lora_style.get_patcher(
+            [sdxl_tileupscaler.pipe, high_res.pipe], task.get_style()
+        )
         lora_patcher.patch()
 
         images, has_nsfw = sdxl_tileupscaler.process(
@@ -555,6 +574,124 @@ def replace_bg(task: Task):
     }
 
 
+@update_db
+@slack.auto_send_alert
+def remove_bg(task: Task):
+    output_image = remove_background_v2.remove(
+        task.get_imageUrl(), model_type=task.get_modelType()
+    )
+
+    output_key = "crecoAI/{}_rmbg.png".format(task.get_taskId())
+    image_url = upload_image(output_image, output_key)
+
+    return {"generated_image_url": image_url}
+
+
+@update_db
+@slack.auto_send_alert
+def upscale_image(task: Task):
+    output_key = "crecoAI/{}_upscale.png".format(task.get_taskId())
+    out_img = None
+    if (
+        task.get_modelType() == ModelType.ANIME
+        or task.get_modelType() == ModelType.COMIC
+    ):
+        print("Using Anime model")
+        out_img = upscaler.upscale_anime(
+            image=task.get_imageUrl(),
+            width=task.get_width(),
+            height=task.get_height(),
+            face_enhance=task.get_face_enhance(),
+            resize_dimension=task.get_resize_dimension(),
+        )
+    else:
+        print("Using Real model")
+        out_img = upscaler.upscale(
+            image=task.get_imageUrl(),
+            width=task.get_width(),
+            height=task.get_height(),
+            face_enhance=task.get_face_enhance(),
+            resize_dimension=task.get_resize_dimension(),
+        )
+
+    image_url = upload_image(BytesIO(out_img), output_key)
+
+    clear_cuda_and_gc()
+
+    return {"generated_image_url": image_url}
+
+
+@update_db
+@slack.auto_send_alert
+def remove_object(task: Task):
+    output_key = "crecoAI/{}_object_remove.png".format(task.get_taskId())
+
+    images = object_removal.process(
+        image_url=task.get_imageUrl(),
+        mask_image_url=task.get_maskImageUrl(),
+        seed=task.get_seed(),
+        width=task.get_width(),
+        height=task.get_height(),
+    )
+    generated_image_urls = upload_image(images[0], output_key)
+
+    clear_cuda()
+
+    return {"generated_image_urls": generated_image_urls}
+
+
+def rt_draw_seg(task: Task):
+    image = task.get_imageUrl()
+    if image.startswith("http"):
+        image = download_image(image)
+    else:  # consider image as base64
+        image = base64_to_image(image)
+
+    img = realtime_draw.process_seg(
+        image=image,
+        prompt=task.get_prompt(),
+        negative_prompt=task.get_negative_prompt(),
+        seed=task.get_seed(),
+    )
+
+    clear_cuda_and_gc()
+
+    base64_image = image_to_base64(img)
+
+    return {"image": base64_image}
+
+
+def rt_draw_img(task: Task):
+    image = task.get_imageUrl()
+    aux_image = task.get_auxilary_imageUrl()
+
+    if image:
+        if image.startswith("http"):
+            image = download_image(image)
+        else:  # consider image as base64
+            image = base64_to_image(image)
+
+    if aux_image:
+        if aux_image.startswith("http"):
+            aux_image = download_image(aux_image)
+        else:  # consider image as base64
+            aux_image = base64_to_image(aux_image)
+
+    img = realtime_draw.process_img(
+        image=image,  # pyright: ignore
+        image2=aux_image,  # pyright: ignore
+        prompt=task.get_prompt(),
+        negative_prompt=task.get_negative_prompt(),
+        seed=task.get_seed(),
+    )
+
+    clear_cuda_and_gc()
+
+    base64_image = image_to_base64(img)
+
+    return {"image": base64_image}
+
+
 def custom_action(task: Task):
     from external.scripts import __scripts__
 
@@ -581,7 +718,7 @@ def custom_action(task: Task):
             return script(task, data)
 
 
-def load_model_by_task(task: Task):
+def load_model_by_task(task_type: TaskType, model_id=-1):
     if not text2img_pipe.is_loaded():
         text2img_pipe.load(get_model_dir())
         img2img_pipe.create(text2img_pipe)
@@ -593,24 +730,30 @@ def load_model_by_task(task: Task):
         safety_checker.apply(text2img_pipe)
         safety_checker.apply(img2img_pipe)
 
-    if task.get_type() == TaskType.INPAINT:
+    if task_type == TaskType.INPAINT:
         inpainter.load()
         safety_checker.apply(inpainter)
-    elif task.get_type() == TaskType.REPLACE_BG:
+    elif task_type == TaskType.REPLACE_BG:
         replace_background.load(base=text2img_pipe, high_res=high_res)
+    elif task_type == TaskType.RT_DRAW_SEG or task_type == TaskType.RT_DRAW_IMG:
+        realtime_draw.load(text2img_pipe)
+    elif task_type == TaskType.OBJECT_REMOVAL:
+        object_removal.load(get_model_dir())
+    elif task_type == TaskType.UPSCALE_IMAGE:
+        upscaler.load()
     else:
-        if task.get_type() == TaskType.TILE_UPSCALE:
+        if task_type == TaskType.TILE_UPSCALE:
             if get_is_sdxl():
-                sdxl_tileupscaler.create(text2img_pipe, task.get_model_id())
+                sdxl_tileupscaler.create(high_res, text2img_pipe, model_id)
             else:
                 controlnet.load_model("tile_upscaler")
-        elif task.get_type() == TaskType.CANNY:
+        elif task_type == TaskType.CANNY:
             controlnet.load_model("canny")
-        elif task.get_type() == TaskType.SCRIBBLE:
+        elif task_type == TaskType.SCRIBBLE:
             controlnet.load_model("scribble")
-        elif task.get_type() == TaskType.LINEARART:
+        elif task_type == TaskType.LINEARART:
             controlnet.load_model("linearart")
-        elif task.get_type() == TaskType.POSE:
+        elif task_type == TaskType.POSE:
             controlnet.load_model("pose")
 
         safety_checker.apply(controlnet)
@@ -629,6 +772,8 @@ def model_fn(model_dir):
 
     lora_style.load(model_dir)
 
+    load_model_by_task(TaskType.TEXT_TO_IMAGE)
+
     print("Logs: model loaded ....")
     return
 
@@ -643,11 +788,21 @@ def predict_fn(data, pipe):
     FailureHandler.handle(task)
 
     try:
+        task_type = task.get_type()
+
         # Set set_environment
         set_configs_from_task(task)
 
         # Load model based on task
-        load_model_by_task(task)
+        load_model_by_task(
+            task.get_type() or TaskType.TEXT_TO_IMAGE, task.get_model_id()
+        )
+
+        # Realtime generation apis
+        if task_type == TaskType.RT_DRAW_SEG:
+            return rt_draw_seg(task)
+        if task_type == TaskType.RT_DRAW_IMG:
+            return rt_draw_img(task)
 
         # Apply arguments
         apply_style_args(data)
@@ -658,8 +813,6 @@ def predict_fn(data, pipe):
         # Fetch avatars
         avatar.fetch_from_network(task.get_model_id())
 
-        task_type = task.get_type()
-
         if task_type == TaskType.TEXT_TO_IMAGE:
             # character sheet
             # if "character sheet" in task.get_prompt().lower():
@@ -684,8 +837,20 @@ def predict_fn(data, pipe):
             return replace_bg(task)
         elif task_type == TaskType.CUSTOM_ACTION:
             return custom_action(task)
+        elif task_type == TaskType.REMOVE_BG:
+            return remove_bg(task)
+        elif task_type == TaskType.UPSCALE_IMAGE:
+            return upscale_image(task)
+        elif task_type == TaskType.OBJECT_REMOVAL:
+            return remove_object(task)
         elif task_type == TaskType.SYSTEM_CMD:
             os.system(task.get_prompt())
+        elif task_type == TaskType.PRELOAD_MODEL:
+            try:
+                task_type = TaskType(task.get_prompt())
+            except:
+                task_type = TaskType.SYSTEM_CMD
+            load_model_by_task(task_type)
         else:
             raise Exception("Invalid task type")
     except Exception as e:

internals/data/task.py

@@ -1,6 +1,6 @@
 from enum import Enum
 from functools import lru_cache
-from typing import Union
+from typing import Optional, Union
 
 import numpy as np
 
@@ -18,6 +18,9 @@ class TaskType(Enum):
     SCRIBBLE = "SCRIBBLE"
     LINEARART = "LINEARART"
     REPLACE_BG = "REPLACE_BG"
+    RT_DRAW_SEG = "RT_DRAW_SEG"
+    RT_DRAW_IMG = "RT_DRAW_IMG"
+    PRELOAD_MODEL = "PRELOAD_MODEL"
     CUSTOM_ACTION = "CUSTOM_ACTION"
     SYSTEM_CMD = "SYSTEM_CMD"
 
@@ -52,7 +55,7 @@ class Task:
     def get_imageUrl(self) -> str:
         return self.__data.get("imageUrl", None)
 
-    def get_auxilary_imageUrl(self) -> str:
+    def get_auxilary_imageUrl(self) -> Optional[str]:
         return self.__data.get("aux_imageUrl", None)
 
     def get_prompt(self) -> str:

internals/pipelines/object_remove.py

@@ -18,7 +18,11 @@ from saicinpainting.training.trainers import load_checkpoint
 
 
 class ObjectRemoval:
+    __loaded = False
+
     def load(self, model_dir):
+        if self.__loaded:
+            return
         print("Downloading LAMA model...")
 
         self.lama_path = Path.home() / ".cache" / "lama"
@@ -36,6 +40,8 @@ class ObjectRemoval:
         self.model.freeze()
         self.model.to("cuda")
 
+        self.__loaded = True
+
     @torch.no_grad()
     def process(
         self,

internals/pipelines/realtime_draw.py

@@ -0,0 +1,97 @@
+from typing import Optional
+
+import torch
+from diffusers import ControlNetModel, StableDiffusionControlNetImg2ImgPipeline
+from PIL import Image
+
+import internals.util.image as ImageUtil
+from internals.pipelines.commons import AbstractPipeline
+from internals.pipelines.controlnets import ControlNet
+from internals.util.config import get_hf_cache_dir
+
+
+class RealtimeDraw(AbstractPipeline):
+    def load(self, pipeline: AbstractPipeline):
+        if hasattr(self, "pipe"):
+            return
+
+        self.__controlnet_scribble = ControlNetModel.from_pretrained(
+            "lllyasviel/control_v11p_sd15_scribble",
+            torch_dtype=torch.float16,
+            cache_dir=get_hf_cache_dir(),
+        )
+
+        self.__controlnet_seg = ControlNetModel.from_pretrained(
+            "lllyasviel/control_v11p_sd15_seg",
+            torch_dtype=torch.float16,
+            cache_dir=get_hf_cache_dir(),
+        )
+
+        kwargs = {**pipeline.pipe.components}  # pyright: ignore
+        kwargs.pop("image_encoder", None)
+        self.pipe = StableDiffusionControlNetImg2ImgPipeline(
+            **kwargs, controlnet=self.__controlnet_seg
+        ).to("cuda")
+        self.pipe.safety_checker = None
+        self.pipe2 = StableDiffusionControlNetImg2ImgPipeline(
+            **kwargs, controlnet=[self.__controlnet_scribble, self.__controlnet_seg]
+        ).to("cuda")
+        self.pipe2.safety_checker = None
+
+    def process_seg(
+        self,
+        image: Image.Image,
+        prompt: str,
+        negative_prompt: str,
+        seed: int,
+    ):
+        torch.manual_seed(seed)
+
+        image = ImageUtil.resize_image(image, 512)
+
+        img = self.pipe.__call__(
+            image=image,
+            control_image=image,
+            prompt=prompt,
+            num_inference_steps=15,
+            negative_prompt=negative_prompt,
+            guidance_scale=10,
+            strength=0.8,
+        ).images[0]
+
+        return img
+
+    def process_img(
+        self,
+        prompt: str,
+        negative_prompt: str,
+        seed: int,
+        image: Optional[Image.Image] = None,
+        image2: Optional[Image.Image] = None,
+    ):
+        torch.manual_seed(seed)
+
+        if not image:
+            image = Image.new("RGB", (512, 512), color=0)
+
+        if not image2:
+            image2 = Image.new("RGB", image.size, color=0)
+
+        image = ImageUtil.resize_image(image, 512)
+
+        scribble = ControlNet.scribble_image(image)
+
+        image2 = ImageUtil.resize_image(image2, 512)
+
+        img = self.pipe2.__call__(
+            image=image,
+            control_image=[scribble, image2],
+            prompt=prompt,
+            num_inference_steps=15,
+            negative_prompt=negative_prompt,
+            guidance_scale=10,
+            strength=0.9,
+            controlnet_conditioning_scale=[1.0, 0.8],
+        ).images[0]
+
+        return img

internals/pipelines/sdxl_tile_upscale.py

@@ -4,10 +4,12 @@ from PIL import Image
 from torchvision import transforms
 
 import internals.util.image as ImageUtils
+from carvekit.api import high
 from internals.data.result import Result
 from internals.pipelines.commons import AbstractPipeline, Text2Img
 from internals.pipelines.controlnets import ControlNet
 from internals.pipelines.demofusion_sdxl import DemoFusionSDXLControlNetPipeline
+from internals.pipelines.high_res import HighRes
 from internals.util.commons import download_image
 from internals.util.config import get_base_dimension
 
@@ -15,7 +17,7 @@ controlnet = ControlNet()
 
 
 class SDXLTileUpscaler(AbstractPipeline):
-    def create(self, pipeline: Text2Img, model_id: int):
+    def create(self, high_res: HighRes, pipeline: Text2Img, model_id: int):
         # temporal hack for upscale model till multicontrolnet support is added
         model = (
             "thibaud/controlnet-openpose-sdxl-1.0"
@@ -32,6 +34,8 @@ class SDXLTileUpscaler(AbstractPipeline):
         pipe.enable_vae_slicing()
         pipe.enable_xformers_memory_efficient_attention()
 
+        self.high_res = high_res
+
         self.pipe = pipe
 
     def process(
@@ -58,18 +62,28 @@ class SDXLTileUpscaler(AbstractPipeline):
 
         image_lr = self.load_and_process_image(img)
         print("img", img2.size, img.size)
-        images = self.pipe.__call__(
-            image_lr=image_lr,
-            prompt=prompt,
-            condition_image=condition_image,
-            negative_prompt="blurry, ugly, duplicate, poorly drawn, deformed, mosaic",
-            guidance_scale=11,
-            sigma=0.8,
-            num_inference_steps=24,
-            width=img2.size[0],
-            height=img2.size[1],
-        )
-        images = images[::-1]
+        if int(model_id) == 2000173:
+            kwargs = {
+                "prompt": prompt,
+                "negative_prompt": negative_prompt,
+                "image": img2,
+                "strength": 0.3,
+                "num_inference_steps": 30,
+            }
+            images = self.high_res.pipe.__call__(**kwargs).images
+        else:
+            images = self.pipe.__call__(
+                image_lr=image_lr,
+                prompt=prompt,
+                condition_image=condition_image,
+                negative_prompt="blurry, ugly, duplicate, poorly drawn, deformed, mosaic",
+                guidance_scale=11,
+                sigma=0.8,
+                num_inference_steps=24,
+                width=img2.size[0],
+                height=img2.size[1],
+            )
+            images = images[::-1]
         return images, False
 
     def load_and_process_image(self, pil_image):

internals/pipelines/upscaler.py

@@ -148,7 +148,7 @@ class Upscaler:
                 model=model,
                 half=False,
                 gpu_id="0",
-                tile=128,
+                tile=320,
                 tile_pad=10,
                 pre_pad=0,
             )

internals/util/commons.py

@@ -1,3 +1,4 @@
+import base64
 import json
 import os
 import pprint
@@ -163,6 +164,18 @@ def download_file(url, out_path: Path):
                 f.write(chunk)
 
 
+def base64_to_image(base64_string):
+    imgdata = base64.b64decode(base64_string)
+    return Image.open(io.BytesIO(imgdata)).convert("RGB")
+
+
+def image_to_base64(image):
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue())
+    return img_str.decode("ascii")
+
+
 def pickPoses():
     random_images = random.sample(characterSheets, 4)
     poses = []