Update app.py

app.py

@@ -53,7 +53,6 @@ def binary_mask_eval(args, model):
             ),
             map_location="cpu",
         )
-        torch.cuda.empty_cache()
 
     # reload the original model
     if model == "sdxl":
@@ -64,7 +63,6 @@ def binary_mask_eval(args, model):
         pipe = FluxPipeline.from_pretrained(
             "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16
         ).to("cpu")
-        torch.cuda.empty_cache()
 
     print("prune complete")
     return pipe, pruned_pipe
@@ -79,7 +77,6 @@ def generate_images(prompt, seed, steps, pipe, pruned_pipe):
     original_image = pipe(
         prompt=prompt, generator=g_cpu, num_inference_steps=steps
     ).images[0]
-    torch.cuda.empty_cache()
     g_cpu = torch.Generator("cuda").manual_seed(seed)
     ecodiff_image = pruned_pipe(
         prompt=prompt, generator=g_cpu, num_inference_steps=steps