update

Author: a-r-r-o-w

scripts/convert_dcae_to_diffusers.py

@@ -0,0 +1,166 @@
+import argparse
+from typing import Any, Dict
+
+import torch
+from safetensors.torch import load_file
+from transformers import T5EncoderModel, T5Tokenizer
+
+from diffusers import AutoencoderDC
+
+
+def remove_keys_(key: str, state_dict: Dict[str, Any]):
+    state_dict.pop(key)
+
+
+TOKENIZER_MAX_LENGTH = 128
+
+TRANSFORMER_KEYS_RENAME_DICT = {}
+
+TRANSFORMER_SPECIAL_KEYS_REMAP = {}
+
+VAE_KEYS_RENAME_DICT = {
+    # common
+    "norm.": "norm.norm.",
+    # encoder
+    "encoder.project_in": "encoder.conv_in",
+    "encoder.project_out.main.op_list.0": "encoder.conv_out",
+    # decoder
+    "decoder.project_in.main": "decoder.conv_in",
+    "decoder.project_out.op_list.0": "decoder.norm_out.norm",
+    "decoder.project_out.op_list.2": "decoder.conv_out",
+}
+
+VAE_SPECIAL_KEYS_REMAP = {}
+
+
+def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]:
+    state_dict = saved_dict
+    if "model" in saved_dict.keys():
+        state_dict = state_dict["model"]
+    if "module" in saved_dict.keys():
+        state_dict = state_dict["module"]
+    if "state_dict" in saved_dict.keys():
+        state_dict = state_dict["state_dict"]
+    return state_dict
+
+
+def update_state_dict_inplace(state_dict: Dict[str, Any], old_key: str, new_key: str) -> Dict[str, Any]:
+    state_dict[new_key] = state_dict.pop(old_key)
+
+
+# def convert_transformer(
+#     ckpt_path: str,
+#     dtype: torch.dtype,
+# ):
+#     PREFIX_KEY = ""
+
+#     original_state_dict = get_state_dict(load_file(ckpt_path))
+#     transformer = LTXTransformer3DModel().to(dtype=dtype)
+
+#     for key in list(original_state_dict.keys()):
+#         new_key = key[len(PREFIX_KEY) :]
+#         for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items():
+#             new_key = new_key.replace(replace_key, rename_key)
+#         update_state_dict_inplace(original_state_dict, key, new_key)
+
+#     for key in list(original_state_dict.keys()):
+#         for special_key, handler_fn_inplace in TRANSFORMER_SPECIAL_KEYS_REMAP.items():
+#             if special_key not in key:
+#                 continue
+#             handler_fn_inplace(key, original_state_dict)
+
+#     transformer.load_state_dict(original_state_dict, strict=True)
+#     return transformer
+
+
+def convert_vae(ckpt_path: str, dtype: torch.dtype):
+    original_state_dict = get_state_dict(load_file(ckpt_path))
+    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],
+        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_block_type=["ResBlock", "ResBlock", "ResBlock", "EViTS5_GLU", "EViTS5_GLU", "EViTS5_GLU"],
+        decoder_norm="rms2d",
+        decoder_act="silu",
+        upsample_block_type="InterpolateConv",
+        scaling_factor=0.41407,
+    ).to(dtype=dtype)
+
+    for key in list(original_state_dict.keys()):
+        new_key = key[:]
+        for replace_key, rename_key in VAE_KEYS_RENAME_DICT.items():
+            new_key = new_key.replace(replace_key, rename_key)
+        update_state_dict_inplace(original_state_dict, key, new_key)
+
+    for key in list(original_state_dict.keys()):
+        for special_key, handler_fn_inplace in VAE_SPECIAL_KEYS_REMAP.items():
+            if special_key not in key:
+                continue
+            handler_fn_inplace(key, original_state_dict)
+
+    vae.load_state_dict(original_state_dict, strict=True)
+    return vae
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--transformer_ckpt_path", type=str, default=None, help="Path to original transformer checkpoint"
+    )
+    parser.add_argument("--vae_ckpt_path", type=str, default=None, help="Path to original vae checkpoint")
+    parser.add_argument(
+        "--text_encoder_cache_dir", type=str, default=None, help="Path to text encoder cache directory"
+    )
+    parser.add_argument(
+        "--typecast_text_encoder",
+        action="store_true",
+        default=False,
+        help="Whether or not to apply fp16/bf16 precision to text_encoder",
+    )
+    parser.add_argument("--save_pipeline", action="store_true")
+    parser.add_argument("--output_path", type=str, required=True, help="Path where converted model should be saved")
+    parser.add_argument("--dtype", default="fp32", help="Torch dtype to save the model in.")
+    return parser.parse_args()
+
+
+DTYPE_MAPPING = {
+    "fp32": torch.float32,
+    "fp16": torch.float16,
+    "bf16": torch.bfloat16,
+}
+
+VARIANT_MAPPING = {
+    "fp32": None,
+    "fp16": "fp16",
+    "bf16": "bf16",
+}
+
+
+if __name__ == "__main__":
+    args = get_args()
+
+    transformer = None
+    dtype = DTYPE_MAPPING[args.dtype]
+    variant = VARIANT_MAPPING[args.dtype]
+
+    if args.save_pipeline:
+        assert args.transformer_ckpt_path is not None and args.vae_ckpt_path is not None
+
+    # if args.transformer_ckpt_path is not None:
+    #     transformer = convert_transformer(args.transformer_ckpt_path, dtype)
+    #     if not args.save_pipeline:
+    #         transformer.save_pretrained(
+    #             args.output_path, safe_serialization=True, max_shard_size="5GB", variant=variant
+    #         )
+
+    if args.vae_ckpt_path is not None:
+        vae = convert_vae(args.vae_ckpt_path, dtype)
+        if not args.save_pipeline:
+            vae.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB", variant=variant)

