flaxdiff 0.2.9__py3-none-any.whl → 0.2.10__py3-none-any.whl

flaxdiff/inference/utils.py

@@ -121,8 +121,8 @@ def parse_config(config, overrides=None):
         'uvit': UViT,
         'diffusers_unet_simple': FlaxUNet2DConditionModel,
         'simple_dit': SimpleDiT,
-        'simple_uvit': SimpleUDiT,
         'simple_mmdit': SimpleMMDiT,
+        'simple_udit': SimpleUDiT,
     }
     
     # Map all the leaves of the model config, converting strings to appropriate types

flaxdiff/trainer/general_diffusion_trainer.py

@@ -427,7 +427,7 @@ class GeneralDiffusionTrainer(DiffusionTrainer):
         process_index = jax.process_index()
         generate_samples = val_step_fn
         
-        val_ds = iter(val_ds) if val_ds else None
+        val_ds = iter(val_ds()) if val_ds else None
         print(f"Validation loop started for process index {process_index} with {global_device_count} devices.")
         # Evaluation step
         try:
@@ -475,10 +475,11 @@ class GeneralDiffusionTrainer(DiffusionTrainer):
                 metrics = {k: np.mean(v) for k, v in metrics.items()}
                 # Update the best validation metrics
                 for key, value in metrics.items():
-                    if key not in self.best_val_metrics:
-                        self.best_val_metrics[key] = value
+                    final_key = f"val/{key}"
+                    if final_key not in self.best_val_metrics:
+                        self.best_val_metrics[final_key] = value
                     else:
-                        self.best_val_metrics[key] = min(self.best_val_metrics[key], value)
+                        self.best_val_metrics[final_key] = min(self.best_val_metrics[final_key], value)
                 # Log the best validation metrics
                 if getattr(self, 'wandb', None) is not None and self.wandb:
                     # Log the metrics
@@ -488,6 +489,13 @@ class GeneralDiffusionTrainer(DiffusionTrainer):
                         self.wandb.log({
                             f"val/{key}": value,
                         }, step=current_step)
+                    # Log the best validation metrics
+                    for key, value in self.best_val_metrics.items():
+                        if isinstance(value, jnp.ndarray):
+                            value = np.array(value)
+                        self.wandb.log({
+                            f"best_{key}": value,
+                        }, step=current_step)
                 print(f"Validation metrics for process index {process_index}: {metrics}")
             
             # Close validation dataset iterator
@@ -622,10 +630,10 @@ class GeneralDiffusionTrainer(DiffusionTrainer):
         if not runs:
             raise ValueError("No runs found in wandb.")
         print(f"Getting best runs from wandb {self.wandb.id}...")
-        runs = sorted(runs, key=lambda x: x.summary.get(metric, float('inf')))
+        runs = sorted(runs, key=lambda x: x.summary.get(f"best_{metric}", float('inf')))
         best_runs = runs[:top_k]
-        lower_bound = best_runs[-1].summary.get(metric, float('inf'))
-        upper_bound = best_runs[0].summary.get(metric, float('inf'))
+        lower_bound = best_runs[-1].summary.get(f"best_{metric}", float('inf'))
+        upper_bound = best_runs[0].summary.get(f"best_{metric}", float('inf'))
         print(f"Best runs from wandb {self.wandb.id}:")
         for run in best_runs:
             print(f"\t\tRun ID: {run.id}, Metric: {run.summary.get(metric, float('inf'))}")
@@ -649,19 +657,21 @@ class GeneralDiffusionTrainer(DiffusionTrainer):
             best_runs, bounds = self.__get_best_general_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()
+        is_lower_better = True
         
         # Check if current run is one of the best
         if metric == "train/best_loss":
             current_run_metric = self.best_loss
         elif metric in self.best_val_metrics:
+            print(f"Fetching best validation metric {metric} from local")
             current_run_metric = self.best_val_metrics[metric]
         else:
             current_run_metric = self.wandb.summary.get(metric, float('inf') if is_lower_better else float('-inf'))
-                
+        
+        print(f"Current run {self.wandb.id} metric: {current_run_metric}, Best bounds: {bounds}")
         # 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.")
+            print(f"Current run {self.wandb.id} meets performance criteria. Current metric: {current_run_metric}, Best bounds: {bounds}")
             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
             

flaxdiff/trainer/simple_trainer.py

@@ -600,7 +600,7 @@ class SimpleTrainer:
 
     def fit(self, data, train_steps_per_epoch, epochs, train_step_args={}, val_steps_per_epoch=5, validation_step_args={}):
         train_ds = iter(data['train']())
