Upload 18 files

data/dataAccessor.py

@@ -0,0 +1,60 @@
+import requests
+from data.task import Task
+from util.slack import Slack
+
+comic_url = "http://internal-k8s-gamma-internal-ea8e32da94-1997933257.ap-south-1.elb.amazonaws.com:80"
+
+
+def updateSource(sourceId, userId, state):
+    print("update source is called")
+    url = comic_url + f"/comic-crecoai/source/{sourceId}"
+    headers = {"Content-Type": "application/json", "user-id": str(userId)}
+
+    data = {"state": state}
+
+    try:
+        response = requests.patch(url, headers=headers, json=data, timeout=10)
+        print("update source response", response)
+    except requests.exceptions.Timeout:
+        print("Request timed out while updating source")
+    except requests.exceptions.RequestException as e:
+        print(f"Error while updating source: {e}")
+
+    return
+
+
+def saveGeneratedImages(sourceId, userId):
+    print("save generation called")
+    url = comic_url + "/comic-crecoai/source/" + str(sourceId) + "/generatedImages"
+    headers = {"Content-Type": "application/json", "user-id": str(userId)}
+    data = {"state": "ACTIVE"}
+
+    try:
+        requests.patch(url, headers=headers, json=data)
+        # print("save generation response", response)
+    except requests.exceptions.Timeout:
+        print("Request timed out while saving image")
+    except requests.exceptions.RequestException as e:
+        print("Failed to mark source as active: ", e)
+        return
+    return
+
+
+def update_db(func):
+    def caller(*args, **kwargs):
+        if type(args[0]) is not Task:
+            raise Exception("First argument must be a Task object")
+        task = args[0]
+        try:
+            updateSource(task.get_sourceId(), task.get_userId(), "INPROGRESS")
+            rargs = func(*args, **kwargs)
+            updateSource(task.get_sourceId(), task.get_userId(), "COMPLETED")
+            saveGeneratedImages(task.get_sourceId(), task.get_userId())
+            return rargs
+        except Exception as e:
+            print("Error processing image: {}".format(str(e)))
+            slack = Slack()
+            slack.error_alert(task, e)
+            updateSource(task.get_sourceId(), task.get_userId(), "FAILED")
+
+    return caller

data/task.py

@@ -0,0 +1,85 @@
+from enum import Enum
+from typing import Union
+
+import numpy as np
+
+
+class TaskType(Enum):
+    TEXT_TO_IMAGE = "GENERATE_AI_IMAGE"
+    IMAGE_TO_IMAGE = "IMAGE_TO_IMAGE"
+    POSE = "POSE"
+    CANNY = "CANNY"
+    REMOVE_BG = "REMOVE_BG"
+    INPAINT = "INPAINT"
+    UPSCALE_IMAGE = "UPSCALE_IMAGE"
+
+
+class ModelType(Enum):
+    REAL = 10000
+    ANIME = 10001
+    COMIC = 10002
+
+
+class Task:
+    def __init__(self, data):
+        self.__data = data
+        if data.get("seed", -1) == None or self.get_seed() == -1:
+            self.__data["seed"] = np.random.randint(0, np.iinfo(np.int64).max)
+
+    def get_taskId(self) -> str:
+        return self.__data.get("task_id")
+
+    def get_sourceId(self) -> str:
+        return self.__data.get("source_id")
+
+    def get_imageUrl(self) -> str:
+        return self.__data.get("imageUrl")
+
+    def get_prompt(self) -> str:
+        return self.__data.get("prompt")
+
+    def get_userId(self) -> str:
+        return self.__data.get("userId", "")
+
+    def get_email(self) -> str:
+        return self.__data.get("email", "")
+
+    def get_style(self) -> str:
+        return self.__data.get("style", None)
+
+    def get_iteration(self) -> float:
+        return float(self.__data.get("iteration", 3.0))
+
+    def get_modelType(self) -> ModelType:
+        id = int(self.__data.get("model_id", 10000))
+        return ModelType(id)
+
+    def get_width(self) -> int:
+        return int(self.__data.get("width", 512))
+
+    def get_height(self) -> int:
+        return int(self.__data.get("height", 512))
+
+    def get_seed(self) -> int:
+        return int(self.__data.get("seed", -1))
+
+    def get_steps(self) -> int:
+        return int(self.__data.get("steps", "75"))
+
+    def get_type(self) -> Union[TaskType, None]:
+        try:
+            return TaskType(self.__data.get("task_type"))
+        except ValueError:
+            return None
+
+    def get_maskImageUrl(self) -> str:
+        return self.__data.get("maskImageUrl")
+
+    def get_negative_prompt(self) -> str:
+        return self.__data.get("negative_prompt", "")
+
+    def is_prompt_engineering(self) -> bool:
+        return self.__data.get("auto_mode", True)
+
+    def get_raw(self) -> dict:
+        return self.__data.copy()

inference.py

