we worked to handle gpu operations

main.py

@@ -15,6 +15,36 @@ from rewards import get_reward_losses
 from training import LatentNoiseTrainer, get_optimizer
 
 
+import torch
+import gc
+
+def clear_gpu():
+    """Clear GPU memory by removing tensors, freeing cache, and moving data to CPU."""
+    # List memory usage before clearing
+    print(f"Memory allocated before clearing: {torch.cuda.memory_allocated() / (1024 ** 2)} MB")
+    print(f"Memory reserved before clearing: {torch.cuda.memory_reserved() / (1024 ** 2)} MB")
+    
+    # Force the garbage collector to free unreferenced objects
+    gc.collect()
+
+    # Move any bound tensors back to CPU if needed
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()  # Free up the cached memory
+        torch.cuda.ipc_collect()  # Clear any cross-process memory
+    
+    print(f"Memory allocated after clearing: {torch.cuda.memory_allocated() / (1024 ** 2)} MB")
+    print(f"Memory reserved after clearing: {torch.cuda.memory_reserved() / (1024 ** 2)} MB")
+
+def unload_previous_model_if_needed(loaded_model_setup):
+    """Unload the current model from the GPU and free resources if a new model is being loaded."""
+    if loaded_model_setup is not None:
+        print("Unloading previous model from GPU to free memory.")
+        previous_model = loaded_model_setup[7]  # Assuming pipe is at position [7] in the setup
+        if hasattr(previous_model, 'to') and loaded_model_setup[0].model != "flux":
+            previous_model.to('cpu')  # Move model to CPU to free GPU memory
+        del previous_model  # Delete the reference to the model
+        clear_gpu()  # Clear all remaining GPU memory
+
 def setup(args, loaded_model_setup=None):
     seed_everything(args.seed)
     bf.makedirs(f"{args.save_dir}/logs/{args.task}")
@@ -52,14 +82,10 @@ def setup(args, loaded_model_setup=None):
         os.environ["CUDA_VISIBLE_DEVICES"] = args.device_id
 
     device = torch.device("cuda")
-    if args.dtype == "float32":
-        dtype = torch.float32
-    elif args.dtype == "float16":
-        dtype = torch.float16
+    dtype = torch.float16 if args.dtype == "float16" else torch.float32
 
     # If args.model is the same as the one in loaded_model_setup, reuse the trainer and pipe
     if loaded_model_setup and args.model == loaded_model_setup[0].model:
-        # Reuse the trainer and pipe from the loaded model setup
         print(f"Reusing model {args.model} from loaded setup.")
         trainer = loaded_model_setup[1]  # Trainer is at position 1 in loaded_model_setup
         
@@ -97,10 +123,13 @@ def setup(args, loaded_model_setup=None):
                 width // trainer.model.vae_scale_factor,
             )
         
-        multi_apply_fn = loaded_model_setup[6]
+        pipe = loaded_model_setup[7]
         enable_grad = not args.no_optim
 
-        return args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings
+        return args, trainer, device, dtype, shape, enable_grad, settings, pipe
+
+    # Unload previous model and clear GPU resources
+    unload_previous_model_if_needed(loaded_model_setup)
 
     # Proceed with full model loading if args.model is different
     print(f"Loading new model: {args.model}")
@@ -113,27 +142,8 @@ def setup(args, loaded_model_setup=None):
         args.model, dtype, device, args.cache_dir, args.memsave, args.cpu_offloading
     )
 
-    # Attempt to move the model to GPU or keep it on CPU if offloading is enabled
-    try:
-        if not args.cpu_offloading:
-            pipe.to(device)
-    except RuntimeError as e:
-        if 'out of memory' in str(e):
-            print("CUDA OOM error. Attempting to handle OOM situation.")
-            # Attempt to clear memory and retry moving to GPU
-            torch.cuda.empty_cache()  # Free up cached memory
-            gc.collect()
-            try:
-                # Retry loading after clearing cache
-                if not args.cpu_offloading:
-                    pipe.to(device)
-            except RuntimeError as e:
-                print("Still facing OOM issues. Keeping model on CPU.")
-                args.cpu_offloading = True  # Force CPU offloading
-        else:
-            raise e  # Re-raise the exception if it's not OOM
-
-    torch.cuda.empty_cache()  # Free up cached memory
+    # Final memory cleanup after model loading
+    torch.cuda.empty_cache()
     gc.collect()
 
     trainer = LatentNoiseTrainer(
@@ -180,28 +190,47 @@ def setup(args, loaded_model_setup=None):
     torch.cuda.empty_cache()  # Free up cached memory
     gc.collect()
 
-    if args.enable_multi_apply:
-        multi_apply_fn = get_multi_apply_fn(
-            model_type=args.multi_step_model,
-            seed=args.seed,
-            pipe=pipe,
-            cache_dir=args.cache_dir,
-            device=device if not args.cpu_offloading else 'cpu',
-            dtype=dtype,
-        )
-    else:
-        multi_apply_fn = None    
+    
 
-    torch.cuda.empty_cache()  # Free up cached memory
-    gc.collect()
+    return args, trainer, device, dtype, shape, enable_grad, settings, pipe
 
-    return args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings
 
 
 
-def execute_task(args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings, progress_callback=None):
+def execute_task(args, trainer, device, dtype, shape, enable_grad, settings, pipe, progress_callback=None):
     
     if args.task == "single":
+        # Attempt to move the model to GPU if model is not Flux
+        if args.model != "flux":
+            if pipe.device != torch.device('cuda'):
+                pipe.to(device, dtype)
+        else:
+            print(f"PIPE:{pipe}")
+            
+            
+            if args.cpu_offloading:
+                pipe.enable_sequential_cpu_offload()
+
+            #if pipe.device != torch.device('cuda'):
+            #    pipe.to(device, dtype)
+        
+        if args.enable_multi_apply:
+            
+            multi_apply_fn = get_multi_apply_fn(
+                model_type=args.multi_step_model,
+                seed=args.seed,
+                pipe=pipe,
+                cache_dir=args.cache_dir,
+                device=device if not args.cpu_offloading else 'cpu',
+                dtype=dtype,
+            )
+        else:
+            multi_apply_fn = None    
+    
+        torch.cuda.empty_cache()  # Free up cached memory
+        gc.collect()
+    
+
         init_latents = torch.randn(shape, device=device, dtype=dtype)
         latents = torch.nn.Parameter(init_latents, requires_grad=enable_grad)
         optimizer = get_optimizer(args.optim, latents, args.lr, args.nesterov)
@@ -383,8 +412,8 @@ def execute_task(args, trainer, device, dtype, shape, enable_grad, multi_apply_f
 
 def main():
     args = parse_args()
-    args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings = setup(args, loaded_model_setup=None)
-    execute_task(args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings)
+    args, trainer, device, dtype, shape, enable_grad, settings, pipe = setup(args, loaded_model_setup=None)
+    execute_task(args, trainer, device, dtype, shape, enable_grad, settings, pipe)
 
 if __name__ == "__main__":
     main()