-        val_ds = data.get('val', data.get('test', None))()
+        val_ds = data.get('val', data.get('test', None))
         train_step = self._define_train_step(**train_step_args)
         val_step = self._define_validation_step(**validation_step_args)
         train_state = self.state
@@ -642,6 +642,19 @@ class SimpleTrainer:
             self.rngstate = rng_state
             total_time = end_time - start_time
             avg_time_per_step = total_time / train_steps_per_epoch
+            
+            if val_steps_per_epoch > 0:
+                print(f"Validation started for process index {process_index}")
+                # Validation step
+                self.validation_loop(
+                    train_state,
+                    val_step,
+                    val_ds,
+                    val_steps_per_epoch,
+                    current_step,
+                )
+                print(colored(f"Validation done on process index {process_index}", PROCESS_COLOR_MAP[process_index]))
+            
             avg_loss = epoch_loss / train_steps_per_epoch
             if avg_loss < self.best_loss:
                 self.best_loss = avg_loss
@@ -659,17 +672,6 @@ class SimpleTrainer:
                     }, step=current_step)
                 print(colored(f"\n\tEpoch {current_epoch} completed. Avg Loss: {avg_loss}, Time: {total_time:.2f}s, Best Loss: {self.best_loss}", 'green'))
                     
-            if val_steps_per_epoch > 0:
-                print(f"Validation started for process index {process_index}")
-                # Validation step
-                self.validation_loop(
-                    train_state,
-                    val_step,
-                    val_ds,
-                    val_steps_per_epoch,
-                    current_step,
-                )
-                print(colored(f"Validation done on process index {process_index}", PROCESS_COLOR_MAP[process_index]))
                 
         self.save(epochs)#
         return self.state

{flaxdiff-0.2.9.dist-info → flaxdiff-0.2.10.dist-info}/METADATA

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

{flaxdiff-0.2.9.dist-info → flaxdiff-0.2.10.dist-info}/RECORD

@@ -15,7 +15,7 @@ flaxdiff/data/sources/videos.py,sha256=NkxwEruNpAwDCM53q4WurQ802gSjQMOqjNLxYOqjo
 flaxdiff/data/sources/voxceleb2.py,sha256=BoKfat_hsw6ObDyyaiQmPbBzuFiqgCGlgAZmf-t5Iz8,18621
 flaxdiff/inference/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 flaxdiff/inference/pipeline.py,sha256=8S30FAlXEjvrDd87H-qdD6biySQZ3cJUflU8gdmPxig,9223
-flaxdiff/inference/utils.py,sha256=JEBZYSgj-0DLJTV-TNmIAllAqqVJMn0KfryHwFO-MFs,12606
+flaxdiff/inference/utils.py,sha256=Dh0KawgvQrZxyqN_9wbsb7gUyvPRendwb-YtAU6zIBE,12606
 flaxdiff/inputs/__init__.py,sha256=ybPjQsFAf5sqRVZG1sRiOl99EnwpI-NQ8HE3y7UbXmU,7197
 flaxdiff/inputs/encoders.py,sha256=pjfbx4Rk7bLoE80MOfThZDm6YtsDncRekmn0Bmg_CwI,2963
 flaxdiff/metrics/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -63,9 +63,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=OtE2spZIBFPpY6q-ijYol5Y-CaP2UHJYIDX3PFBiPtg,29492
-flaxdiff/trainer/simple_trainer.py,sha256=Hdltuo3lgF61N04Lxc7L3z6NLveW4_h1ff7_5mu3Wbg,28730
-flaxdiff-0.2.9.dist-info/METADATA,sha256=a8btxHRkAZVieuZfTyXgPkJbEG9fZRknEhq2Ti3_7m4,24057
-flaxdiff-0.2.9.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
-flaxdiff-0.2.9.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
-flaxdiff-0.2.9.dist-info/RECORD,,
+flaxdiff/trainer/general_diffusion_trainer.py,sha256=gMo0OOz8EFKGfiqZnDwhVSxtk_IUMGUvyt5TTr_Hk8g,30168
+flaxdiff/trainer/simple_trainer.py,sha256=nXYy9tadteG8N0RovpevPPEs6oeFvbr2gVq7Zot9l78,28754
+flaxdiff-0.2.10.dist-info/METADATA,sha256=xsqksvLSps2a9nNdvZkguWvsC07kX8A3Z26DPTq-tGI,24058
+flaxdiff-0.2.10.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
+flaxdiff-0.2.10.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
+flaxdiff-0.2.10.dist-info/RECORD,,