Update app.py

app.py

@@ -6,29 +6,24 @@ import gradio as gr
 import numpy as np
 import PIL.Image
 import torch
-from diffusers import DiffusionPipeline, UNet2DConditionModel, LCMScheduler
+from diffusers import LCMScheduler, AutoPipelineForText2Image
 
 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')
 if torch.cuda.is_available():
-    unet = UNet2DConditionModel.from_pretrained(
-        "latent-consistency/lcm-ssd-1b",
-        torch_dtype=torch.float16,
-        variant="fp16"
-    )
-
-    pipe = DiffusionPipeline.from_pretrained(
-        "segmind/SSD-1B",
-        unet=unet,
-        torch_dtype=torch.float16,
-        variant="fp16"
-    )
-
+    pipe = AutoPipelineForText2Image.from_pretrained(MODEL_ID, torch_dtype=torch.float16, variant="fp16")
     pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
-    pipe.to(device)
+    pipe.to("cuda")
+
+    # load and fuse
+    pipe.load_lora_weights(ADAPTER_ID)
+    pipe.fuse_lora()
 else:
     pipe = None
     
@@ -44,8 +39,8 @@ def generate(prompt: str,
              seed: int = 0,
              width: int = 1024,
              height: int = 1024,
-             guidance_scale: float = 1.0,
-             num_inference_steps: int = 6,
+             guidance_scale: float = 0.0,
+             num_inference_steps: int = 4,
              secret_token: str = '') -> PIL.Image.Image:
     if secret_token != SECRET_TOKEN:
         raise gr.Error(
@@ -69,7 +64,7 @@ with gr.Blocks() as demo:
     gr.HTML("""
     <div style="z-index: 100; position: fixed; top: 0px; right: 0px; left: 0px; bottom: 0px; width: 100%; height: 100%; background: white; display: flex; align-items: center; justify-content: center; color: black;">
         <div style="text-align: center; color: black;">
-            <p style="color: black;">This space is a REST API to programmatically generate images using LCM-SSD-1B.</p>
+            <p style="color: black;">This space is a REST API to programmatically generate images using LCM SDXL LoRA.</p>
             <p style="color: black;">It is not meant to be directly used through a user interface, but using code and an access key.</p>
         </div>
     </div>""")
@@ -117,16 +112,16 @@ with gr.Blocks() as demo:
     )
     guidance_scale = gr.Slider(
         label='Guidance scale',
-        minimum=1,
-        maximum=20,
+        minimum=0,
+        maximum=2,
         step=0.1,
-        value=1.0)
+        value=0.0)
     num_inference_steps = gr.Slider(
         label='Number of inference steps',
-        minimum=2,
-        maximum=40,
+        minimum=1,
+        maximum=8,
         step=1,
-        value=6)
+        value=4)
     
     use_negative_prompt.change(
         fn=lambda x: gr.update(visible=x),