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

flaxdiff/models/vit_common.py

@@ -0,0 +1,262 @@
+import jax
+import jax.numpy as jnp
+from flax import linen as nn
+from typing import Any, Optional
+import einops
+from flax.typing import Dtype, PrecisionLike
+
+from .attention import NormalAttention
+
+def unpatchify(x, channels=3):
+    patch_size = int((x.shape[2] // channels) ** 0.5)
+    h = w = int(x.shape[1] ** .5)
+    assert h * w == x.shape[1] and patch_size ** 2 * \
+        channels == x.shape[2], f"Invalid shape: {x.shape}, should be {h*w}, {patch_size**2*channels}"
+    x = einops.rearrange(
+        x, 'B (h w) (p1 p2 C) -> B (h p1) (w p2) C', h=h, p1=patch_size, p2=patch_size)
+    return x
+
+
+class PatchEmbedding(nn.Module):
+    patch_size: int
+    embedding_dim: int
+    dtype: Any = jnp.float32
+    precision: Any = jax.lax.Precision.HIGH
+
+    @nn.compact
+    def __call__(self, x):
+        batch, height, width, channels = x.shape
+        assert height % self.patch_size == 0 and width % self.patch_size == 0, "Image dimensions must be divisible by patch size"
+
+        x = nn.Conv(features=self.embedding_dim,
+                    kernel_size=(self.patch_size, self.patch_size),
+                    strides=(self.patch_size, self.patch_size),
+                    dtype=self.dtype,
+                    precision=self.precision)(x)
+        x = jnp.reshape(x, (batch, -1, self.embedding_dim))
+        return x
+
+
+class PositionalEncoding(nn.Module):
+    max_len: int
+    embedding_dim: int
+
+    @nn.compact
+    def __call__(self, x):
+        pe = self.param('pos_encoding',
+                        jax.nn.initializers.zeros,
+                        (1, self.max_len, self.embedding_dim))
+        return x + pe[:, :x.shape[1], :]
+
+
+# --- Rotary Positional Embedding (RoPE) ---
+# Adapted from https://github.com/google-deepmind/ring_attention/blob/main/ring_attention/layers/rotary.py
+
+
+def _rotate_half(x: jax.Array) -> jax.Array:
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2:]
+    return jnp.concatenate((-x2, x1), axis=-1)
+
+def apply_rotary_embedding(
+    x: jax.Array, freqs_cos: jax.Array, freqs_sin: jax.Array
+) -> jax.Array:
+    """Applies rotary embedding to the input tensor using rotate_half method."""
+    # x shape: [..., Sequence, Dimension] e.g. [B, H, S, D] or [B, S, D]
+    # freqs_cos/sin shape: [Sequence, Dimension / 2]
+
+    # Expand dims for broadcasting: [1, 1, S, D/2] or [1, S, D/2]
+    if x.ndim == 4:  # [B, H, S, D]
+        cos_freqs = jnp.expand_dims(freqs_cos, axis=(0, 1))
+        sin_freqs = jnp.expand_dims(freqs_sin, axis=(0, 1))
+    elif x.ndim == 3:  # [B, S, D]
+        cos_freqs = jnp.expand_dims(freqs_cos, axis=0)
+        sin_freqs = jnp.expand_dims(freqs_sin, axis=0)
+
+    # Duplicate cos and sin for the full dimension D
+    # Shape becomes [..., S, D]
+    cos_freqs = jnp.concatenate([cos_freqs, cos_freqs], axis=-1)
+    sin_freqs = jnp.concatenate([sin_freqs, sin_freqs], axis=-1)
+
+    # Apply rotation: x * cos + rotate_half(x) * sin
+    x_rotated = x * cos_freqs + _rotate_half(x) * sin_freqs
+    return x_rotated.astype(x.dtype)
+
+class RotaryEmbedding(nn.Module):
+    dim: int
+    max_seq_len: int = 4096  # Increased default based on SimpleDiT
+    base: int = 10000
+    dtype: Dtype = jnp.float32
+
+    def setup(self):
+        inv_freq = 1.0 / (
+            self.base ** (jnp.arange(0, self.dim, 2,
+                          dtype=jnp.float32) / self.dim)
+        )
+        t = jnp.arange(self.max_seq_len, dtype=jnp.float32)
+        freqs = jnp.outer(t, inv_freq)
+        self.freqs_cos = jnp.cos(freqs)
+        self.freqs_sin = jnp.sin(freqs)
+
+    def __call__(self, seq_len: int):
+        if seq_len > self.max_seq_len:
+            # Dynamically extend frequencies if needed (more robust)
+            t = jnp.arange(seq_len, dtype=jnp.float32)
+            inv_freq = 1.0 / (
+                self.base ** (jnp.arange(0, self.dim, 2,
+                              dtype=jnp.float32) / self.dim)
+            )
+            freqs = jnp.outer(t, inv_freq)
+            freqs_cos = jnp.cos(freqs)
+            freqs_sin = jnp.sin(freqs)
+            # Consider caching extended freqs if this happens often
+            return freqs_cos, freqs_sin
+            # Or raise error like before:
+            # raise ValueError(f"Sequence length {seq_len} exceeds max_seq_len {self.max_seq_len}")
+        return self.freqs_cos[:seq_len, :], self.freqs_sin[:seq_len, :]
+
+    
+# --- Attention with RoPE ---
+
+
+class RoPEAttention(NormalAttention):
+    rope_emb: RotaryEmbedding = None
+
+    @nn.compact
+    def __call__(self, x, context=None, freqs_cis=None):
+        orig_x_shape = x.shape
+        is_4d = len(orig_x_shape) == 4
+        if is_4d:
+            B, H, W, C = x.shape
+            seq_len = H * W
+            x = x.reshape((B, seq_len, C))
+        else:
+            B, seq_len, C = x.shape
+
+        context = x if context is None else context
+        if len(context.shape) == 4:
+            _B, _H, _W, _C = context.shape
+            context_seq_len = _H * _W
+            context = context.reshape((B, context_seq_len, _C))
+        # else: # context is already [B, S_ctx, C]
+
+        query = self.query(x)      # [B, S, H, D]
+        key = self.key(context)    # [B, S_ctx, H, D]
+        value = self.value(context)  # [B, S_ctx, H, D]
+
+        if freqs_cis is None and self.rope_emb is not None:
+            seq_len_q = query.shape[1]  # Use query's sequence length
+            freqs_cos, freqs_sin = self.rope_emb(seq_len_q)
+        elif freqs_cis is not None:
+            freqs_cos, freqs_sin = freqs_cis
+        else:
+            # Should not happen if rope_emb is provided or freqs_cis are passed
+            raise ValueError("RoPE frequencies not provided.")
+
+        # Apply RoPE to query and key
+        # Permute to [B, H, S, D] for RoPE application
+        query = einops.rearrange(query, 'b s h d -> b h s d')
+        key = einops.rearrange(key, 'b s h d -> b h s d')
+
+        # Apply RoPE only up to the context sequence length for keys if different
+        # Assuming self-attention or context has same seq len for simplicity here
+        query = apply_rotary_embedding(query, freqs_cos, freqs_sin)
+        key = apply_rotary_embedding(
+            key, freqs_cos, freqs_sin)  # Apply same freqs to key
+
+        # Permute back to [B, S, H, D] for dot_product_attention
+        query = einops.rearrange(query, 'b h s d -> b s h d')
+        key = einops.rearrange(key, 'b h s d -> b s h d')
+
+        hidden_states = nn.dot_product_attention(
+            query, key, value, dtype=self.dtype, broadcast_dropout=False,
+            dropout_rng=None, precision=self.precision, force_fp32_for_softmax=self.force_fp32_for_softmax,
+            deterministic=True
+        )
+
+        proj = self.proj_attn(hidden_states)
+
+        if is_4d:
+            proj = proj.reshape(orig_x_shape)
+
+        return proj
+
+
+# --- adaLN-Zero ---
+
+
+class AdaLNZero(nn.Module):
+    features: int
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
+    norm_epsilon: float = 1e-5  # Standard LayerNorm epsilon
+
+    @nn.compact
+    def __call__(self, x, conditioning):
+        # Project conditioning signal to get scale and shift parameters
+        # Conditioning shape: [B, D_cond] -> [B, 1, ..., 1, 6 * features] for broadcasting
+        # Or [B, 1, 6*features] if x is [B, S, F]
+
+        # Ensure conditioning has seq dim if x does
+        # x=[B,S,F], cond=[B,D_cond]
+        if x.ndim == 3 and conditioning.ndim == 2:
+            conditioning = jnp.expand_dims(
+                conditioning, axis=1)  # cond=[B,1,D_cond]
+
+        # Project conditioning to get 6 params per feature (scale_mlp, shift_mlp, gate_mlp, scale_attn, shift_attn, gate_attn)
+        # Using nn.DenseGeneral for flexibility if needed, but nn.Dense is fine if cond is [B, D_cond] or [B, 1, D_cond]
+        ada_params = nn.Dense(
+            features=6 * self.features,
+            dtype=self.dtype,
+            precision=self.precision,
+            # Initialize projection to zero (Zero init)
+            kernel_init=nn.initializers.zeros,
+            name="ada_proj"
+        )(conditioning)
+
+        # Split into scale, shift, gate for MLP and Attention
+        scale_mlp, shift_mlp, gate_mlp, scale_attn, shift_attn, gate_attn = jnp.split(
+            ada_params, 6, axis=-1)
+
+        scale_mlp = jnp.clip(scale_mlp, -10.0, 10.0)
+        shift_mlp = jnp.clip(shift_mlp, -10.0, 10.0)
+        # Apply Layer Normalization
+        norm = nn.LayerNorm(epsilon=self.norm_epsilon,
+                            use_scale=False, use_bias=False, dtype=self.dtype)
+        # norm = nn.RMSNorm(epsilon=self.norm_epsilon, dtype=self.dtype) # Alternative: RMSNorm
+
+        norm_x = norm(x)
+
+        # Modulate for Attention path
+        x_attn = norm_x * (1 + scale_attn) + shift_attn
+
+        # Modulate for MLP path
+        x_mlp = norm_x * (1 + scale_mlp) + shift_mlp
+
+        # Return modulated outputs and gates
+        return x_attn, gate_attn, x_mlp, gate_mlp
+
+class AdaLNParams(nn.Module): # Renamed for clarity
+    features: int
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
+
+    @nn.compact
+    def __call__(self, conditioning):
+        # Ensure conditioning is broadcastable if needed (e.g., [B, 1, D_cond])
+        if conditioning.ndim == 2:
+             conditioning = jnp.expand_dims(conditioning, axis=1)
+
+        # Project conditioning to get 6 params per feature
+        ada_params = nn.Dense(
+            features=6 * self.features,
+            dtype=self.dtype,
+            precision=self.precision,
+            kernel_init=nn.initializers.zeros,
+            name="ada_proj"
+        )(conditioning)
+        # Return all params (or split if preferred, but maybe return tuple/dict)
+        # Shape: [B, 1, 6*F]
+        return ada_params # Or split and return tuple: jnp.split(ada_params, 6, axis=-1)
+    

flaxdiff/trainer/general_diffusion_trainer.py

@@ -129,6 +129,7 @@ class GeneralDiffusionTrainer(DiffusionTrainer):
                  frames_per_sample: int = None,
                  wandb_config: Dict[str, Any] = None,
                  eval_metrics: List[EvaluationMetric] = None,
+                 best_tracker_metric: str = "train/best_loss",
                  **kwargs
                  ):
         """
@@ -196,6 +197,8 @@ class GeneralDiffusionTrainer(DiffusionTrainer):
             **kwargs
         )
         
+        self.best_tracker_metric = best_tracker_metric
+        
         # Store video-specific parameters
         self.frames_per_sample = frames_per_sample
         
@@ -203,6 +206,7 @@ class GeneralDiffusionTrainer(DiffusionTrainer):
         self.conditional_inputs = input_config.conditions
         # Determine if we're working with video or images
         self.is_video = self._is_video_data()
+        self.best_val_metrics = {}
     
     def _is_video_data(self):
         sample_data_shape = self.input_config.sample_data_shape
@@ -423,7 +427,8 @@ 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:
             metrics = {metric.name: [] for metric in self.eval_metrics} if self.eval_metrics else {}
@@ -465,11 +470,17 @@ class GeneralDiffusionTrainer(DiffusionTrainer):
                         else:  # [B,H,W,C] - Image data
                             self._log_image_samples(samples, current_step)
                     
-            if getattr(self, 'wandb', None) is not None and self.wandb:
-                # metrics is a dict of metrics
-                if metrics and type(metrics) == dict:
-                    # Flatten the metrics
-                    metrics = {k: np.mean(v) for k, v in metrics.items()}
+            # Flatten the metrics
+            if metrics:
+                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
+                    else:
+                        self.best_val_metrics[key] = min(self.best_val_metrics[key], value)
+                # Log the best validation metrics
+                if getattr(self, 'wandb', None) is not None and self.wandb:
                     # Log the metrics
                     for key, value in metrics.items():
                         if isinstance(value, jnp.ndarray):
@@ -477,7 +488,10 @@ class GeneralDiffusionTrainer(DiffusionTrainer):
                         self.wandb.log({
                             f"val/{key}": value,
                         }, step=current_step)
-                        
+                print(f"Validation metrics for process index {process_index}: {metrics}")
+            
+            # Close validation dataset iterator
+            del val_ds
         except StopIteration:
             print(f"Validation dataset exhausted for process index {process_index}")
         except Exception as e:
@@ -602,9 +616,13 @@ class GeneralDiffusionTrainer(DiffusionTrainer):
         """
         if self.wandb is None:
             raise ValueError("Wandb is not initialized. Cannot get best runs.")
