eykarim commited on
Commit
136fb9a
·
1 Parent(s): 112093a

Delete handler.py

Browse files
Files changed (1) hide show
  1. handler.py +0 -42
handler.py DELETED
@@ -1,42 +0,0 @@
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()}