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
+from diffusers import StableDiffusionPipeline,StableDiffusionXLPipeline
 
 from diffusion_webui.utils.model_list import stable_model_list
 from diffusion_webui.utils.scheduler_list import (
@@ -19,9 +19,14 @@ class StableDiffusionText2ImageGenerator:
         scheduler,
     ):
         if self.pipe is None or self.pipe.model_name != stable_model_path or self.pipe.scheduler_name != scheduler:
-            self.pipe = StableDiffusionPipeline.from_pretrained(
-                stable_model_path, safety_checker=None, torch_dtype=torch.float16
-            )
+            if stable_model_path == "stabilityai/stable-diffusion-xl-base-0.9":
+                self.pipe = StableDiffusionXLPipeline.from_pretrained(
+                    stable_model_path, safety_checker=None, torch_dtype=torch.float16
+                )
+            else:
+                self.pipe = StableDiffusionPipeline.from_pretrained(
+                    stable_model_path, safety_checker=None, torch_dtype=torch.float16
+                )
 
         self.pipe = get_scheduler(pipe=self.pipe, scheduler=scheduler)
         self.pipe.to("cuda")