Update handler.py

handler.py

@@ -0,0 +1,54 @@
+from typing import Dict, List, Any
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+import torch
+import os
+from huggingface_hub import HfApi
+from pathlib import Path
+from diffusers.utils import load_image
+from PIL import Image
+import numpy as np
+from controlnet_aux import PidiNetDetector, HEDdetector
+from diffusers import (
+    ControlNetModel,
+    StableDiffusionControlNetPipeline,
+    UniPCMultistepScheduler,
+)
+from io import BytesIO
+import base64
+
+checkpoint = "lllyasviel/control_v11p_sd15_scribble"
+
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        # load model and processor from path
+        self.model =  ControlNetModel.from_pretrained(checkpoint, torch_dtype=torch.float16)
+        self.processor = HEDdetector.from_pretrained('lllyasviel/Annotators')
+)
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
+        """
+        Args:
+            data (:dict:):
+                The payload with the text prompt and generation parameters.
+        """
+        # process input
+        image_base64 = data.pop("image_base64", None)
+        prompt = data.pop("prompt", None)
+
+        # preprocess
+        image = Image.open(BytesIO(base64.b64decode(image_base64)))
+        control_image = self.processor(image, scribble=True)
+        pipe = StableDiffusionControlNetPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5", controlnet=self.model, torch_dtype=torch.float16
+)
+        pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+        pipe.enable_model_cpu_offload()
+
+        generator = torch.manual_seed(0)
+        image = pipe(prompt, num_inference_steps=30, generator=generator, image=control_image).images[0]
+
+        # postprocess the prediction
+        res_base64 = base64.b64encode(image)
+
+        return [{"result": res_base64}]