Update diffusion_webui/diffusion_models/text2img_app.py

diffusion_webui/diffusion_models/text2img_app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from diffusers import StableDiffusionPipeline,StableDiffusionXLPipeline
+from diffusers import StableDiffusionPipeline,DiffusionPipeline
 
 from diffusion_webui.utils.model_list import stable_model_list
 from diffusion_webui.utils.scheduler_list import (
@@ -20,7 +20,7 @@ class StableDiffusionText2ImageGenerator:
     ):
         if self.pipe is None or self.pipe.model_name != stable_model_path or self.pipe.scheduler_name != scheduler:
             if stable_model_path == "stabilityai/stable-diffusion-xl-base-0.9":
-                self.pipe = StableDiffusionXLPipeline.from_pretrained(
+                self.pipe = DiffusionPipeline.from_pretrained(
                     stable_model_path, safety_checker=None, torch_dtype=torch.float16
                 )
             else: