add generate video during training

Author: susanbao

src/maxdiffusion/generate_wan.py

@@ -19,9 +19,34 @@
 from maxdiffusion import pyconfig, max_logging, max_utils
 from absl import app
 from maxdiffusion.utils import export_to_video
+import os
 
 jax.config.update("jax_use_shardy_partitioner", True)
 
+def inference_generate_video(config, pipeline, filename_prefix=""):
+  s0 = time.perf_counter()
+  prompt = [config.prompt] * config.global_batch_size_to_train_on
+  negative_prompt = [config.negative_prompt] * config.global_batch_size_to_train_on
+
+  max_logging.log(
+      f"Num steps: {config.num_inference_steps}, height: {config.height}, width: {config.width}, frames: {config.num_frames}, video: {filename_prefix}"
+  )
+
+  videos = pipeline(
+      prompt=prompt,
+      negative_prompt=negative_prompt,
+      height=config.height,
+      width=config.width,
+      num_frames=config.num_frames,
+      num_inference_steps=config.num_inference_steps,
+      guidance_scale=config.guidance_scale,
+  )
+
+  print(f"video {filename_prefix}, compile time: {(time.perf_counter() - s0)}")
+  for i in range(len(videos)):
+    video_path = os.path.join(config.output_dir, "videos", f"{filename_prefix}wan_output_{config.seed}_{i}.mp4")
+    export_to_video(videos[i], video_path, fps=config.fps)
+  return
 
 def run(config, pipeline=None, filename_prefix=""):
   print("seed: ", config.seed)
@@ -54,7 +79,7 @@ def run(config, pipeline=None, filename_prefix=""):
   print("compile time: ", (time.perf_counter() - s0))
   saved_video_path = []
   for i in range(len(videos)):
-    video_path = f"{filename_prefix}wan_output_{config.seed}_{i}.mp4"
+    video_path = os.path.join(config.output_dir, "videos", f"{filename_prefix}wan_output_{config.seed}_{i}.mp4")
     export_to_video(videos[i], video_path, fps=config.fps)
     saved_video_path.append(video_path)
 

src/maxdiffusion/trainers/wan_trainer.py

@@ -37,6 +37,7 @@
 from skimage.metrics import structural_similarity as ssim
 from flax.training import train_state
 from maxdiffusion.pipelines.wan.wan_pipeline import WanPipeline
+from maxdiffusion.generate_wan import inference_generate_video
 
 
 class TrainState(train_state.TrainState):
@@ -251,6 +252,7 @@ def training_loop(self, pipeline, optimizer, learning_rate_scheduler, train_data
         if self.config.eval_every > 0 and (step + 1) % self.config.eval_every == 0:
           # Re-create the iterator each time you start evaluation to reset it
           # This assumes your data loading logic can be called to get a fresh iterator.
+          inference_generate_video(self.config, pipeline, filename_prefix=f"{step+1}-train_steps-")
           eval_data_iterator = self.load_dataset(mesh, is_training=False)
           eval_rng = jax.random.key(self.config.seed + step)
           eval_metrics = []