Update app.py

app.py

@@ -16,20 +16,24 @@ opts = {
 }
 
 pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to(device)
+last_step = None
 
 # Function 
 @spaces.GPU(enable_queue=True)
 def generate_image(prompt, option):
     ckpt, step = opts[option]
-    pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", prediction_type="sample" if step == 1 else "epsilon")
-    pipe.unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device))
+    if last_step != step:
+        pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", prediction_type="sample" if step == 1 else "epsilon")
+        pipe.unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device))
+        last_step = step
     image = pipe(prompt, num_inference_steps=step, guidance_scale=0).images[0]
     return image
 
 
 with gr.Blocks() as demo:
-    gr.HTML("<h1><center>SDXL-Lightning ⚡</center></h1>")
-    gr.Markdown("Lightning-fast text-to-image generation! https://huggingface.co/ByteDance/SDXL-Lightning")
+    gr.HTML("<h1><center>SDXL-Lightning</center></h1>")
+    gr.HTML("<p><center>Lightning-fast text-to-image generation.</center></p>")
+    gr.HTML("<p><center><a href='https://huggingface.co/ByteDance/SDXL-Lightning'>https://huggingface.co/ByteDance/SDXL-Lightning</a></center></p>")
     
     with gr.Group():
         with gr.Row():
@@ -40,7 +44,7 @@ with gr.Blocks() as demo:
             option = gr.Dropdown(
                 label="Inference steps",
                 choices=["1 Step", "2 Steps", "4 Steps", "8 Steps"],
-                value="4-Step",
+                value="4 Steps",
                 interactive=True
             )
             submit = gr.Button(