Update app.py

app.py

@@ -1,3 +1,4 @@
+import spaces
 import base64
 from io import BytesIO
 
@@ -7,18 +8,25 @@ import torch
 
 from diffusers import StableDiffusionPipeline, AutoencoderKL, AutoencoderTiny
 
-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 = StableDiffusionPipeline.from_pretrained(
+    "IDKiro/sdxs-512-dreamshaper", torch_dtype=weight_type
+)
 pipe.to(torch_device=device, torch_dtype=weight_type)
 
-vae_tiny = AutoencoderTiny.from_pretrained("IDKiro/sdxs-512-dreamshaper", subfolder="vae")
+vae_tiny = AutoencoderTiny.from_pretrained(
+    "IDKiro/sdxs-512-dreamshaper", subfolder="vae"
+)
 vae_tiny.to(device, dtype=weight_type)
 
-vae_large = AutoencoderKL.from_pretrained("IDKiro/sdxs-512-dreamshaper", subfolder="vae_large")
+vae_large = AutoencoderKL.from_pretrained(
+    "IDKiro/sdxs-512-dreamshaper", subfolder="vae_large"
+)
 vae_tiny.to(device, dtype=weight_type)
 
+
 def pil_image_to_data_url(img, format="PNG"):
     buffered = BytesIO()
     img.save(buffered, format=format)
@@ -26,11 +34,12 @@ 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",
     vae_type=None,
-    param_dtype='torch.float16',
+    param_dtype="torch.float16",
 ) -> PIL.Image.Image:
     if vae_type == "tiny vae":
         pipe.vae = vae_tiny
@@ -38,12 +47,15 @@ def run(
         pipe.vae = vae_large
 
     if device_type == "CPU":
-        device = "cpu" 
-        param_dtype = 'torch.float32'
+        device = "cpu"
+        param_dtype = "torch.float32"
     else:
         device = "cuda"
-    
-    pipe.to(torch_device=device, torch_dtype=torch.float16 if param_dtype == 'torch.float16' else torch.float32)
+
+    pipe.to(
+        torch_device=device,
+        torch_dtype=torch.float16 if param_dtype == "torch.float16" else torch.float32,
+    )
 
     result = pipe(
         prompt=prompt,
@@ -62,7 +74,7 @@ examples = [
 ]
 
 with gr.Blocks(css="style.css") as demo:
-    gr.Markdown("# SDXS-512-DreamShaper (only CPU now)")
+    gr.Markdown("# SDXS-512-DreamShaper")
     with gr.Group():
         with gr.Row():
             with gr.Column(min_width=685):
@@ -75,38 +87,51 @@ with gr.Blocks(css="style.css") as demo:
                         container=False,
                     )
                     run_button = gr.Button("Run", scale=0)
-                    
-                device_choices = ['GPU','CPU']
-                device_type = gr.Radio(device_choices, label='Device',  
-                                            value=device_choices[1],
-                                            interactive=False,
-                                            info='Only CPU now.')
-
-                vae_choices = ['tiny vae','large vae']
-                vae_type = gr.Radio(vae_choices, label='Image Decoder Type',  
-                                            value=vae_choices[0],
-                                            interactive=True,
-                                            info='To save GPU memory, use tiny vae. For better quality, use large vae.')
-
-                dtype_choices = ['torch.float16','torch.float32']
-                param_dtype = gr.Radio(dtype_choices,label='torch.weight_type',  
-                                            value=dtype_choices[0],
-                                            interactive=True,
-                                            info='To save GPU memory, use torch.float16. For better quality, use torch.float32.')
-                
-                download_output = gr.Button("Download output", elem_id="download_output")
 
-            with gr.Column(min_width=512):
-                result = gr.Image(label="Result", height=512, width=512, elem_id="output_image", show_label=False, show_download_button=True)
+                device_choices = ["GPU", "CPU"]
+                device_type = gr.Radio(
+                    device_choices,
+                    label="Device",
+                    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",
+                    value=vae_choices[0],
+                    interactive=True,
+                    info="To save GPU memory, use tiny vae. For better quality, use large vae.",
+                )
+
+                dtype_choices = ["torch.float16", "torch.float32"]
+                param_dtype = gr.Radio(
+                    dtype_choices,
+                    label="torch.weight_type",
+                    value=dtype_choices[0],
+                    interactive=True,
+                    info="To save GPU memory, use torch.float16. For better quality, use torch.float32.",
+                )
+
+                download_output = gr.Button(
+                    "Download output", elem_id="download_output"
+                )
 
-    gr.Examples(
-        examples=examples,
-        inputs=prompt,
-        outputs=result,
-        fn=run
-    )
-
-    demo.load(None,None,None)
+            with gr.Column(min_width=512):
+                result = gr.Image(
+                    label="Result",
+                    height=512,
+                    width=512,
+                    elem_id="output_image",
+                    show_label=False,
+                    show_download_button=True,
+                )
+
+    gr.Examples(examples=examples, inputs=prompt, outputs=result, fn=run)
+
+    demo.load(None, None, None)
 
     inputs = [prompt, device_type, vae_type, param_dtype]
     outputs = [result, download_output]