Update app.py

app.py

@@ -8,24 +8,34 @@ import PIL
 import base64
 import io
 import torch
-from diffusers import LCMScheduler, AutoPipelineForText2Image
+
+# SSD-1B
+#from diffusers import LCMScheduler, AutoPipelineForText2Image
+
+# SDXL
+from diffusers import UNet2DConditionModel, DiffusionPipeline, LCMScheduler
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = int(os.getenv('MAX_IMAGE_SIZE', '1024'))
 SECRET_TOKEN = os.getenv('SECRET_TOKEN', 'default_secret')
 
-MODEL_ID = "segmind/SSD-1B"
-ADAPTER_ID = "latent-consistency/lcm-lora-ssd-1b"
-
-device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+#device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
 if torch.cuda.is_available():
-    pipe = AutoPipelineForText2Image.from_pretrained(MODEL_ID, torch_dtype=torch.float16, variant="fp16")
-    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
-    pipe.to("cuda")
+
+    #pipe = AutoPipelineForText2Image.from_pretrained("segmind/SSD-1B", torch_dtype=torch.float16, variant="fp16")
+    #pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+    #pipe.to("cuda")
 
     # load and fuse
-    pipe.load_lora_weights(ADAPTER_ID)
-    pipe.fuse_lora()
+    #pipe.load_lora_weights("latent-consistency/lcm-lora-ssd-1b")
+    #pipe.fuse_lora()
+
+    unet = UNet2DConditionModel.from_pretrained("latent-consistency/lcm-sdxl", torch_dtype=torch.float16, variant="fp16")
+    pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", unet=unet, torch_dtype=torch.float16, variant="fp16")
+
+    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+    pipe.to('cuda')
+    
 else:
     pipe = None