sdxl pipeline

app.py

@@ -1,4 +1,6 @@
 from diffusers import AutoencoderKL, UNet2DConditionModel, StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, UniPCMultistepScheduler
+from diffusers import StableDiffusionXLPipeline
+from diffusers import StableDiffusionXLImg2ImgPipeline
 import gradio as gr
 import torch
 from PIL import Image
@@ -43,9 +45,16 @@ if torch.cuda.is_available():
   vae = AutoencoderKL.from_pretrained(current_model.path, subfolder="vae", torch_dtype=torch.float16, use_auth_token=auth_token)
   for model in models:
     try:
+        if 'XL'in model.name:
+            PipeClass  = StableDiffusionXLPipeline
+            PipeI2IClass = StableDiffusionXLImg2ImgPipeline
+        else:
+            PipeClass  = StableDIffusionXPipeline
+            PipeI2IClass = StableDiffusionImg2ImgPipeline
+
         unet = UNet2DConditionModel.from_pretrained(model.path, subfolder="unet", torch_dtype=torch.float16, use_auth_token=auth_token)
-        model.pipe_t2i = StableDiffusionPipeline.from_pretrained(model.path, unet=unet, vae=vae, torch_dtype=torch.float16, scheduler=scheduler, use_auth_token=auth_token)
-        model.pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(model.path, unet=unet, vae=vae, torch_dtype=torch.float16, scheduler=scheduler, use_auth_token=auth_token)
+        model.pipe_t2i = PipeClass.from_pretrained(model.path, unet=unet, vae=vae, torch_dtype=torch.float16, scheduler=scheduler, use_auth_token=auth_token)
+        model.pipe_i2i = PipeI2IClass.from_pretrained(model.path, unet=unet, vae=vae, torch_dtype=torch.float16, scheduler=scheduler, use_auth_token=auth_token)
     except:
         models.remove(model)
   pipe = models[0].pipe_t2i