-        
+        import wandb
         # Get the sweep runs
-        runs = sorted(self.wandb.runs, key=lambda x: x.summary.get(metric, float('inf')))
+        runs = [i for i in wandb.Api().runs(path=f"{self.wandb.entity}/{self.wandb.project}", filters={"config.dataset.name": self.wandb.config['dataset']['name']})]
+        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')))
         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'))
@@ -636,6 +654,8 @@ class GeneralDiffusionTrainer(DiffusionTrainer):
         # 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:
+            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'))
                 
@@ -653,7 +673,7 @@ class GeneralDiffusionTrainer(DiffusionTrainer):
         if self.wandb is not None:
             checkpoint = get_latest_checkpoint(self.checkpoint_path())
             try:
-                is_good, is_best = self.__compare_run_against_best__(top_k=5, metric="train/best_loss", from_sweeps=hasattr(self, "wandb_sweep"))
+                is_good, is_best = self.__compare_run_against_best__(top_k=5, metric=self.best_tracker_metric, from_sweeps=hasattr(self, "wandb_sweep"))
                 if is_good:
                     # Push to registry with appropriate aliases
                     aliases = []

flaxdiff/trainer/simple_trainer.py

@@ -26,6 +26,7 @@ from flax.training.dynamic_scale import DynamicScale
 from flaxdiff.utils import RandomMarkovState
 from flax.training import dynamic_scale as dynamic_scale_lib
 from dataclasses import dataclass
