model setup optimizations

main.py

@@ -15,12 +15,11 @@ from rewards import get_reward_losses
 from training import LatentNoiseTrainer, get_optimizer
 
 
-def setup(args):
-    
+def setup(args, loaded_model_setup=None):
     seed_everything(args.seed)
     bf.makedirs(f"{args.save_dir}/logs/{args.task}")
+    
     # Set up logging and name settings
-    # Get the root logger and clear existing handlers
     logger = logging.getLogger()
     logger.handlers.clear()  # Clear existing handlers
     settings = (
@@ -34,6 +33,7 @@ def setup(args):
         f"{'_imagereward' + str(args.imagereward_weighting) if args.enable_imagereward else ''}"
         f"{'_aesthetic' + str(args.aesthetic_weighting) if args.enable_aesthetic else ''}"
     )
+    
     file_stream = open(f"{args.save_dir}/logs/{args.task}/{settings}.txt", "w")
     handler = logging.StreamHandler(file_stream)
     formatter = logging.Formatter("%(asctime)s - %(message)s")
@@ -43,16 +43,68 @@ def setup(args):
     consoleHandler = logging.StreamHandler()
     consoleHandler.setFormatter(formatter)
     logger.addHandler(consoleHandler)
+    
     logging.info(args)
+    
     if args.device_id is not None:
         logging.info(f"Using CUDA device {args.device_id}")
         os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
         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
+
+    # 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
+        
+        # Update trainer with the new arguments
+        trainer.n_iters = args.n_iters
+        trainer.n_inference_steps = args.n_inference_steps
+        trainer.seed = args.seed
+        trainer.save_all_images = args.save_all_images
+        trainer.no_optim = args.no_optim
+        trainer.regularize = args.enable_reg
+        trainer.regularization_weight = args.reg_weight
+        trainer.grad_clip = args.grad_clip
+        trainer.log_metrics = args.task == "single" or not args.no_optim
+        trainer.imageselect = args.imageselect
+        
+        # Get latents (this step is still required)
+        if args.model == "flux":
+            shape = (1, 16 * 64, 64)
+        elif args.model != "pixart":
+            height = trainer.model.unet.config.sample_size * trainer.model.vae_scale_factor
+            width = trainer.model.unet.config.sample_size * trainer.model.vae_scale_factor
+            shape = (
+                1,
+                trainer.model.unet.in_channels,
+                height // trainer.model.vae_scale_factor,
+                width // trainer.model.vae_scale_factor,
+            )
+        else:
+            height = trainer.model.transformer.config.sample_size * trainer.model.vae_scale_factor
+            width = trainer.model.transformer.config.sample_size * trainer.model.vae_scale_factor
+            shape = (
+                1,
+                trainer.model.transformer.config.in_channels,
+                height // trainer.model.vae_scale_factor,
+                width // trainer.model.vae_scale_factor,
+            )
+        
+        multi_apply_fn = loaded_model_setup[6]
+        enable_grad = not args.no_optim
+
+        return args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings
+
+    # Proceed with full model loading if args.model is different
+    print(f"Loading new model: {args.model}")
+    
     # Get reward losses
     reward_losses = get_reward_losses(args, dtype, device, args.cache_dir)
 
@@ -63,7 +115,7 @@ def setup(args):
     
     torch.cuda.empty_cache()  # Free up cached memory
     gc.collect()
-    
+
     trainer = LatentNoiseTrainer(
         reward_losses=reward_losses,
         model=pipe,
@@ -85,7 +137,6 @@ def setup(args):
 
     # Create latents
     if args.model == "flux":
-        # currently only support 512x512 generation
         shape = (1, 16 * 64, 64)
     elif args.model != "pixart":
         height = pipe.unet.config.sample_size * pipe.vae_scale_factor
@@ -107,6 +158,9 @@ def setup(args):
         )
     
     enable_grad = not args.no_optim
+
+    torch.cuda.empty_cache()  # Free up cached memory
+    gc.collect()
     
     if args.enable_multi_apply:
         multi_apply_fn = get_multi_apply_fn(
@@ -121,6 +175,7 @@ def setup(args):
         multi_apply_fn = None    
 
     torch.cuda.empty_cache()  # Free up cached memory
+    gc.collect()
 
     return args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings
 
@@ -308,7 +363,7 @@ 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)
+    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)
 
 if __name__ == "__main__":