Update config

transformer/config.json

@@ -8,9 +8,11 @@
   "clip_channels": null,
   "dropout": 0.0,
   "enable_clip_in_inpaint": false,
+  "enable_text_attention_mask": false,
   "flip_sin_to_cos": true,
   "freq_shift": 0,
   "in_channels": 33,
+  "mmdit_layers": 48,
   "norm_elementwise_affine": true,
   "norm_eps": 1e-05,
   "num_attention_heads": 48,

vae/config.json

@@ -1,43 +1,58 @@
 {
-    "_class_name": "AutoencoderKL",
-    "_diffusers_version": "0.31.0.dev0",
+    "_class_name": "AutoencoderKLMagvit",
+    "_diffusers_version": "0.30.1",
     "act_fn": "silu",
     "block_out_channels": [
-        128,
-        256,
-        512,
-        512
+      128,
+      256,
+      512,
+      512
+    ],
+    "cache_compression_vae": false,
+    "cache_mag_vae": true,
+    "ch": 128,
+    "ch_mult": [
+      1,
+      2,
+      4,
+      4
     ],
     "down_block_types": [
-        "SpatialDownBlock3D", 
-        "SpatialTemporalDownBlock3D", 
-        "SpatialTemporalDownBlock3D",
-        "SpatialTemporalDownBlock3D"
+      "SpatialDownBlock3D",
+      "SpatialTemporalDownBlock3D",
+      "SpatialTemporalDownBlock3D",
+      "SpatialTemporalDownBlock3D"
     ],
     "force_upcast": true,
     "in_channels": 3,
     "latent_channels": 16,
     "layers_per_block": 2,
+    "mid_block_attention_type": "spatial",
+    "mid_block_num_attention_heads": 1,
+    "mid_block_type": "MidBlock3D",
+    "mid_block_use_attention": true,
+    "mini_batch_decoder": 1,
+    "mini_batch_encoder": 4,
     "norm_num_groups": 32,
+    "num_attention_heads": 1,
     "out_channels": 3,
     "sample_size": 256,
     "scaling_factor": 0.7125,
-    "slice_mag_vae": false,
     "slice_compression_vae": false,
-    "cache_compression_vae": false,
-    "cache_mag_vae": true,
-    "use_tiling": false,
-    "use_tiling_encoder": false,
-    "use_tiling_decoder": false,
-    "mini_batch_encoder": 4,
-    "mini_batch_decoder": 1,
+    "slice_mag_vae": false,
     "spatial_group_norm": true,
-    "mid_block_attention_type": "spatial",
+    "tile_overlap_factor": 0.25,
+    "tile_sample_min_size": 384,
     "up_block_types": [
-        "SpatialUpBlock3D", 
-        "SpatialTemporalUpBlock3D", 
-        "SpatialTemporalUpBlock3D",
-        "SpatialTemporalUpBlock3D"
-    ]
-}
+      "SpatialUpBlock3D",
+      "SpatialTemporalUpBlock3D",
+      "SpatialTemporalUpBlock3D",
+      "SpatialTemporalUpBlock3D"
+    ],
+    "upcast_vae": false,
+    "use_gc_blocks": null,
+    "use_tiling": false,
+    "use_tiling_decoder": false,
+    "use_tiling_encoder": false
+  }