+import shutil
 import gc
 
 PROCESS_COLOR_MAP = {
@@ -73,6 +74,76 @@ class SimpleTrainState(train_state.TrainState):
     metrics: Metrics
     dynamic_scale: dynamic_scale_lib.DynamicScale
 
+def move_contents_to_subdir(target_dir, new_subdir_name):
+    # --- 1. Validate Target Directory ---
+    if not os.path.isdir(target_dir):
+        print(f"Error: Target directory '{target_dir}' not found or is not a directory.")
+        return
+    # --- 2. Define Paths ---
+    # Construct the full path for the new subdirectory
+    new_subdir_path = os.path.join(target_dir, new_subdir_name)
+    # --- 3. Create New Subdirectory ---
+    try:
+        # Create the subdirectory.
+        # exist_ok=True prevents an error if the directory already exists.
+        os.makedirs(new_subdir_path, exist_ok=True)
+        print(f"Subdirectory '{new_subdir_path}' created or already exists.")
+    except OSError as e:
+        print(f"Error creating subdirectory '{new_subdir_path}': {e}")
+        return # Stop execution if subdirectory creation fails
+    # --- 4. List Contents of Target Directory ---
+    try:
+        items_to_move = os.listdir(target_dir)
+    except OSError as e:
+        print(f"Error listing contents of '{target_dir}': {e}")
+        return # Stop if we can't list directory contents
+    # --- 5. Move Items ---
+    print(f"Moving items from '{target_dir}' to '{new_subdir_path}'...")
+    moved_count = 0
+    error_count = 0
+    for item_name in items_to_move:
+        # Construct the full path of the item in the target directory
+        source_path = os.path.join(target_dir, item_name)
+        # IMPORTANT: Skip the newly created subdirectory itself!
+        if source_path == new_subdir_path:
+            continue
+        # Construct the destination path inside the new subdirectory
+        destination_path = os.path.join(new_subdir_path, item_name)
+        # Move the item
+        try:
+            shutil.move(source_path, destination_path)
+            # print(f"  Moved: '{item_name}'") # Uncomment for verbose output
+            moved_count += 1
+        except Exception as e:
+            print(f"  Error moving '{item_name}': {e}")
+            error_count += 1
+    print(f"\nOperation complete.")
+    print(f"  Successfully moved: {moved_count} item(s).")
+    if error_count > 0:
+        print(f"  Errors encountered: {error_count} item(s).")
+
+def load_from_checkpoint(
+    checkpoint_dir: str,
+):
+    try:
+        checkpointer = orbax.checkpoint.PyTreeCheckpointer()
+        options = orbax.checkpoint.CheckpointManagerOptions(create=False)
+        # Convert checkpoint_dir to absolute path
+        checkpoint_dir = os.path.abspath(checkpoint_dir)
+        manager = orbax.checkpoint.CheckpointManager(checkpoint_dir, checkpointer, options)
+        ckpt = manager.restore(checkpoint_dir)
+        # Extract as above
+        state, best_state = None, None
+        if 'state' in ckpt:
+            state = ckpt['state']
+        if 'best_state' in ckpt:
+            best_state = ckpt['best_state']
+        print(f"Loaded checkpoint from local dir {checkpoint_dir}")
+        return state, best_state
+    except Exception as e:
+        print(f"Warning: Failed to load checkpoint from local dir: {e}")
+        return None, None
+    
 @dataclass
 class SimpleTrainer:
     state: SimpleTrainState
@@ -97,6 +168,7 @@ class SimpleTrainer:
                  checkpoint_step: int = None,
                  use_dynamic_scale: bool = False,
                  max_checkpoints_to_keep: int = 2,
+                 train_start_step_override: int = None,
                  ):
         if distributed_training is None or distributed_training is True:
             # Auto-detect if we are running on multiple devices