@@ -0,0 +1,212 @@
+from typing import List, Optional
+
+import torch
+from data.dataAccessor import update_db
+from data.task import Task, TaskType
+from pipelines.commons import Img2Img, Text2Img
+from pipelines.controlnets import ControlNet
+from pipelines.prompt_modifier import PromptModifier
+from util.cache import auto_clear_cuda_and_gc, clear_cuda
+from util.commons import add_code_names, pickPoses, upload_images
+from util.lora_style import LoraStyle
+from util.slack import Slack
+
+torch.backends.cudnn.benchmark = True
+torch.backends.cuda.matmul.allow_tf32 = True
+
+num_return_sequences = 4  # the number of results to generate
+auto_mode = False
+
+prompt_modifier = PromptModifier(num_of_sequences=num_return_sequences)
+controlnet = ControlNet()
+lora_style = LoraStyle()
+text2img_pipe = Text2Img()
+img2img_pipe = Img2Img()
+slack = Slack()
+
+
+def get_patched_prompt(task: Task):
+    def add_style_and_character(prompt: List[str]):
+        for i in range(len(prompt)):
+            prompt[i] = add_code_names(prompt[i])
+            prompt[i] = lora_style.prepend_style_to_prompt(prompt[i], task.get_style())
+
+    prompt = task.get_prompt()
+
+    if task.is_prompt_engineering():
+        prompt = prompt_modifier.modify(prompt)
+    else:
+        prompt = [prompt] * num_return_sequences
+
+    ori_prompt = [task.get_prompt()] * num_return_sequences
+
+    add_style_and_character(ori_prompt)
+    add_style_and_character(prompt)
+
+    print({"prompts": prompt})
+
+    return (prompt, ori_prompt)
+
+
+@update_db
+@auto_clear_cuda_and_gc(controlnet)
+@slack.auto_send_alert
+def canny(task: Task):
+    prompt, _ = get_patched_prompt(task)
+
+    controlnet.load_canny()
+
+    lora_patcher = lora_style.get_patcher(controlnet.pipe, task.get_style())
+    lora_patcher.patch()
+
+    images = controlnet.process_canny(
+        prompt=prompt,
+        imageUrl=task.get_imageUrl(),
+        seed=task.get_seed(),
+        steps=task.get_steps(),
+        width=task.get_width(),
+        height=task.get_height(),
+        negative_prompt=[
+            f"monochrome, neon, x-ray, negative image, oversaturated, {task.get_negative_prompt()}"
+        ]
+        * num_return_sequences,
+        **lora_patcher.kwargs(),
+    )
+
+    generated_image_urls = upload_images(images, "_canny", task.get_taskId())
+
+    lora_patcher.cleanup()
+    controlnet.cleanup()
+
+    return {"modified_prompts": prompt, "generated_image_urls": generated_image_urls}
+
+
+@update_db
+@auto_clear_cuda_and_gc(controlnet)
+@slack.auto_send_alert
+def pose(task: Task, s3_outkey: str = "_pose", poses: Optional[list] = None):
+    prompt, _ = get_patched_prompt(task)
+
+    controlnet.load_pose()
+
+    lora_patcher = lora_style.get_patcher(controlnet.pipe, task.get_style())
+    lora_patcher.patch()
+
+    if poses is None:
+        poses = [controlnet.detect_pose(task.get_imageUrl())] * num_return_sequences
+
+    images = controlnet.process_pose(
+        prompt=prompt,
+        image=poses,
+        seed=task.get_seed(),
+        steps=task.get_steps(),
+        negative_prompt=[task.get_negative_prompt()] * num_return_sequences,
+        width=task.get_width(),
+        height=task.get_height(),
+        **lora_patcher.kwargs(),
+    )
+
+    generated_image_urls = upload_images(images, s3_outkey, task.get_taskId())
+
+    lora_patcher.cleanup()
+    controlnet.cleanup()
+
+    return {"modified_prompts": prompt, "generated_image_urls": generated_image_urls}
+
+
+@update_db
+@auto_clear_cuda_and_gc(controlnet)
+@slack.auto_send_alert
+def text2img(task: Task):
+    prompt, ori_prompt = get_patched_prompt(task)
+
+    lora_patcher = lora_style.get_patcher(text2img_pipe.pipe, task.get_style())
+    lora_patcher.patch()
+
+    torch.manual_seed(task.get_seed())
+
+    images = text2img_pipe.process(
+        prompt=ori_prompt,
+        modified_prompts=prompt,
+        num_inference_steps=task.get_steps(),
+        guidance_scale=7.5,
+        height=task.get_height(),
+        width=task.get_width(),
+        negative_prompt=[task.get_negative_prompt()] * num_return_sequences,
+        iteration=task.get_iteration(),
+        **lora_patcher.kwargs(),
+    )
+
+    generated_image_urls = upload_images(images, "", task.get_taskId())
+
+    lora_patcher.cleanup()
+
+    return {"modified_prompts": prompt, "generated_image_urls": generated_image_urls}
+
+
+@update_db
+@auto_clear_cuda_and_gc(controlnet)
+@slack.auto_send_alert
+def img2img(task: Task):
+    prompt, _ = get_patched_prompt(task)
+
+    lora_patcher = lora_style.get_patcher(img2img_pipe.pipe, task.get_style())
+    lora_patcher.patch()
+
+    torch.manual_seed(task.get_seed())
+
+    images = img2img_pipe.process(
+        prompt=prompt,
+        imageUrl=task.get_imageUrl(),
+        negative_prompt=[task.get_negative_prompt()] * num_return_sequences,
+        steps=task.get_steps(),
+        **lora_patcher.kwargs(),
+    )
+
+    generated_image_urls = upload_images(images, "_imgtoimg", task.get_taskId())
+
+    lora_patcher.cleanup()
+
+    return {"modified_prompts": prompt, "generated_image_urls": generated_image_urls}
+
+
+def model_fn(model_dir):
+    print("Logs: model loaded .... starts")
+
+    prompt_modifier.load()
+
+    lora_style.load(model_dir)
+    controlnet.load(model_dir)
+
+    text2img_pipe.load(model_dir)
+    img2img_pipe.load(model_dir)
+
+    print("Logs: model loaded ....")
+    return
+
+
+def predict_fn(data, pipe):
+    task = Task(data)
+    print("task is ", data)
+
+    try:
+        task_type = task.get_type()
+
+        if task_type == TaskType.TEXT_TO_IMAGE:
+            # character sheet
+            if "character sheet" in task.get_prompt().lower():
+                return pose(task, s3_outkey="", poses=pickPoses())
+            else:
+                return text2img(task)
+        elif task_type == TaskType.IMAGE_TO_IMAGE:
+            return img2img(task)
+        elif task_type == TaskType.CANNY:
+            return canny(task)
+        elif task_type == TaskType.POSE:
+            return pose(task)
+        else:
+            raise Exception("Invalid task type")
+    except Exception as e:
+        print(f"Error: {e}")
+        slack.error_alert(task, e)
+        return None

