julien-c HF staff commited on
Commit
ad11739
1 Parent(s): 8f663d4

Copy template from https://huggingface.co/philschmid/stable-diffusion-v1-4-endpoints/blob/main/handler.py

Browse files
Files changed (2) hide show
  1. .gitignore +1 -0
  2. handler.py +42 -0
.gitignore ADDED
@@ -0,0 +1 @@
 
 
1
+ .env/
handler.py ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List, Any
2
+ import torch
3
+ from torch import autocast
4
+ from diffusers import StableDiffusionPipeline
5
+ import base64
6
+ from io import BytesIO
7
+
8
+
9
+ # set device
10
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
11
+
12
+ if device.type != 'cuda':
13
+ raise ValueError("need to run on GPU")
14
+
15
+ class EndpointHandler():
16
+ def __init__(self, path=""):
17
+ # load the optimized model
18
+ self.pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16)
19
+ self.pipe = self.pipe.to(device)
20
+
21
+
22
+ def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
23
+ """
24
+ Args:
25
+ data (:obj:):
26
+ includes the input data and the parameters for the inference.
27
+ Return:
28
+ A :obj:`dict`:. base64 encoded image
29
+ """
30
+ inputs = data.pop("inputs", data)
31
+
32
+ # run inference pipeline
33
+ with autocast(device.type):
34
+ image = self.pipe(inputs, guidance_scale=7.5)["sample"][0]
35
+
36
+ # encode image as base 64
37
+ buffered = BytesIO()
38
+ image.save(buffered, format="JPEG")
39
+ img_str = base64.b64encode(buffered.getvalue())
40
+
41
+ # postprocess the prediction
42
+ return {"image": img_str.decode()}