fix

Author: prishajain1

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -1123,27 +1123,44 @@ def _encode(self, x: jax.Array, feat_cache: AutoencoderKLWanCache):
       x = jnp.transpose(x, (0, 2, 3, 4, 1))
       assert x.shape[-1] == 3, f"Expected input shape (N, D, H, W, 3), got {x.shape}"
 
-    # Swap to (T, B, H, W, C) for scanning over time
-    x_scan = jnp.swapaxes(x, 0, 1)
+    t = x.shape[1]
     enc_feat_map = feat_cache._enc_feat_map
 
-    def scan_fn(carry_cache, input_frame):
-      # Expand time dimension to 1 for the encoder
-      input_frame = jnp.expand_dims(input_frame, 1)
-      # idx is restarted at 0 for each chunk/frame conceptually
-      enc_conv_idx = 0
-      out_frame, new_cache, _ = self.encoder(input_frame, feat_cache=carry_cache, feat_idx=enc_conv_idx)
-      out_frame = jnp.squeeze(out_frame, 1)
-      return new_cache, out_frame
-
-    # Perform JAX scan
-    final_enc_feat_map, encoded_frames = jax.lax.scan(scan_fn, enc_feat_map, x_scan)
-    
-    # Swap back to (B, T, ... )
-    out = jnp.swapaxes(encoded_frames, 0, 1)
+    # 1. Evaluate the first frame manually to establish the initial cache with JAX Arrays.
+    # This prevents jax.lax.scan from crashing on type mismatch between None and ShapedArray.
+    out_0, enc_feat_map, _ = self.encoder(x[:, :1, :, :, :], feat_cache=enc_feat_map, feat_idx=0)
+    out = out_0
 
-    # Update back to the wrapper object if needed
-    feat_cache._enc_feat_map = final_enc_feat_map
+    # 2. Process remaining frames in chunks of 4 using jax.lax.scan
+    if t > 1:
+      x_rest = x[:, 1:, :, :, :]
+      B, T_rest, H, W, C = x_rest.shape
+      num_chunks = T_rest // 4
+
+      # Reshape to (B, num_chunks, 4, H, W, C)
+      x_chunks = jnp.reshape(x_rest, (B, num_chunks, 4, H, W, C))
+      
+      # Swap axes for scan traversal: (num_chunks, B, 4, H, W, C)
+      x_scan = jnp.swapaxes(x_chunks, 0, 1)
+      
+      def scan_fn(carry_cache, input_chunk):
+        # input_chunk shape: (B, 4, H, W, C)
+        out_chunk, new_cache, _ = self.encoder(input_chunk, feat_cache=carry_cache, feat_idx=0)
+        # out_chunk shape: (B, 1, H', W', C')
+        return new_cache, out_chunk
+        
+      enc_feat_map, scanned_out_chunks = jax.lax.scan(scan_fn, enc_feat_map, x_scan)
+      
+      # scanned_out_chunks shape: (num_chunks, B, 1, H', W', C')
+      scanned_out_chunks = jnp.swapaxes(scanned_out_chunks, 0, 1)
+      
+      B_out, _, _, H_out, W_out, C_out = scanned_out_chunks.shape
+      scanned_out_chunks = jnp.reshape(scanned_out_chunks, (B_out, num_chunks, H_out, W_out, C_out))
+      
+      out = jnp.concatenate([out_0, scanned_out_chunks], axis=1)
+
+    # 3. Update back to the wrapper object if needed
+    feat_cache._enc_feat_map = enc_feat_map
 
     enc = self.quant_conv(out)
     mu, logvar = enc[:, :, :, :, : self.z_dim], enc[:, :, :, :, self.z_dim :]
