Upload folder using huggingface_hub

external/scripts/__init__.py

@@ -0,0 +1,13 @@
+import importlib
+import os
+
+from internals.util import getcwd
+
+path = os.path.join(getcwd(), "external/scripts")
+
+__scripts__ = []
+for name in os.listdir(path):
+    name = name.split("/")[-1].replace(".py", "")
+    imp = importlib.import_module(f"external.scripts.{name}")
+    if hasattr(imp, "Script") and imp not in __scripts__:
+        __scripts__.append(imp)

external/scripts/day_night_ip2p.py

@@ -0,0 +1,63 @@
+import torch
+from diffusers import StableDiffusionInstructPix2PixPipeline
+
+import internals.util.image as ImageUtil
+from internals.data.dataAccessor import update_db
+from internals.data.task import Task
+from internals.util.cache import clear_cuda_and_gc
+from internals.util.commons import download_image, upload_images
+from internals.util.config import get_hf_token
+from internals.util.slack import Slack
+
+slack = Slack()
+
+
+class Script:
+    def __init__(self, **kwargs):
+        self.__name__ = "day_night_ip2p"
+
+    @update_db
+    @slack.auto_send_alert
+    def __call__(self, task: Task, args: dict):
+        clear_cuda_and_gc()
+
+        model_id = args.get("model_id", None)
+        steps = args.get("steps", 50)
+        image_guidance_scale = args.get("image_guidance_scale", 1.5)
+        guidance_scale = args.get("guidance_scale", 7.5)
+
+        pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
+            model_id,
+            use_auth_token=get_hf_token(),
+            torch_dtype=torch.float16,
+            safety_checker=None,
+        ).to("cuda")
+        pipe.enable_xformers_memory_efficient_attention()
+
+        prompt = ["convert to night", "convert to evening", "convert to midnight"]
+        image = download_image(task.get_imageUrl())
+        image = ImageUtil.resize_image(image, 1024)
+
+        images = []
+        for p in prompt:
+            print("Generating: ", p)
+            image = pipe.__call__(
+                prompt=p,
+                num_inference_steps=steps,
+                image=image,
+                guidance_scale=guidance_scale,
+                num_images_per_prompt=1,
+                image_guidance_scale=image_guidance_scale,
+            ).images[0]
+            images.append(image)
+
+        generated_image_urls = upload_images(
+            images, "_" + self.__name__, task.get_taskId()
+        )
+
+        pipe = None
+        del pipe
+
+        clear_cuda_and_gc()
+
+        return {"generated_image_urls": generated_image_urls}

inference.py

@@ -2,7 +2,9 @@ import os
 import traceback
 from typing import List, Optional
 
+import pydash as _
 import torch
+from numpy import who
 
 import internals.util.prompt as prompt_util
 from internals.data.dataAccessor import update_db, update_db_source_failed
@@ -54,6 +56,8 @@ safety_checker = SafetyChecker()
 slack = Slack()
 avatar = Avatar()
 
+custom_scripts: List = []
+
 
 def get_patched_prompt(task: Task):
     return prompt_util.get_patched_prompt(task, avatar, lora_style, prompt_modifier)
@@ -533,6 +537,32 @@ def replace_bg(task: Task):
     }
 
 
+def custom_action(task: Task):
+    from external.scripts import __scripts__
+
+    global custom_scripts
+    kwargs = {
+        "CONTROLNET": controlnet,
+        "LORASTYLE": lora_style,
+    }
+
+    torch.manual_seed(task.get_seed())
+
+    for script in __scripts__:
+        script = script.Script(**kwargs)
+        existing_script = _.find(
+            custom_scripts, lambda x: x.__name__ == script.__name__
+        )
+        if existing_script:
+            script = existing_script
+        else:
+            custom_scripts.append(script)
+
+        data = task.get_action_data()
+        if data["name"] == script.__name__:
+            return script(task, data)
+
+
 def load_model_by_task(task: Task):
     if not text2img_pipe.is_loaded():
         text2img_pipe.load(get_model_dir())
@@ -587,6 +617,8 @@ def predict_fn(data, pipe):
     task = Task(data)
     print("task is ", data)
 
