Upload handler and requirements

handler.py

@@ -0,0 +1,87 @@
+from typing import Dict, List, Any
+import torch
+import base64
+from PIL import Image
+from io import BytesIO
+from diffusers import T2IAdapter, StableDiffusionXLAdapterPipeline, AutoencoderKL
+from controlnet_aux.pidi import PidiNetDetector
+
+# set device
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+if device.type != 'cuda':
+    raise ValueError("need to run on GPU")
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        # Preload all the elements you are going to need at inference.
+        # pseudo:
+        # self.model= load_model(path)
+
+        adapter = T2IAdapter.from_pretrained(
+            "Adapter/t2iadapter",
+            subfolder="sketch_sdxl_1.0",
+            torch_dtype=torch.float16,
+            adapter_type="full_adapter_xl"
+        )
+
+        vae = AutoencoderKL.from_pretrained(
+            "madebyollin/sdxl-vae-fp16-fix",
+            torch_dtype=torch.float16,
+            use_safetensors=True
+        )
+
+        self.pipeline = StableDiffusionXLAdapterPipeline.from_pretrained(
+            "stabilityai/stable-diffusion-xl-base-1.0",
+            adapter=adapter,
+            vae=vae,
+            torch_dtype=torch.float16,
+            variant="fp16"
+        ).to("cuda")
+        self.pipeline.enable_sequential_cpu_offload()
+
+        self.pidinet = PidiNetDetector.from_pretrained("lllyasviel/Annotators").to("cuda")
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        data args:
+            inputs (:obj: `str` | `PIL.Image` | `np.array`)
+            kwargs
+        Return:
+            A :obj:`list` | `dict`: will be serialized and returned
+        """
+
+        # pseudo
+        # self.model(input)
+
+        # get inputs
+        inputs = data.pop("inputs", "")
+        encoded_image = data.pop("image", None)
+
+        # Decode image and convert to black and white sketch
+        decoded_image = self.decode_base64_image(encoded_image).convert('RGB')
+        sketch_image = self.pidinet(
+            decoded_image,
+            detect_resolution=1024,
+            image_resolution=1024,
+            apply_filter=True
+        ).convert('L')
+
+        # sketch_image.save("./output1.png")
+
+        output_image = self.pipeline(
+            prompt=inputs,
+            negative_prompt="extra digit, fewer digits, cropped, worst quality, low quality",
+            image=sketch_image,
+            guidance_scale=7.5,
+        ).images[0]
+
+        # output_image.save("./output2.png")
+        return output_image
+
+    # helper to decode input image
+    def decode_base64_image(self, image_string):
+        base64_image = base64.b64decode(image_string)
+        buffer = BytesIO(base64_image)
+        image = Image.open(buffer)
+        return image

requirements.txt

@@ -0,0 +1 @@
+controlnet-aux