@@ -1170,22 +1187,35 @@ def _decode(
 
     dec_feat_map = feat_cache._feat_map
 
-    # Evaluate the first frame manually to establish the initial cache.
-    # The decoder returns 1 frame on the first step, and 4 frames on subsequent steps due to temporal upsampling.
+    # 1. Evaluate the first frame manually (Cache: None -> RepSentinel/ShapedArray)
+    # The decoder returns 1 frame on the first step.
     out_0, dec_feat_map, _ = self.decoder(x[:, 0:1, :, :, :], feat_cache=dec_feat_map, feat_idx=0)
     out = out_0
 
-    # Process remaining frames using jax.lax.scan (requires homogenous output shapes)
+    # 2. Evaluate the second frame manually (Cache: RepSentinel -> ShapedArray)
+    # This ensures that ALL cache components are ShapredArrays before entering jax.lax.scan,
+    # preventing TraceContext errors due to type mismatches.
     if iter_ > 1:
-      x_rest = x[:, 1:, :, :, :]
-      x_scan = jnp.swapaxes(x_rest, 0, 1) # (T-1, B, H, W, C)
+      out_1, dec_feat_map, _ = self.decoder(x[:, 1:2, :, :, :], feat_cache=dec_feat_map, feat_idx=0)
+      
+      # Bypass an issue where frame[1] should be frame[2] and vice versa.
+      fm1 = out_1[:, 0:1, ...]
+      fm2 = out_1[:, 1:2, ...]
+      fm3 = out_1[:, 2:3, ...]
+      fm4 = out_1[:, 3:4, ...]
+      out_1_fixed = jnp.concatenate([fm1, fm3, fm2, fm4], axis=1)
+      out = jnp.concatenate([out_0, out_1_fixed], axis=1)
+
+    # 3. Process remaining frames using jax.lax.scan (requires homogenous output and carry shapes)
+    if iter_ > 2:
+      x_rest = x[:, 2:, :, :, :]
+      x_scan = jnp.swapaxes(x_rest, 0, 1) # (T-2, B, H, W, C)
 
       def scan_fn(carry_cache, input_frame):
         input_frame = jnp.expand_dims(input_frame, 1) # (B, 1, H, W, C)
         out_frames, new_cache, _ = self.decoder(input_frame, feat_cache=carry_cache, feat_idx=0)
 
         # Bypass an issue where frame[1] should be frame[2] and vice versa.
-        # Ensure dimensionality allows straightforward slicing:
         fm1 = out_frames[:, 0:1, ...]
         fm2 = out_frames[:, 1:2, ...]
         fm3 = out_frames[:, 2:3, ...]
@@ -1196,17 +1226,17 @@ def scan_fn(carry_cache, input_frame):
 
       dec_feat_map, scanned_out_frames = jax.lax.scan(scan_fn, dec_feat_map, x_scan)
 
-      # scanned_out_frames is (T-1, B, 4, H, W, C)
+      # scanned_out_frames is (T-2, B, 4, H, W, C)
       B = scanned_out_frames.shape[1]
-      T_minus_1 = scanned_out_frames.shape[0]
+      T_minus_2 = scanned_out_frames.shape[0]
       H, W, C = scanned_out_frames.shape[3], scanned_out_frames.shape[4], scanned_out_frames.shape[5]
 
-      # Swap back to (B, T-1, 4, H, W, C)
+      # Swap back to (B, T-2, 4, H, W, C)
       scanned_out_frames = jnp.swapaxes(scanned_out_frames, 0, 1)
-      # Flatten the temporal axes to (B, (T-1)*4, H, W, C)
-      scanned_out_frames = jnp.reshape(scanned_out_frames, (B, T_minus_1 * 4, H, W, C))
+      # Flatten the temporal axes to (B, (T-2)*4, H, W, C)
+      scanned_out_frames = jnp.reshape(scanned_out_frames, (B, T_minus_2 * 4, H, W, C))
 
-      out = jnp.concatenate([out_0, scanned_out_frames], axis=1)
+      out = jnp.concatenate([out, scanned_out_frames], axis=1)
 
     feat_cache._feat_map = dec_feat_map