fix in audio vae

prishajain1 · prishajain1 · commit 8861b92ba765 · 2026-03-10T23:03:23.000+05:30
diff --git a/src/maxdiffusion/models/ltx2/autoencoder_kl_ltx2_audio.py b/src/maxdiffusion/models/ltx2/autoencoder_kl_ltx2_audio.py
@@ -506,7 +506,7 @@ def __init__(
         )
         curr_res = curr_res // 2
 
-      self.down_stages.append({"blocks": stage_blocks, "attns": stage_attns, "downsample": downsample})
+      self.down_stages.append(nnx.Dict({"blocks": stage_blocks, "attns": stage_attns, "downsample": downsample}))
 
     self.mid_block1 = FlaxLTX2AudioResnetBlock(
         block_in,
@@ -671,7 +671,7 @@ def __init__(
         )
         curr_res *= 2
 
-      self.up_stages.append({"blocks": stage_blocks, "attns": stage_attns, "upsample": upsample})
+      self.up_stages.append(nnx.Dict({"blocks": stage_blocks, "attns": stage_attns, "upsample": upsample}))
 
     if self.norm_type == "group":
       self.norm_out = nnx.GroupNorm(num_groups=32, num_channels=block_in, epsilon=1e-6, dtype=dtype, rngs=rngs)

Original file line number	Diff line number	Diff line change
`@@ -506,7 +506,7 @@ def __init__(`
`506`	`506`	`)`
`507`	`507`	`curr_res = curr_res // 2`
`508`	`508`
`509`		`- self.down_stages.append({"blocks": stage_blocks, "attns": stage_attns, "downsample": downsample})`
	`509`	`+ self.down_stages.append(nnx.Dict({"blocks": stage_blocks, "attns": stage_attns, "downsample": downsample}))`
`510`	`510`
`511`	`511`	`self.mid_block1 = FlaxLTX2AudioResnetBlock(`
`512`	`512`	`block_in,`
`@@ -671,7 +671,7 @@ def __init__(`
`671`	`671`	`)`
`672`	`672`	`curr_res *= 2`
`673`	`673`
`674`		`- self.up_stages.append({"blocks": stage_blocks, "attns": stage_attns, "upsample": upsample})`
	`674`	`+ self.up_stages.append(nnx.Dict({"blocks": stage_blocks, "attns": stage_attns, "upsample": upsample}))`
`675`	`675`
`676`	`676`	`if self.norm_type == "group":`
`677`	`677`	`self.norm_out = nnx.GroupNorm(num_groups=32, num_channels=block_in, epsilon=1e-6, dtype=dtype, rngs=rngs)`