we can load model an handle oom errors on model switch

main.py

@@ -112,7 +112,27 @@ def setup(args, loaded_model_setup=None):
     pipe = get_model(
         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
     gc.collect()
 
@@ -123,7 +143,7 @@ def setup(args, loaded_model_setup=None):
         n_inference_steps=args.n_inference_steps,
         seed=args.seed,
         save_all_images=args.save_all_images,
-        device=device,
+        device=device if not args.cpu_offloading else 'cpu',  # Use CPU if offloading is enabled
         no_optim=args.no_optim,
         regularize=args.enable_reg,
         regularization_weight=args.reg_weight,
@@ -132,9 +152,6 @@ def setup(args, loaded_model_setup=None):
         imageselect=args.imageselect,
     )
 
-    torch.cuda.empty_cache()  # Free up cached memory
-    gc.collect()
-
     # Create latents
     if args.model == "flux":
         shape = (1, 16 * 64, 64)
@@ -159,16 +176,17 @@ def setup(args, loaded_model_setup=None):
     
     enable_grad = not args.no_optim
 
+    # Final memory cleanup
     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,
+            device=device if not args.cpu_offloading else 'cpu',
             dtype=dtype,
         )
     else:
@@ -179,6 +197,8 @@ def setup(args, loaded_model_setup=None):
 
     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):
     
     if args.task == "single":