dmingod commited on
Commit
6a8f037
·
1 Parent(s): e8cc959
Files changed (1) hide show
  1. handler.py +7 -6
handler.py CHANGED
@@ -4,17 +4,18 @@ import torch
4
  import base64
5
  from io import BytesIO
6
  from PIL import Image
7
- from diffusers import StableDiffusionXLImg2ImgPipeline
8
- from diffusers.utils import load_image
9
  import numpy as np
 
 
10
 
11
 
12
  class EndpointHandler():
13
- def __init__(self, path=""):
14
- self.pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16")
15
  self.pipe.to("cuda")
16
- self.pipe.upcast_vae()
17
-
18
 
19
  def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
20
  """
 
4
  import base64
5
  from io import BytesIO
6
  from PIL import Image
7
+ # from diffusers import StableDiffusionXLImg2ImgPipeline
8
+ # from diffusers.utils import load_image
9
  import numpy as np
10
+ from diffusers import AutoPipelineForImage2Image
11
+ from diffusers.utils import load_image
12
 
13
 
14
  class EndpointHandler():
15
+ def __init__(self, path=""):
16
+ pipe = AutoPipelineForImage2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16")
17
  self.pipe.to("cuda")
18
+
 
19
 
20
  def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
21
  """