modify

app.py

@@ -14,6 +14,7 @@ from diffusers import StableDiffusionXLPipeline
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+
 def get_model_param_summary(model, verbose=False):
     params_dict = dict()
     overall_params = 0
@@ -50,7 +51,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,18 +92,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("cpu")
 
     torch_dtype = torch.bfloat16 if args.mix_precision == "bf16" else torch.float32
     mask_pipe, hookers = create_pipeline(
         pipe,
         args.model,
-        device,
+        "cpu",
         torch_dtype,
         args.ckpt,
         binary=args.binary,
@@ -132,7 +133,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("cpu")
 
     # get model param summary
     print(f"original model param: {get_model_param_summary(pipe.unet)['overall']}")
@@ -140,12 +141,15 @@ def binary_mask_eval(args):
     print("prune complete")
     return pipe, pruned_pipe
 
+
 @spaces.GPU
 def generate_images(prompt, seed, steps, pipe, pruned_pipe):
+    pipe.to("cuda")
+    pruned_pipe.to("cuda")
     # Run the model and return images directly
-    g_cpu = torch.Generator(device).manual_seed(seed)
+    g_cpu = torch.Generator("cuda").manual_seed(seed)
     original_image = pipe(prompt=prompt, generator=g_cpu, num_inference_steps=steps).images[0]
-    g_cpu = torch.Generator(device).manual_seed(seed)
+    g_cpu = torch.Generator("cuda").manual_seed(seed)
     ecodiff_image = pruned_pipe(prompt=prompt, generator=g_cpu, num_inference_steps=steps).images[0]
     return original_image, ecodiff_image
 
@@ -177,8 +181,8 @@ def create_demo():
         with gr.Row():
             model_choice = gr.Dropdown(choices=["SDXL"], value="SDXL", label="Model", scale=1.2)
             pruning_ratio = gr.Dropdown(choices=["20%"], value="20%", label="Pruning Ratio", scale=1.2)
-            prune_btn = gr.Button("Initialize Original and Pruned Models", variant="primary", scale=1)
             status_label = gr.HighlightedText(label="Model Status", value=[("Model Not Initialized", "red")], scale=1)
+            prune_btn = gr.Button("Initialize Original and Pruned Models", variant="primary", scale=1)
         with gr.Row():
             prompt = gr.Textbox(label="Prompt", value="A clock tower floating in a sea of clouds", scale=3)
             seed = gr.Number(label="Seed", value=44, precision=0, scale=1)