flaxdiff 0.2.9__tar.gz → 0.2.11__tar.gz

{flaxdiff-0.2.9 → flaxdiff-0.2.11}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: flaxdiff
-Version: 0.2.9
+Version: 0.2.11
 Summary: A versatile and easy to understand Diffusion library
 Author-email: Ashish Kumar Singh <ashishkmr472@gmail.com>
-License-Expression: MIT
+License: MIT
 Description-Content-Type: text/markdown
 Requires-Dist: flax>=0.8.4
 Requires-Dist: jax>=0.4.28
@@ -22,8 +22,6 @@ Requires-Dist: python-dotenv
 
 # ![](images/logo.jpeg "FlaxDiff")
 
-**This project is being used for the UMD Course project MSML 605: MLOps**
-
 **This project is partially supported by [Google TPU Research Cloud](https://sites.research.google/trc/about/). I would like to thank the Google Cloud TPU team for providing me with the resources to train the bigger text-conditional models in multi-host distributed settings.**
 
 ## A Versatile and simple Diffusion Library

{flaxdiff-0.2.9 → flaxdiff-0.2.11}/README.md

@@ -1,7 +1,5 @@
 # ![](images/logo.jpeg "FlaxDiff")
 
-**This project is being used for the UMD Course project MSML 605: MLOps**
-
 **This project is partially supported by [Google TPU Research Cloud](https://sites.research.google/trc/about/). I would like to thank the Google Cloud TPU team for providing me with the resources to train the bigger text-conditional models in multi-host distributed settings.**
 
 ## A Versatile and simple Diffusion Library

{flaxdiff-0.2.9 → flaxdiff-0.2.11}/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-0.2.9 → flaxdiff-0.2.11}/flaxdiff/models/hilbert.py

@@ -241,7 +241,6 @@ def hilbert_unpatchify(x: jnp.ndarray, inv_idx: jnp.ndarray, patch_size: int, H:
     target_shape = (B, total_patches_expected, patch_dim)
 
     # Create indices for gathering from x (Hilbert order h) based on inv_idx (map k -> h)
-    # inv_idx contains the 'h' index for each 'k' index.
     # Clamp invalid indices (-1) to 0; we'll mask these results later.
     # Values must be < N (the actual number of patches in x).
     h_indices_for_gather = jnp.maximum(inv_idx, 0) # Shape [total_patches_expected]
@@ -320,26 +319,37 @@ def visualize_hilbert_curve(H: int, W: int, patch_size: int, figsize=(12, 5)):
     orig_grid = np.arange(H_P * W_P).reshape((H_P, W_P))
     im0 = axes[0].imshow(orig_grid, cmap='viridis', aspect='auto')
     axes[0].set_title(f"Original Grid ({H_P}x{W_P})\n(Row-Major Order)")
-    # Add text labels for indices
-    for r in range(H_P):
-        for c in range(W_P):
-            axes[0].text(c, r, f'{orig_grid[r, c]}', ha='center', va='center', color='white' if orig_grid[r,c] < (H_P*W_P)/2 else 'black', fontsize=8)
+    # Remove text labels for indices
     axes[0].set_xticks(np.arange(W_P))
     axes[0].set_yticks(np.arange(H_P))
     axes[0].set_xticklabels(np.arange(W_P))
     axes[0].set_yticklabels(np.arange(H_P))
     plt.colorbar(im0, ax=axes[0], fraction=0.046, pad=0.04, label="Row-Major Index")
 
+    # Draw the row-major scanning line
+    if H_P * W_P > 1:
+        row_major_coords_y = []
+        row_major_coords_x = []
+        for r_idx in range(H_P * W_P):
+            r = r_idx // W_P
+            c = r_idx % W_P
+            row_major_coords_y.append(r + 0.5) # Cell center
+            row_major_coords_x.append(c + 0.5) # Cell center
+        axes[0].plot(row_major_coords_x, row_major_coords_y, color='red', linestyle='-', linewidth=1.5, alpha=0.8)
+        axes[0].plot(row_major_coords_x[0], row_major_coords_y[0], 'go', markersize=5, label='Start') # Smaller marker
+        axes[0].plot(row_major_coords_x[-1], row_major_coords_y[-1], 'mo', markersize=5, label='End') # Smaller marker
+        axes[0].legend(fontsize='x-small')
+
     # --- Plot 2: Hilbert Curve Ordering ---
     # Mask unmapped cells for visualization
     masked_grid = np.ma.masked_where(grid == -1, grid)
     im1 = axes[1].imshow(masked_grid, cmap=cmap, aspect='auto', vmin=0, vmax=max(0, len(idx)-1))
     axes[1].set_title(f"Hilbert Curve Ordering ({len(idx)} points)")
     # Add text labels for Hilbert indices
-    for r in range(H_P):
-        for c in range(W_P):
-            if grid[r,c] != -1:
-                axes[1].text(c, r, f'{int(grid[r, c])}', ha='center', va='center', color='black', fontsize=8)
+    # for r in range(H_P):
+    #     for c in range(W_P):
+    #         if grid[r,c] != -1:
+    #             axes[1].text(c, r, f'{int(grid[r, c])}', ha='center', va='center', color='black', fontsize=8)
     axes[1].set_xticks(np.arange(W_P))
     axes[1].set_yticks(np.arange(H_P))
     axes[1].set_xticklabels(np.arange(W_P))
@@ -355,22 +365,17 @@ def visualize_hilbert_curve(H: int, W: int, patch_size: int, figsize=(12, 5)):
         for i in range(len(idx)):
              if i in row_col_map:
                  coords.append(row_col_map[i])
-             # Fallback (slower):
-             # row_indices, col_indices = np.where(grid == i)
-             # if len(row_indices) > 0:
-             #     coords.append((row_indices[0], col_indices[0]))
 
         if coords:
              # Get coordinates for plotting (centers of cells)
-             y_coords = [r + 0.5 for r, c in coords]
-             x_coords = [c + 0.5 for r, c in coords]
-             axes[1].plot(x_coords, y_coords, color='black', linestyle='-', linewidth=1.5, alpha=0.8)
+             y_coords = [r + 0.5 for r, c in coords] # Cell center
+             x_coords = [c + 0.5 for r, c in coords] # Cell center
+             axes[1].plot(x_coords, y_coords, color='red', linestyle='-', linewidth=1.5, alpha=0.8) # Ensure Hilbert curve is red
              # Mark start point
-             axes[1].plot(x_coords[0], y_coords[0], 'go', markersize=8, label='Start (Idx 0)') # Green circle
+             axes[1].plot(x_coords[0], y_coords[0], 'go', markersize=5, label='Start') # Smaller marker
              # Mark end point
-             axes[1].plot(x_coords[-1], y_coords[-1], 'mo', markersize=8, label=f'End (Idx {len(idx)-1})') # Magenta circle
-             axes[1].legend(fontsize='small')
-
+             axes[1].plot(x_coords[-1], y_coords[-1], 'mo', markersize=5, label='End') # Smaller marker
+             axes[1].legend(fontsize='x-small')
 
     plt.tight_layout()
     return fig

flaxdiff-0.2.11/flaxdiff/models/ssm_dit.py

@@ -0,0 +1,560 @@
+"""
+SSM-based DiT Block and Hybrid SSM-Attention DiT for diffusion models.
+
+Implements S5 (Simplified State Space) layers as a drop-in replacement for
+attention in DiT blocks. The S5 layer uses diagonal state spaces with
+JAX associative_scan for efficient parallel computation on TPUs.
+
+References:
+    - Smith et al., "Simplified State Space Layers for Sequence Modeling", ICLR 2023
+    - Peebles & Xie, "Scalable Diffusion Models with Transformers", ICCV 2023
+"""
+
+import jax
+import jax.numpy as jnp
+from flax import linen as nn
+from typing import Callable, Any, Optional, Tuple, Sequence, Union
+import einops
+from functools import partial
+
+from .vit_common import PatchEmbedding, unpatchify, RotaryEmbedding, RoPEAttention, AdaLNParams
+from .common import kernel_init, FourierEmbedding, TimeProjection
+from .attention import NormalAttention
+from flax.typing import Dtype, PrecisionLike
+
+from .hilbert import hilbert_indices, inverse_permutation, hilbert_patchify, hilbert_unpatchify
+from .simple_dit import DiTBlock
+
+
+# =============================================================================
+# S5 SSM Layer - Diagonal State Space with Parallel Scan
+# =============================================================================
+
+def hippo_initializer(state_dim):
+    """Initialize A matrix using HiPPO-LegS (Legendre State Space) framework.
+
+    For diagonal S5, we use the diagonal approximation of HiPPO.
+    Returns complex diagonal elements that capture long-range dependencies.
+    """
+    def init(key, shape, dtype=jnp.float32):
+        # Diagonal HiPPO initialization: lambda_n = -(2n+1)/2 + i*pi*n
+        n = jnp.arange(state_dim)
+        real = -(n + 0.5)
+        imag = jnp.pi * n
+        # Normalize to unit circle for stability
+        lambda_init = real + 1j * imag
+        # Discretize with dt=1 for initialization, actual dt is learned
+        return lambda_init.astype(jnp.complex64)
+    return init
+
+
+class S5Layer(nn.Module):
+    """S5 (Simplified State Space) layer with diagonal state matrix.
+
+    Processes sequences using state space recurrence:
+        x_k = A * x_{k-1} + B * u_k
+        y_k = Re(C * x_k) + D * u_k
+
+    Where A is diagonal (complex), enabling efficient parallel scan via
+    jax.lax.associative_scan. This is TPU-friendly as it maps to
+    prefix-sum operations on the systolic array.
+
+    Args:
+        features: Output dimension (must match input for residual)
+        state_dim: Dimension of the hidden state (per feature)
+        dt_min: Minimum discretization step (log scale)
+        dt_max: Maximum discretization step (log scale)
+        dtype: Computation dtype
+        precision: JAX precision level
+    """
+    features: int
+    state_dim: int = 64
+    dt_min: float = 0.001
+    dt_max: float = 0.1
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
+
+    @nn.compact
+    def __call__(self, u):
+        """
+        Args:
+            u: Input sequence [B, S, F]
+        Returns:
+            y: Output sequence [B, S, F]
+        """
+        B, S, F = u.shape
+
+        # --- Learnable SSM Parameters ---
+        # A: Diagonal state matrix (complex) - initialized with HiPPO
+        # We parameterize as log of negative real part for stability
+        log_A_real = self.param(
+            'log_A_real',
+            nn.initializers.normal(stddev=0.5),
+            (self.state_dim,)
+        )
+        A_imag = self.param(
+            'A_imag',
+            nn.initializers.normal(stddev=0.5),
+            (self.state_dim,)
+        )
+
+        # B: Input-to-state projection [state_dim, F]
+        B_re = self.param(
+            'B_re',
+            nn.initializers.lecun_normal(),
+            (self.state_dim, F)
+        )
+        B_im = self.param(
+            'B_im',
+            nn.initializers.lecun_normal(),
+            (self.state_dim, F)
+        )
+
+        # C: State-to-output projection [F, state_dim]
+        C_re = self.param(
+            'C_re',
+            nn.initializers.lecun_normal(),
+            (F, self.state_dim)
+        )
+        C_im = self.param(
+            'C_im',
+            nn.initializers.lecun_normal(),
+            (F, self.state_dim)
+        )
+
+        # D: Skip connection (direct input-to-output)
+        D = self.param('D', nn.initializers.ones, (F,))
+
+        # dt: Discretization timestep (learned, per-feature)
+        log_dt = self.param(
+            'log_dt',
+            lambda key, shape: jax.random.uniform(
+                key, shape,
+                minval=jnp.log(self.dt_min),
+                maxval=jnp.log(self.dt_max)
+            ),
+            (F,)
+        )
+        dt = jnp.exp(log_dt)  # [F]
+
+        # --- Construct complex A and discretize ---
+        # A_diag: complex diagonal [state_dim]
+        A_real = -jnp.exp(log_A_real)  # Ensure negative real part for stability
+        A_diag = A_real + 1j * A_imag  # [state_dim]
+
+        # ZOH discretization: A_bar = exp(A * dt), B_bar = (A_bar - I) * A^{-1} * B
+        # For diagonal A, this simplifies element-wise
+        # dt broadcasts: [F] -> we use mean dt for state transition
+        dt_mean = jnp.mean(dt)
+        A_bar = jnp.exp(A_diag * dt_mean)  # [state_dim], complex
+
+        # B as complex: [state_dim, F]
+        B_complex = B_re + 1j * B_im
+        # Discretized B: element-wise
+        B_bar = ((A_bar[:, None] - 1.0) / (A_diag[:, None] + 1e-8)) * B_complex  # [state_dim, F]
+
+        # C as complex: [F, state_dim]
+        C_complex = C_re + 1j * C_im
+
+        # --- Parallel Scan (associative scan for TPU efficiency) ---
+        # For diagonal SSM: x_k = A_bar * x_{k-1} + B_bar @ u_k
+        # This can be computed via associative scan with the binary operator:
+        #   (a1, b1) * (a2, b2) = (a1 * a2, a2 * b1 + b2)
+
+        # Compute per-step inputs: B_bar @ u_k for each timestep
+        # u: [B, S, F], B_bar: [state_dim, F]
+        # Bu: [B, S, state_dim] (complex)
+        u_float = u.astype(jnp.float32)
+        Bu = jnp.einsum('bsf,nf->bsn', u_float, B_bar)  # [B, S, state_dim]
+
+        # A_bar is the same for all timesteps: [state_dim]
+        # Expand for scan: [B, S, state_dim]
+        A_bar_expanded = jnp.broadcast_to(A_bar[None, None, :], (B, S, self.state_dim))
+
+        # Associative scan elements: (a, b) where a=A_bar, b=Bu
+        def binary_operator(e1, e2):
+            a1, b1 = e1
+            a2, b2 = e2
+            return a1 * a2, a2 * b1 + b2
+
+        # Run parallel scan along sequence dimension (axis=1)
+        _, x_states = jax.lax.associative_scan(
+            binary_operator,
+            (A_bar_expanded, Bu),
+            axis=1
+        )
+        # x_states: [B, S, state_dim] (complex) - hidden states at each step
+
+        # --- Output computation ---
+        # y_k = Re(C @ x_k) + D * u_k
+        # C: [F, state_dim], x_states: [B, S, state_dim]
+        y_complex = jnp.einsum('fn,bsn->bsf', C_complex, x_states)  # [B, S, F]
+        y = y_complex.real  # Take real part: [B, S, F]
+
+        # Add skip connection
+        y = y + D[None, None, :] * u_float  # [B, S, F]
+
+        # Cast back to input dtype
+        if self.dtype is not None:
+            y = y.astype(self.dtype)
+        else:
+            y = y.astype(u.dtype)
+
+        return y
+
+
+# =============================================================================
+# Bidirectional S5 - processes sequence in both directions
+# =============================================================================
+
+class BidirectionalS5Layer(nn.Module):
+    """Bidirectional S5 layer - runs forward and backward scans.
+
+    For diffusion models, spatial patches have no inherent directionality,
+    so bidirectional processing captures dependencies in both directions
+    along the serialization curve (Hilbert, raster, etc.).
+
+    Output is the sum of forward and backward scans, projected to features.
+    """
+    features: int
+    state_dim: int = 64
+    dt_min: float = 0.001
+    dt_max: float = 0.1
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
+
+    @nn.compact
+    def __call__(self, u):
+        """
+        Args:
+            u: [B, S, F]
+        Returns:
+            y: [B, S, F]
+        """
+        # Forward scan
+        y_fwd = S5Layer(
+            features=self.features,
+            state_dim=self.state_dim,
+            dt_min=self.dt_min,
+            dt_max=self.dt_max,
+            dtype=self.dtype,
+            precision=self.precision,
+            name="s5_forward"
+        )(u)
+
+        # Backward scan (reverse input, scan, reverse output)
+        u_rev = jnp.flip(u, axis=1)
+        y_bwd_rev = S5Layer(
+            features=self.features,
+            state_dim=self.state_dim,
+            dt_min=self.dt_min,
+            dt_max=self.dt_max,
+            dtype=self.dtype,
+            precision=self.precision,
+            name="s5_backward"
+        )(u_rev)
+        y_bwd = jnp.flip(y_bwd_rev, axis=1)
+
+        # Combine forward and backward
+        y = y_fwd + y_bwd
+
+        # Output projection to mix directions
+        y = nn.Dense(
+            features=self.features,
+            dtype=self.dtype,
+            precision=self.precision,
+            name="out_proj"
+        )(y)
+
+        return y
+
+
+# =============================================================================
+# SSMDiTBlock - Drop-in replacement for DiTBlock
+# =============================================================================
+
+class SSMDiTBlock(nn.Module):
+    """SSM-based DiT Block that replaces attention with bidirectional S5.
+
+    Maintains the EXACT same interface as DiTBlock:
+        __call__(self, x, conditioning, freqs_cis) -> x
+
+    The AdaLN modulation, gating, residual connections, and MLP path
+    are IDENTICAL to DiTBlock. Only the attention is replaced with S5.
+
+    Note: freqs_cis (RoPE frequencies) is accepted for interface compatibility
+    but not used by the SSM - positional information is implicitly captured
+    by the sequential scan along the serialization order (Hilbert/raster/etc).
+    """
+    features: int
+    num_heads: int  # Not used by SSM, kept for interface compat
+    rope_emb: RotaryEmbedding  # Not used by SSM, kept for interface compat
+    state_dim: int = 64
+    mlp_ratio: int = 4
+    dropout_rate: float = 0.0
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
+    use_flash_attention: bool = False  # Ignored, interface compat
+    force_fp32_for_softmax: bool = True  # Ignored, interface compat
+    norm_epsilon: float = 1e-5
+    use_gating: bool = True
+    bidirectional: bool = True
+
+    def setup(self):
+        hidden_features = int(self.features * self.mlp_ratio)
+
+        # AdaLN modulation - IDENTICAL to DiTBlock
+        self.ada_params_module = AdaLNParams(
+            self.features, dtype=self.dtype, precision=self.precision)
+
+        # Layer Norms - IDENTICAL to DiTBlock
+        self.norm1 = nn.LayerNorm(
+            epsilon=self.norm_epsilon, use_scale=False, use_bias=False,
+            dtype=self.dtype, name="norm1")
+        self.norm2 = nn.LayerNorm(
+            epsilon=self.norm_epsilon, use_scale=False, use_bias=False,
+            dtype=self.dtype, name="norm2")
+
+        # S5 SSM layer (replaces attention)
+        ssm_cls = BidirectionalS5Layer if self.bidirectional else S5Layer
+        self.ssm = ssm_cls(
+            features=self.features,
+            state_dim=self.state_dim,
+            dtype=self.dtype,
+            precision=self.precision,
+            name="ssm"
+        )
+
+        # MLP - IDENTICAL to DiTBlock
+        self.mlp = nn.Sequential([
+            nn.Dense(features=hidden_features, dtype=self.dtype, precision=self.precision),
+            nn.gelu,
+            nn.Dense(features=self.features, dtype=self.dtype, precision=self.precision)
+        ])
+
+    @nn.compact
+    def __call__(self, x, conditioning, freqs_cis):
+        """Exact same signature as DiTBlock.__call__"""
+        # Get scale/shift/gate parameters - IDENTICAL to DiTBlock
+        scale_mlp, shift_mlp, gate_mlp, scale_attn, shift_attn, gate_attn = jnp.split(
+            self.ada_params_module(conditioning), 6, axis=-1
+        )
+
+        # --- SSM Path (replaces Attention Path) ---
+        residual = x
+        norm_x = self.norm1(x)
+        x_modulated = norm_x * (1 + scale_attn) + shift_attn
+        ssm_output = self.ssm(x_modulated)
+
+        if self.use_gating:
+            x = residual + gate_attn * ssm_output
+        else:
+            x = residual + ssm_output
+
+        # --- MLP Path --- IDENTICAL to DiTBlock
+        residual = x
+        norm_x_mlp = self.norm2(x)
+        x_mlp_modulated = norm_x_mlp * (1 + scale_mlp) + shift_mlp
+        mlp_output = self.mlp(x_mlp_modulated)
+
+        if self.use_gating:
+            x = residual + gate_mlp * mlp_output
+        else:
+            x = residual + mlp_output
+
+        return x
+
+
+# =============================================================================
+# HybridSSMAttentionDiT - The proposed novel architecture
+# =============================================================================
+
+class HybridSSMAttentionDiT(nn.Module):
+    """Hybrid SSM-Attention Diffusion Transformer.
+
+    Interleaves SSM blocks (for O(n) local processing along Hilbert curve)
+    with attention blocks (for O(n^2) global composition) in a configurable
+    ratio. Both block types share the same AdaLN conditioning interface.
+
+    Args:
+        block_pattern: List of 'ssm' or 'attn' strings defining the block sequence.
+            Example: ['ssm','ssm','ssm','attn'] for 3:1 ratio repeated.
+            If None, uses ssm_attention_ratio to auto-generate pattern.
+        ssm_attention_ratio: Shorthand ratio like '3:1' (3 SSM per 1 attention).
+            Only used if block_pattern is None.
+    """
+    output_channels: int = 3
+    patch_size: int = 16
+    emb_features: int = 768
+    num_layers: int = 12
+    num_heads: int = 12
+    mlp_ratio: int = 4
+    ssm_state_dim: int = 64
+    dropout_rate: float = 0.0
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
+    use_flash_attention: bool = False
+    force_fp32_for_softmax: bool = True
+    norm_epsilon: float = 1e-5
+    learn_sigma: bool = False
+    use_hilbert: bool = False
+    norm_groups: int = 0
+    activation: Callable = jax.nn.swish
+    block_pattern: Optional[Sequence[str]] = None  # e.g., ['ssm','ssm','ssm','attn']
+    ssm_attention_ratio: str = "3:1"  # e.g., "3:1", "1:1", "all-ssm", "all-attn"
+    bidirectional_ssm: bool = True
+
+    def _build_block_pattern(self):
+        """Generate block pattern from ratio string."""
+        if self.block_pattern is not None:
+            pattern = list(self.block_pattern)
+        elif self.ssm_attention_ratio == "all-ssm":
+            pattern = ['ssm'] * self.num_layers
+        elif self.ssm_attention_ratio == "all-attn":
+            pattern = ['attn'] * self.num_layers
+        else:
+            parts = self.ssm_attention_ratio.split(':')
+            n_ssm, n_attn = int(parts[0]), int(parts[1])
+            unit = ['ssm'] * n_ssm + ['attn'] * n_attn
+            pattern = (unit * (self.num_layers // len(unit) + 1))[:self.num_layers]
+        return pattern
+
+    def setup(self):
+        self.patch_embed = PatchEmbedding(
+            patch_size=self.patch_size,
+            embedding_dim=self.emb_features,
+            dtype=self.dtype,
+            precision=self.precision
+        )
+
+        if self.use_hilbert:
+            self.hilbert_proj = nn.Dense(
+                features=self.emb_features,
+                dtype=self.dtype,
+                precision=self.precision,
+                name="hilbert_projection"
+            )
+
+        # Time embedding
+        self.time_embed = nn.Sequential([
+            FourierEmbedding(features=self.emb_features),
+            TimeProjection(features=self.emb_features * self.mlp_ratio),
+            nn.Dense(features=self.emb_features, dtype=self.dtype, precision=self.precision)
+        ])
+
+        # Text context projection
+        self.text_proj = nn.Dense(
+            features=self.emb_features, dtype=self.dtype,
+            precision=self.precision, name="text_context_proj")
+
+        # RoPE (used by attention blocks, passed through SSM blocks)
+        self.rope = RotaryEmbedding(
+            dim=self.emb_features // self.num_heads,
+            max_seq_len=4096, dtype=self.dtype)
+
+        # Build hybrid block sequence
+        pattern = self._build_block_pattern()
+        blocks = []
+        for i, block_type in enumerate(pattern):
+            if block_type == 'ssm':
+                blocks.append(SSMDiTBlock(
+                    features=self.emb_features,
+                    num_heads=self.num_heads,
+                    rope_emb=self.rope,
+                    state_dim=self.ssm_state_dim,
+                    mlp_ratio=self.mlp_ratio,
+                    dropout_rate=self.dropout_rate,
+                    dtype=self.dtype,
+                    precision=self.precision,
+                    norm_epsilon=self.norm_epsilon,
+                    bidirectional=self.bidirectional_ssm,
+                    name=f"ssm_block_{i}"
+                ))
+            else:  # 'attn'
+                blocks.append(DiTBlock(
+                    features=self.emb_features,
+                    num_heads=self.num_heads,
+                    rope_emb=self.rope,
+                    mlp_ratio=self.mlp_ratio,
+                    dropout_rate=self.dropout_rate,
+                    dtype=self.dtype,
+                    precision=self.precision,
+                    use_flash_attention=self.use_flash_attention,
+                    force_fp32_for_softmax=self.force_fp32_for_softmax,
+                    norm_epsilon=self.norm_epsilon,
+                    name=f"dit_block_{i}"
+                ))
+        self.blocks = blocks
+
+        # Final layer
+        self.final_norm = nn.LayerNorm(
+            epsilon=self.norm_epsilon, dtype=self.dtype, name="final_norm")
+
+        output_dim = self.patch_size * self.patch_size * self.output_channels
+        if self.learn_sigma:
+            output_dim *= 2
+
+        self.final_proj = nn.Dense(
+            features=output_dim,
+            dtype=self.dtype,
+            precision=self.precision,
+            kernel_init=nn.initializers.zeros,
+            name="final_proj"
+        )
+
+    @nn.compact
+    def __call__(self, x, temb, textcontext=None):
+        """Exact same signature as SimpleDiT.__call__"""
+        B, H, W, C = x.shape
+        assert H % self.patch_size == 0 and W % self.patch_size == 0
+
+        H_P = H // self.patch_size
+        W_P = W // self.patch_size
+
+        # 1. Patch Embedding (identical to SimpleDiT)
+        if self.use_hilbert:
+            patches_raw, hilbert_inv_idx = hilbert_patchify(x, self.patch_size)
+            patches = self.hilbert_proj(patches_raw)
+        else:
+            patches = self.patch_embed(x)
+            hilbert_inv_idx = None
+
+        num_patches = patches.shape[1]
+        x_seq = patches
+
+        # 2. Conditioning (identical to SimpleDiT)
+        t_emb = self.time_embed(temb)
+        cond_emb = t_emb
+        if textcontext is not None:
+            text_emb = self.text_proj(textcontext)
+            text_emb_pooled = jnp.mean(text_emb, axis=1)
+            cond_emb = cond_emb + text_emb_pooled
+
+        # 3. RoPE frequencies
+        freqs_cos, freqs_sin = self.rope(seq_len=num_patches)
+
+        # 4. Hybrid blocks (SSM and attention interleaved)
+        for block in self.blocks:
+            x_seq = block(x_seq, conditioning=cond_emb, freqs_cis=(freqs_cos, freqs_sin))
+
+        # 5. Final output (identical to SimpleDiT)
+        x_out = self.final_norm(x_seq)
+        x_out = self.final_proj(x_out)
+
+        # 6. Unpatchify
+        if self.use_hilbert:
+            if self.learn_sigma:
+                x_mean, x_logvar = jnp.split(x_out, 2, axis=-1)
+                x_image = hilbert_unpatchify(x_mean, hilbert_inv_idx, self.patch_size, H, W, self.output_channels)
+                return x_image
+            else:
+                x_image = hilbert_unpatchify(x_out, hilbert_inv_idx, self.patch_size, H, W, self.output_channels)
+                return x_image
+        else:
+            if self.learn_sigma:
+                x_mean, x_logvar = jnp.split(x_out, 2, axis=-1)
+                x = unpatchify(x_mean, channels=self.output_channels)
+                return x
+            else:
+                x = unpatchify(x_out, channels=self.output_channels)
+                return x

{flaxdiff-0.2.9 → flaxdiff-0.2.11}/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-0.2.9 → flaxdiff-0.2.11}/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 → flaxdiff-0.2.11}/flaxdiff.egg-info/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: flaxdiff
-Version: 0.2.9
+Version: 0.2.11
 Summary: A versatile and easy to understand Diffusion library
 Author-email: Ashish Kumar Singh <ashishkmr472@gmail.com>
-License-Expression: MIT
+License: MIT
 Description-Content-Type: text/markdown
 Requires-Dist: flax>=0.8.4
 Requires-Dist: jax>=0.4.28
@@ -22,8 +22,6 @@ Requires-Dist: python-dotenv
 
 # ![](images/logo.jpeg "FlaxDiff")
 
-**This project is being used for the UMD Course project MSML 605: MLOps**
-
 **This project is partially supported by [Google TPU Research Cloud](https://sites.research.google/trc/about/). I would like to thank the Google Cloud TPU team for providing me with the resources to train the bigger text-conditional models in multi-host distributed settings.**
 
 ## A Versatile and simple Diffusion Library

{flaxdiff-0.2.9 → flaxdiff-0.2.11}/flaxdiff.egg-info/SOURCES.txt

@@ -42,6 +42,7 @@ flaxdiff/models/simple_dit.py
 flaxdiff/models/simple_mmdit.py
 flaxdiff/models/simple_unet.py
 flaxdiff/models/simple_vit.py
+flaxdiff/models/ssm_dit.py
 flaxdiff/models/unet_3d.py
 flaxdiff/models/unet_3d_blocks.py
 flaxdiff/models/vit_common.py

{flaxdiff-0.2.9 → flaxdiff-0.2.11}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "flaxdiff"
-version = "0.2.9"
+version = "0.2.11"
 description = "A versatile and easy to understand Diffusion library"
 readme = "README.md"
 authors = [
@@ -26,7 +26,7 @@ dependencies = [
     "rich",
     "python-dotenv",
 ]
-license = "MIT"
+license = {text = "MIT"}
 
 [tool.setuptools.packages.find]
 include = ["flaxdiff*"]