debug

Author: prishajain1

src/maxdiffusion/models/ltx2/autoencoder_kl_ltx2.py

@@ -145,11 +145,6 @@ def __call__(self, hidden_states: jax.Array, causal: bool = True) -> jax.Array:
     # 3. Conv
     hidden_states = self.conv(hidden_states)
 
-        # LTX-2 specific output expansion
-    last_channel = hidden_states[..., -1:]
-    repeats = 127 # 256 - 129
-    last_channel_repeated = jnp.repeat(last_channel, repeats, axis=-1)
-    hidden_states = jnp.concatenate([hidden_states, last_channel_repeated], axis=-1)
 
     return hidden_states
 
@@ -296,11 +291,6 @@ def __call__(
         inputs = self.conv_shortcut(inputs)
 
     hidden_states = hidden_states + inputs
-        # LTX-2 specific output expansion
-    last_channel = hidden_states[..., -1:]
-    repeats = 127 # 256 - 129
-    last_channel_repeated = jnp.repeat(last_channel, repeats, axis=-1)
-    hidden_states = jnp.concatenate([hidden_states, last_channel_repeated], axis=-1)
 
     return hidden_states
 
@@ -442,11 +432,6 @@ def __call__(self, hidden_states: jax.Array, causal: bool = True) -> jax.Array:
     if self.residual:
         hidden_states = hidden_states + residual
 
-        # LTX-2 specific output expansion
-    last_channel = hidden_states[..., -1:]
-    repeats = 127 # 256 - 129
-    last_channel_repeated = jnp.repeat(last_channel, repeats, axis=-1)
-    hidden_states = jnp.concatenate([hidden_states, last_channel_repeated], axis=-1)
 
     return hidden_states
 
@@ -641,11 +626,6 @@ def __call__(
             deterministic=deterministic
         )
 
-        # LTX-2 specific output expansion
-    last_channel = hidden_states[..., -1:]
-    repeats = 127 # 256 - 129
-    last_channel_repeated = jnp.repeat(last_channel, repeats, axis=-1)
-    hidden_states = jnp.concatenate([hidden_states, last_channel_repeated], axis=-1)
 
     return hidden_states
 
@@ -770,11 +750,6 @@ def __call__(
             deterministic=deterministic
         )
 
-        # LTX-2 specific output expansion
-    last_channel = hidden_states[..., -1:]
-    repeats = 127 # 256 - 129
-    last_channel_repeated = jnp.repeat(last_channel, repeats, axis=-1)
-    hidden_states = jnp.concatenate([hidden_states, last_channel_repeated], axis=-1)
 
     return hidden_states
 
@@ -914,12 +889,13 @@ def __call__(
     hidden_states = self.conv_act(hidden_states)
     hidden_states = self.conv_out(hidden_states, causal=causal)
 
-        # LTX-2 specific output expansion
+    # LTX-2 specific output expansion
     last_channel = hidden_states[..., -1:]
     repeats = 127 # 256 - 129
     last_channel_repeated = jnp.repeat(last_channel, repeats, axis=-1)
     hidden_states = jnp.concatenate([hidden_states, last_channel_repeated], axis=-1)
 
+
     return hidden_states
 
 
@@ -1095,6 +1071,12 @@ def __call__(
 
     hidden_states = self.conv_act(hidden_states)
     hidden_states = self.conv_out(hidden_states, causal=causal)
+
+    # LTX-2 specific output expansion
+    last_channel = hidden_states[..., -1:]
+    repeats = 127 # 256 - 129
+    last_channel_repeated = jnp.repeat(last_channel, repeats, axis=-1)
+    hidden_states = jnp.concatenate([hidden_states, last_channel_repeated], axis=-1)
 
     # Unpatchify
     B, T, H, W, C = hidden_states.shape
@@ -1108,11 +1090,6 @@ def __call__(
     hidden_states = hidden_states.transpose(0, 1, 5, 2, 7, 3, 6, 4)
     hidden_states = hidden_states.reshape(B, T * p_t, H * p, W * p, C_out_final)
 
-        # LTX-2 specific output expansion
-    last_channel = hidden_states[..., -1:]
-    repeats = 127 # 256 - 129
-    last_channel_repeated = jnp.repeat(last_channel, repeats, axis=-1)
-    hidden_states = jnp.concatenate([hidden_states, last_channel_repeated], axis=-1)
 
     return hidden_states