Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1,6 @@
+*.pyc
+.ipynb_checkpoints                            █
+env
+test.py
+*.jpeg
+__pycache__

inference.py

@@ -8,11 +8,12 @@ from internals.pipelines.commons import Img2Img, Text2Img
 from internals.pipelines.controlnets import ControlNet
 from internals.pipelines.img_classifier import ImageClassifier
 from internals.pipelines.img_to_text import Image2Text
+from internals.pipelines.inpainter import InPainter
 from internals.pipelines.prompt_modifier import PromptModifier
 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
+from internals.util.cache import auto_clear_cuda_and_gc, clear_cuda
 from internals.util.commons import pickPoses, upload_image, upload_images
 from internals.util.config import set_configs_from_task, set_root_dir
 from internals.util.failure_hander import FailureHandler
@@ -26,6 +27,7 @@ num_return_sequences = 4  # the number of results to generate
 auto_mode = False
 
 prompt_modifier = PromptModifier(num_of_sequences=num_return_sequences)
+inpainter = InPainter()
 img2text = Image2Text()
 img_classifier = ImageClassifier()
 controlnet = ControlNet()
@@ -269,6 +271,29 @@ def img2img(task: Task):
     }
 
 
+@update_db
+@slack.auto_send_alert
+def inpaint(task: Task):
+    prompt, _ = get_patched_prompt(task)
+
+    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}
+
+
 def model_fn(model_dir):
     print("Logs: model loaded .... starts")
 
@@ -288,6 +313,7 @@ def model_fn(model_dir):
     controlnet.load(model_dir)
     text2img_pipe.load(model_dir)
     img2img_pipe.create(text2img_pipe)
+    inpainter.create(text2img_pipe)
 
     safety_checker.apply(text2img_pipe)
     safety_checker.apply(img2img_pipe)
@@ -333,6 +359,8 @@ def predict_fn(data, pipe):
             return pose(task)
         elif task_type == TaskType.TILE_UPSCALE:
             return tile_upscale(task)
+        elif task_type == TaskType.INPAINT:
+            return inpaint(task)
         else:
             raise Exception("Invalid task type")
     except Exception as e:

internals/pipelines/inpainter.py

@@ -15,6 +15,12 @@ class InPainter(AbstractPipeline):
         ).to("cuda")
         disable_safety_checker(self.pipe)
 
+    def create(self, pipeline: AbstractPipeline):
+        self.pipe = StableDiffusionInpaintPipeline(**pipeline.pipe.components).to(
+            "cuda"
+        )
+        disable_safety_checker(self.pipe)
+
     @torch.inference_mode()
     def process(
         self,