[WIP][Fix] order of sigmas, dataloader issue to debug

Author: Eigensystem

fastvideo/distill/solver.py

@@ -242,11 +242,16 @@ class EulerSolver:
 
     def __init__(self, sigmas, timesteps=1000, euler_timesteps=50):
         self.step_ratio = timesteps // euler_timesteps
+    
         self.euler_timesteps = (np.arange(1, euler_timesteps + 1) * self.step_ratio).round().astype(np.int64) - 1
         self.euler_timesteps_prev = np.asarray([0] + self.euler_timesteps[:-1].tolist())
         self.sigmas = sigmas[self.euler_timesteps]
         self.sigmas_prev = np.asarray([sigmas[0]] +
                                       sigmas[self.euler_timesteps[:-1]].tolist())  # either use sigma0 or 0
+        print(f"sigmas: {sigmas}")
+        print(f"euler_timesteps: {self.euler_timesteps}")
+        print(f"sigmas: {self.sigmas}")
+        print(f"sigmas_prev: {self.sigmas_prev}")
 
         self.euler_timesteps = torch.from_numpy(self.euler_timesteps).long()
         self.euler_timesteps_prev = torch.from_numpy(self.euler_timesteps_prev).long()

fastvideo/v1/training/distillation_pipeline.py

@@ -150,7 +150,7 @@ def initialize_distillation_pipeline(self, fastvideo_args: TrainingArgs):
             sigmas = noise_scheduler.sigmas
 
         self.solver = EulerSolver(
-            sigmas.numpy()[::-1],
+            sigmas.numpy(),
             noise_scheduler.config.num_train_timesteps,
             euler_timesteps=fastvideo_args.num_euler_timesteps,
         )
@@ -324,7 +324,19 @@ def log_validation(self, transformer, fastvideo_args, global_step):
                 x = torchvision.utils.make_grid(x, nrow=6)
                 x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
                 frames.append((x * 255).numpy().astype(np.uint8))
-            videos.append(frames)
+            # videos.append(frames)
+            videos = [frames]
+
+            video_filenames = []
+            video_captions = []
+            for i, video in enumerate(videos):
+                caption = captions[i]
+                filename = os.path.join(
+                    fastvideo_args.output_dir,
+                    f"validation_step_{global_step}_video_{i}.mp4")
+                imageio.mimsave(filename, video, fps=sampling_param.fps)
+                video_filenames.append(filename)
+                video_captions.append(caption)
 
         # Log validation results
         if self.rank == 0: