include changes to track iteration callbacks

main.py

@@ -14,7 +14,7 @@ from rewards import get_reward_losses
 from training import LatentNoiseTrainer, get_optimizer
 
 
-def main(args):
+def main(args, progress_callback=None):
     seed_everything(args.seed)
     bf.makedirs(f"{args.save_dir}/logs/{args.task}")
     # Set up logging and name settings
@@ -99,9 +99,10 @@ def main(args):
         save_dir = f"{args.save_dir}/{args.task}/{settings}/{args.prompt}"
         os.makedirs(f"{save_dir}", exist_ok=True)
         best_image, total_init_rewards, total_best_rewards = trainer.train(
-            latents, args.prompt, optimizer, save_dir
+            latents, args.prompt, optimizer, save_dir, progress_callback=progress_callback
         )
         best_image.save(f"{save_dir}/best_image.png")
+        return best_image, total_init_rewards, total_best_rewards
     elif args.task == "example-prompts":
         fo = open("assets/example_prompts.txt", "r")
         prompts = fo.readlines()
@@ -271,4 +272,4 @@ def main(args):
     
 if __name__ == "__main__":
     args = parse_args()
-    main(args)
+    main(args)