Commit
•
d280d4b
1
Parent(s):
61bb2bf
Update handler.py
Browse files- handler.py +54 -0
handler.py
CHANGED
@@ -0,0 +1,54 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from typing import Dict, List, Any
|
2 |
+
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
|
3 |
+
import torch
|
4 |
+
import os
|
5 |
+
from huggingface_hub import HfApi
|
6 |
+
from pathlib import Path
|
7 |
+
from diffusers.utils import load_image
|
8 |
+
from PIL import Image
|
9 |
+
import numpy as np
|
10 |
+
from controlnet_aux import PidiNetDetector, HEDdetector
|
11 |
+
from diffusers import (
|
12 |
+
ControlNetModel,
|
13 |
+
StableDiffusionControlNetPipeline,
|
14 |
+
UniPCMultistepScheduler,
|
15 |
+
)
|
16 |
+
from io import BytesIO
|
17 |
+
import base64
|
18 |
+
|
19 |
+
checkpoint = "lllyasviel/control_v11p_sd15_scribble"
|
20 |
+
|
21 |
+
|
22 |
+
class EndpointHandler:
|
23 |
+
def __init__(self, path=""):
|
24 |
+
# load model and processor from path
|
25 |
+
self.model = ControlNetModel.from_pretrained(checkpoint, torch_dtype=torch.float16)
|
26 |
+
self.processor = HEDdetector.from_pretrained('lllyasviel/Annotators')
|
27 |
+
)
|
28 |
+
|
29 |
+
def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
|
30 |
+
"""
|
31 |
+
Args:
|
32 |
+
data (:dict:):
|
33 |
+
The payload with the text prompt and generation parameters.
|
34 |
+
"""
|
35 |
+
# process input
|
36 |
+
image_base64 = data.pop("image_base64", None)
|
37 |
+
prompt = data.pop("prompt", None)
|
38 |
+
|
39 |
+
# preprocess
|
40 |
+
image = Image.open(BytesIO(base64.b64decode(image_base64)))
|
41 |
+
control_image = self.processor(image, scribble=True)
|
42 |
+
pipe = StableDiffusionControlNetPipeline.from_pretrained(
|
43 |
+
"runwayml/stable-diffusion-v1-5", controlnet=self.model, torch_dtype=torch.float16
|
44 |
+
)
|
45 |
+
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
|
46 |
+
pipe.enable_model_cpu_offload()
|
47 |
+
|
48 |
+
generator = torch.manual_seed(0)
|
49 |
+
image = pipe(prompt, num_inference_steps=30, generator=generator, image=control_image).images[0]
|
50 |
+
|
51 |
+
# postprocess the prediction
|
52 |
+
res_base64 = base64.b64encode(image)
|
53 |
+
|
54 |
+
return [{"result": res_base64}]
|