enable zero GPU

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from dataclasses import dataclass
-
+import spaces
 import torch
 from tqdm import tqdm
 
@@ -91,7 +91,7 @@ def prune_model(pipe, hookers):
     ffn_hook.clear_hooks()
     return pipe
 
-
+@spaces.GPU
 def binary_mask_eval(args):
     # load sdxl model
     pipe = StableDiffusionXLPipeline.from_pretrained(
@@ -141,7 +141,7 @@ def binary_mask_eval(args):
     print("prune complete")
     return pipe, pruned_pipe
 
-
+@spaces.GPU
 def generate_images(prompt, seed, steps, pipe, pruned_pipe):
     # Run the model and return images directly
     g_cpu = torch.Generator("cuda:0").manual_seed(seed)
@@ -221,4 +221,4 @@ def create_demo():
 
 if __name__ == "__main__":
     demo = create_demo()
-    demo.launch(share=True)
+    demo.launch()