inference2.py

@@ -0,0 +1,116 @@
+from io import BytesIO
+
+import torch
+from data.dataAccessor import update_db
+from data.task import ModelType, Task, TaskType
+from pipelines.inpainter import InPainter
+from pipelines.prompt_modifier import PromptModifier
+from pipelines.remove_background import RemoveBackground
+from pipelines.upscaler import Upscaler
+from util.cache import clear_cuda
+from util.commons import (
+    add_code_names,
+    construct_default_s3_url,
+    upload_image,
+    upload_images,
+)
+from util.slack import Slack
+
+torch.backends.cudnn.benchmark = True
+torch.backends.cuda.matmul.allow_tf32 = True
+
+num_return_sequences = 4
+auto_mode = False
+
+slack = Slack()
+
+prompt_modifier = PromptModifier(num_of_sequences=num_return_sequences)
+upscaler = Upscaler()
+inpainter = InPainter()
+
+
+@update_db
+@slack.auto_send_alert
+def remove_bg(task: Task):
+    remove_background = RemoveBackground()
+    output_image = remove_background.remove(task.get_imageUrl())
+
+    output_key = "crecoAI/{}_rmbg.png".format(task.get_taskId())
+    upload_image(output_image, output_key)
+
+    return {"generated_image_url": construct_default_s3_url(output_key)}
+
+
+@update_db
+@slack.auto_send_alert
+def inpaint(task: Task):
+    prompt = add_code_names(task.get_prompt())
+    if task.is_prompt_engineering():
+        prompt = prompt_modifier.modify(prompt)
+    else:
+        prompt = [prompt] * num_return_sequences
+
+    print({"prompts": prompt})
+
+    images = inpainter.process(
+        prompt=prompt,
+        image_url=task.get_imageUrl(),
+        mask_image_url=task.get_maskImageUrl(),
+        width=task.get_width(),
+        height=task.get_height(),
+        seed=task.get_seed(),
+        negative_prompt=[task.get_negative_prompt()] * num_return_sequences,
+    )
+    generated_image_urls = upload_images(images, "_inpaint", task.get_taskId())
+
+    clear_cuda()
+
+    return {"modified_prompts": prompt, "generated_image_urls": generated_image_urls}
+
+
+@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:
+        print("Using Anime model")
+        out_img = upscaler.upscale_anime(task.get_imageUrl())
+    else:
+        print("Using Real model")
+        out_img = upscaler.upscale(task.get_imageUrl())
+
+    upload_image(BytesIO(out_img), output_key)
+    return {"generated_image_url": construct_default_s3_url(output_key)}
+
+
+def model_fn(model_dir):
+    print("Logs: model loaded .... starts")
+
+    prompt_modifier.load()
+    upscaler.load()
+    inpainter.load()
+
+    print("Logs: model loaded ....")
+    return
+
+
+def predict_fn(data, pipe):
+    task = Task(data)
+    print("task is ", data)
+
+    try:
+        task_type = task.get_type()
+
+        if task_type == TaskType.REMOVE_BG:
+            return remove_bg(task)
+        elif task_type == TaskType.INPAINT:
+            return inpaint(task)
+        elif task_type == TaskType.UPSCALE_IMAGE:
+            return upscale_image(task)
+        else:
+            raise Exception("Invalid task type")
+    except Exception as e:
+        print(f"Error: {e}")
+        slack.error_alert(task, e)
+        return None

pipelines/commons.py

@@ -0,0 +1,85 @@
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import torch
+from diffusers import StableDiffusionImg2ImgPipeline
+from pipelines.twoStepPipeline import two_step_pipeline
+from util.commons import disable_safety_checker, download_image
+
+
+class Text2Img:
+    def load(self, model_dir: str):
+        self.pipe = two_step_pipeline.from_pretrained(
+            model_dir, torch_dtype=torch.float16
+        ).to("cuda")
+        self.pipe.enable_xformers_memory_efficient_attention()
+        disable_safety_checker(self.pipe)
+
+    @torch.inference_mode()
+    def process(
+        self,
+        prompt: Union[str, List[str]] = None,
+        modified_prompts: Union[str, List[str]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: int = 1,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        iteration: float = 3.0,
+    ):
+        return self.pipe.two_step_pipeline(
+            prompt=prompt,
+            modified_prompts=modified_prompts,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            negative_prompt=negative_prompt,
+            num_images_per_prompt=num_images_per_prompt,
+            eta=eta,
+            generator=generator,
+            latents=latents,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            output_type=output_type,
+            return_dict=return_dict,
+            callback=callback,
+            callback_steps=callback_steps,
+            cross_attention_kwargs=cross_attention_kwargs,
+            iteration=iteration,
+        ).images
+
+
+class Img2Img:
+    def load(self, model_dir: str):
+        self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+            model_dir, torch_dtype=torch.float16
+        ).to("cuda")
+        self.pipe.enable_xformers_memory_efficient_attention()
+        disable_safety_checker(self.pipe)
+
+    @torch.inference_mode()
+    def process(
+        self, prompt: List[str], imageUrl: str, negative_prompt: List[str], steps: int
+    ):
+        image = download_image(imageUrl)
+
+        return self.pipe.__call__(
+            prompt=prompt,
+            image=image,
+            strength=0.75,
+            negative_prompt=negative_prompt,
+            guidance_scale=7.5,
+            num_images_per_prompt=1,
+            num_inference_steps=steps,
+        ).images

