Update app.py

app.py

@@ -6,24 +6,29 @@ import gradio as gr
 import numpy as np
 import PIL.Image
 import torch
-from diffusers import LCMScheduler, AutoPipelineForText2Image
+from diffusers import DiffusionPipeline, UNet2DConditionModel, 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')
 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")
+    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"
+    )
 
-    # load and fuse
-    pipe.load_lora_weights(ADAPTER_ID)
-    pipe.fuse_lora()
+    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+    pipe.to(device)
 else:
     pipe = None
     
@@ -39,8 +44,8 @@ def generate(prompt: str,
              seed: int = 0,
              width: int = 1024,
              height: int = 1024,
-             guidance_scale: float = 0.0,
-             num_inference_steps: int = 4,
+             guidance_scale: float = 1.0,
+             num_inference_steps: int = 6,
              secret_token: str = '') -> PIL.Image.Image:
     if secret_token != SECRET_TOKEN:
         raise gr.Error(
@@ -64,7 +69,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 SDXL LoRA.</p>
+            <p style="color: black;">This space is a REST API to programmatically generate images using LCM-SSD-1B.</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>""")
@@ -112,16 +117,16 @@ with gr.Blocks() as demo:
     )
     guidance_scale = gr.Slider(
         label='Guidance scale',
-        minimum=0,
-        maximum=2,
+        minimum=1,
+        maximum=20,
         step=0.1,
-        value=0.0)
+        value=1.0)
     num_inference_steps = gr.Slider(
         label='Number of inference steps',
-        minimum=1,
-        maximum=8,
+        minimum=2,
+        maximum=40,
         step=1,
-        value=4)
+        value=6)
     
     use_negative_prompt.change(
         fn=lambda x: gr.update(visible=x),