src/diffusers/__init__.py

@@ -80,6 +80,7 @@
             "AllegroTransformer3DModel",
             "AsymmetricAutoencoderKL",
             "AuraFlowTransformer2DModel",
+            "AutoencoderDC",
             "AutoencoderKL",
             "AutoencoderKLAllegro",
             "AutoencoderKLCogVideoX",
@@ -572,6 +573,7 @@
             AsymmetricAutoencoderKL,
             AuraFlowTransformer2DModel,
             AutoencoderKL,
+            AutoencoderDC,
             AutoencoderKLAllegro,
             AutoencoderKLCogVideoX,
             AutoencoderKLMochi,

src/diffusers/models/__init__.py

@@ -27,6 +27,7 @@
 if is_torch_available():
     _import_structure["adapter"] = ["MultiAdapter", "T2IAdapter"]
     _import_structure["autoencoders.autoencoder_asym_kl"] = ["AsymmetricAutoencoderKL"]
+    _import_structure["autoencoders.autoencoder_dc"] = ["AutoencoderDC"]
     _import_structure["autoencoders.autoencoder_kl"] = ["AutoencoderKL"]
     _import_structure["autoencoders.autoencoder_kl_allegro"] = ["AutoencoderKLAllegro"]
     _import_structure["autoencoders.autoencoder_kl_cogvideox"] = ["AutoencoderKLCogVideoX"]
