fffiloni commited on
Commit
1a6b2a7
1 Parent(s): 94f8ab2

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +1 -19
main.py CHANGED
@@ -300,25 +300,7 @@ def execute_task(args, trainer, device, dtype, shape, enable_grad, settings, pip
300
  best_image.save(f"{save_dir}/best_image.png")
301
  #init_image.save(f"{save_dir}/init_image.png")
302
 
303
- # Move the pipe back to CPU
304
- if args.model != "flux":
305
- if args.model == "hyper-sd":
306
- if pipe.device == torch.device('cuda'):
307
- print("Moving the entire pipe back to CPU.")
308
- # Transfer UNet to GPU
309
- pipe.unet = pipe.unet.to("cpu")
310
- pipe.to('cpu') # Move all components of the pipe back to CPU
311
- # Delete the pipe to free resources
312
- del pipe
313
- print("Pipe deleted to free resources.")
314
-
315
- else:
316
- if pipe.device == torch.device('cuda'):
317
- print("Moving the entire pipe back to CPU.")
318
- pipe.to("cpu")
319
- # Delete the pipe to free resources
320
- del pipe
321
- print("Pipe deleted to free resources.")
322
 
323
  clear_gpu()
324
 
 
300
  best_image.save(f"{save_dir}/best_image.png")
301
  #init_image.save(f"{save_dir}/init_image.png")
302
 
303
+ find_and_move_object_to_cpu()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
304
 
305
  clear_gpu()
306