fix

Author: a-r-r-o-w

scripts/convert_dcae_to_diffusers.py

@@ -28,7 +28,7 @@ def remove_keys_(key: str, state_dict: Dict[str, Any]):
     "conv.conv.": "conv.",
     "conv1.conv": "conv1",
     "conv2.conv": "conv2",
-    "conv1.norm": "norm2",
+    "conv1.norm": "norm1",
     "conv2.norm": "norm2",
     "qkv.conv": "qkv",
     "proj.conv": "proj_out",
@@ -90,14 +90,11 @@ def convert_vae(ckpt_path: str, dtype: torch.dtype):
     vae = AutoencoderDC(
         in_channels=3,
         latent_channels=32,
-        encoder_width_list=[128, 256, 512, 512, 1024, 1024],
-        encoder_depth_list=[2, 2, 2, 3, 3, 3],
+        block_out_channels=[128, 256, 512, 512, 1024, 1024],
+        encoder_layers_per_block=[2, 2, 2, 3, 3, 3],
         encoder_block_type=["ResBlock", "ResBlock", "ResBlock", "EViTS5_GLU", "EViTS5_GLU", "EViTS5_GLU"],
-        encoder_norm="rms2d",
-        encoder_act="silu",
         downsample_block_type="Conv",
-        decoder_width_list=[128, 256, 512, 512, 1024, 1024],
-        decoder_depth_list=[3, 3, 3, 3, 3, 3],
+        decoder_layers_per_block=[3, 3, 3, 3, 3, 3],
         decoder_block_type=["ResBlock", "ResBlock", "ResBlock", "EViTS5_GLU", "EViTS5_GLU", "EViTS5_GLU"],
         decoder_norm="rms2d",
         decoder_act="silu",