@@ -112,11 +184,32 @@ class SimpleTrainer:
         self.input_shapes = input_shapes
         self.checkpoint_base_path = checkpoint_base_path
         
+        load_directly_from_dir = False
+        
         if wandb_config is not None and jax.process_index() == 0:
             import wandb
             run = wandb.init(resume='allow', **wandb_config)
             self.wandb = run
             
+            if 'id' in wandb_config:
+                # If resuming from a previous run, and train_start_step_override is not set, 
+                # set the start step to the last step of the previous run
+                if train_start_step_override is None:
+                    train_start_step_override = run.summary['train/step'] + 1
+                print(f"Resuming from previous run {wandb_config['id']} with start step {train_start_step_override}")
+                
+                # If load_from_checkpoint is not set, and an artifact is found, load the artifact
+                if load_from_checkpoint is None:
+                    api_run = wandb.Api().run(f"{wandb_config['entity']}/{wandb_config['project']}/{wandb_config['id']}")
+                    model_artifacts = [i for i in api_run.logged_artifacts() if i.type == 'model']
+                    if model_artifacts:
+                        artifact = model_artifacts[0]
+                        artifact_dir = artifact.download()
+                        print(f"Loading model from artifact {artifact.name} at {artifact_dir}")
+                        # Move the artifact's contents
+                        load_from_checkpoint = artifact_dir
+                        load_directly_from_dir = True
+            
             # define our custom x axis metric
             self.wandb.define_metric("train/step")
             self.wandb.define_metric("train/epoch")