pipelines/controlnets.py

@@ -0,0 +1,133 @@
+from typing import List
+
+import cv2
+import numpy as np
+import torch
+from controlnet_aux import OpenposeDetector
+from diffusers import (ControlNetModel, StableDiffusionControlNetPipeline,
+                       UniPCMultistepScheduler)
+from PIL import Image
+from util.cache import clear_cuda_and_gc
+from util.commons import disable_safety_checker, download_image
+
+
+class ControlNet:
+    __current_task_name = ""
+
+    def load(self, model_dir: str):
+        # we will load canny by default
+        self.load_canny()
+
+        pipe = StableDiffusionControlNetPipeline.from_pretrained(
+            model_dir, controlnet=self.controlnet, torch_dtype=torch.float16
+        )
+        pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+        pipe.enable_model_cpu_offload()
+        pipe.enable_xformers_memory_efficient_attention()
+        disable_safety_checker(pipe)
+        self.pipe = pipe
+
+    def load_canny(self):
+        if self.__current_task_name == "canny":
+            return
+        canny = ControlNetModel.from_pretrained(
+            "lllyasviel/control_v11p_sd15_canny", torch_dtype=torch.float16
+        ).to("cuda")
+        self.__current_task_name = "canny"
+        self.controlnet = canny
+        if hasattr(self, "pipe"):
+            self.pipe.controlnet = canny
+        clear_cuda_and_gc()
+
+    def load_pose(self):
+        if self.__current_task_name == "pose":
+            return
+        pose = ControlNetModel.from_pretrained(
+            "lllyasviel/control_v11p_sd15_openpose", torch_dtype=torch.float16
+        ).to("cuda")
+        self.__current_task_name = "pose"
+        self.controlnet = pose
+        if hasattr(self, "pipe"):
+            self.pipe.controlnet = pose
+        clear_cuda_and_gc()
+
+    def cleanup(self):
+        self.pipe.controlnet = None
+        self.controlnet = None
+        self.__current_task_name = ""
+
+        clear_cuda_and_gc()
+
+    @torch.inference_mode()
+    def process_canny(
+        self,
+        prompt: List[str],
+        imageUrl: str,
+        seed: int,
+        steps: int,
+        negative_prompt: List[str],
+        height: int,
+        width: int,
+    ):
+        if self.__current_task_name != "canny":
+            raise Exception("ControlNet is not loaded with canny model")
+
+        torch.manual_seed(seed)
+
+        init_image = download_image(imageUrl)
+        init_image = self.__canny_detect_edge(init_image)
+
+        return self.pipe.__call__(
+            prompt=prompt,
+            image=init_image,
+            guidance_scale=9,
+            num_images_per_prompt=1,
+            negative_prompt=negative_prompt,
+            num_inference_steps=steps,
+            height=height,
+            width=width,
+        ).images
+
+    @torch.inference_mode()
+    def process_pose(
+        self,
+        prompt: List[str],
+        image: List[Image.Image],
+        seed: int,
+        steps: int,
+        negative_prompt: List[str],
+        height: int,
+        width: int,
+    ):
+        if self.__current_task_name != "pose":
+            raise Exception("ControlNet is not loaded with pose model")
+
+        torch.manual_seed(seed)
+
+        return self.pipe.__call__(
+            prompt=prompt,
+            image=image,
+            num_images_per_prompt=1,
+            num_inference_steps=steps,
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+        ).images
+
+    def detect_pose(self, imageUrl: str) -> Image.Image:
+        detector = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
+        image = download_image(imageUrl)
+        image = detector.__call__(image)
+        return image
+
+    def __canny_detect_edge(self, image: Image.Image) -> Image.Image:
+        image_array = np.array(image)
+
+        low_threshold = 100
+        high_threshold = 200
+
+        image_array = cv2.Canny(image_array, low_threshold, high_threshold)
+        image_array = image_array[:, :, None]
+        image_array = np.concatenate([image_array, image_array, image_array], axis=2)
+        canny_image = Image.fromarray(image_array)
+        return canny_image

pipelines/inpainter.py

@@ -0,0 +1,40 @@
+from typing import List, Union
+
+import torch
+from diffusers import StableDiffusionInpaintPipeline
+from util.commons import disable_safety_checker, download_image
+
+
+class InPainter:
+    def load(self):
+        self.pipe = StableDiffusionInpaintPipeline.from_pretrained(
+            "runwayml/stable-diffusion-inpainting",
+            torch_dtype=torch.float16,
+            revision="fp16",
+        ).to("cuda")
+        disable_safety_checker(self.pipe)
+
+    @torch.inference_mode()
+    def process(
+        self,
+        image_url: str,
+        mask_image_url: str,
+        width: int,
+        height: int,
+        seed: int,
+        prompt: Union[str, List[str]],
+        negative_prompt: Union[str, List[str]],
+    ):
+        torch.manual_seed(seed)
+
+        input_img = download_image(image_url).resize((width, height))
+        mask_img = download_image(mask_image_url).resize((width, height))
+
+        return self.pipe.__call__(
+            prompt=prompt,
+            image=input_img,
+            mask_image=mask_img,
+            height=height,
+            width=width,
+            negative_prompt=negative_prompt,
+        ).images

