modify

app.py

@@ -50,7 +50,7 @@ class GradioArgs:
         if self.ratio is None:
             self.ratio = [0.68, 0.88]
 
-
+@spaces.GPU
 def prune_model(pipe, hookers):
     # remove parameters in attention blocks
     cross_attn_hooker = hookers[0]
@@ -91,17 +91,18 @@ 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(
         "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
-    )#.to(device)
+    ).to(device)
 
     torch_dtype = torch.bfloat16 if args.mix_precision == "bf16" else torch.float32
     mask_pipe, hookers = create_pipeline(
         pipe,
         args.model,
-        "cpu",
+        device,
         torch_dtype,
         args.ckpt,
         binary=args.binary,
@@ -131,7 +132,7 @@ def binary_mask_eval(args):
     # reload the original model
     pipe = StableDiffusionXLPipeline.from_pretrained(
         "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
-    )#.to(device)
+    ).to(device)
 
     # get model param summary
     print(f"original model param: {get_model_param_summary(pipe.unet)['overall']}")