@@ -142,12 +235,16 @@ class SimpleTrainer:
             self.checkpoint_path(), async_checkpointer, options)
 
         if load_from_checkpoint is not None:
-            latest_epoch, latest_step, old_state, old_best_state, rngstate = self.load(load_from_checkpoint, checkpoint_step)
+            latest_step, old_state, old_best_state, rngstate = self.load(load_from_checkpoint, checkpoint_step, load_directly_from_dir)
         else:
-            latest_epoch, latest_step, old_state, old_best_state, rngstate = 0, 0, None, None, None
+            latest_step, old_state, old_best_state, rngstate = 0, None, None, None
 
         self.latest_step = latest_step
         
+        if train_start_step_override is not None:
+            self.latest_step = train_start_step_override
+            print(f"Overriding start step to {self.latest_step}")
+        
         if rngstate:
             self.rngstate = RandomMarkovState(**rngstate)
         else:
@@ -239,15 +336,12 @@ class SimpleTrainer:
             os.makedirs(path)
         return path
 
-    def load(self, checkpoint_path=None, checkpoint_step=None):
-        if checkpoint_path is None:
-            checkpointer = self.checkpointer
-        else:
-            checkpointer = orbax.checkpoint.PyTreeCheckpointer()
-            options = orbax.checkpoint.CheckpointManagerOptions(
-                max_to_keep=4, create=False)
-            checkpointer = orbax.checkpoint.CheckpointManager(
-                checkpoint_path, checkpointer, options)    
+    def load(self, checkpoint_path, checkpoint_step=None, load_directly_from_dir=False):
+        checkpointer = orbax.checkpoint.PyTreeCheckpointer()
+        options = orbax.checkpoint.CheckpointManagerOptions(
+            max_to_keep=4, create=False)
+        checkpointer = orbax.checkpoint.CheckpointManager(
+            checkpoint_path, checkpointer, options)    
         
         if checkpoint_step is None:
             step = checkpointer.latest_step()
