match new flax version

Author: susanbao

src/maxdiffusion/models/attention_flax.py

@@ -795,8 +795,8 @@ def __init__(
 
     self.drop_out = nnx.Dropout(dropout)
 
-    self.norm_q = None
-    self.norm_k = None
+    self.norm_q = nnx.data(None)
+    self.norm_k = nnx.data(None)
     if qk_norm is not None:
       self.norm_q = nnx.RMSNorm(
           num_features=self.inner_dim,

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -225,7 +225,7 @@ def __init__(
   ):
     self.dim = dim
     self.mode = mode
-    self.time_conv = None
+    self.time_conv = nnx.data(None)
 
     if mode == "upsample2d":
       self.resample = nnx.Sequential(
@@ -554,8 +554,8 @@ def __init__(
               precision=precision,
           )
       )
-    self.attentions = attentions
-    self.resnets = resnets
+    self.attentions = nnx.data(attentions)
+    self.resnets = nnx.data(resnets)
 
   def __call__(self, x: jax.Array, feat_cache=None, feat_idx=[0]):
     x = self.resnets[0](x, feat_cache, feat_idx)
@@ -601,10 +601,10 @@ def __init__(
           )
       )
       current_dim = out_dim
-    self.resnets = resnets
+    self.resnets = nnx.data(resnets)
 
     # Add upsampling layer if needed.
-    self.upsamplers = None
+    self.upsamplers = nnx.data(None)
     if upsample_mode is not None:
       self.upsamplers = [
           WanResample(
@@ -710,6 +710,7 @@ def __init__(
             )
         )
         scale /= 2.0
+    self.down_blocks = nnx.data(self.down_blocks)
 
     # middle_blocks
     self.mid_block = WanMidBlock(
@@ -873,6 +874,7 @@ def __init__(
       # Update scale for next iteration
       if upsample_mode is not None:
         scale *= 2.0
+    self.up_blocks = nnx.data(self.up_blocks)
 
     # output blocks
     self.norm_out = WanRMS_norm(dim=out_dim, images=False, rngs=rngs, channel_first=False)

src/maxdiffusion/models/wan/transformers/transformer_wan.py

@@ -209,7 +209,7 @@ def __init__(
       inner_dim = int(dim * mult)
     dim_out = dim_out if dim_out is not None else dim
 
-    self.act_fn = None
+    self.act_fn = nnx.data(None)
     if activation_fn == "gelu-approximate":
       self.act_fn = ApproximateGELU(
           rngs=rngs, dim_in=dim, dim_out=inner_dim, bias=bias, dtype=dtype, weights_dtype=weights_dtype, precision=precision