add wan mid block vae test

Author: jfacevedo-google

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -414,8 +414,20 @@ def __init__(
     num_layers: int = 1
   ):
     self.dim = dim
+    resnets = [WanResidualBlock(in_dim=dim, out_dim=dim, rngs=rngs,dropout=dropout, non_linearity=non_linearity)]
+    attentions = []
+    for _ in range(num_layers):
+      attentions.append(WanAttentionBlock(dim=dim, rngs=rngs))
+      resnets.append(WanResidualBlock(in_dim=dim, out_dim=dim, rngs=rngs,dropout=dropout, non_linearity=non_linearity))
+    self.attentions = attentions
+    self.resnets = resnets
 
   def __call__(self, x: jax.Array, feat_cache=None, feat_idx=[0]):
+    x = self.resnets[0](x)
+    for attn, resnet in zip(self.attentions, self.resnets[1:]):
+      if attn is not None:
+        x = attn(x)
+      x = resnet(x)
     return x
 
 class WanUpBlock(nnx.Module):
@@ -519,6 +531,8 @@ def __call__(self, x: jax.Array, feat_cache=None, feat_idx=[0]):
     # (1, 1, 480, 720, 96)
     for layer in self.down_blocks:
       x = layer(x)
+    
+    x = self.mid_block(x)
     breakpoint()
     return x
 

src/maxdiffusion/tests/wan_vae_test.py

@@ -30,6 +30,7 @@
   WanUpsample,
   AutoencoderKLWan,
   WanEncoder3d,
+  WanMidBlock,
   WanResidualBlock,
   WanRMS_norm,
   WanResample,
@@ -373,6 +374,22 @@ def test_wan_attention(self):
     output = wan_attention(dummy_input)
     assert output.shape == input_shape
 
+  def test_wan_midblock(self):
+    key = jax.random.key(0)
+    rngs = nnx.Rngs(key)
+    batch = 1
+    t = 1
+    dim = 384
+    height = 60
+    width = 90
+    input_shape = (batch, t, height, width, dim)
+    wan_midblock = WanMidBlock(
+      dim=dim, rngs=rngs
+    )
+    dummy_input = jnp.ones(input_shape)
+    output = wan_midblock(dummy_input)
+    assert output.shape == input_shape
+
   def test_wan_encode(self):
     key = jax.random.key(0)
     rngs = nnx.Rngs(key)
@@ -392,16 +409,6 @@ def test_wan_encode(self):
       attn_scales=attn_scales,
       temperal_downsample=temperal_downsample,
     )
-    # wan_encoder = WanEncoder3d(
-    #    rngs=rngs,
-    #    dim=dim,
-    #    z_dim=z_dim,
-    #    dim_mult=dim_mult,
-    #    num_res_blocks=num_res_blocks,
-    #    attn_scales=attn_scales,
-    #    temperal_downsample=temperal_downsample,
-    #    non_linearity=nonlinearity
-    # )
     batch = 1
     channels = 3
     t = 49