Upload folder using huggingface_hub

inference.py

@@ -19,6 +19,7 @@ from internals.pipelines.pose_detector import PoseDetector
 from internals.pipelines.prompt_modifier import PromptModifier
 from internals.pipelines.replace_background import ReplaceBackground
 from internals.pipelines.safety_checker import SafetyChecker
+from internals.pipelines.sdxl_tile_upscale import SDXLTileUpscaler
 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
@@ -55,6 +56,8 @@ img2img_pipe = Img2Img()
 safety_checker = SafetyChecker()
 slack = Slack()
 avatar = Avatar()
+sdxl_tileupscaler = SDXLTileUpscaler()
+
 
 custom_scripts: List = []
 
@@ -145,28 +148,42 @@ def tile_upscale(task: Task):
 
     prompt = get_patched_prompt_tile_upscale(task)
 
-    controlnet.load_model("tile_upscaler")
+    if get_is_sdxl():
+        lora_patcher = lora_style.get_patcher(sdxl_tileupscaler.pipe, task.get_style())
+        lora_patcher.patch()
+
+        images, has_nsfw = sdxl_tileupscaler.process(
+            prompt=prompt,
+            imageUrl=task.get_imageUrl(),
+            resize_dimension=task.get_resize_dimension(),
+            negative_prompt=task.get_negative_prompt(),
+            width=task.get_width(),
+            height=task.get_height(),
+        )
 
-    lora_patcher = lora_style.get_patcher(controlnet.pipe, task.get_style())
-    lora_patcher.patch()
+        lora_patcher.cleanup()
+    else:
+        controlnet.load_model("tile_upscaler")
 
-    kwargs = {
-        "imageUrl": task.get_imageUrl(),
-        "seed": task.get_seed(),
-        "num_inference_steps": task.get_steps(),
-        "negative_prompt": task.get_negative_prompt(),
-        "width": task.get_width(),
-        "height": task.get_height(),
-        "prompt": prompt,
-        "resize_dimension": task.get_resize_dimension(),
-        **task.cnt_kwargs(),
-    }
-    images, has_nsfw = controlnet.process(**kwargs)
+        lora_patcher = lora_style.get_patcher(controlnet.pipe, task.get_style())
+        lora_patcher.patch()
 
-    generated_image_url = upload_image(images[0], output_key)
+        kwargs = {
+            "imageUrl": task.get_imageUrl(),
+            "seed": task.get_seed(),
+            "num_inference_steps": task.get_steps(),
+            "negative_prompt": task.get_negative_prompt(),
+            "width": task.get_width(),
+            "height": task.get_height(),
+            "prompt": prompt,
+            "resize_dimension": task.get_resize_dimension(),
+            **task.cnt_kwargs(),
+        }
+        images, has_nsfw = controlnet.process(**kwargs)
+        lora_patcher.cleanup()
+        controlnet.cleanup()
 
-    lora_patcher.cleanup()
-    controlnet.cleanup()
+    generated_image_url = upload_image(images[0], output_key)
 
     return {
         "modified_prompts": prompt,
@@ -582,7 +599,10 @@ def load_model_by_task(task: Task):
         replace_background.load(base=text2img_pipe, high_res=high_res)
     else:
         if task.get_type() == TaskType.TILE_UPSCALE:
-            controlnet.load_model("tile_upscaler")
+            if get_is_sdxl():
+                sdxl_tileupscaler.create(text2img_pipe)
+            else:
+                controlnet.load_model("tile_upscaler")
         elif task.get_type() == TaskType.CANNY:
             controlnet.load_model("canny")
         elif task.get_type() == TaskType.SCRIBBLE:

internals/pipelines/commons.py

@@ -3,15 +3,16 @@ from typing import Any, Callable, Dict, List, Optional, Union
 
 import torch
 from diffusers import (
+    AutoencoderKL,
     StableDiffusionImg2ImgPipeline,
-    StableDiffusionXLPipeline,
     StableDiffusionXLImg2ImgPipeline,
+    StableDiffusionXLPipeline,
 )
 
 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, get_is_sdxl
+from internals.util.config import get_hf_token, get_is_sdxl, num_return_sequences
 
 
 class AbstractPipeline:
@@ -32,12 +33,18 @@ class Text2Img(AbstractPipeline):
 
     def load(self, model_dir: str):
         if get_is_sdxl():
-            self.pipe = StableDiffusionXLPipeline.from_pretrained(
+            vae = AutoencoderKL.from_pretrained(
+                "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
+            )
+            pipe = StableDiffusionXLPipeline.from_pretrained(
                 model_dir,
                 torch_dtype=torch.float16,
                 use_auth_token=get_hf_token(),
                 use_safetensors=True,
-            ).to("cuda")
+            )
+            pipe.vae = vae
+            pipe.to("cuda")
+            self.pipe = pipe
         else:
             self.pipe = two_step_pipeline.from_pretrained(
                 model_dir, torch_dtype=torch.float16, use_auth_token=get_hf_token()

internals/pipelines/sdxl_tile_upscale.py

@@ -0,0 +1,87 @@
+import torch
+from diffusers import ControlNetModel
+from PIL import Image
+from torchvision import transforms
+
+import internals.util.image as ImageUtils
+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.util.commons import download_image
+from internals.util.config import get_base_dimension
+
+controlnet = ControlNet()
+
+
+class SDXLTileUpscaler(AbstractPipeline):
+    def create(self, pipeline: Text2Img):
+        controlnet = ControlNetModel.from_pretrained(
+            "thibaud/controlnet-openpose-sdxl-1.0", torch_dtype=torch.float16
+        )
+        pipe = DemoFusionSDXLControlNetPipeline(
+            **pipeline.pipe.components, controlnet=controlnet
+        )
+        pipe = pipe.to("cuda")
+        pipe.enable_vae_tiling()
+        pipe.enable_vae_slicing()
+        pipe.enable_xformers_memory_efficient_attention()
+
+        self.pipe = pipe
+
+    def process(
+        self,
+        prompt: str,
+        imageUrl: str,
+        resize_dimension: int,
+        negative_prompt: str,
+        width: int,
+        height: int,
+    ):
+        pose_image = controlnet.detect_pose(imageUrl)
+        img = download_image(imageUrl).resize((width, height))
+
+        img = ImageUtils.resize_image(img, get_base_dimension())
+        pose_image = pose_image.resize(img.size)
+
+        img2 = self.__resize_for_condition_image(img, resize_dimension)
+
+        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=pose_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):
+        transform = transforms.Compose(
+            [
+                transforms.Resize((1024, 1024)),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+            ]
+        )
+        image = transform(pil_image)
+        image = image.unsqueeze(0).half()
+        image = image.to("cuda")
+        return image
+
+    def __resize_for_condition_image(self, image: Image.Image, resolution: int):
+        input_image = image.convert("RGB")
+        W, H = input_image.size
+        k = float(resolution) / max(W, H)
+        H *= k
+        W *= k
+        H = int(round(H / 64.0)) * 64
+        W = int(round(W / 64.0)) * 64
+        img = input_image.resize((W, H), resample=Image.LANCZOS)
+        return img