Update app.py

app.py

@@ -19,22 +19,24 @@ opts = {
     "8 Steps" : ("sdxl_lightning_8step_unet.safetensors", 8),
 }
 
-step_loaded = None
-# unet = UNet2DConditionModel.from_config(base, subfolder="unet").to(device, dtype)
-# unet.load_state_dict(load_file(hf_hub_download(repo, opts["4 Steps"][0]), device=device))
-pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=dtype, variant="fp16").to(device)
+step_loaded = 4
+unet = UNet2DConditionModel.from_config(base, subfolder="unet").to(device, dtype)
+unet.load_state_dict(load_file(hf_hub_download(repo, opts["4 Steps"][0]), device=device))
+pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=dtype, variant="fp16").to(device, dtype)
 
 @spaces.GPU(enable_queue=True)
 def generate_image(prompt, option):
     global step_loaded
+    print(prompt, option)
     ckpt, step = opts[option]
     if step != step_loaded:
+        print(f"Switching checkpoint from {step_loaded} to {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))
         step_loaded = step
     return pipe(prompt, num_inference_steps=step, guidance_scale=0).images[0]
 
-with gr.Blocks() as demo:
+with gr.Blocks(css="style.css") as demo:
     gr.HTML(
         "<h1><center>SDXL-Lightning</center></h1>" +
         "<p><center>Lightning-fast text-to-image generation</center></p>" +
@@ -57,7 +59,7 @@ with gr.Blocks() as demo:
             variant="primary"
         )
     
-    img = gr.Image(label="SDXL-Lightening Generated Image")
+    img = gr.Image(label="SDXL-Lighting Generated Image")
 
     prompt.submit(
         fn=generate_image,