Update app.py

app.py

@@ -16,7 +16,7 @@ opts = {
 }
 
 step_loaded = 4
-unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
+unet = UNet2DConditionModel.from_config(base, subfolder="unet").to(device, torch.float16)
 unet.load_state_dict(load_file(hf_hub_download(repo, opts["4 Steps"][0]), device=device))
 pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to(device)
 
@@ -36,22 +36,21 @@ with gr.Blocks() as demo:
         "<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():
-            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 Steps",
-                interactive=True
-            )
-            submit = gr.Button(
-                scale=1,
-                variant="primary"
-            )
+    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 Steps",
+            interactive=True
+        )
+        submit = gr.Button(
+            scale=1,
+            variant="primary"
+        )
     
     img = gr.Image(label="SDXL-Lightening Generated Image")