Suraj Narayanan Sasikumar commited on
Commit
4ccc92a
1 Parent(s): bba78d3

trying a diff pipe

Browse files
Files changed (1) hide show
  1. handler.py +3 -3
handler.py CHANGED
@@ -1,6 +1,6 @@
1
  from typing import Dict, List, Any
2
  import torch
3
- from diffusers import DPMSolverMultistepScheduler, StableDiffusionXLPipeline
4
  from PIL import Image
5
  import base64
6
  from io import BytesIO
@@ -16,7 +16,7 @@ if device.type != "cuda":
16
  class EndpointHandler:
17
  def __init__(self, path=""):
18
  # load StableDiffusionInpaintPipeline pipeline
19
- self.base = StableDiffusionXLPipeline.from_pretrained(
20
  path, torch_dtype=torch.float16, variant="fp16", use_safetensors=True
21
  )
22
  # use DPMSolverMultistepScheduler
@@ -27,7 +27,7 @@ class EndpointHandler:
27
  self.base = self.base.to(device)
28
  self.base.unet = torch.compile(self.base.unet, mode="reduce-overhead", fullgraph=True)
29
 
30
- self.refiner = StableDiffusionXLPipeline.from_pretrained(
31
  "stabilityai/stable-diffusion-xl-refiner-1.0",
32
  text_encoder_2=self.base.text_encoder_2,
33
  vae=self.base.vae,
 
1
  from typing import Dict, List, Any
2
  import torch
3
+ from diffusers import DPMSolverMultistepScheduler, DiffusionPipeline
4
  from PIL import Image
5
  import base64
6
  from io import BytesIO
 
16
  class EndpointHandler:
17
  def __init__(self, path=""):
18
  # load StableDiffusionInpaintPipeline pipeline
19
+ self.base = DiffusionPipeline.from_pretrained(
20
  path, torch_dtype=torch.float16, variant="fp16", use_safetensors=True
21
  )
22
  # use DPMSolverMultistepScheduler
 
27
  self.base = self.base.to(device)
28
  self.base.unet = torch.compile(self.base.unet, mode="reduce-overhead", fullgraph=True)
29
 
30
+ self.refiner = DiffusionPipeline.from_pretrained(
31
  "stabilityai/stable-diffusion-xl-refiner-1.0",
32
  text_encoder_2=self.base.text_encoder_2,
33
  vae=self.base.vae,