Update app.py

app.py

@@ -6,28 +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'))
 
+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
     
@@ -43,7 +39,6 @@ def generate(prompt: str,
              seed: int = 0,
              width: int = 1024,
              height: int = 1024,
-             guidance_scale: float = 1.0,
              num_inference_steps: int = 6) -> PIL.Image.Image:
 
     generator = torch.Generator().manual_seed(seed)
@@ -55,7 +50,6 @@ def generate(prompt: str,
                 negative_prompt=negative_prompt,
                 width=width,
                 height=height,
-                guidance_scale=guidance_scale,
                 num_inference_steps=num_inference_steps,
                 generator=generator,
                 output_type='pil').images[0]
@@ -105,12 +99,6 @@ with gr.Blocks() as demo:
                     value=1024,
                 )
             with gr.Row():
-                guidance_scale = gr.Slider(
-                    label='Guidance scale',
-                    minimum=1,
-                    maximum=20,
-                    step=0.1,
-                    value=5.0)
                 num_inference_steps = gr.Slider(
                     label='Number of inference steps',
                     minimum=2,
@@ -133,7 +121,6 @@ with gr.Blocks() as demo:
         seed,
         width,
         height,
-        guidance_scale,
         num_inference_steps,
     ]
     prompt.submit(