new fix

Author: prishajain1

src/maxdiffusion/models/ltx2/autoencoder_kl_ltx2.py

@@ -655,14 +655,16 @@ def __init__(
   ):
     out_channels = out_channels or in_channels
 
-    self.time_embedder = nnx.data(NNXPixArtAlphaCombinedTimestepSizeEmbeddings(
-        rngs=rngs,
-        embedding_dim=in_channels * 4,
-        size_emb_dim=0,
-        use_additional_conditions=False,
-        dtype=dtype,
-        weights_dtype=weights_dtype
-    ))
+    self.time_embedder = None
+    if timestep_conditioning:
+        self.time_embedder = nnx.data(NNXPixArtAlphaCombinedTimestepSizeEmbeddings(
+            rngs=rngs,
+            embedding_dim=in_channels * 4,
+            size_emb_dim=0,
+            use_additional_conditions=False,
+            dtype=dtype,
+            weights_dtype=weights_dtype
+        ))
 
     if in_channels != out_channels:
         self.conv_in = nnx.data(LTX2VideoResnetBlock3d(
@@ -1068,6 +1070,60 @@ def __call__(
     return hidden_states
 
 
+
+class LTX2DiagonalGaussianDistribution(nnx.Module):
+    def __init__(self, parameters: jax.Array, deterministic: bool = False):
+        self.parameters = parameters
+        # Split into mean and logvar
+        # LTX-2 specific: 128 channels for mean, 1 channel for logvar
+        self.mean, self.logvar = jnp.split(parameters, [128], axis=-1)
+        self.logvar = jnp.clip(self.logvar, -30.0, 20.0)
+        self.deterministic = deterministic
+        self.std = jnp.exp(0.5 * self.logvar)
+        self.var = jnp.exp(self.logvar)
+        if self.deterministic:
+            self.var = self.std = jnp.zeros_like(
+                self.mean, dtype=self.parameters.dtype
+            )
+
+    def sample(self, key: jax.Array) -> jax.Array:
+        # make sure sample is on the same device as the parameters and has same dtype
+        sample = jax.random.normal(key, self.mean.shape, dtype=self.parameters.dtype)
+        x = self.mean + self.std * sample
+        return x
+
+    def kl(self, other: "LTX2DiagonalGaussianDistribution" = None) -> jax.Array:
+        if self.deterministic:
+            return jnp.array([0.0])
+        else:
+            if other is None:
+                return 0.5 * jnp.sum(
+                    jnp.power(self.mean, 2) + self.var - 1.0 - self.logvar,
+                    axis=[1, 2, 3],
+                )
+            else:
+                return 0.5 * jnp.sum(
+                    jnp.power(self.mean - other.mean, 2) / other.var
+                    + self.var / other.var
+                    - 1.0
+                    - self.logvar
+                    + other.logvar,
+                    axis=[1, 2, 3],
+                )
+
+    def nll(self, sample: jax.Array, dims: Tuple[int, ...] = (1, 2, 3)) -> jax.Array:
+        if self.deterministic:
+            return jnp.array([0.0])
+        logtwopi = jnp.log(2.0 * jnp.pi)
+        return 0.5 * jnp.sum(
+            logtwopi + self.logvar + jnp.power(sample - self.mean, 2) / self.var,
+            axis=dims,
+        )
+
+    def mode(self) -> jax.Array:
+        return self.mean
+
+
 class LTX2VideoAutoencoderKL(nnx.Module, ConfigMixin):
   _supports_gradient_checkpointing = True
   config_name = "config.json"
@@ -1510,7 +1566,8 @@ def encode(
     else:
         moments = self._encode(sample, key=key, causal=causal)
 
-    posterior = FlaxDiagonalGaussianDistribution(moments)
+        
+    posterior = LTX2DiagonalGaussianDistribution(moments)
 
     if not return_dict:
       return (posterior,)

src/maxdiffusion/tests/ltx2_vae_parity_test.py

@@ -130,5 +130,19 @@ def fix_keys(d):
     print(f"Latents Mean: {latents.mean():.6f}")
     print(f"Latents Std:  {latents.std():.6f}")
 
+    # Assertions
+    # 1. Check Output Shape
+    assert jax_recon.shape == jax_input.shape, f"Output shape mismatch: {jax_recon.shape} vs {jax_input.shape}"
+    
+    # 2. Check Latents Shape (Mean should be 128 channels)
+    assert latents.shape[-1] == 128, f"Latents channel mismatch: {latents.shape[-1]} vs 128"
+    
+    # 3. Check Encoder Output Channels (should be 129 before split)
+    # We can check parameters of the distribution indirectly via moments if accessible, 
+    # but here we checked latents (mode) which is derived from mean (128).
+    # The fact that it ran without error implies the split worked.
+
+    print("\nTest Passed!")
+
 if __name__ == "__main__":
     test_ltx2_vae_parity()