@@ -88,6 +89,7 @@
         from .adapter import MultiAdapter, T2IAdapter
         from .autoencoders import (
             AsymmetricAutoencoderKL,
+            AutoencoderDC,
             AutoencoderKL,
             AutoencoderKLAllegro,
             AutoencoderKLCogVideoX,

src/diffusers/models/attention.py

@@ -22,7 +22,7 @@
 from .activations import GEGLU, GELU, ApproximateGELU, FP32SiLU, SwiGLU, get_activation
 from .attention_processor import Attention, JointAttnProcessor2_0
 from .embeddings import SinusoidalPositionalEmbedding
-from .normalization import AdaLayerNorm, AdaLayerNormContinuous, AdaLayerNormZero, RMSNorm, SD35AdaLayerNormZeroX, RMSNorm2d
+from .normalization import AdaLayerNorm, AdaLayerNormContinuous, AdaLayerNormZero, RMSNorm, SD35AdaLayerNormZeroX, RMSNormNd
 
 
 logger = logging.get_logger(__name__)
@@ -1241,160 +1241,3 @@ def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
         for module in self.net:
             hidden_states = module(hidden_states)
         return hidden_states
-
-
-class DCAELiteMLA(nn.Module):
-    r"""Lightweight multi-scale linear attention used in DC-AE"""
-
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        heads: Optional[int] = None,
-        heads_ratio: float = 1.0,
-        dim=8,
-        use_bias=(False, False),
-        norm=(None, "bn2d"),
-        act_func=(None, None),
-        kernel_func="relu",
-        scales: Tuple[int, ...] = (5,),
-        eps=1.0e-15,
-    ):
-        super().__init__()
-        self.eps = eps
-        heads = int(in_channels // dim * heads_ratio) if heads is None else heads
-
-        total_dim = heads * dim
-
-        self.dim = dim
-        
-        qkv = [nn.Conv2d(in_channels=in_channels, out_channels=3 * total_dim, kernel_size=1, bias=use_bias[0])]
-        if norm[0] is None:
-            pass
-        elif norm[0] == "rms2d":
-            qkv.append(RMSNorm2d(num_features=3 * total_dim))
-        elif norm[0] == "bn2d":
-            qkv.append(nn.BatchNorm2d(num_features=3 * total_dim))
-        else:
-            raise ValueError(f"norm {norm[0]} is not supported")
-        if act_func[0] is not None:
-            qkv.append(get_activation(act_func[0]))
-        self.qkv = nn.Sequential(*qkv)
-
-        self.aggreg = nn.ModuleList(
-            [
-                nn.Sequential(
-                    nn.Conv2d(
-                        3 * total_dim,
-                        3 * total_dim,
-                        scale,
-                        padding=scale // 2,
-                        groups=3 * total_dim,
-                        bias=use_bias[0],
-                    ),
-                    nn.Conv2d(3 * total_dim, 3 * total_dim, 1, groups=3 * heads, bias=use_bias[0]),
-                )
-                for scale in scales
-            ]
-        )
-        self.kernel_func = get_activation(kernel_func)
-
-        proj = [nn.Conv2d(in_channels=total_dim * (1 + len(scales)), out_channels=out_channels, kernel_size=1, bias=use_bias[1])]
-        if norm[1] is None:
-            pass
-        elif norm[1] == "rms2d":
-            proj.append(RMSNorm2d(num_features=out_channels))
-        elif norm[1] == "bn2d":
-            proj.append(nn.BatchNorm2d(num_features=out_channels))
-        else:
-            raise ValueError(f"norm {norm[1]} is not supported")
-        if act_func[1] is not None:
-            proj.append(get_activation(act_func[1]))
-        self.proj = nn.Sequential(*proj)
-
-    def relu_linear_att(self, qkv: torch.Tensor) -> torch.Tensor:
-        B, _, H, W = list(qkv.size())
-
-        if qkv.dtype == torch.float16:
-            qkv = qkv.float()
-
-        qkv = torch.reshape(
-            qkv,
-            (
-                B,
-                -1,
-                3 * self.dim,
-                H * W,
-            ),
-        )
-        q, k, v = (
-            qkv[:, :, 0 : self.dim],
-            qkv[:, :, self.dim : 2 * self.dim],
-            qkv[:, :, 2 * self.dim :],
-        )
-
-        # lightweight linear attention
-        q = self.kernel_func(q)
-        k = self.kernel_func(k)
-
-        # linear matmul
-        trans_k = k.transpose(-1, -2)
-
-        v = F.pad(v, (0, 0, 0, 1), mode="constant", value=1)
-        vk = torch.matmul(v, trans_k)
-        out = torch.matmul(vk, q)
-        if out.dtype == torch.bfloat16:
-            out = out.float()
-        out = out[:, :, :-1] / (out[:, :, -1:] + self.eps)
-
-        out = torch.reshape(out, (B, -1, H, W))
-        return out
-
-    def relu_quadratic_att(self, qkv: torch.Tensor) -> torch.Tensor:
-        B, _, H, W = list(qkv.size())
-
-        qkv = torch.reshape(
-            qkv,
-            (
-                B,
-                -1,
-                3 * self.dim,
-                H * W,
-            ),
-        )
-        q, k, v = (
-            qkv[:, :, 0 : self.dim],
-            qkv[:, :, self.dim : 2 * self.dim],
-            qkv[:, :, 2 * self.dim :],
-        )
-
-        q = self.kernel_func(q)
-        k = self.kernel_func(k)
-
-        att_map = torch.matmul(k.transpose(-1, -2), q)  # b h n n
-        original_dtype = att_map.dtype
-        if original_dtype in [torch.float16, torch.bfloat16]:
-            att_map = att_map.float()
-        att_map = att_map / (torch.sum(att_map, dim=2, keepdim=True) + self.eps)  # b h n n
-        att_map = att_map.to(original_dtype)
-        out = torch.matmul(v, att_map)  # b h d n
-
-        out = torch.reshape(out, (B, -1, H, W))
-        return out
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # generate multi-scale q, k, v
-        qkv = self.qkv(x)
-        multi_scale_qkv = [qkv]
-        for op in self.aggreg:
-            multi_scale_qkv.append(op(qkv))
-        qkv = torch.cat(multi_scale_qkv, dim=1)
-
-        H, W = list(qkv.size())[-2:]
-        if H * W > self.dim:
-            out = self.relu_linear_att(qkv).to(qkv.dtype)
-        else:
-            out = self.relu_quadratic_att(qkv)
-        out = self.proj(out)
-
-        return x + out

src/diffusers/models/autoencoders/__init__.py

@@ -1,4 +1,5 @@
 from .autoencoder_asym_kl import AsymmetricAutoencoderKL
+from .autoencoder_dc import AutoencoderDC
 from .autoencoder_kl import AutoencoderKL
 from .autoencoder_kl_allegro import AutoencoderKLAllegro
 from .autoencoder_kl_cogvideox import AutoencoderKLCogVideoX