update LatentNoiseTrainer class to include iteration callbacks

training/trainer.py

@@ -51,6 +51,7 @@ class LatentNoiseTrainer:
         prompt: str,
         optimizer: torch.optim.Optimizer,
         save_dir: Optional[str] = None,
+        progress_callback=None,
     ) -> Tuple[PIL.Image.Image, Dict[str, float], Dict[str, float]]:
         logging.info(f"Optimizing latents for prompt '{prompt}'.")
         best_loss = torch.inf
@@ -120,6 +121,8 @@ class LatentNoiseTrainer:
                 image_numpy = image.detach().cpu().permute(0, 2, 3, 1).float().numpy()
                 image_pil = DiffusionPipeline.numpy_to_pil(image_numpy)[0]
                 image_pil.save(f"{save_dir}/{iteration}.png")
+            if progress_callback:
+                progress_callback(iteration + 1)
         image_numpy = best_image.detach().cpu().permute(0, 2, 3, 1).float().numpy()
         image_pil = DiffusionPipeline.numpy_to_pil(image_numpy)[0]
-        return image_pil, initial_rewards, best_rewards
+        return image_pil, initial_rewards, best_rewards