Update app.py

app.py

@@ -6,6 +6,8 @@ from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
 base = "stabilityai/stable-diffusion-xl-base-1.0"
 repo = "ByteDance/SDXL-Lightning"
 opts = {
@@ -16,9 +18,9 @@ opts = {
 }
 
 step_loaded = 4
-unet = UNet2DConditionModel.from_config(base, subfolder="unet").to(device, torch.float16)
+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=torch.float16, variant="fp16").to(device)
+pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=dtype, variant="fp16").to(device)
 
 @spaces.GPU(enable_queue=True)
 def generate_image(prompt, option):
@@ -33,7 +35,7 @@ def generate_image(prompt, option):
 with gr.Blocks() as demo:
     gr.HTML(
         "<h1><center>SDXL-Lightning</center></h1>" +
-        "<p><center>Lightning-fast text-to-image generation.</center></p>" +
+        "<p><center>Lightning-fast text-to-image generation</center></p>" +
         "<p><center><a href='https://huggingface.co/ByteDance/SDXL-Lightning'>https://huggingface.co/ByteDance/SDXL-Lightning</a></center></p>"
     )