Update app.py

app.py

@@ -0,0 +1,64 @@
+import gradio as gr
+import torch
+import spaces
+from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+base = "stabilityai/stable-diffusion-xl-base-1.0"
+repo = "ByteDance/SDXL-Lightning"
+opts = {
+    "1 Step"  : ["sdxl_lightning_1step_unet_x0.safetensors", 1],
+    "2 Steps" : ["sdxl_lightning_2step_unet.safetensors", 2],
+    "4 Steps" : ["sdxl_lightning_4step_unet.safetensors", 4],
+    "8 Steps" : ["sdxl_lightning_8step_unet.safetensors", 8],
+}
+
+pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to(device)
+
+# 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))
+    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")
+    
+    with gr.Group():
+        with gr.Row():
+            prompt = gr.Textbox(
+                label="Text prompt",
+                scale=8
+            )
+            option = gr.Dropdown(
+                label="Inference steps",
+                choices=["1 Step", "2 Steps", "4 Steps", "8 Steps"],
+                value="4-Step",
+                interactive=True
+            )
+            submit = gr.Button(
+                scale=1,
+                variant="primary"
+            )
+    
+    img = gr.Image(label="SDXL-Lightening Generated Image")
+
+    prompt.submit(
+        fn=generate_image,
+        inputs=[prompt, option],
+        outputs=img,
+    )
+    submit.click(
+        fn=generate_image,
+        inputs=[prompt, option],
+        outputs=img,
+    )
+    
+demo.queue().launch()