@@ -259,7 +353,8 @@ class SimpleTrainer:
             checkpoint_path if checkpoint_path else self.checkpoint_path(),
             f"{step}")
         self.loaded_checkpoint_path = loaded_checkpoint_path
-        ckpt = checkpointer.restore(step)
+        ckpt = checkpointer.restore(step) if not load_directly_from_dir else checkpointer.restore(checkpoint_path)
+        
         state = ckpt['state']
         best_state = ckpt['best_state']
         rngstate = ckpt['rngs']
@@ -268,10 +363,8 @@ class SimpleTrainer:
         if self.best_loss == 0:
             # It cant be zero as that must have been some problem
             self.best_loss = 1e9
-        current_epoch = ckpt.get('epoch', step) # Must be a checkpoint from an older version which used epochs instead of steps
-        print(
-            f"Loaded model from checkpoint at epoch {current_epoch} step {step}", ckpt['best_loss'])
-        return current_epoch, step, state, best_state, rngstate
+        print(f"Loaded model from checkpoint at step {step}", ckpt['best_loss'])
+        return step, state, best_state, rngstate
 
     def save(self, epoch=0, step=0, state=None, rngstate=None):
         print(f"Saving model at epoch {epoch} step {step}")
@@ -507,6 +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))()
         train_step = self._define_train_step(**train_step_args)
         val_step = self._define_validation_step(**validation_step_args)
         train_state = self.state