pipelines/prompt_modifier.py

@@ -0,0 +1,54 @@
+from typing import List, Optional
+
+from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
+
+
+class PromptModifier:
+    def __init__(self, num_of_sequences: Optional[int] = 4):
+        self.__blacklist = {"alphonse mucha": "", "adolphe bouguereau": ""}
+        self.__num_of_sequences = num_of_sequences
+
+    def load(self):
+        self.prompter_model = AutoModelForCausalLM.from_pretrained(
+            "Gustavosta/MagicPrompt-Stable-Diffusion"
+        )
+        self.prompter_tokenizer = AutoTokenizer.from_pretrained(
+            "Gustavosta/MagicPrompt-Stable-Diffusion"
+        )
+        self.prompter_tokenizer.pad_token = self.prompter_tokenizer.eos_token
+        self.prompter_tokenizer.padding_side = "left"
+
+    def modify(self, text: str) -> List[str]:
+        eos_id = self.prompter_tokenizer.eos_token_id
+        # restricted_words_list = ["octane", "cyber"]
+        # restricted_words_token_ids = prompter_tokenizer(
+        #     restricted_words_list, add_special_tokens=False
+        # ).input_ids
+
+        generation_config = GenerationConfig(
+            do_sample=False,
+            max_new_tokens=75,
+            num_beams=4,
+            num_return_sequences=self.__num_of_sequences,
+            eos_token_id=eos_id,
+            pad_token_id=eos_id,
+            length_penalty=-1.0,
+        )
+
+        input_ids = self.prompter_tokenizer(text.strip(), return_tensors="pt").input_ids
+        outputs = self.prompter_model.generate(
+            input_ids, generation_config=generation_config
+        )
+        output_texts = self.prompter_tokenizer.batch_decode(
+            outputs, skip_special_tokens=True
+        )
+        output_texts = self.__patch_blacklist_words(output_texts)
+        return output_texts
+
+    def __patch_blacklist_words(self, texts: List[str]):
+        def replace_all(text, dic):
+            for i, j in dic.items():
+                text = text.replace(i, j)
+            return text
+
+        return [replace_all(text, self.__blacklist) for text in texts]

pipelines/remove_background.py

@@ -0,0 +1,15 @@
+import io
+from typing import Union
+
+from PIL import Image
+from rembg import remove
+from util.commons import read_url
+
+
+class RemoveBackground:
+    def remove(self, image: Union[str, Image.Image]) -> Image.Image:
+        if type(image) is str:
+            image = Image.open(io.BytesIO(read_url(image)))
+
+        output = remove(image)
+        return output

pipelines/twoStepPipeline.py

