flaxdiff 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

flaxdiff/data/sources/av_utils.py

@@ -7,7 +7,6 @@ import shutil
 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 @@ class GeneralDiffusionTrainer(DiffusionTrainer):
     def push_to_registry(
         self,
         registry_name: str = 'wandb-registry-model',
+        aliases: List[str] = [],
     ):
         """
         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 @@ class GeneralDiffusionTrainer(DiffusionTrainer):
             artifact_or_path=latest_checkpoint_path,
             name=modelname,
             type="model",
+            aliases=['latest'] + aliases,
         )
         
         target_path = f"{registry_name}/{modelname}"
@@ -509,6 +512,7 @@ class GeneralDiffusionTrainer(DiffusionTrainer):
         self.wandb.link_artifact(
             artifact=logged_artifact,
             target_path=target_path,
+            aliases=aliases,
         )
         print(f"Model pushed to registry at {target_path}")
         return logged_artifact
@@ -541,6 +545,15 @@ class GeneralDiffusionTrainer(DiffusionTrainer):
         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)
         
@@ -548,20 +561,18 @@ class GeneralDiffusionTrainer(DiffusionTrainer):
         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)
@@ -569,9 +580,14 @@ class GeneralDiffusionTrainer(DiffusionTrainer):
         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 = []
+                    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.")
                 

{flaxdiff-0.2.0.dist-info → flaxdiff-0.2.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flaxdiff
-Version: 0.2.0
+Version: 0.2.2
 Summary: A versatile and easy to understand Diffusion library
 Author-email: Ashish Kumar Singh <ashishkmr472@gmail.com>
 License-Expression: MIT

{flaxdiff-0.2.0.dist-info → flaxdiff-0.2.2.dist-info}/RECORD

@@ -7,7 +7,7 @@ flaxdiff/data/dataset_map.py,sha256=NrLG1XtIxy8GcCsZ-e6eascjgsP0Xq5lVA1z3HIIYyI,
 flaxdiff/data/online_loader.py,sha256=t1jEhdB6gWTlwx68ehj1ol_PrImbwXYiRlrJPCmNgCM,35701
 flaxdiff/data/sources/audio_utils.py,sha256=X27gG1yQt_abVOYgMtruYmZD7-8_uQCRhhTSpn4clkI,4514
 flaxdiff/data/sources/av_example.py,sha256=RIcbVKqckFqbfnV65NQotzIBxjdDuM67kD1nY8fqw5Q,3826
-flaxdiff/data/sources/av_utils.py,sha256=n2qwMBQGouoBH025vdE7gitWC6RduUommUrs-SPdWe4,24041
+flaxdiff/data/sources/av_utils.py,sha256=LCr9MJNurOaoxY-sjzkLqJS_MlX0x3gRSlKAVIglAU0,24045
 flaxdiff/data/sources/base.py,sha256=uhF0odJSYRy0SLw1xnI9Q_q_xiVht2DmEYcX1j9AWT4,4246
 flaxdiff/data/sources/images.py,sha256=WpH4ywZhNol26peX3m6m5NrmDJ1K2s6fRcYHvOFlOk8,11102
 flaxdiff/data/sources/utils.py,sha256=kFzM4_kPoThbAu54ulABmEDAR33tR50NgzXIpC0Dzjk,7316
@@ -56,9 +56,9 @@ flaxdiff/schedulers/sqrt.py,sha256=mCd_szmOqF6vqQKiAiEOqV_3eBIPGYrW3VxK0o4rBuo,4
 flaxdiff/trainer/__init__.py,sha256=xSoierfi26gxfgxlNnwvyyPmuPAJ--5i3mEHxt3S-AE,215
 flaxdiff/trainer/autoencoder_trainer.py,sha256=2FP2P-k9c0n_k3eT0trkq73dQrHRdBj9ObK1idcyhSw,6996
 flaxdiff/trainer/diffusion_trainer.py,sha256=reQEVWKTqKAeyCMQ-curPOfSRmBKxKooK8EVtUuorcM,14599
-flaxdiff/trainer/general_diffusion_trainer.py,sha256=VQ5p2ZaTv2R1LM0Epz4e719_EfK2dh1eoKK3WIysIW0,24040
+flaxdiff/trainer/general_diffusion_trainer.py,sha256=7VAeT3TzCDUyns8wdZbIwXJqDKx_FYSzq8toOkaeQMI,24802
 flaxdiff/trainer/simple_trainer.py,sha256=CF2mMcc6AtBgcR1XiqKevRL0paGS0S9ZJofCns32nRM,24214
-flaxdiff-0.2.0.dist-info/METADATA,sha256=1WLpd9RQy_mJE2E2uOdXptY5Fm3n_MTNcgZyBD7YmGw,23982
-flaxdiff-0.2.0.dist-info/WHEEL,sha256=lTU6B6eIfYoiQJTZNc-fyaR6BpL6ehTzU3xGYxn2n8k,91
-flaxdiff-0.2.0.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
-flaxdiff-0.2.0.dist-info/RECORD,,
+flaxdiff-0.2.2.dist-info/METADATA,sha256=pzYYdy1zK7lbaqSRdpopZHHYx7q3BP0DL11hGTOO7h4,23982
+flaxdiff-0.2.2.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
+flaxdiff-0.2.2.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
+flaxdiff-0.2.2.dist-info/RECORD,,

{flaxdiff-0.2.0.dist-info → flaxdiff-0.2.2.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.1.1)
+Generator: setuptools (79.0.0)
 Root-Is-Purelib: true
 Tag: py3-none-any