@@ -520,7 +614,7 @@ class SimpleTrainer:
             self.validation_loop(
                 train_state,
                 val_step,
-                data.get('val', data.get('test', None)),
+                val_ds,
                 val_steps_per_epoch,
                 self.latest_step,
             )
@@ -571,11 +665,11 @@ class SimpleTrainer:
                 self.validation_loop(
                     train_state,
                     val_step,
-                    data.get('val', data.get('test', None)),
+                    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)
+        self.save(epochs)#
         return self.state

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

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

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

@@ -2,20 +2,20 @@ flaxdiff/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 flaxdiff/utils.py,sha256=DmlWUY1FGz4ESxIHaPQJf92CHjsdMjyDd651wFUtyNg,8838
 flaxdiff/data/__init__.py,sha256=8W5y7NyAOWtpLi8WRawk4VYeE3DMDnM3B_jKPD8BoFQ,143
 flaxdiff/data/benchmark_decord.py,sha256=x56Db1VPmziv_9KJvWdfS0O7cffsYkF5tt5WvldOKc0,13720
-flaxdiff/data/dataloaders.py,sha256=LV8ugqoB86yihfYeOJZHHdRZJNmZ63A2NQkdILMR9QA,23564
-flaxdiff/data/dataset_map.py,sha256=_6SYnmrYO-URDd8vPAmALTV6r0eMGWWmwUtsdjKGXnA,5072
+flaxdiff/data/dataloaders.py,sha256=k_3YGJhiY2Wt_-7qK0Yjl4pmF2QJjX_-BlSFuXbH5-M,23628
+flaxdiff/data/dataset_map.py,sha256=p30U23RkfgMbR8kfPBDIjrjfzDBszWQ9Q1ff2BvDYZk,5116
 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=LCr9MJNurOaoxY-sjzkLqJS_MlX0x3gRSlKAVIglAU0,24045
-flaxdiff/data/sources/base.py,sha256=uhF0odJSYRy0SLw1xnI9Q_q_xiVht2DmEYcX1j9AWT4,4246
-flaxdiff/data/sources/images.py,sha256=RFLtKW1xzw6ZPVXtCMmnTg1MPb8dc7rP77rZWbK7qpo,11796
+flaxdiff/data/sources/base.py,sha256=4Rm9pCtXxzoB8FO0lkDHsrX3ULoU_PNNcid978e6ir0,4610
+flaxdiff/data/sources/images.py,sha256=ZHBmZ2fnPN75Hc2kiog-Wcs_NZJZOiqw4WcSH5WZJHA,16572
 flaxdiff/data/sources/utils.py,sha256=kFzM4_kPoThbAu54ulABmEDAR33tR50NgzXIpC0Dzjk,7316
-flaxdiff/data/sources/videos.py,sha256=CVpOH6A4P2D8iv3gZIhd2GB5ATUD8Vsm_wVYbbugWD4,9359
+flaxdiff/data/sources/videos.py,sha256=NkxwEruNpAwDCM53q4WurQ802gSjQMOqjNLxYOqjoNE,9545
 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=MVnWl0LnC-1ILk0SsLd1YFu6igaQFR7mGhzo0jE797E,12323
+flaxdiff/inference/utils.py,sha256=JEBZYSgj-0DLJTV-TNmIAllAqqVJMn0KfryHwFO-MFs,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
@@ -27,14 +27,17 @@ flaxdiff/metrics/ssim.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 flaxdiff/metrics/utils.py,sha256=YuuOfqvqgIjsceupwNeJ59vQ2TnGeNMIyKdkIqOmoNg,1702
 flaxdiff/models/__init__.py,sha256=amtDF07DfiAdnZsvWX4eaW79nwNEU1s8Zb4PB3ewtg4,118
 flaxdiff/models/attention.py,sha256=YkED3_MRTjI9aTMTTQdsuReHhG8MK0Z4OVuU2j8ZAHs,13524
