trying a diff pipe

handler.py

@@ -1,6 +1,6 @@
 from typing import Dict, List, Any
 import torch
-from diffusers import DPMSolverMultistepScheduler, StableDiffusionXLPipeline
+from diffusers import DPMSolverMultistepScheduler, DiffusionPipeline
 from PIL import Image
 import base64
 from io import BytesIO
@@ -16,7 +16,7 @@ if device.type != "cuda":
 class EndpointHandler:
     def __init__(self, path=""):
         # load StableDiffusionInpaintPipeline pipeline
-        self.base = StableDiffusionXLPipeline.from_pretrained(
+        self.base = DiffusionPipeline.from_pretrained(
             path, torch_dtype=torch.float16, variant="fp16", use_safetensors=True
         )
         # use DPMSolverMultistepScheduler
@@ -27,7 +27,7 @@ class EndpointHandler:
         self.base = self.base.to(device)
         self.base.unet = torch.compile(self.base.unet, mode="reduce-overhead", fullgraph=True)
 
-        self.refiner = StableDiffusionXLPipeline.from_pretrained(
+        self.refiner = DiffusionPipeline.from_pretrained(
             "stabilityai/stable-diffusion-xl-refiner-1.0",
             text_encoder_2=self.base.text_encoder_2,
             vae=self.base.vae,