Update app.py

app.py

@@ -1,3 +1,4 @@
+import space
 import base64
 from io import BytesIO
 
@@ -8,8 +9,8 @@ import torch
 from diffusers import StableDiffusionPipeline, AutoencoderKL, AutoencoderTiny
 from peft import PeftModel
 
-device = "cpu"   # Linux & Windows
-weight_type = torch.float32  # torch.float16 works as well, but pictures seem to be a bit worse
+device = "cuda"
+weight_type = torch.float16
 
 pipe = StableDiffusionPipeline.from_pretrained("IDKiro/sdxs-512-dreamshaper", torch_dtype=weight_type)
 pipe.unet = PeftModel.from_pretrained(pipe.unet, "IDKiro/sdxs-512-dreamshaper-anime")
@@ -28,6 +29,7 @@ def pil_image_to_data_url(img, format="PNG"):
     return f"data:image/{format.lower()};base64,{img_str}"
 
 
+@spaces.GPU
 def run(
     prompt: str,
     device_type="GPU",
@@ -64,7 +66,8 @@ examples = [
 ]
 
 with gr.Blocks(css="style.css") as demo:
-    gr.Markdown("# SDXS-512-DreamShaper-Anime (only CPU now)")
+    gr.Markdown("# SDXS-512-DreamShaper-Anime")
+    gr.Markdown("[SDXS: Real-Time One-Step Latent Diffusion Models with Image Conditions](https://arxiv.org/abs/2403.16627) | [GitHub](https://github.com/IDKiro/sdxs")
     with gr.Group():
         with gr.Row():
             with gr.Column(min_width=685):
@@ -80,9 +83,9 @@ with gr.Blocks(css="style.css") as demo:
                     
                 device_choices = ['GPU','CPU']
                 device_type = gr.Radio(device_choices, label='Device',  
-                                            value=device_choices[1],
-                                            interactive=False,
-                                            info='Only CPU now.')
+                                            value=device_choices[0],
+                                            interactive=True,
+                                            info='Thanks to the community for the GPU!')
 
                 vae_choices = ['tiny vae','large vae']
                 vae_type = gr.Radio(vae_choices, label='Image Decoder Type',