+    clear_cuda_and_gc()
+
     FailureHandler.handle(task)
 
     try:
@@ -629,6 +661,8 @@ def predict_fn(data, pipe):
             return linearart(task)
         elif task_type == TaskType.REPLACE_BG:
             return replace_bg(task)
+        elif task_type == TaskType.CUSTOM_ACTION:
+            return custom_action(task)
         elif task_type == TaskType.SYSTEM_CMD:
             os.system(task.get_prompt())
         else:

internals/data/dataAccessor.py

@@ -1,9 +1,9 @@
 import traceback
 from typing import Dict, List, Optional
 
-from requests.adapters import Retry, HTTPAdapter
 import requests
 from pydash import includes
+from requests.adapters import HTTPAdapter, Retry
 
 from internals.data.task import Task
 from internals.util.config import api_endpoint, api_headers
@@ -104,9 +104,13 @@ def update_db_source_failed(sourceId, userId):
 
 def update_db(func):
     def caller(*args, **kwargs):
-        if type(args[0]) is not Task:
+        task = None
+        for arg in args:
+            if type(arg) is Task:
+                task = arg
+                break
+        if task is None:
             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)

internals/data/task.py

@@ -18,6 +18,7 @@ class TaskType(Enum):
     SCRIBBLE = "SCRIBBLE"
     LINEARART = "LINEARART"
     REPLACE_BG = "REPLACE_BG"
+    CUSTOM_ACTION = "CUSTOM_ACTION"
     SYSTEM_CMD = "SYSTEM_CMD"
 
 
@@ -148,6 +149,10 @@ class Task:
     def get_base_dimension(self):
         return self.__data.get("base_dimension", None)
 
+    def get_action_data(self) -> dict:
+        "If task_type is CUSTOM_ACTION, then this will return the action data with 'name' as key"
+        return self.__data.get("action_data", {})
+
     def get_raw(self) -> dict:
         return self.__data.copy()
 

internals/pipelines/controlnets.py

@@ -151,7 +151,6 @@ class ControlNet(AbstractPipeline):
 
         self.__load_pipeline(model, pipeline_type)
 
-        self.network_model = model
         self.__current_task_name = task_name
 
         clear_cuda_and_gc()
@@ -247,6 +246,11 @@ class ControlNet(AbstractPipeline):
             if hasattr(self, "pipe2"):
                 setattr(self.pipe2, "adapter", network_model)
 
+        if hasattr(self, "pipe"):
+            self.pipe = self.pipe.to("cuda")
+        if hasattr(self, "pipe2"):
+            self.pipe2 = self.pipe2.to("cuda")
+
         clear_cuda_and_gc()
 
     def process(self, **kwargs):

internals/pipelines/high_res.py

@@ -5,7 +5,8 @@ from PIL import Image
 
 from internals.data.result import Result
 from internals.pipelines.commons import AbstractPipeline, Img2Img
-from internals.util.config import get_model_dir, get_base_dimension
+from internals.util.cache import clear_cuda_and_gc
+from internals.util.config import get_base_dimension, get_model_dir
 
 
 class HighRes(AbstractPipeline):
@@ -32,6 +33,8 @@ class HighRes(AbstractPipeline):
         guidance_scale: int = 9,
         **kwargs,
     ):
+        clear_cuda_and_gc()
+
         images = [image.resize((width, height)) for image in images]
         kwargs = {
             "prompt": prompt,

internals/util/__init__.py

@@ -0,0 +1,7 @@
+import os
+
+from internals.util.config import get_root_dir
+
+
+def getcwd():
+    return get_root_dir()

internals/util/lora_style.py

@@ -10,8 +10,8 @@ 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
+from internals.util.config import get_is_sdxl
 
 
 class LoraStyle:
@@ -113,9 +113,6 @@ 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]

internals/util/slack.py

@@ -55,7 +55,12 @@ class Slack:
     def auto_send_alert(self, func):
         def inner(*args, **kwargs):
             rargs = func(*args, **kwargs)
-            self.send_alert(args[0], rargs)
+            task = Task({})
+            for arg in args:
+                if type(arg) is Task:
+                    task = arg
+                    break
+            self.send_alert(task, rargs)
             return rargs
 
         return inner