feat: push best alias as well

Author: AshishKumar4

flaxdiff/data/sources/av_utils.py

@@ -7,7 +7,6 @@
 import subprocess
 import numpy as np
 from typing import Tuple, Optional, Union, List
-from video_reader import PyVideoReader
 from .audio_utils import read_audio
 
 def get_video_fps(video_path: str):
@@ -113,6 +112,7 @@ def read_av_improved(
     Returns:
         Tuple of (audio_data, video_frames) where video_frames is a numpy array.
     """
+    from video_reader import PyVideoReader
     # Calculate time information for audio extraction
     start_time = start / fps if start > 0 else 0
     duration = None

flaxdiff/trainer/general_diffusion_trainer.py

@@ -484,11 +484,13 @@ def _log_image_samples(self, samples, current_step):
     def push_to_registry(
         self,
         registry_name: str = 'wandb-registry-model',
+        aliases: List[str] = ['latest'],
     ):
         """
         Push the model to wandb registry.
         Args:
             registry_name: Name of the model registry.
+            aliases: List of aliases for the model.
         """
         if self.wandb is None:
             raise ValueError("Wandb is not initialized. Cannot push to registry.")
@@ -502,6 +504,7 @@ def push_to_registry(
             artifact_or_path=latest_checkpoint_path,
             name=modelname,
             type="model",
+            aliases=aliases,
         )
 
         target_path = f"{registry_name}/{modelname}"
@@ -541,37 +544,49 @@ def __get_best_sweep_runs__(
         return best_runs, (min(lower_bound, upper_bound), max(lower_bound, upper_bound))
 
     def __compare_run_against_best__(self, top_k=2, metric="train/best_loss"):
+        """
+        Compare the current run against the best runs from the sweep.
+        Args:
+            top_k: Number of top runs to consider.
+            metric: Metric to compare against.
+        Returns:
+            is_good: Whether the current run is among the best.
+            is_best: Whether the current run is the best.
+        """
         # Get best runs
         best_runs, bounds = self.__get_best_sweep_runs__(metric=metric, top_k=top_k)
 
         # Determine if lower or higher values are better (for loss, lower is better)
         is_lower_better = "loss" in metric.lower()
 
         # Check if current run is one of the best
-        current_run_metric = self.wandb.summary.get(metric, float('inf') if is_lower_better else float('-inf'))
-        
-        # Direct check if current run is in best runs
-        for run in best_runs:
-            if run.id == self.wandb.id:
-                print(f"Current run {self.wandb.id} is one of the best runs.")
-                return True
+        if metric == "train/best_loss":
+            current_run_metric = self.best_loss
+        else:
+            current_run_metric = self.wandb.summary.get(metric, float('inf') if is_lower_better else float('-inf'))
 
-        # Backup check based on metric value
+        # Check based on bounds
         if (is_lower_better and current_run_metric < bounds[1]) or (not is_lower_better and current_run_metric > bounds[0]):
             print(f"Current run {self.wandb.id} meets performance criteria.")
-            return True
+            is_best = (is_lower_better and current_run_metric < bounds[0]) or (not is_lower_better and current_run_metric > bounds[1])
+            return True, is_best
 
-        return False
+        return False, False
 
     def save(self, epoch=0, step=0, state=None, rngstate=None):
         super().save(epoch=epoch, step=step, state=state, rngstate=rngstate)
 
         if self.wandb is not None and hasattr(self, "wandb_sweep"):
             checkpoint = get_latest_checkpoint(self.checkpoint_path())
             try:
-                if self.__compare_run_against_best__(top_k=5, metric="train/best_loss"):
-                    self.push_to_registry()
-                    print("Model pushed to registry successfully")
+                is_good, is_best = self.__compare_run_against_best__(top_k=5, metric="train/best_loss")
+                if is_good:
+                    # Push to registry with appropriate aliases
+                    aliases = ["latest"]
+                    if is_best:
+                        aliases.append("best")
+                    self.push_to_registry(aliases=aliases)
+                    print("Model pushed to registry successfully with aliases:", aliases)
                 else:
                     print("Current run is not one of the best runs. Not saving model.")