asoderznik commited on
Commit
dfa9c7c
1 Parent(s): d188461

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +2 -2
handler.py CHANGED
@@ -2,7 +2,7 @@ from typing import Dict, List, Any
2
  from PIL import Image
3
  import torch
4
  from torch import autocast
5
- from diffusers import StableDiffusionUpscalePipeline
6
  import base64
7
  from io import BytesIO
8
  from transformers.utils import logging
@@ -24,7 +24,7 @@ class EndpointHandler():
24
  self.path = path
25
  # load the optimized model
26
  #model_id = "stabilityai/stable-diffusion-x4-upscaler"
27
- self.pipe = StableDiffusionUpscalePipeline.from_pretrained(path, torch_dtype=torch.float16)
28
  self.pipe = self.pipe.to(device)
29
 
30
 
 
2
  from PIL import Image
3
  import torch
4
  from torch import autocast
5
+ from diffusers import StableDiffusionLatentUpscalePipeline
6
  import base64
7
  from io import BytesIO
8
  from transformers.utils import logging
 
24
  self.path = path
25
  # load the optimized model
26
  #model_id = "stabilityai/stable-diffusion-x4-upscaler"
27
+ self.pipe = StableDiffusionLatentUpscalePipeline.from_pretrained(path, torch_dtype=torch.float16)
28
  self.pipe = self.pipe.to(device)
29
 
30