@@ -0,0 +1,252 @@
+import torch
+from diffusers import StableDiffusionPipeline
+
+torch.backends.cudnn.benchmark = True
+torch.backends.cuda.matmul.allow_tf32 = True
+
+from typing import Any, Callable, Dict, List, Optional, Union
+
+from diffusers import StableDiffusionPipeline
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+
+
+class two_step_pipeline(StableDiffusionPipeline):
+    @torch.no_grad()
+    def two_step_pipeline(
+        self,
+        prompt: Union[str, List[str]] = None,
+        modified_prompts: Union[str, List[str]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: int = 1,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        iteration: float = 3.0,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            guidance_scale (`float`, *optional*, defaults to 7.5):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
+                [`schedulers.DDIMScheduler`], will be ignored for others.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
+                plain tuple.
+            callback (`Callable`, *optional*):
+                A function that will be called every `callback_steps` steps during inference. The function will be
+                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
+            callback_steps (`int`, *optional*, defaults to 1):
+                The frequency at which the `callback` function will be called. If not specified, the callback will be
+                called at every step.
+            cross_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).
+        Examples:
+        Returns:
+            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
+            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
+            When returning a tuple, the first element is a list with the generated images, and the second element is a
+            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
+            (nsfw) content, according to the `safety_checker`.
+        """
+        # 0. Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            height,
+            width,
+            callback_steps,
+            negative_prompt,
+            prompt_embeds,
+            negative_prompt_embeds,
+        )
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode input prompt
+        modified_embeds = self._encode_prompt(
+            modified_prompts,
+            device,
+            num_images_per_prompt,
+            do_classifier_free_guidance,
+            negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+        )
+        print("mod prompt size : ", modified_embeds.size(), modified_embeds.dtype)
+
+        prompt_embeds = self._encode_prompt(
+            prompt,
+            device,
+            num_images_per_prompt,
+            do_classifier_free_guidance,
+            negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+        )
+
+        print("prompt size : ", prompt_embeds.size(), prompt_embeds.dtype)
+
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 7. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = (
+                    torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                )
+                latent_model_input = self.scheduler.scale_model_input(
+                    latent_model_input, t
+                )
+
+                # predict the noise residual
+                noise_pred = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=prompt_embeds,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                ).sample
+
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (
+                        noise_pred_text - noise_pred_uncond
+                    )
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(
+                    noise_pred, t, latents, **extra_step_kwargs
+                ).prev_sample
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or (
+                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
+                ):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        callback(i, t, latents)
+
+                if i == int(len(timesteps) / iteration):
+                    print("modified prompts")
+                    prompt_embeds = modified_embeds
+
+        if output_type == "latent":
+            image = latents
+            has_nsfw_concept = None
+        elif output_type == "pil":
+            # 8. Post-processing
+            image = self.decode_latents(latents)
+
+            # 9. Run safety checker
+            image, has_nsfw_concept = self.run_safety_checker(
+                image, device, prompt_embeds.dtype
+            )
+
+            # 10. Convert to PIL
+            image = self.numpy_to_pil(image)
+        else:
+            # 8. Post-processing
+            image = self.decode_latents(latents)
+
+            # 9. Run safety checker
+            image, has_nsfw_concept = self.run_safety_checker(
+                image, device, prompt_embeds.dtype
+            )
+
+        # Offload last model to CPU
+        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
+            self.final_offload_hook.offload()
+
+        if not return_dict:
+            return (image, has_nsfw_concept)
+
+        return StableDiffusionPipelineOutput(
+            images=image, nsfw_content_detected=has_nsfw_concept
+        )

pipelines/upscaler.py

@@ -0,0 +1,77 @@
+import os
+from pathlib import Path
+from typing import Union
+
+import cv2
+import numpy as np
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from basicsr.utils.download_util import load_file_from_url
+from PIL import Image
+from realesrgan import RealESRGANer
+from util.commons import read_url
+
+
+class Upscaler:
+    __model_esrgan_url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
+    __model_esrgan_anime_url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth"
+
+    def load(self):
+        download_dir = Path(Path.home() / ".cache" / "realesrgan")
+        download_dir.mkdir(parents=True, exist_ok=True)
+
+        self.__model_path = self.__preload_model(self.__model_esrgan_url, download_dir)
+        self.__model_path_anime = self.__preload_model(
+            self.__model_esrgan_anime_url, download_dir
+        )
+
+    def upscale(self, image: Union[str, bytes]) -> bytes:
+        model = RRDBNet(
+            num_in_ch=3,
+            num_out_ch=3,
+            num_feat=64,
+            num_block=23,
+            num_grow_ch=32,
+            scale=4,
+        )
+        return self.__internal_upscale(image, self.__model_path, model)
+
+    def upscale_anime(self, image: Union[str, bytes]) -> bytes:
+        model = RRDBNet(
+            num_in_ch=3,
+            num_out_ch=3,
+            num_feat=64,
+            num_block=23,
+            num_grow_ch=32,
+            scale=4,
+        )
+        return self.__internal_upscale(image, self.__model_path_anime, model)
+
+    def __preload_model(self, url: str, download_dir: Path):
+        name = url.split("/")[-1]
+        if not os.path.exists(str(download_dir / name)):
+            return load_file_from_url(
+                url=url,
+                model_dir=str(download_dir),
+                progress=True,
+                file_name=None,
+            )
+        else:
+            return str(download_dir / name)
+
+    def __internal_upscale(
+        self,
+        image: Union[str, bytes],
+        model_path: str,
+        rrbdnet: RRDBNet,
+    ) -> bytes:
+        if type(image) is str:
+            image = read_url(image)
+
+        upsampler = RealESRGANer(
+            scale=4, model_path=model_path, model=rrbdnet, half="fp16", gpu_id="0"
+        )
+        image_array = np.frombuffer(image, dtype=np.uint8)
+        input_image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
+        output, _ = upsampler.enhance(input_image, outscale=4)
+        out_bytes = cv2.imencode(".png", output)[1].tobytes()
+        return out_bytes

requirements.txt

@@ -0,0 +1,19 @@
+aioredis==1.3.1
+boto3==1.24.61
+triton==2.0.0
+diffusers==0.14.0
+fastapi==0.87.0
+Pillow==9.3.0
+redis==4.3.4
+requests==2.28.1
+transformers
+rembg==2.0.30
+accelerate==0.17.0
+gfpgan==1.3.8
+rembg==2.0.30
+controlnet-aux==0.0.1
+realesrgan==0.3.0
+compel==1.0.4
+xformers
+torchvision
+git+https://github.com/cloneofsimo/lora.git

util/cache.py

@@ -0,0 +1,31 @@
+import gc
+
+import torch
+
+
+def clear_cuda_and_gc():
+    clear_cuda()
+    clear_gc()
+
+
+def clear_cuda():
+    torch.cuda.empty_cache()
+
+
+def clear_gc():
+    gc.collect()
+
+
+def auto_clear_cuda_and_gc(controlnet):
+    def auto_clear_cuda_and_gc_wrapper(func):
+        def wrapper(*args, **kwargs):
+            try:
+                return func(*args, **kwargs)
+            except Exception as e:
+                controlnet.cleanup()
+                clear_cuda_and_gc()
+                raise e
+
+        return wrapper
+
+    return auto_clear_cuda_and_gc_wrapper

util/commons.py

@@ -0,0 +1,220 @@
+import pprint
+import random
+import re
+from io import BytesIO
+from typing import Union
+
+import boto3
+import requests
+
+s3 = boto3.client("s3")
+import io
+import urllib.request
+
+from PIL import Image
+
+black_list = {"alphonse mucha": "", "adolphe bouguereau": ""}
+pp = pprint.PrettyPrinter(indent=4)
+Avatar = [
+    {
+        "avatarName": "niomi",
+        "codename": "1jMGp1kFkG",
+        "avatarImage": "https://comic-assets.s3.ap-south-1.amazonaws.com/7_char_assets/niyomi_1jMGp1kFkG.png",
+        "extraPrompt": "1jMGp1kFkG  girl",
+    },
+    {
+        "avatarName": "riya",
+        "codename": "vW6AUQtoaY",
+        "avatarImage": "https://comic-assets.s3.ap-south-1.amazonaws.com/12_char_assets/riya_vW6AUQtoaY.png",
+        "extraPrompt": "vW6AUQtoaY girl",
+    },
+    {
+        "avatarName": "rajveer",
+        "codename": "fSLF0OPkBw",
+        "avatarImage": "https://comic-assets.s3.ap-south-1.amazonaws.com/12_character_assets/rajveer_fSLF0OPkBw.png",
+        "extraPrompt": "fSLF0OPkBw guy",
+    },
+    {
+        "avatarName": "bheem",
+        "codename": "HL79CB3ODZ",
+        "avatarImage": "https://comic-assets.s3.ap-south-1.amazonaws.com/7_char_assets/scene01001_1.png",
+        "extraPrompt": "HL79CB3ODZ boy",
+    },
+    {
+        "avatarName": "chutki",
+        "codename": "SJ7JVIS9M7",
+        "avatarImage": "https://comic-assets.s3.ap-south-1.amazonaws.com/7_char_assets/14.png",
+        "extraPrompt": "SJ7JVIS9M7 girl",
+    },
+]
+
+webhook_url = (
+    "https://hooks.slack.com/services/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"
+)
+error_webhook = (
+    "https://hooks.slack.com/services/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"
+)
+
+characterSheets = [
+    "character+sheets/1.1.png",
+    "character+sheets/10.1.png",
+    "character+sheets/11.1.png",
+    "character+sheets/12.1.png",
+    "character+sheets/13.1.png",
+    "character+sheets/14.1.png",
+    "character+sheets/16.1.png",
+    "character+sheets/17.1.png",
+    "character+sheets/18.1.png",
+    "character+sheets/19.1.png",
+    "character+sheets/2.1.png",
+    "character+sheets/20.1.png",
+    "character+sheets/21.1.png",
+    "character+sheets/22.1.png",
+    "character+sheets/23.1.png",
+    "character+sheets/24.1.png",
+    "character+sheets/25.1.png",
+    "character+sheets/26.1.png",
+    "character+sheets/27.1.png",
+    "character+sheets/28.1.png",
+    "character+sheets/29.1.png",
+    "character+sheets/3.1.png",
+    "character+sheets/30.1.png",
+    "character+sheets/31.1.png",
+    "character+sheets/32.1.png",
+    "character+sheets/33.1.png",
+    "character+sheets/34.1.png",
+    "character+sheets/35.1.png",
+    "character+sheets/36.1.png",
+    "character+sheets/38.1.png",
+    "character+sheets/39.1.png",
+    "character+sheets/4.1.png",
+    "character+sheets/40.1.png",
+    "character+sheets/42.1.png",
+    "character+sheets/43.1.png",
+    "character+sheets/44.1.png",
+    "character+sheets/45.1.png",
+    "character+sheets/46.1.png",
+    "character+sheets/47.1.png",
+    "character+sheets/48.1.png",
+    "character+sheets/49.1.png",
+    "character+sheets/5.1.png",
+    "character+sheets/50.1.png",
+    "character+sheets/51.1.png",
+    "character+sheets/52.1.png",
+    "character+sheets/53.1.png",
+    "character+sheets/54.1.png",
+    "character+sheets/55.1.png",
+    "character+sheets/56.1.png",
+    "character+sheets/57.1.png",
+    "character+sheets/58.1.png",
+    "character+sheets/59.1.png",
+    "character+sheets/60.1.png",
+    "character+sheets/61.1.png",
+    "character+sheets/62.1.png",
+    "character+sheets/63.1.png",
+    "character+sheets/64.1.png",
+    "character+sheets/65.1.png",
+    "character+sheets/66.1.png",
+    "character+sheets/7.1.png",
+    "character+sheets/8.1.png",
+    "character+sheets/9.1.png",
+]
+
+
+def add_code_names(sentence):
+    array_of_objects = Avatar
+
+    for obj in array_of_objects:
+        sentence = (
+            re.sub(
+                r"\b" + obj["avatarName"] + r"\b",
+                obj["extraPrompt"],
+                sentence,
+                flags=re.IGNORECASE,
+            )
+            + " "
+        )
+    print(sentence)
+    return sentence
+
+
+def upload_images(images, processName: str, taskId: str):
+    imageUrls = []
+    for i, image in enumerate(images):
+        # img_io = BytesIO()
+        # image.save(img_io, "JPEG", quality=70)
+        # img_io.seek(0)
+        # key = "crecoAI/{}{}_{}.png".format(taskId, processName, i)
+        # t = s3.put_object(
+        #     Bucket="comic-assets", Key=key, Body=img_io.getvalue(), ACL="public-read"
+        # )
+        # print("uploading done to s3", key, t)
+        imageUrls.append(
+            "https://comic-assets.s3.ap-south-1.amazonaws.com/crecoAI/{}{}_{}.png".format(
+                taskId, processName, i
+            )
+        )
+
+    print({"promptImages": imageUrls})
+
+    return imageUrls
+
+
+# def upload_image(image: Union[Image.Image, BytesIO], out_path):
+    # if type(image) is Image.Image:
+    #     buffer = io.BytesIO()
+    #     image.save(buffer, format="PNG")
+    #     image = buffer
+
+    # image.seek(0)
+    # s3.upload_fileobj(image, "comic-assets", out_path, ExtraArgs={"ACL": "public-read"})
+    # image.close()
+
+
+def download_image(url) -> Image.Image:
+    response = requests.get(url)
+    return Image.open(BytesIO(response.content)).convert("RGB")
+
+
+def pickPoses():
+    random_images = random.sample(characterSheets, 4)
+    poses = []
+    prefix = "https://comic-assets.s3.ap-south-1.amazonaws.com/"
+
+    # Use list comprehension to add prefix to all elements in the array
+    random_images_with_prefix = [prefix + img for img in random_images]
+
+    print(random_images_with_prefix)
+    for imageUrl in random_images_with_prefix:
+        # Download and resize the image
+        init_image = download_image(imageUrl).resize((512, 512))
+
+        # Open the pose image
+        imageUrlPose = imageUrl
+        # print(imageUrl)
+        input_image_bytes = read_url(imageUrlPose)
+        # print(input_image_bytes)
+        pose_image = Image.open(io.BytesIO(input_image_bytes)).convert("RGB")
+        # print(pose_image)
+        pose_image = pose_image.resize((512, 512))
+        # print(pose_image)
+        # Append the result to the poses array
+        poses.append(pose_image)
+
+    return poses
+
+
+def construct_default_s3_url(key):
+    return "https://comic-assets.s3.ap-south-1.amazonaws.com/" + key
+
+
+def read_url(url: str):
+    with urllib.request.urlopen(url) as u:
+        return u.read()
+
+
+def disable_safety_checker(pipe):
+    def dummy(images, **kwargs):
+        return images, False
+
+    pipe.safety_checker = None

util/lora_style.py

@@ -0,0 +1,92 @@
+import os
+from typing import Any, Dict, Union
+
+from lora_diffusion import patch_pipe, tune_lora_scale
+
+
+class LoraStyle:
+    class LoraPatcher:
+        def __init__(self, pipe, style: Dict[str, Any]):
+            self.__style = style
+            self.pipe = pipe
+
+        def patch(self):
+            patch_pipe(self.pipe, self.__style["path"])
+            tune_lora_scale(self.pipe.unet, self.__style["weight"])
+            tune_lora_scale(self.pipe.text_encoder, self.__style["weight"])
+
+        def kwargs(self):
+            return {}
+
+        def cleanup(self):
+            tune_lora_scale(self.pipe.unet, 0.0)
+            tune_lora_scale(self.pipe.text_encoder, 0.0)
+            pass
+
+    class EmptyLoraPatcher:
+        def patch(self):
+            pass
+
+        def kwargs(self):
+            return {}
+
+        def cleanup(self):
+            pass
+
+    def load(self, model_dir: str):
+        self.__styles = {
+            "nq6akX1CIp": {
+                "path": model_dir + "/laur_style/nq6akX1CIp/final_lora.safetensors",
+                "weight": 0.5,
+                "negativePrompt": [""],
+                "type": "custom",
+            },
+            "ghibli": {
+                "path": model_dir + "/laur_style/nq6akX1CIp/ghibli.bin",
+                "weight": 1,
+                "negativePrompt": [""],
+                "type": "custom",
+            },
+            "eQAmnK2kB2": {
+                "path": model_dir + "/laur_style/eQAmnK2kB2/final_lora.safetensors",
+                "weight": 0.5,
+                "negativePrompt": [""],
+                "type": "custom",
+            },
+            "to8contrast": {
+                "path": model_dir + "/laur_style/rpjgusOgqD/final_lora.bin",
+                "weight": 0.5,
+                "negativePrompt": [""],
+                "type": "custom",
+            },
+            "jim lee": {
+                "path": model_dir + "/laur_style/e2j9mz0jqj/final_lora.bin",
+                "weight": 0.8,
+                "negativePrompt": [""],
+                "type": "custom",
+            },
+        }
+        self.__verify()
+
+    def prepend_style_to_prompt(self, prompt: str, key: str) -> str:
+        if key in self.__styles:
+            return f"{key} style {prompt}"
+        return prompt
+
+    def get_patcher(self, pipe, key: str) -> Union[LoraPatcher, EmptyLoraPatcher]:
+        if key in self.__styles:
+            style = self.__styles[key]
+            return self.LoraPatcher(pipe, style)
+        return self.EmptyLoraPatcher()
+
+    def __verify(self):
+        "A method to verify if lora exists within the required path otherwise throw error"
+
+        for item in self.__styles.keys():
+            if not os.path.exists(self.__styles[item]["path"]):
+                raise Exception(
+                    "Lora style model "
+                    + item
+                    + " not found at path: "
+                    + self.__styles[item]["path"]
+                )

util/slack.py

@@ -0,0 +1,55 @@
+from time import sleep
+from typing import Optional
+
+import requests
+from data.task import Task
+
+
+class Slack:
+    def __init__(self):
+        # self.webhook_url = "https://hooks.slack.com/services/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"
+        self.webhook_url = "https://hooks.slack.com/services/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"
+        self.error_webhook = "https://hooks.slack.com/services/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"
+
+    def send_alert(self, task: Task, args: Optional[dict]):
+        raw = task.get_raw().copy()
+
+        raw.pop("queue_name", None)
+        raw.pop("attempt", None)
+        raw.pop("timestamp", None)
+        raw.pop("task_id", None)
+        raw.pop("maskImageUrl", None)
+
+        if args is not None:
+            raw.update(args.items())
+
+        message = ""
+        for key, value in raw.items():
+            if value:
+                if type(value) == list:
+                    message += f"*{key}*: {', '.join(value)}\n"
+                else:
+                    message += f"*{key}*: {value}\n"
+
+        requests.post(
+            self.webhook_url,
+            headers={"Content-Type": "application/json"},
+            json={"text": message},
+        )
+
+    def error_alert(self, task: Task, e: Exception):
+        requests.post(
+            self.error_webhook,
+            headers={"Content-Type": "application/json"},
+            json={
+                "text": "Task failed:\n{} \n error is: \n {}".format(task.get_raw(), e)
+            },
+        )
+
+    def auto_send_alert(self, func):
+        def inner(*args, **kwargs):
+            rargs = func(*args, **kwargs)
+            self.send_alert(args[0], rargs)
+            return rargs
+
+        return inner