al-livefront
commited on
Commit
•
f040cb4
1
Parent(s):
e80d746
Read inputs from base64
Browse files- handler.py +8 -2
handler.py
CHANGED
@@ -1,6 +1,9 @@
|
|
1 |
from typing import Dict, List, Any
|
2 |
from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler
|
3 |
import torch
|
|
|
|
|
|
|
4 |
|
5 |
class EndpointHandler():
|
6 |
def __init__(self, path=""):
|
@@ -10,8 +13,11 @@ class EndpointHandler():
|
|
10 |
self.pipeline.to("cuda")
|
11 |
|
12 |
def __call__(self, data: Dict[str, any]) -> List[Dict[str, Any]]:
|
13 |
-
|
14 |
-
|
|
|
|
|
|
|
15 |
|
16 |
prompt = "a picture of a person with a nice haircut"
|
17 |
|
|
|
1 |
from typing import Dict, List, Any
|
2 |
from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler
|
3 |
import torch
|
4 |
+
import base64
|
5 |
+
from PIL import Image
|
6 |
+
import io
|
7 |
|
8 |
class EndpointHandler():
|
9 |
def __init__(self, path=""):
|
|
|
13 |
self.pipeline.to("cuda")
|
14 |
|
15 |
def __call__(self, data: Dict[str, any]) -> List[Dict[str, Any]]:
|
16 |
+
mask_bytes = base64.b64decode(data["inputs"]["mask"])
|
17 |
+
mask = Image.open(io.BytesIO(mask_bytes)).convert('RGB')
|
18 |
+
|
19 |
+
image_bytes = base64.b64decode(data["inputs"]["image"])
|
20 |
+
image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
|
21 |
|
22 |
prompt = "a picture of a person with a nice haircut"
|
23 |
|