-flaxdiff/models/better_uvit.py,sha256=wPxvYBjuWQH6-OqW79VedzN6_WRY1f2mysPxaciWLww,15598
-flaxdiff/models/common.py,sha256=0j9AAjGPgkBLHo2DlYj0R6OsUNw2QaoDjaXSKq2mqkA,12647
+flaxdiff/models/common.py,sha256=QpciwuJldvLUwyAyWBQqiPPGVI-c9qLR7h7C1YoRX7w,10510
 flaxdiff/models/favor_fastattn.py,sha256=79Ew1nqarsNLPzZaBSd1ILORzJr74CupYeqGiCQK5E4,27689
 flaxdiff/models/general.py,sha256=7xMME6KVKQY8sScyHYH4f-Kek4j1pRfplKShFXwVZd4,587
+flaxdiff/models/hilbert.py,sha256=AjlAv49dL6UAYWslMJfCMLiFqY4kTgpiUWr2nc1mk34,24823
+flaxdiff/models/simple_dit.py,sha256=l238MYHRTArv_pS57aY24C2PTfxeL8EmzJ24iQqdoWI,11702
+flaxdiff/models/simple_mmdit.py,sha256=ARk0juopn2k7giln5BAUrnYD1pTFwgTJoSzrhozQ0A8,31356
 flaxdiff/models/simple_unet.py,sha256=pjeixszG_6gEY5PNFbQ7KbOyg4z5bfn4RUbINCJexOM,10758
-flaxdiff/models/simple_vit.py,sha256=no0o3os8dEKGU5I4PMBJlXq6XKjhUex8S8uZ9BDPZS4,7971
+flaxdiff/models/simple_vit.py,sha256=J9s3hBF87_iVrJDBe2cs9a56N7ect6pux_f_ge07XXc,17357
 flaxdiff/models/unet_3d.py,sha256=LF0PMxBKGU-_lAMtO_Coxy1yRE02yKKdgb7i6YZxI_4,20163
 flaxdiff/models/unet_3d_blocks.py,sha256=lRYDc9X1VEu54Kg7xEEphXYiQ09tabPXKi-hEcKFYug,19687
+flaxdiff/models/vit_common.py,sha256=1OGu4ezY3uzKinTnw3p8YkQAslHDqEbN78JheXnTleI,9831
 flaxdiff/models/autoencoder/__init__.py,sha256=qY-7MldZpsfkF-_T2LqlRK7VHbqfmosz0NmvzDlBkOk,78
 flaxdiff/models/autoencoder/autoencoder.py,sha256=8XWdsWvsPsyWGtzpCT8w0KXi_ZLGpRuQpn4oXo1gHKw,6039
 flaxdiff/models/autoencoder/diffusers.py,sha256=tPz77YuctrT--jF2AOL8G6vr0NiIr3RXANNrZCxe0bg,5921
@@ -60,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=BeDpJzgR8bUClJI4epQXlAul27MwiSfRW0lIBZSiPWk,28342
-flaxdiff/trainer/simple_trainer.py,sha256=raLS1shwpjJBT_bYXLAB2E4kA9MbwasDTzDTUqfCCUc,24312
-flaxdiff-0.2.7.dist-info/METADATA,sha256=nwglJYeF2lH_MNq5PeFLR8TSPU-I9tzJUcBbTaLYxRM,24057
-flaxdiff-0.2.7.dist-info/WHEEL,sha256=wXxTzcEDnjrTwFYjLPcsW_7_XihufBwmpiBeiXNBGEA,91
-flaxdiff-0.2.7.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
-flaxdiff-0.2.7.dist-info/RECORD,,
+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-0.2.7.dist-info → flaxdiff-0.2.9.dist-info}/WHEEL

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