Refactor

Author: prishajain1

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -756,7 +756,7 @@ def __init__(
         precision=precision,
     )
 
-  @nnx.jit
+  @nnx.jit(static_argnames="feat_idx")
   def __call__(self, x: jax.Array, feat_cache=None, feat_idx=0):
     if feat_cache is not None:
       idx = feat_idx
@@ -787,7 +787,7 @@ def __call__(self, x: jax.Array, feat_cache=None, feat_idx=0):
       feat_idx += 1
     else:
       x = self.conv_out(x)
-    return x, feat_cache, feat_idx
+    return x, feat_cache, jnp.array(feat_idx, dtype=jnp.int32)
 
 
 class WanDecoder3d(nnx.Module):
@@ -905,7 +905,7 @@ def __init__(
         precision=precision,
     )
 
-  @nnx.jit
+  @nnx.jit(static_argnames="feat_idx")
   def __call__(self, x: jax.Array, feat_cache=None, feat_idx=0):
     if feat_cache is not None:
       idx = feat_idx
@@ -939,7 +939,7 @@ def __call__(self, x: jax.Array, feat_cache=None, feat_idx=0):
       feat_idx += 1
     else:
       x = self.conv_out(x)
-    return x, feat_cache, feat_idx
+    return x, feat_cache, jnp.array(feat_idx, dtype=jnp.int32)
 
 
 class WanDiagonalGaussianDistribution(FlaxDiagonalGaussianDistribution):