Update main.py

main.py

@@ -300,25 +300,7 @@ def execute_task(args, trainer, device, dtype, shape, enable_grad, settings, pip
         best_image.save(f"{save_dir}/best_image.png")
         #init_image.save(f"{save_dir}/init_image.png")
 
-        # Move the pipe back to CPU
-        if args.model != "flux":
-            if args.model == "hyper-sd":
-                if pipe.device == torch.device('cuda'):
-                    print("Moving the entire pipe back to CPU.")
-                    # Transfer UNet to GPU
-                    pipe.unet = pipe.unet.to("cpu") 
-                    pipe.to('cpu')  # Move all components of the pipe back to CPU
-                    # Delete the pipe to free resources
-                    del pipe
-                    print("Pipe deleted to free resources.")
-                            
-            else: 
-                if pipe.device == torch.device('cuda'):
-                    print("Moving the entire pipe back to CPU.")
-                    pipe.to("cpu")
-                    # Delete the pipe to free resources
-                    del pipe
-                    print("Pipe deleted to free resources.")
+        find_and_move_object_to_cpu()
         
         clear_gpu()