refactor

Author: a-r-r-o-w

scripts/convert_dcae_to_diffusers.py

@@ -136,10 +136,10 @@ def get_vae_config(name: str):
     elif name in ["dc-ae-f32c32-in-1.0", "dc-ae-f32c32-mix-1.0"]:
         config = {
             "latent_channels": 32,
-            "encoder_block_type": ["ResBlock", "ResBlock", "ResBlock", "EViT_GLU", "EViT_GLU", "EViT_GLU"],
+            "encoder_block_type": ["ResBlock", "ResBlock", "ResBlock", "EfficientViTBlock", "EfficientViTBlock", "EfficientViTBlock"],
             "block_out_channels": [128, 256, 512, 512, 1024, 1024],
             "encoder_layers_per_block": [0, 4, 8, 2, 2, 2],
-            "decoder_block_type": ["ResBlock", "ResBlock", "ResBlock", "EViT_GLU", "EViT_GLU", "EViT_GLU"],
+            "decoder_block_type": ["ResBlock", "ResBlock", "ResBlock", "EfficientViTBlock", "EfficientViTBlock", "EfficientViTBlock"],
             "decoder_layers_per_block": [0, 5, 10, 2, 2, 2],
             "decoder_norm": ["bn2d", "bn2d", "bn2d", "rms2d", "rms2d", "rms2d"],
             "decoder_act": ["relu", "relu", "relu", "silu", "silu", "silu"],
@@ -151,23 +151,23 @@ def get_vae_config(name: str):
                 "ResBlock",
                 "ResBlock",
                 "ResBlock",
-                "EViT_GLU",
-                "EViT_GLU",
-                "EViT_GLU",
-                "EViT_GLU",
-                "EViT_GLU",
+                "EfficientViTBlock",
+                "EfficientViTBlock",
+                "EfficientViTBlock",
+                "EfficientViTBlock",
+                "EfficientViTBlock",
             ],
             "block_out_channels": [128, 256, 512, 512, 1024, 1024, 2048, 2048],
             "encoder_layers_per_block": [0, 4, 8, 2, 2, 2, 2, 2],
             "decoder_block_type": [
                 "ResBlock",
                 "ResBlock",
                 "ResBlock",
-                "EViT_GLU",
-                "EViT_GLU",
-                "EViT_GLU",
-                "EViT_GLU",
-                "EViT_GLU",
+                "EfficientViTBlock",
+                "EfficientViTBlock",
+                "EfficientViTBlock",
+                "EfficientViTBlock",
+                "EfficientViTBlock",
             ],
             "decoder_layers_per_block": [0, 5, 10, 2, 2, 2, 2, 2],
             "decoder_norm": ["bn2d", "bn2d", "bn2d", "rms2d", "rms2d", "rms2d", "rms2d", "rms2d"],
@@ -176,10 +176,10 @@ def get_vae_config(name: str):
     elif name in ["dc-ae-f64c128-in-1.0", "dc-ae-f64c128-mix-1.0"]:
         config = {
             "latent_channels": 128,
-            "encoder_block_type": ["ResBlock", "ResBlock", "ResBlock", "EViT_GLU", "EViT_GLU", "EViT_GLU", "EViT_GLU"],
+            "encoder_block_type": ["ResBlock", "ResBlock", "ResBlock", "EfficientViTBlock", "EfficientViTBlock", "EfficientViTBlock", "EfficientViTBlock"],
             "block_out_channels": [128, 256, 512, 512, 1024, 1024, 2048],
             "encoder_layers_per_block": [0, 4, 8, 2, 2, 2, 2],
-            "decoder_block_type": ["ResBlock", "ResBlock", "ResBlock", "EViT_GLU", "EViT_GLU", "EViT_GLU", "EViT_GLU"],
+            "decoder_block_type": ["ResBlock", "ResBlock", "ResBlock", "EfficientViTBlock", "EfficientViTBlock", "EfficientViTBlock", "EfficientViTBlock"],
             "decoder_layers_per_block": [0, 5, 10, 2, 2, 2, 2],
             "decoder_norm": ["bn2d", "bn2d", "bn2d", "rms2d", "rms2d", "rms2d", "rms2d"],
             "decoder_act": ["relu", "relu", "relu", "silu", "silu", "silu", "silu"],

src/diffusers/models/autoencoders/autoencoder_dc.py

@@ -80,13 +80,11 @@ def __init__(
         in_channels: int,
         out_channels: int,
         norm_type: str = "bn2d",
-        act_func=("relu6", None),
+        act_fn: str = "relu6",
     ) -> None:
         super().__init__()
 
-        act_func = val2tuple(act_func, 2)
-
-        self.nonlinearity = get_activation(act_func[0]) if act_func[0] is not None else nn.Identity()
+        self.nonlinearity = get_activation(act_fn) if act_fn is not None else nn.Identity()
 
         self.conv1 = nn.Conv2d(in_channels, in_channels, 3, 1, 1)
         self.conv2 = nn.Conv2d(in_channels, out_channels, 3, 1, 1, bias=False)
@@ -259,6 +257,23 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x
 
 
+def get_block_from_block_type(
+    block_type: str,
+    in_channels: int,
+    out_channels: int,
+    norm_type: str,
+    act_fn: str,
+):
+    if block_type == "ResBlock":
+        block = ResBlock(in_channels, out_channels, norm_type, act_fn)
+    
+    elif block_type == "EfficientViTBlock":
+        block = EfficientViTBlock(in_channels, norm=norm_type, scales=())
+
+    else:
+        raise ValueError(f"Block with {block_type=} is not supported.")
+
+
 def build_stage_main(
     width: int, depth: int, block_type: str | List[str], norm: str, act: str, input_width: int
 ) -> list[nn.Module]:
@@ -278,7 +293,7 @@ def build_stage_main(
                 norm_type=norm,
                 act_func=(act, None),
             )
-        elif current_block_type == "EViT_GLU":
+        elif current_block_type == "EfficientViTBlock":
             assert in_channels == out_channels
             block = EfficientViTBlock(in_channels, norm=norm, scales=())
         elif current_block_type == "EViTS5_GLU":