Upload folder using huggingface_hub

AFM-B-2-1NFE-guided/generator/config.json

@@ -1,6 +1,22 @@
 {
-  "_class_name": "AFMGenerator2DModel",
+  "_class_name": "DiTTransformer2DModel",
   "_diffusers_version": "0.36.0",
+  "sample_size": 32,
+  "num_layers": 12,
+  "num_attention_heads": 12,
+  "attention_head_dim": 64,
+  "in_channels": 4,
+  "out_channels": 4,
+  "patch_size": 2,
+  "attention_bias": true,
+  "activation_fn": "gelu-approximate",
+  "num_embeds_ada_norm": 1000,
+  "norm_type": "ada_norm_zero",
+  "norm_elementwise_affine": false,
+  "dropout": 0.0,
+  "norm_num_groups": 32,
+  "norm_eps": 1e-05,
+  "upcast_attention": false,
   "model_type": "AFM-B/2",
   "architecture": "standard",
   "repeat": 1,
@@ -12,9 +28,7 @@
   "class_dropout_prob": 0.0,
   "input_size": 32,
   "num_classes": 1000,
-  "in_channels": 4,
   "depth": 12,
   "hidden_size": 768,
-  "patch_size": 2,
   "num_heads": 12
 }

AFM-B-2-1NFE-guided/generator/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8288ea86cfca336987ea264c9ec315c683a39fdd7ec9740efb32bbcd463a4875
-size 589740400
+oid sha256:72d904dc41423e59c3e70f429a5a12eb01613e04b104f3cfb875b13ba77186c1
+size 589737592

AFM-B-2-1NFE-guided/generator/modeling_afm.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import copy
 import math
 from typing import Any, Dict, Optional
 
@@ -10,15 +9,10 @@ import torch
 import torch.nn.functional as F
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
 from torch import nn
 
 
-def get_dit_transformer_class():
-    from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
-
-    return DiTTransformer2DModel
-
-
 def build_dit_config(
     *,
     depth: int = 28,
@@ -54,110 +48,15 @@ def build_dit_config(
     }
 
 
-def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]:
-    converted = copy.deepcopy(state_dict)
-
-    converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous()
-    converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous()
-
-    if "t_embedder.mlp.0.weight" in converted:
-        timestep_weights = {
-            "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"),
-            "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"),
-            "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"),
-            "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"),
-        }
-    else:
-        timestep_weights = None
-    class_embedding = converted.pop("y_embedder.embedding_table.weight")
-    if class_embedding.shape[0] == 1000:
-        null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype)
-        class_embedding = torch.cat([class_embedding, null_class], dim=0)
-
-    for block_idx in range(depth):
-        if timestep_weights is not None:
-            for key, tensor in timestep_weights.items():
-                converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone()
-        converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = (
-            class_embedding.clone()
-        )
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.bias"
-        ]
-
-        q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0)
-        q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0)
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[
-            f"blocks.{block_idx}.attn.proj.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[
-            f"blocks.{block_idx}.attn.proj.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.bias"
-        ]
-
-        for suffix in (
-            "attn.qkv.weight",
-            "attn.qkv.bias",
-            "attn.proj.weight",
-            "attn.proj.bias",
-            "mlp.fc1.weight",
-            "mlp.fc1.bias",
-            "mlp.fc2.weight",
-            "mlp.fc2.bias",
-            "adaLN_modulation.1.weight",
-            "adaLN_modulation.1.bias",
-        ):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight")
-    converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias")
-    converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight")
-    converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias")
-
-    converted.pop("pos_embed", None)
-    for block_idx in range(depth):
-        for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()}
-
-
-def load_legacy_dit_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    depth: int,
-    patch_size: int,
-    hidden_size: int,
-    num_heads: int,
-) -> Dict[str, torch.Tensor]:
-    if "transformer_blocks.0.attn1.to_q.weight" in state_dict:
-        return state_dict
-    if any(key.startswith("blocks.") for key in state_dict):
-        return _convert_legacy_dit_state_dict(state_dict, depth)
-    return state_dict
+_PREFIX_KEYS = ("transformer.", "t_tgt_embedder.")
+
+
+def _wrap_dit_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    return {k if k.startswith(_PREFIX_KEYS) else f"transformer.{k}": v for k, v in state_dict.items()}
 
 
 def forward_dit_output(
-    transformer,
+    transformer: DiTTransformer2DModel,
     hidden_states: torch.Tensor,
     timestep: Optional[torch.Tensor],
     class_labels: torch.Tensor,
@@ -207,23 +106,6 @@ def forward_dit_output(
     )
 
 
-def remap_transformer_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    legacy_kwargs: Dict[str, int],
-    extra_prefixes: tuple[str, ...] = (),
-) -> Dict[str, torch.Tensor]:
-    state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs)
-    remapped: Dict[str, torch.Tensor] = {}
-    for key, value in state_dict.items():
-        if any(key.startswith(prefix) for prefix in extra_prefixes):
-            remapped[key] = value
-        elif key.startswith("transformer."):
-            remapped[key] = value
-        else:
-            remapped[f"transformer.{key}"] = value
-    return remapped
-
-
 def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor:
     if output.shape[1] == in_channels:
         return output
@@ -259,6 +141,14 @@ class TimestepEmbedder(nn.Module):
         return self.mlp(emb.to(dtype=weight_dtype))
 
 
+class _AFMCheckpointMixin(ModelMixin):
+    def _fix_state_dict_keys_on_load(self, state_dict):
+        wrapped = _wrap_dit_state_dict(state_dict)
+        state_dict.clear()
+        state_dict.update(wrapped)
+        super()._fix_state_dict_keys_on_load(state_dict)
+
+
 class Generator(nn.Module):
     def __init__(
         self,
@@ -279,14 +169,7 @@ class Generator(nn.Module):
         super().__init__()
         self.use_t_src = use_t_src
         self.use_t_tgt = use_t_tgt
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -306,20 +189,9 @@ class Generator(nn.Module):
             nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02)
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        extra = ("t_tgt_embedder.",)
-        has_extra = any(key.startswith(extra) for key in state_dict)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra)
-        if self.use_t_tgt and not has_extra:
-            strict = False
-        if not has_t_embedder:
+        if self.use_t_tgt and not any(key.startswith("t_tgt_embedder.") for key in state_dict):
             strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, t_src=None, t_tgt=None):
         if self.use_t_src and t_src is not None:
@@ -353,14 +225,7 @@ class GeneratorDeep(nn.Module):
     ):
         super().__init__()
         self.repeat = repeat
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -376,16 +241,7 @@ class GeneratorDeep(nn.Module):
         self.in_channels = in_channels
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        if not has_t_embedder:
-            strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, *args, **kwargs):
         hidden_states = self.transformer.pos_embed(x)
@@ -429,7 +285,7 @@ class GeneratorDeep(nn.Module):
         return split_output_channels(output, self.in_channels)
 
 
-class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
+class AFMGenerator2DModel(_AFMCheckpointMixin, Generator, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config
@@ -470,7 +326,7 @@ class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
         )
 
 
-class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin):
+class AFMGeneratorDeep2DModel(_AFMCheckpointMixin, GeneratorDeep, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config

AFM-B-2-1NFE-noguide/generator/config.json

@@ -1,6 +1,22 @@
 {
-  "_class_name": "AFMGenerator2DModel",
+  "_class_name": "DiTTransformer2DModel",
   "_diffusers_version": "0.36.0",
+  "sample_size": 32,
+  "num_layers": 12,
+  "num_attention_heads": 12,
+  "attention_head_dim": 64,
+  "in_channels": 4,
+  "out_channels": 4,
+  "patch_size": 2,
+  "attention_bias": true,
+  "activation_fn": "gelu-approximate",
+  "num_embeds_ada_norm": 1000,
+  "norm_type": "ada_norm_zero",
+  "norm_elementwise_affine": false,
+  "dropout": 0.0,
+  "norm_num_groups": 32,
+  "norm_eps": 1e-05,
+  "upcast_attention": false,
   "model_type": "AFM-B/2",
   "architecture": "standard",
   "repeat": 1,
@@ -12,9 +28,7 @@
   "class_dropout_prob": 0.0,
   "input_size": 32,
   "num_classes": 1000,
-  "in_channels": 4,
   "depth": 12,
   "hidden_size": 768,
-  "patch_size": 2,
   "num_heads": 12
 }

AFM-B-2-1NFE-noguide/generator/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:eb5297181aebfeb15fd7bf69781c80a2ae15aef5c24a2479991cba2fd17ba14d
-size 589740400
+oid sha256:af0b3881166119fbca961d253c9e96a0414e56096531dd34046961c934977051
+size 589737592

AFM-B-2-1NFE-noguide/generator/modeling_afm.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import copy
 import math
 from typing import Any, Dict, Optional
 
@@ -10,15 +9,10 @@ import torch
 import torch.nn.functional as F
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
 from torch import nn
 
 
-def get_dit_transformer_class():
-    from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
-
-    return DiTTransformer2DModel
-
-
 def build_dit_config(
     *,
     depth: int = 28,
@@ -54,110 +48,15 @@ def build_dit_config(
     }
 
 
-def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]:
-    converted = copy.deepcopy(state_dict)
-
-    converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous()
-    converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous()
-
-    if "t_embedder.mlp.0.weight" in converted:
-        timestep_weights = {
-            "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"),
-            "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"),
-            "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"),
-            "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"),
-        }
-    else:
-        timestep_weights = None
-    class_embedding = converted.pop("y_embedder.embedding_table.weight")
-    if class_embedding.shape[0] == 1000:
-        null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype)
-        class_embedding = torch.cat([class_embedding, null_class], dim=0)
-
-    for block_idx in range(depth):
-        if timestep_weights is not None:
-            for key, tensor in timestep_weights.items():
-                converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone()
-        converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = (
-            class_embedding.clone()
-        )
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.bias"
-        ]
-
-        q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0)
-        q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0)
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[
-            f"blocks.{block_idx}.attn.proj.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[
-            f"blocks.{block_idx}.attn.proj.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.bias"
-        ]
-
-        for suffix in (
-            "attn.qkv.weight",
-            "attn.qkv.bias",
-            "attn.proj.weight",
-            "attn.proj.bias",
-            "mlp.fc1.weight",
-            "mlp.fc1.bias",
-            "mlp.fc2.weight",
-            "mlp.fc2.bias",
-            "adaLN_modulation.1.weight",
-            "adaLN_modulation.1.bias",
-        ):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight")
-    converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias")
-    converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight")
-    converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias")
-
-    converted.pop("pos_embed", None)
-    for block_idx in range(depth):
-        for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()}
-
-
-def load_legacy_dit_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    depth: int,
-    patch_size: int,
-    hidden_size: int,
-    num_heads: int,
-) -> Dict[str, torch.Tensor]:
-    if "transformer_blocks.0.attn1.to_q.weight" in state_dict:
-        return state_dict
-    if any(key.startswith("blocks.") for key in state_dict):
-        return _convert_legacy_dit_state_dict(state_dict, depth)
-    return state_dict
+_PREFIX_KEYS = ("transformer.", "t_tgt_embedder.")
+
+
+def _wrap_dit_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    return {k if k.startswith(_PREFIX_KEYS) else f"transformer.{k}": v for k, v in state_dict.items()}
 
 
 def forward_dit_output(
-    transformer,
+    transformer: DiTTransformer2DModel,
     hidden_states: torch.Tensor,
     timestep: Optional[torch.Tensor],
     class_labels: torch.Tensor,
@@ -207,23 +106,6 @@ def forward_dit_output(
     )
 
 
-def remap_transformer_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    legacy_kwargs: Dict[str, int],
-    extra_prefixes: tuple[str, ...] = (),
-) -> Dict[str, torch.Tensor]:
-    state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs)
-    remapped: Dict[str, torch.Tensor] = {}
-    for key, value in state_dict.items():
-        if any(key.startswith(prefix) for prefix in extra_prefixes):
-            remapped[key] = value
-        elif key.startswith("transformer."):
-            remapped[key] = value
-        else:
-            remapped[f"transformer.{key}"] = value
-    return remapped
-
-
 def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor:
     if output.shape[1] == in_channels:
         return output
@@ -259,6 +141,14 @@ class TimestepEmbedder(nn.Module):
         return self.mlp(emb.to(dtype=weight_dtype))
 
 
+class _AFMCheckpointMixin(ModelMixin):
+    def _fix_state_dict_keys_on_load(self, state_dict):
+        wrapped = _wrap_dit_state_dict(state_dict)
+        state_dict.clear()
+        state_dict.update(wrapped)
+        super()._fix_state_dict_keys_on_load(state_dict)
+
+
 class Generator(nn.Module):
     def __init__(
         self,
@@ -279,14 +169,7 @@ class Generator(nn.Module):
         super().__init__()
         self.use_t_src = use_t_src
         self.use_t_tgt = use_t_tgt
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -306,20 +189,9 @@ class Generator(nn.Module):
             nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02)
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        extra = ("t_tgt_embedder.",)
-        has_extra = any(key.startswith(extra) for key in state_dict)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra)
-        if self.use_t_tgt and not has_extra:
-            strict = False
-        if not has_t_embedder:
+        if self.use_t_tgt and not any(key.startswith("t_tgt_embedder.") for key in state_dict):
             strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, t_src=None, t_tgt=None):
         if self.use_t_src and t_src is not None:
@@ -353,14 +225,7 @@ class GeneratorDeep(nn.Module):
     ):
         super().__init__()
         self.repeat = repeat
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -376,16 +241,7 @@ class GeneratorDeep(nn.Module):
         self.in_channels = in_channels
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        if not has_t_embedder:
-            strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, *args, **kwargs):
         hidden_states = self.transformer.pos_embed(x)
@@ -429,7 +285,7 @@ class GeneratorDeep(nn.Module):
         return split_output_channels(output, self.in_channels)
 
 
-class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
+class AFMGenerator2DModel(_AFMCheckpointMixin, Generator, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config
@@ -470,7 +326,7 @@ class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
         )
 
 
-class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin):
+class AFMGeneratorDeep2DModel(_AFMCheckpointMixin, GeneratorDeep, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config

AFM-L-2-1NFE-guided/generator/config.json

@@ -1,6 +1,22 @@
 {
-  "_class_name": "AFMGenerator2DModel",
+  "_class_name": "DiTTransformer2DModel",
   "_diffusers_version": "0.36.0",
+  "sample_size": 32,
+  "num_layers": 24,
+  "num_attention_heads": 16,
+  "attention_head_dim": 64,
+  "in_channels": 4,
+  "out_channels": 4,
+  "patch_size": 2,
+  "attention_bias": true,
+  "activation_fn": "gelu-approximate",
+  "num_embeds_ada_norm": 1000,
+  "norm_type": "ada_norm_zero",
+  "norm_elementwise_affine": false,
+  "dropout": 0.0,
+  "norm_num_groups": 32,
+  "norm_eps": 1e-05,
+  "upcast_attention": false,
   "model_type": "AFM-L/2",
   "architecture": "standard",
   "repeat": 1,
@@ -12,9 +28,7 @@
   "class_dropout_prob": 0.0,
   "input_size": 32,
   "num_classes": 1000,
-  "in_channels": 4,
   "depth": 24,
   "hidden_size": 1024,
-  "patch_size": 2,
   "num_heads": 16
 }

AFM-L-2-1NFE-guided/generator/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f95d199f5759ef66f561d7d9f08f31489b1794272f450713a844551c2c816682
-size 2046434688
+oid sha256:193cda8648b3cd31432402b249df6bd67ee7f0bc59ac00080f3f2d811919e871
+size 2046429144

AFM-L-2-1NFE-guided/generator/modeling_afm.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import copy
 import math
 from typing import Any, Dict, Optional
 
@@ -10,15 +9,10 @@ import torch
 import torch.nn.functional as F
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
 from torch import nn
 
 
-def get_dit_transformer_class():
-    from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
-
-    return DiTTransformer2DModel
-
-
 def build_dit_config(
     *,
     depth: int = 28,
@@ -54,110 +48,15 @@ def build_dit_config(
     }
 
 
-def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]:
-    converted = copy.deepcopy(state_dict)
-
-    converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous()
-    converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous()
-
-    if "t_embedder.mlp.0.weight" in converted:
-        timestep_weights = {
-            "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"),
-            "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"),
-            "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"),
-            "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"),
-        }
-    else:
-        timestep_weights = None
-    class_embedding = converted.pop("y_embedder.embedding_table.weight")
-    if class_embedding.shape[0] == 1000:
-        null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype)
-        class_embedding = torch.cat([class_embedding, null_class], dim=0)
-
-    for block_idx in range(depth):
-        if timestep_weights is not None:
-            for key, tensor in timestep_weights.items():
-                converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone()
-        converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = (
-            class_embedding.clone()
-        )
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.bias"
-        ]
-
-        q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0)
-        q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0)
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[
-            f"blocks.{block_idx}.attn.proj.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[
-            f"blocks.{block_idx}.attn.proj.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.bias"
-        ]
-
-        for suffix in (
-            "attn.qkv.weight",
-            "attn.qkv.bias",
-            "attn.proj.weight",
-            "attn.proj.bias",
-            "mlp.fc1.weight",
-            "mlp.fc1.bias",
-            "mlp.fc2.weight",
-            "mlp.fc2.bias",
-            "adaLN_modulation.1.weight",
-            "adaLN_modulation.1.bias",
-        ):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight")
-    converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias")
-    converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight")
-    converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias")
-
-    converted.pop("pos_embed", None)
-    for block_idx in range(depth):
-        for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()}
-
-
-def load_legacy_dit_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    depth: int,
-    patch_size: int,
-    hidden_size: int,
-    num_heads: int,
-) -> Dict[str, torch.Tensor]:
-    if "transformer_blocks.0.attn1.to_q.weight" in state_dict:
-        return state_dict
-    if any(key.startswith("blocks.") for key in state_dict):
-        return _convert_legacy_dit_state_dict(state_dict, depth)
-    return state_dict
+_PREFIX_KEYS = ("transformer.", "t_tgt_embedder.")
+
+
+def _wrap_dit_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    return {k if k.startswith(_PREFIX_KEYS) else f"transformer.{k}": v for k, v in state_dict.items()}
 
 
 def forward_dit_output(
-    transformer,
+    transformer: DiTTransformer2DModel,
     hidden_states: torch.Tensor,
     timestep: Optional[torch.Tensor],
     class_labels: torch.Tensor,
@@ -207,23 +106,6 @@ def forward_dit_output(
     )
 
 
-def remap_transformer_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    legacy_kwargs: Dict[str, int],
-    extra_prefixes: tuple[str, ...] = (),
-) -> Dict[str, torch.Tensor]:
-    state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs)
-    remapped: Dict[str, torch.Tensor] = {}
-    for key, value in state_dict.items():
-        if any(key.startswith(prefix) for prefix in extra_prefixes):
-            remapped[key] = value
-        elif key.startswith("transformer."):
-            remapped[key] = value
-        else:
-            remapped[f"transformer.{key}"] = value
-    return remapped
-
-
 def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor:
     if output.shape[1] == in_channels:
         return output
@@ -259,6 +141,14 @@ class TimestepEmbedder(nn.Module):
         return self.mlp(emb.to(dtype=weight_dtype))
 
 
+class _AFMCheckpointMixin(ModelMixin):
+    def _fix_state_dict_keys_on_load(self, state_dict):
+        wrapped = _wrap_dit_state_dict(state_dict)
+        state_dict.clear()
+        state_dict.update(wrapped)
+        super()._fix_state_dict_keys_on_load(state_dict)
+
+
 class Generator(nn.Module):
     def __init__(
         self,
@@ -279,14 +169,7 @@ class Generator(nn.Module):
         super().__init__()
         self.use_t_src = use_t_src
         self.use_t_tgt = use_t_tgt
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -306,20 +189,9 @@ class Generator(nn.Module):
             nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02)
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        extra = ("t_tgt_embedder.",)
-        has_extra = any(key.startswith(extra) for key in state_dict)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra)
-        if self.use_t_tgt and not has_extra:
-            strict = False
-        if not has_t_embedder:
+        if self.use_t_tgt and not any(key.startswith("t_tgt_embedder.") for key in state_dict):
             strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, t_src=None, t_tgt=None):
         if self.use_t_src and t_src is not None:
@@ -353,14 +225,7 @@ class GeneratorDeep(nn.Module):
     ):
         super().__init__()
         self.repeat = repeat
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -376,16 +241,7 @@ class GeneratorDeep(nn.Module):
         self.in_channels = in_channels
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        if not has_t_embedder:
-            strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, *args, **kwargs):
         hidden_states = self.transformer.pos_embed(x)
@@ -429,7 +285,7 @@ class GeneratorDeep(nn.Module):
         return split_output_channels(output, self.in_channels)
 
 
-class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
+class AFMGenerator2DModel(_AFMCheckpointMixin, Generator, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config
@@ -470,7 +326,7 @@ class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
         )
 
 
-class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin):
+class AFMGeneratorDeep2DModel(_AFMCheckpointMixin, GeneratorDeep, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config

AFM-L-2-1NFE-noguide/generator/config.json

@@ -1,6 +1,22 @@
 {
-  "_class_name": "AFMGenerator2DModel",
+  "_class_name": "DiTTransformer2DModel",
   "_diffusers_version": "0.36.0",
+  "sample_size": 32,
+  "num_layers": 24,
+  "num_attention_heads": 16,
+  "attention_head_dim": 64,
+  "in_channels": 4,
+  "out_channels": 4,
+  "patch_size": 2,
+  "attention_bias": true,
+  "activation_fn": "gelu-approximate",
+  "num_embeds_ada_norm": 1000,
+  "norm_type": "ada_norm_zero",
+  "norm_elementwise_affine": false,
+  "dropout": 0.0,
+  "norm_num_groups": 32,
+  "norm_eps": 1e-05,
+  "upcast_attention": false,
   "model_type": "AFM-L/2",
   "architecture": "standard",
   "repeat": 1,
@@ -12,9 +28,7 @@
   "class_dropout_prob": 0.0,
   "input_size": 32,
   "num_classes": 1000,
-  "in_channels": 4,
   "depth": 24,
   "hidden_size": 1024,
-  "patch_size": 2,
   "num_heads": 16
 }

AFM-L-2-1NFE-noguide/generator/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f8519fdc60ab76ee443981ec7a2601933ae54e13b67a247a0cc03e706c239c22
-size 2046434688
+oid sha256:e8fb9aadfdae6bbcfa0334d7769414bc9cecba5bfd106db0aeaa082a34f9b329
+size 2046429144

AFM-L-2-1NFE-noguide/generator/modeling_afm.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import copy
 import math
 from typing import Any, Dict, Optional
 
@@ -10,15 +9,10 @@ import torch
 import torch.nn.functional as F
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
 from torch import nn
 
 
-def get_dit_transformer_class():
-    from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
-
-    return DiTTransformer2DModel
-
-
 def build_dit_config(
     *,
     depth: int = 28,
@@ -54,110 +48,15 @@ def build_dit_config(
     }
 
 
-def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]:
-    converted = copy.deepcopy(state_dict)
-
-    converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous()
-    converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous()
-
-    if "t_embedder.mlp.0.weight" in converted:
-        timestep_weights = {
-            "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"),
-            "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"),
-            "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"),
-            "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"),
-        }
-    else:
-        timestep_weights = None
-    class_embedding = converted.pop("y_embedder.embedding_table.weight")
-    if class_embedding.shape[0] == 1000:
-        null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype)
-        class_embedding = torch.cat([class_embedding, null_class], dim=0)
-
-    for block_idx in range(depth):
-        if timestep_weights is not None:
-            for key, tensor in timestep_weights.items():
-                converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone()
-        converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = (
-            class_embedding.clone()
-        )
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.bias"
-        ]
-
-        q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0)
-        q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0)
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[
-            f"blocks.{block_idx}.attn.proj.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[
-            f"blocks.{block_idx}.attn.proj.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.bias"
-        ]
-
-        for suffix in (
-            "attn.qkv.weight",
-            "attn.qkv.bias",
-            "attn.proj.weight",
-            "attn.proj.bias",
-            "mlp.fc1.weight",
-            "mlp.fc1.bias",
-            "mlp.fc2.weight",
-            "mlp.fc2.bias",
-            "adaLN_modulation.1.weight",
-            "adaLN_modulation.1.bias",
-        ):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight")
-    converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias")
-    converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight")
-    converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias")
-
-    converted.pop("pos_embed", None)
-    for block_idx in range(depth):
-        for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()}
-
-
-def load_legacy_dit_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    depth: int,
-    patch_size: int,
-    hidden_size: int,
-    num_heads: int,
-) -> Dict[str, torch.Tensor]:
-    if "transformer_blocks.0.attn1.to_q.weight" in state_dict:
-        return state_dict
-    if any(key.startswith("blocks.") for key in state_dict):
-        return _convert_legacy_dit_state_dict(state_dict, depth)
-    return state_dict
+_PREFIX_KEYS = ("transformer.", "t_tgt_embedder.")
+
+
+def _wrap_dit_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    return {k if k.startswith(_PREFIX_KEYS) else f"transformer.{k}": v for k, v in state_dict.items()}
 
 
 def forward_dit_output(
-    transformer,
+    transformer: DiTTransformer2DModel,
     hidden_states: torch.Tensor,
     timestep: Optional[torch.Tensor],
     class_labels: torch.Tensor,
@@ -207,23 +106,6 @@ def forward_dit_output(
     )
 
 
-def remap_transformer_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    legacy_kwargs: Dict[str, int],
-    extra_prefixes: tuple[str, ...] = (),
-) -> Dict[str, torch.Tensor]:
-    state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs)
-    remapped: Dict[str, torch.Tensor] = {}
-    for key, value in state_dict.items():
-        if any(key.startswith(prefix) for prefix in extra_prefixes):
-            remapped[key] = value
-        elif key.startswith("transformer."):
-            remapped[key] = value
-        else:
-            remapped[f"transformer.{key}"] = value
-    return remapped
-
-
 def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor:
     if output.shape[1] == in_channels:
         return output
@@ -259,6 +141,14 @@ class TimestepEmbedder(nn.Module):
         return self.mlp(emb.to(dtype=weight_dtype))
 
 
+class _AFMCheckpointMixin(ModelMixin):
+    def _fix_state_dict_keys_on_load(self, state_dict):
+        wrapped = _wrap_dit_state_dict(state_dict)
+        state_dict.clear()
+        state_dict.update(wrapped)
+        super()._fix_state_dict_keys_on_load(state_dict)
+
+
 class Generator(nn.Module):
     def __init__(
         self,
@@ -279,14 +169,7 @@ class Generator(nn.Module):
         super().__init__()
         self.use_t_src = use_t_src
         self.use_t_tgt = use_t_tgt
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -306,20 +189,9 @@ class Generator(nn.Module):
             nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02)
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        extra = ("t_tgt_embedder.",)
-        has_extra = any(key.startswith(extra) for key in state_dict)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra)
-        if self.use_t_tgt and not has_extra:
-            strict = False
-        if not has_t_embedder:
+        if self.use_t_tgt and not any(key.startswith("t_tgt_embedder.") for key in state_dict):
             strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, t_src=None, t_tgt=None):
         if self.use_t_src and t_src is not None:
@@ -353,14 +225,7 @@ class GeneratorDeep(nn.Module):
     ):
         super().__init__()
         self.repeat = repeat
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -376,16 +241,7 @@ class GeneratorDeep(nn.Module):
         self.in_channels = in_channels
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        if not has_t_embedder:
-            strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, *args, **kwargs):
         hidden_states = self.transformer.pos_embed(x)
@@ -429,7 +285,7 @@ class GeneratorDeep(nn.Module):
         return split_output_channels(output, self.in_channels)
 
 
-class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
+class AFMGenerator2DModel(_AFMCheckpointMixin, Generator, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config
@@ -470,7 +326,7 @@ class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
         )
 
 
-class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin):
+class AFMGeneratorDeep2DModel(_AFMCheckpointMixin, GeneratorDeep, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config

AFM-M-2-1NFE-guided/generator/config.json

@@ -1,6 +1,22 @@
 {
-  "_class_name": "AFMGenerator2DModel",
+  "_class_name": "DiTTransformer2DModel",
   "_diffusers_version": "0.36.0",
+  "sample_size": 32,
+  "num_layers": 16,
+  "num_attention_heads": 16,
+  "attention_head_dim": 64,
+  "in_channels": 4,
+  "out_channels": 4,
+  "patch_size": 2,
+  "attention_bias": true,
+  "activation_fn": "gelu-approximate",
+  "num_embeds_ada_norm": 1000,
+  "norm_type": "ada_norm_zero",
+  "norm_elementwise_affine": false,
+  "dropout": 0.0,
+  "norm_num_groups": 32,
+  "norm_eps": 1e-05,
+  "upcast_attention": false,
   "model_type": "AFM-M/2",
   "architecture": "standard",
   "repeat": 1,
@@ -12,9 +28,7 @@
   "class_dropout_prob": 0.0,
   "input_size": 32,
   "num_classes": 1000,
-  "in_channels": 4,
   "depth": 16,
   "hidden_size": 1024,
-  "patch_size": 2,
   "num_heads": 16
 }

AFM-M-2-1NFE-guided/generator/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:12c57e80b3d69a132a46821571db1d7559f7b5df75fc74454b99c394c5c3385c
-size 1367133768
+oid sha256:4350f61ca9852882a8123ad337705fe805a7f419acbc715c3aae2d9d8c36684c
+size 1367130048

AFM-M-2-1NFE-guided/generator/modeling_afm.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import copy
 import math
 from typing import Any, Dict, Optional
 
@@ -10,15 +9,10 @@ import torch
 import torch.nn.functional as F
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
 from torch import nn
 
 
-def get_dit_transformer_class():
-    from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
-
-    return DiTTransformer2DModel
-
-
 def build_dit_config(
     *,
     depth: int = 28,
@@ -54,110 +48,15 @@ def build_dit_config(
     }
 
 
-def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]:
-    converted = copy.deepcopy(state_dict)
-
-    converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous()
-    converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous()
-
-    if "t_embedder.mlp.0.weight" in converted:
-        timestep_weights = {
-            "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"),
-            "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"),
-            "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"),
-            "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"),
-        }
-    else:
-        timestep_weights = None
-    class_embedding = converted.pop("y_embedder.embedding_table.weight")
-    if class_embedding.shape[0] == 1000:
-        null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype)
-        class_embedding = torch.cat([class_embedding, null_class], dim=0)
-
-    for block_idx in range(depth):
-        if timestep_weights is not None:
-            for key, tensor in timestep_weights.items():
-                converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone()
-        converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = (
-            class_embedding.clone()
-        )
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.bias"
-        ]
-
-        q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0)
-        q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0)
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[
-            f"blocks.{block_idx}.attn.proj.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[
-            f"blocks.{block_idx}.attn.proj.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.bias"
-        ]
-
-        for suffix in (
-            "attn.qkv.weight",
-            "attn.qkv.bias",
-            "attn.proj.weight",
-            "attn.proj.bias",
-            "mlp.fc1.weight",
-            "mlp.fc1.bias",
-            "mlp.fc2.weight",
-            "mlp.fc2.bias",
-            "adaLN_modulation.1.weight",
-            "adaLN_modulation.1.bias",
-        ):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight")
-    converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias")
-    converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight")
-    converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias")
-
-    converted.pop("pos_embed", None)
-    for block_idx in range(depth):
-        for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()}
-
-
-def load_legacy_dit_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    depth: int,
-    patch_size: int,
-    hidden_size: int,
-    num_heads: int,
-) -> Dict[str, torch.Tensor]:
-    if "transformer_blocks.0.attn1.to_q.weight" in state_dict:
-        return state_dict
-    if any(key.startswith("blocks.") for key in state_dict):
-        return _convert_legacy_dit_state_dict(state_dict, depth)
-    return state_dict
+_PREFIX_KEYS = ("transformer.", "t_tgt_embedder.")
+
+
+def _wrap_dit_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    return {k if k.startswith(_PREFIX_KEYS) else f"transformer.{k}": v for k, v in state_dict.items()}
 
 
 def forward_dit_output(
-    transformer,
+    transformer: DiTTransformer2DModel,
     hidden_states: torch.Tensor,
     timestep: Optional[torch.Tensor],
     class_labels: torch.Tensor,
@@ -207,23 +106,6 @@ def forward_dit_output(
     )
 
 
-def remap_transformer_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    legacy_kwargs: Dict[str, int],
-    extra_prefixes: tuple[str, ...] = (),
-) -> Dict[str, torch.Tensor]:
-    state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs)
-    remapped: Dict[str, torch.Tensor] = {}
-    for key, value in state_dict.items():
-        if any(key.startswith(prefix) for prefix in extra_prefixes):
-            remapped[key] = value
-        elif key.startswith("transformer."):
-            remapped[key] = value
-        else:
-            remapped[f"transformer.{key}"] = value
-    return remapped
-
-
 def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor:
     if output.shape[1] == in_channels:
         return output
@@ -259,6 +141,14 @@ class TimestepEmbedder(nn.Module):
         return self.mlp(emb.to(dtype=weight_dtype))
 
 
+class _AFMCheckpointMixin(ModelMixin):
+    def _fix_state_dict_keys_on_load(self, state_dict):
+        wrapped = _wrap_dit_state_dict(state_dict)
+        state_dict.clear()
+        state_dict.update(wrapped)
+        super()._fix_state_dict_keys_on_load(state_dict)
+
+
 class Generator(nn.Module):
     def __init__(
         self,
@@ -279,14 +169,7 @@ class Generator(nn.Module):
         super().__init__()
         self.use_t_src = use_t_src
         self.use_t_tgt = use_t_tgt
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -306,20 +189,9 @@ class Generator(nn.Module):
             nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02)
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        extra = ("t_tgt_embedder.",)
-        has_extra = any(key.startswith(extra) for key in state_dict)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra)
-        if self.use_t_tgt and not has_extra:
-            strict = False
-        if not has_t_embedder:
+        if self.use_t_tgt and not any(key.startswith("t_tgt_embedder.") for key in state_dict):
             strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, t_src=None, t_tgt=None):
         if self.use_t_src and t_src is not None:
@@ -353,14 +225,7 @@ class GeneratorDeep(nn.Module):
     ):
         super().__init__()
         self.repeat = repeat
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -376,16 +241,7 @@ class GeneratorDeep(nn.Module):
         self.in_channels = in_channels
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        if not has_t_embedder:
-            strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, *args, **kwargs):
         hidden_states = self.transformer.pos_embed(x)
@@ -429,7 +285,7 @@ class GeneratorDeep(nn.Module):
         return split_output_channels(output, self.in_channels)
 
 
-class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
+class AFMGenerator2DModel(_AFMCheckpointMixin, Generator, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config
@@ -470,7 +326,7 @@ class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
         )
 
 
-class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin):
+class AFMGeneratorDeep2DModel(_AFMCheckpointMixin, GeneratorDeep, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config

AFM-M-2-1NFE-noguide/generator/config.json

@@ -1,6 +1,22 @@
 {
-  "_class_name": "AFMGenerator2DModel",
+  "_class_name": "DiTTransformer2DModel",
   "_diffusers_version": "0.36.0",
+  "sample_size": 32,
+  "num_layers": 16,
+  "num_attention_heads": 16,
+  "attention_head_dim": 64,
+  "in_channels": 4,
+  "out_channels": 4,
+  "patch_size": 2,
+  "attention_bias": true,
+  "activation_fn": "gelu-approximate",
+  "num_embeds_ada_norm": 1000,
+  "norm_type": "ada_norm_zero",
+  "norm_elementwise_affine": false,
+  "dropout": 0.0,
+  "norm_num_groups": 32,
+  "norm_eps": 1e-05,
+  "upcast_attention": false,
   "model_type": "AFM-M/2",
   "architecture": "standard",
   "repeat": 1,
@@ -12,9 +28,7 @@
   "class_dropout_prob": 0.0,
   "input_size": 32,
   "num_classes": 1000,
-  "in_channels": 4,
   "depth": 16,
   "hidden_size": 1024,
-  "patch_size": 2,
   "num_heads": 16
 }

AFM-M-2-1NFE-noguide/generator/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ccdd512d0451e0fb435e8e0c16add4b4e949ecdb8996d6abf4e0e933b43d6e1c
-size 1367133768
+oid sha256:42a43983d07f79b9bb9a7d6de1ef1a3f3f201310f7fb30245973189e1a90dbf8
+size 1367130048

AFM-M-2-1NFE-noguide/generator/modeling_afm.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import copy
 import math
 from typing import Any, Dict, Optional
 
@@ -10,15 +9,10 @@ import torch
 import torch.nn.functional as F
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
 from torch import nn
 
 
-def get_dit_transformer_class():
-    from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
-
-    return DiTTransformer2DModel
-
-
 def build_dit_config(
     *,
     depth: int = 28,
@@ -54,110 +48,15 @@ def build_dit_config(
     }
 
 
-def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]:
-    converted = copy.deepcopy(state_dict)
-
-    converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous()
-    converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous()
-
-    if "t_embedder.mlp.0.weight" in converted:
-        timestep_weights = {
-            "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"),
-            "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"),
-            "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"),
-            "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"),
-        }
-    else:
-        timestep_weights = None
-    class_embedding = converted.pop("y_embedder.embedding_table.weight")
-    if class_embedding.shape[0] == 1000:
-        null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype)
-        class_embedding = torch.cat([class_embedding, null_class], dim=0)
-
-    for block_idx in range(depth):
-        if timestep_weights is not None:
-            for key, tensor in timestep_weights.items():
-                converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone()
-        converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = (
-            class_embedding.clone()
-        )
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.bias"
-        ]
-
-        q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0)
-        q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0)
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[
-            f"blocks.{block_idx}.attn.proj.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[
-            f"blocks.{block_idx}.attn.proj.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.bias"
-        ]
-
-        for suffix in (
-            "attn.qkv.weight",
-            "attn.qkv.bias",
-            "attn.proj.weight",
-            "attn.proj.bias",
-            "mlp.fc1.weight",
-            "mlp.fc1.bias",
-            "mlp.fc2.weight",
-            "mlp.fc2.bias",
-            "adaLN_modulation.1.weight",
-            "adaLN_modulation.1.bias",
-        ):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight")
-    converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias")
-    converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight")
-    converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias")
-
-    converted.pop("pos_embed", None)
-    for block_idx in range(depth):
-        for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()}
-
-
-def load_legacy_dit_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    depth: int,
-    patch_size: int,
-    hidden_size: int,
-    num_heads: int,
-) -> Dict[str, torch.Tensor]:
-    if "transformer_blocks.0.attn1.to_q.weight" in state_dict:
-        return state_dict
-    if any(key.startswith("blocks.") for key in state_dict):
-        return _convert_legacy_dit_state_dict(state_dict, depth)
-    return state_dict
+_PREFIX_KEYS = ("transformer.", "t_tgt_embedder.")
+
+
+def _wrap_dit_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    return {k if k.startswith(_PREFIX_KEYS) else f"transformer.{k}": v for k, v in state_dict.items()}
 
 
 def forward_dit_output(
-    transformer,
+    transformer: DiTTransformer2DModel,
     hidden_states: torch.Tensor,
     timestep: Optional[torch.Tensor],
     class_labels: torch.Tensor,
@@ -207,23 +106,6 @@ def forward_dit_output(
     )
 
 
-def remap_transformer_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    legacy_kwargs: Dict[str, int],
-    extra_prefixes: tuple[str, ...] = (),
-) -> Dict[str, torch.Tensor]:
-    state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs)
-    remapped: Dict[str, torch.Tensor] = {}
-    for key, value in state_dict.items():
-        if any(key.startswith(prefix) for prefix in extra_prefixes):
-            remapped[key] = value
-        elif key.startswith("transformer."):
-            remapped[key] = value
-        else:
-            remapped[f"transformer.{key}"] = value
-    return remapped
-
-
 def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor:
     if output.shape[1] == in_channels:
         return output
@@ -259,6 +141,14 @@ class TimestepEmbedder(nn.Module):
         return self.mlp(emb.to(dtype=weight_dtype))
 
 
+class _AFMCheckpointMixin(ModelMixin):
+    def _fix_state_dict_keys_on_load(self, state_dict):
+        wrapped = _wrap_dit_state_dict(state_dict)
+        state_dict.clear()
+        state_dict.update(wrapped)
+        super()._fix_state_dict_keys_on_load(state_dict)
+
+
 class Generator(nn.Module):
     def __init__(
         self,
@@ -279,14 +169,7 @@ class Generator(nn.Module):
         super().__init__()
         self.use_t_src = use_t_src
         self.use_t_tgt = use_t_tgt
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -306,20 +189,9 @@ class Generator(nn.Module):
             nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02)
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        extra = ("t_tgt_embedder.",)
-        has_extra = any(key.startswith(extra) for key in state_dict)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra)
-        if self.use_t_tgt and not has_extra:
-            strict = False
-        if not has_t_embedder:
+        if self.use_t_tgt and not any(key.startswith("t_tgt_embedder.") for key in state_dict):
             strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, t_src=None, t_tgt=None):
         if self.use_t_src and t_src is not None:
@@ -353,14 +225,7 @@ class GeneratorDeep(nn.Module):
     ):
         super().__init__()
         self.repeat = repeat
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -376,16 +241,7 @@ class GeneratorDeep(nn.Module):
         self.in_channels = in_channels
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        if not has_t_embedder:
-            strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, *args, **kwargs):
         hidden_states = self.transformer.pos_embed(x)
@@ -429,7 +285,7 @@ class GeneratorDeep(nn.Module):
         return split_output_channels(output, self.in_channels)
 
 
-class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
+class AFMGenerator2DModel(_AFMCheckpointMixin, Generator, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config
@@ -470,7 +326,7 @@ class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
         )
 
 
-class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin):
+class AFMGeneratorDeep2DModel(_AFMCheckpointMixin, GeneratorDeep, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config

AFM-XL-2-112layer-1NFE-guided/generator/config.json

@@ -1,6 +1,22 @@
 {
-  "_class_name": "AFMGeneratorDeep2DModel",
+  "_class_name": "DiTTransformer2DModel",
   "_diffusers_version": "0.36.0",
+  "sample_size": 32,
+  "num_layers": 28,
+  "num_attention_heads": 16,
+  "attention_head_dim": 72,
+  "in_channels": 4,
+  "out_channels": 4,
+  "patch_size": 2,
+  "attention_bias": true,
+  "activation_fn": "gelu-approximate",
+  "num_embeds_ada_norm": 1000,
+  "norm_type": "ada_norm_zero",
+  "norm_elementwise_affine": false,
+  "dropout": 0.0,
+  "norm_num_groups": 32,
+  "norm_eps": 1e-05,
+  "upcast_attention": false,
   "model_type": "AFM-XL/2",
   "architecture": "deep",
   "repeat": 4,
@@ -12,9 +28,7 @@
   "class_dropout_prob": 0.0,
   "input_size": 32,
   "num_classes": 1000,
-  "in_channels": 4,
   "depth": 28,
   "hidden_size": 1152,
-  "patch_size": 2,
   "num_heads": 16
 }

AFM-XL-2-112layer-1NFE-guided/generator/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f49b04fcd6f77d6fb7039781d0a872bb80ed4a07001b11dc3cd340492965c1a3
-size 2999303040
+oid sha256:4cbdfe1714c45d0d2245a3de495a709b4c45d4f87bc01bfd0e734474b85afdd9
+size 2999296584

AFM-XL-2-112layer-1NFE-guided/generator/modeling_afm.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import copy
 import math
 from typing import Any, Dict, Optional
 
@@ -10,15 +9,10 @@ import torch
 import torch.nn.functional as F
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
 from torch import nn
 
 
-def get_dit_transformer_class():
-    from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
-
-    return DiTTransformer2DModel
-
-
 def build_dit_config(
     *,
     depth: int = 28,
@@ -54,110 +48,15 @@ def build_dit_config(
     }
 
 
-def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]:
-    converted = copy.deepcopy(state_dict)
-
-    converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous()
-    converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous()
-
-    if "t_embedder.mlp.0.weight" in converted:
-        timestep_weights = {
-            "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"),
-            "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"),
-            "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"),
-            "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"),
-        }
-    else:
-        timestep_weights = None
-    class_embedding = converted.pop("y_embedder.embedding_table.weight")
-    if class_embedding.shape[0] == 1000:
-        null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype)
-        class_embedding = torch.cat([class_embedding, null_class], dim=0)
-
-    for block_idx in range(depth):
-        if timestep_weights is not None:
-            for key, tensor in timestep_weights.items():
-                converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone()
-        converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = (
-            class_embedding.clone()
-        )
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.bias"
-        ]
-
-        q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0)
-        q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0)
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[
-            f"blocks.{block_idx}.attn.proj.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[
-            f"blocks.{block_idx}.attn.proj.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.bias"
-        ]
-
-        for suffix in (
-            "attn.qkv.weight",
-            "attn.qkv.bias",
-            "attn.proj.weight",
-            "attn.proj.bias",
-            "mlp.fc1.weight",
-            "mlp.fc1.bias",
-            "mlp.fc2.weight",
-            "mlp.fc2.bias",
-            "adaLN_modulation.1.weight",
-            "adaLN_modulation.1.bias",
-        ):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight")
-    converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias")
-    converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight")
-    converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias")
-
-    converted.pop("pos_embed", None)
-    for block_idx in range(depth):
-        for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()}
-
-
-def load_legacy_dit_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    depth: int,
-    patch_size: int,
-    hidden_size: int,
-    num_heads: int,
-) -> Dict[str, torch.Tensor]:
-    if "transformer_blocks.0.attn1.to_q.weight" in state_dict:
-        return state_dict
-    if any(key.startswith("blocks.") for key in state_dict):
-        return _convert_legacy_dit_state_dict(state_dict, depth)
-    return state_dict
+_PREFIX_KEYS = ("transformer.", "t_tgt_embedder.")
+
+
+def _wrap_dit_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    return {k if k.startswith(_PREFIX_KEYS) else f"transformer.{k}": v for k, v in state_dict.items()}
 
 
 def forward_dit_output(
-    transformer,
+    transformer: DiTTransformer2DModel,
     hidden_states: torch.Tensor,
     timestep: Optional[torch.Tensor],
     class_labels: torch.Tensor,
@@ -207,23 +106,6 @@ def forward_dit_output(
     )
 
 
-def remap_transformer_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    legacy_kwargs: Dict[str, int],
-    extra_prefixes: tuple[str, ...] = (),
-) -> Dict[str, torch.Tensor]:
-    state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs)
-    remapped: Dict[str, torch.Tensor] = {}
-    for key, value in state_dict.items():
-        if any(key.startswith(prefix) for prefix in extra_prefixes):
-            remapped[key] = value
-        elif key.startswith("transformer."):
-            remapped[key] = value
-        else:
-            remapped[f"transformer.{key}"] = value
-    return remapped
-
-
 def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor:
     if output.shape[1] == in_channels:
         return output
@@ -259,6 +141,14 @@ class TimestepEmbedder(nn.Module):
         return self.mlp(emb.to(dtype=weight_dtype))
 
 
+class _AFMCheckpointMixin(ModelMixin):
+    def _fix_state_dict_keys_on_load(self, state_dict):
+        wrapped = _wrap_dit_state_dict(state_dict)
+        state_dict.clear()
+        state_dict.update(wrapped)
+        super()._fix_state_dict_keys_on_load(state_dict)
+
+
 class Generator(nn.Module):
     def __init__(
         self,
@@ -279,14 +169,7 @@ class Generator(nn.Module):
         super().__init__()
         self.use_t_src = use_t_src
         self.use_t_tgt = use_t_tgt
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -306,20 +189,9 @@ class Generator(nn.Module):
             nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02)
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        extra = ("t_tgt_embedder.",)
-        has_extra = any(key.startswith(extra) for key in state_dict)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra)
-        if self.use_t_tgt and not has_extra:
-            strict = False
-        if not has_t_embedder:
+        if self.use_t_tgt and not any(key.startswith("t_tgt_embedder.") for key in state_dict):
             strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, t_src=None, t_tgt=None):
         if self.use_t_src and t_src is not None:
@@ -353,14 +225,7 @@ class GeneratorDeep(nn.Module):
     ):
         super().__init__()
         self.repeat = repeat
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -376,16 +241,7 @@ class GeneratorDeep(nn.Module):
         self.in_channels = in_channels
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        if not has_t_embedder:
-            strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, *args, **kwargs):
         hidden_states = self.transformer.pos_embed(x)
@@ -429,7 +285,7 @@ class GeneratorDeep(nn.Module):
         return split_output_channels(output, self.in_channels)
 
 
-class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
+class AFMGenerator2DModel(_AFMCheckpointMixin, Generator, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config
@@ -470,7 +326,7 @@ class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
         )
 
 
-class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin):
+class AFMGeneratorDeep2DModel(_AFMCheckpointMixin, GeneratorDeep, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config

AFM-XL-2-1NFE-guided/generator/config.json

@@ -1,6 +1,22 @@
 {
-  "_class_name": "AFMGenerator2DModel",
+  "_class_name": "DiTTransformer2DModel",
   "_diffusers_version": "0.36.0",
+  "sample_size": 32,
+  "num_layers": 28,
+  "num_attention_heads": 16,
+  "attention_head_dim": 72,
+  "in_channels": 4,
+  "out_channels": 4,
+  "patch_size": 2,
+  "attention_bias": true,
+  "activation_fn": "gelu-approximate",
+  "num_embeds_ada_norm": 1000,
+  "norm_type": "ada_norm_zero",
+  "norm_elementwise_affine": false,
+  "dropout": 0.0,
+  "norm_num_groups": 32,
+  "norm_eps": 1e-05,
+  "upcast_attention": false,
   "model_type": "AFM-XL/2",
   "architecture": "standard",
   "repeat": 1,
@@ -12,9 +28,7 @@
   "class_dropout_prob": 0.0,
   "input_size": 32,
   "num_classes": 1000,
-  "in_channels": 4,
   "depth": 28,
   "hidden_size": 1152,
-  "patch_size": 2,
   "num_heads": 16
 }

AFM-XL-2-1NFE-guided/generator/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c4b03219ec7f002d176494286d7bed760ff916e5b1518b3f9ce3e2bb9d0c25d1
-size 2999303040
+oid sha256:e8c3fc7246258cac7b4b16b8c95e5fa08a64a2f853e17689783a46cc4bccad9b
+size 2999296584

AFM-XL-2-1NFE-guided/generator/modeling_afm.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import copy
 import math
 from typing import Any, Dict, Optional
 
@@ -10,15 +9,10 @@ import torch
 import torch.nn.functional as F
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
 from torch import nn
 
 
-def get_dit_transformer_class():
-    from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
-
-    return DiTTransformer2DModel
-
-
 def build_dit_config(
     *,
     depth: int = 28,
@@ -54,110 +48,15 @@ def build_dit_config(
     }
 
 
-def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]:
-    converted = copy.deepcopy(state_dict)
-
-    converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous()
-    converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous()
-
-    if "t_embedder.mlp.0.weight" in converted:
-        timestep_weights = {
-            "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"),
-            "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"),
-            "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"),
-            "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"),
-        }
-    else:
-        timestep_weights = None
-    class_embedding = converted.pop("y_embedder.embedding_table.weight")
-    if class_embedding.shape[0] == 1000:
-        null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype)
-        class_embedding = torch.cat([class_embedding, null_class], dim=0)
-
-    for block_idx in range(depth):
-        if timestep_weights is not None:
-            for key, tensor in timestep_weights.items():
-                converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone()
-        converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = (
-            class_embedding.clone()
-        )
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.bias"
-        ]
-
-        q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0)
-        q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0)
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[
-            f"blocks.{block_idx}.attn.proj.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[
-            f"blocks.{block_idx}.attn.proj.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.bias"
-        ]
-
-        for suffix in (
-            "attn.qkv.weight",
-            "attn.qkv.bias",
-            "attn.proj.weight",
-            "attn.proj.bias",
-            "mlp.fc1.weight",
-            "mlp.fc1.bias",
-            "mlp.fc2.weight",
-            "mlp.fc2.bias",
-            "adaLN_modulation.1.weight",
-            "adaLN_modulation.1.bias",
-        ):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight")
-    converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias")
-    converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight")
-    converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias")
-
-    converted.pop("pos_embed", None)
-    for block_idx in range(depth):
-        for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()}
-
-
-def load_legacy_dit_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    depth: int,
-    patch_size: int,
-    hidden_size: int,
-    num_heads: int,
-) -> Dict[str, torch.Tensor]:
-    if "transformer_blocks.0.attn1.to_q.weight" in state_dict:
-        return state_dict
-    if any(key.startswith("blocks.") for key in state_dict):
-        return _convert_legacy_dit_state_dict(state_dict, depth)
-    return state_dict
+_PREFIX_KEYS = ("transformer.", "t_tgt_embedder.")
+
+
+def _wrap_dit_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    return {k if k.startswith(_PREFIX_KEYS) else f"transformer.{k}": v for k, v in state_dict.items()}
 
 
 def forward_dit_output(
-    transformer,
+    transformer: DiTTransformer2DModel,
     hidden_states: torch.Tensor,
     timestep: Optional[torch.Tensor],
     class_labels: torch.Tensor,
@@ -207,23 +106,6 @@ def forward_dit_output(
     )
 
 
-def remap_transformer_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    legacy_kwargs: Dict[str, int],
-    extra_prefixes: tuple[str, ...] = (),
-) -> Dict[str, torch.Tensor]:
-    state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs)
-    remapped: Dict[str, torch.Tensor] = {}
-    for key, value in state_dict.items():
-        if any(key.startswith(prefix) for prefix in extra_prefixes):
-            remapped[key] = value
-        elif key.startswith("transformer."):
-            remapped[key] = value
-        else:
-            remapped[f"transformer.{key}"] = value
-    return remapped
-
-
 def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor:
     if output.shape[1] == in_channels:
         return output
@@ -259,6 +141,14 @@ class TimestepEmbedder(nn.Module):
         return self.mlp(emb.to(dtype=weight_dtype))
 
 
+class _AFMCheckpointMixin(ModelMixin):
+    def _fix_state_dict_keys_on_load(self, state_dict):
+        wrapped = _wrap_dit_state_dict(state_dict)
+        state_dict.clear()
+        state_dict.update(wrapped)
+        super()._fix_state_dict_keys_on_load(state_dict)
+
+
 class Generator(nn.Module):
     def __init__(
         self,
@@ -279,14 +169,7 @@ class Generator(nn.Module):
         super().__init__()
         self.use_t_src = use_t_src
         self.use_t_tgt = use_t_tgt
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -306,20 +189,9 @@ class Generator(nn.Module):
             nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02)
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        extra = ("t_tgt_embedder.",)
-        has_extra = any(key.startswith(extra) for key in state_dict)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra)
-        if self.use_t_tgt and not has_extra:
-            strict = False
-        if not has_t_embedder:
+        if self.use_t_tgt and not any(key.startswith("t_tgt_embedder.") for key in state_dict):
             strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, t_src=None, t_tgt=None):
         if self.use_t_src and t_src is not None:
@@ -353,14 +225,7 @@ class GeneratorDeep(nn.Module):
     ):
         super().__init__()
         self.repeat = repeat
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -376,16 +241,7 @@ class GeneratorDeep(nn.Module):
         self.in_channels = in_channels
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        if not has_t_embedder:
-            strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, *args, **kwargs):
         hidden_states = self.transformer.pos_embed(x)
@@ -429,7 +285,7 @@ class GeneratorDeep(nn.Module):
         return split_output_channels(output, self.in_channels)
 
 
-class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
+class AFMGenerator2DModel(_AFMCheckpointMixin, Generator, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config
@@ -470,7 +326,7 @@ class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
         )
 
 
-class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin):
+class AFMGeneratorDeep2DModel(_AFMCheckpointMixin, GeneratorDeep, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config

AFM-XL-2-1NFE-noguide/generator/config.json

@@ -1,6 +1,22 @@
 {
-  "_class_name": "AFMGenerator2DModel",
+  "_class_name": "DiTTransformer2DModel",
   "_diffusers_version": "0.36.0",
+  "sample_size": 32,
+  "num_layers": 28,
+  "num_attention_heads": 16,
+  "attention_head_dim": 72,
+  "in_channels": 4,
+  "out_channels": 4,
+  "patch_size": 2,
+  "attention_bias": true,
+  "activation_fn": "gelu-approximate",
+  "num_embeds_ada_norm": 1000,
+  "norm_type": "ada_norm_zero",
+  "norm_elementwise_affine": false,
+  "dropout": 0.0,
+  "norm_num_groups": 32,
+  "norm_eps": 1e-05,
+  "upcast_attention": false,
   "model_type": "AFM-XL/2",
   "architecture": "standard",
   "repeat": 1,
@@ -12,9 +28,7 @@
   "class_dropout_prob": 0.0,
   "input_size": 32,
   "num_classes": 1000,
-  "in_channels": 4,
   "depth": 28,
   "hidden_size": 1152,
-  "patch_size": 2,
   "num_heads": 16
 }

AFM-XL-2-1NFE-noguide/generator/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:acb2f8302ed8e316e57a73a6fd34dd4da9b14d752e94658dfc1d6f86f580ebfc
-size 2999303040
+oid sha256:1582a0f05105962729018de245acd95d74651ab3b5bdec9c0aed52ec6c7bf5cd
+size 2999296584

AFM-XL-2-1NFE-noguide/generator/modeling_afm.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import copy
 import math
 from typing import Any, Dict, Optional
 
@@ -10,15 +9,10 @@ import torch
 import torch.nn.functional as F
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
 from torch import nn
 
 
-def get_dit_transformer_class():
-    from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
-
-    return DiTTransformer2DModel
-
-
 def build_dit_config(
     *,
     depth: int = 28,
@@ -54,110 +48,15 @@ def build_dit_config(
     }
 
 
-def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]:
-    converted = copy.deepcopy(state_dict)
-
-    converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous()
-    converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous()
-
-    if "t_embedder.mlp.0.weight" in converted:
-        timestep_weights = {
-            "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"),
-            "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"),
-            "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"),
-            "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"),
-        }
-    else:
-        timestep_weights = None
-    class_embedding = converted.pop("y_embedder.embedding_table.weight")
-    if class_embedding.shape[0] == 1000:
-        null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype)
-        class_embedding = torch.cat([class_embedding, null_class], dim=0)
-
-    for block_idx in range(depth):
-        if timestep_weights is not None:
-            for key, tensor in timestep_weights.items():
-                converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone()
-        converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = (
-            class_embedding.clone()
-        )
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.bias"
-        ]
-
-        q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0)
-        q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0)
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[
-            f"blocks.{block_idx}.attn.proj.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[
-            f"blocks.{block_idx}.attn.proj.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.bias"
-        ]
-
-        for suffix in (
-            "attn.qkv.weight",
-            "attn.qkv.bias",
-            "attn.proj.weight",
-            "attn.proj.bias",
-            "mlp.fc1.weight",
-            "mlp.fc1.bias",
-            "mlp.fc2.weight",
-            "mlp.fc2.bias",
-            "adaLN_modulation.1.weight",
-            "adaLN_modulation.1.bias",
-        ):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight")
-    converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias")
-    converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight")
-    converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias")
-
-    converted.pop("pos_embed", None)
-    for block_idx in range(depth):
-        for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()}
-
-
-def load_legacy_dit_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    depth: int,
-    patch_size: int,
-    hidden_size: int,
-    num_heads: int,
-) -> Dict[str, torch.Tensor]:
-    if "transformer_blocks.0.attn1.to_q.weight" in state_dict:
-        return state_dict
-    if any(key.startswith("blocks.") for key in state_dict):
-        return _convert_legacy_dit_state_dict(state_dict, depth)
-    return state_dict
+_PREFIX_KEYS = ("transformer.", "t_tgt_embedder.")
+
+
+def _wrap_dit_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    return {k if k.startswith(_PREFIX_KEYS) else f"transformer.{k}": v for k, v in state_dict.items()}
 
 
 def forward_dit_output(
-    transformer,
+    transformer: DiTTransformer2DModel,
     hidden_states: torch.Tensor,
     timestep: Optional[torch.Tensor],
     class_labels: torch.Tensor,
@@ -207,23 +106,6 @@ def forward_dit_output(
     )
 
 
-def remap_transformer_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    legacy_kwargs: Dict[str, int],
-    extra_prefixes: tuple[str, ...] = (),
-) -> Dict[str, torch.Tensor]:
-    state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs)
-    remapped: Dict[str, torch.Tensor] = {}
-    for key, value in state_dict.items():
-        if any(key.startswith(prefix) for prefix in extra_prefixes):
-            remapped[key] = value
-        elif key.startswith("transformer."):
-            remapped[key] = value
-        else:
-            remapped[f"transformer.{key}"] = value
-    return remapped
-
-
 def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor:
     if output.shape[1] == in_channels:
         return output
@@ -259,6 +141,14 @@ class TimestepEmbedder(nn.Module):
         return self.mlp(emb.to(dtype=weight_dtype))
 
 
+class _AFMCheckpointMixin(ModelMixin):
+    def _fix_state_dict_keys_on_load(self, state_dict):
+        wrapped = _wrap_dit_state_dict(state_dict)
+        state_dict.clear()
+        state_dict.update(wrapped)
+        super()._fix_state_dict_keys_on_load(state_dict)
+
+
 class Generator(nn.Module):
     def __init__(
         self,
@@ -279,14 +169,7 @@ class Generator(nn.Module):
         super().__init__()
         self.use_t_src = use_t_src
         self.use_t_tgt = use_t_tgt
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -306,20 +189,9 @@ class Generator(nn.Module):
             nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02)
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        extra = ("t_tgt_embedder.",)
-        has_extra = any(key.startswith(extra) for key in state_dict)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra)
-        if self.use_t_tgt and not has_extra:
-            strict = False
-        if not has_t_embedder:
+        if self.use_t_tgt and not any(key.startswith("t_tgt_embedder.") for key in state_dict):
             strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, t_src=None, t_tgt=None):
         if self.use_t_src and t_src is not None:
@@ -353,14 +225,7 @@ class GeneratorDeep(nn.Module):
     ):
         super().__init__()
         self.repeat = repeat
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -376,16 +241,7 @@ class GeneratorDeep(nn.Module):
         self.in_channels = in_channels
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        if not has_t_embedder:
-            strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, *args, **kwargs):
         hidden_states = self.transformer.pos_embed(x)
@@ -429,7 +285,7 @@ class GeneratorDeep(nn.Module):
         return split_output_channels(output, self.in_channels)
 
 
-class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
+class AFMGenerator2DModel(_AFMCheckpointMixin, Generator, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config
@@ -470,7 +326,7 @@ class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
         )
 
 
-class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin):
+class AFMGeneratorDeep2DModel(_AFMCheckpointMixin, GeneratorDeep, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config

AFM-XL-2-2NFE-guided/generator/config.json

@@ -1,6 +1,22 @@
 {
-  "_class_name": "AFMGenerator2DModel",
+  "_class_name": "DiTTransformer2DModel",
   "_diffusers_version": "0.36.0",
+  "sample_size": 32,
+  "num_layers": 28,
+  "num_attention_heads": 16,
+  "attention_head_dim": 72,
+  "in_channels": 4,
+  "out_channels": 4,
+  "patch_size": 2,
+  "attention_bias": true,
+  "activation_fn": "gelu-approximate",
+  "num_embeds_ada_norm": 1000,
+  "norm_type": "ada_norm_zero",
+  "norm_elementwise_affine": false,
+  "dropout": 0.0,
+  "norm_num_groups": 32,
+  "norm_eps": 1e-05,
+  "upcast_attention": false,
   "model_type": "AFM-XL/2",
   "architecture": "standard",
   "repeat": 1,
@@ -12,9 +28,7 @@
   "class_dropout_prob": 0.0,
   "input_size": 32,
   "num_classes": 1000,
-  "in_channels": 4,
   "depth": 28,
   "hidden_size": 1152,
-  "patch_size": 2,
   "num_heads": 16
 }

AFM-XL-2-2NFE-guided/generator/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c4ba62222809142771f48d4281e98299d6c11b83ff2b9b75163f98c7a23a9167
-size 2999303040
+oid sha256:d5b5c2e9fceff4d97eaad3e150878d765b4121b44b0c782c7c62145f2e2219bb
+size 2999296584

AFM-XL-2-2NFE-guided/generator/modeling_afm.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import copy
 import math
 from typing import Any, Dict, Optional
 
@@ -10,15 +9,10 @@ import torch
 import torch.nn.functional as F
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
 from torch import nn
 
 
-def get_dit_transformer_class():
-    from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
-
-    return DiTTransformer2DModel
-
-
 def build_dit_config(
     *,
     depth: int = 28,
@@ -54,110 +48,15 @@ def build_dit_config(
     }
 
 
-def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]:
-    converted = copy.deepcopy(state_dict)
-
-    converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous()
-    converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous()
-
-    if "t_embedder.mlp.0.weight" in converted:
-        timestep_weights = {
-            "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"),
-            "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"),
-            "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"),
-            "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"),
-        }
-    else:
-        timestep_weights = None
-    class_embedding = converted.pop("y_embedder.embedding_table.weight")
-    if class_embedding.shape[0] == 1000:
-        null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype)
-        class_embedding = torch.cat([class_embedding, null_class], dim=0)
-
-    for block_idx in range(depth):
-        if timestep_weights is not None:
-            for key, tensor in timestep_weights.items():
-                converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone()
-        converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = (
-            class_embedding.clone()
-        )
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.bias"
-        ]
-
-        q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0)
-        q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0)
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[
-            f"blocks.{block_idx}.attn.proj.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[
-            f"blocks.{block_idx}.attn.proj.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.bias"
-        ]
-
-        for suffix in (
-            "attn.qkv.weight",
-            "attn.qkv.bias",
-            "attn.proj.weight",
-            "attn.proj.bias",
-            "mlp.fc1.weight",
-            "mlp.fc1.bias",
-            "mlp.fc2.weight",
-            "mlp.fc2.bias",
-            "adaLN_modulation.1.weight",
-            "adaLN_modulation.1.bias",
-        ):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight")
-    converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias")
-    converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight")
-    converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias")
-
-    converted.pop("pos_embed", None)
-    for block_idx in range(depth):
-        for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()}
-
-
-def load_legacy_dit_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    depth: int,
-    patch_size: int,
-    hidden_size: int,
-    num_heads: int,
-) -> Dict[str, torch.Tensor]:
-    if "transformer_blocks.0.attn1.to_q.weight" in state_dict:
-        return state_dict
-    if any(key.startswith("blocks.") for key in state_dict):
-        return _convert_legacy_dit_state_dict(state_dict, depth)
-    return state_dict
+_PREFIX_KEYS = ("transformer.", "t_tgt_embedder.")
+
+
+def _wrap_dit_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    return {k if k.startswith(_PREFIX_KEYS) else f"transformer.{k}": v for k, v in state_dict.items()}
 
 
 def forward_dit_output(
-    transformer,
+    transformer: DiTTransformer2DModel,
     hidden_states: torch.Tensor,
     timestep: Optional[torch.Tensor],
     class_labels: torch.Tensor,
@@ -207,23 +106,6 @@ def forward_dit_output(
     )
 
 
-def remap_transformer_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    legacy_kwargs: Dict[str, int],
-    extra_prefixes: tuple[str, ...] = (),
-) -> Dict[str, torch.Tensor]:
-    state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs)
-    remapped: Dict[str, torch.Tensor] = {}
-    for key, value in state_dict.items():
-        if any(key.startswith(prefix) for prefix in extra_prefixes):
-            remapped[key] = value
-        elif key.startswith("transformer."):
-            remapped[key] = value
-        else:
-            remapped[f"transformer.{key}"] = value
-    return remapped
-
-
 def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor:
     if output.shape[1] == in_channels:
         return output
@@ -259,6 +141,14 @@ class TimestepEmbedder(nn.Module):
         return self.mlp(emb.to(dtype=weight_dtype))
 
 
+class _AFMCheckpointMixin(ModelMixin):
+    def _fix_state_dict_keys_on_load(self, state_dict):
+        wrapped = _wrap_dit_state_dict(state_dict)
+        state_dict.clear()
+        state_dict.update(wrapped)
+        super()._fix_state_dict_keys_on_load(state_dict)
+
+
 class Generator(nn.Module):
     def __init__(
         self,
@@ -279,14 +169,7 @@ class Generator(nn.Module):
         super().__init__()
         self.use_t_src = use_t_src
         self.use_t_tgt = use_t_tgt
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -306,20 +189,9 @@ class Generator(nn.Module):
             nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02)
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        extra = ("t_tgt_embedder.",)
-        has_extra = any(key.startswith(extra) for key in state_dict)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra)
-        if self.use_t_tgt and not has_extra:
-            strict = False
-        if not has_t_embedder:
+        if self.use_t_tgt and not any(key.startswith("t_tgt_embedder.") for key in state_dict):
             strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, t_src=None, t_tgt=None):
         if self.use_t_src and t_src is not None:
@@ -353,14 +225,7 @@ class GeneratorDeep(nn.Module):
     ):
         super().__init__()
         self.repeat = repeat
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -376,16 +241,7 @@ class GeneratorDeep(nn.Module):
         self.in_channels = in_channels
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        if not has_t_embedder:
-            strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, *args, **kwargs):
         hidden_states = self.transformer.pos_embed(x)
@@ -429,7 +285,7 @@ class GeneratorDeep(nn.Module):
         return split_output_channels(output, self.in_channels)
 
 
-class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
+class AFMGenerator2DModel(_AFMCheckpointMixin, Generator, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config
@@ -470,7 +326,7 @@ class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
         )
 
 
-class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin):
+class AFMGeneratorDeep2DModel(_AFMCheckpointMixin, GeneratorDeep, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config

AFM-XL-2-2NFE-noguide/generator/config.json

@@ -1,6 +1,22 @@
 {
-  "_class_name": "AFMGenerator2DModel",
+  "_class_name": "DiTTransformer2DModel",
   "_diffusers_version": "0.36.0",
+  "sample_size": 32,
+  "num_layers": 28,
+  "num_attention_heads": 16,
+  "attention_head_dim": 72,
+  "in_channels": 4,
+  "out_channels": 4,
+  "patch_size": 2,
+  "attention_bias": true,
+  "activation_fn": "gelu-approximate",
+  "num_embeds_ada_norm": 1000,
+  "norm_type": "ada_norm_zero",
+  "norm_elementwise_affine": false,
+  "dropout": 0.0,
+  "norm_num_groups": 32,
+  "norm_eps": 1e-05,
+  "upcast_attention": false,
   "model_type": "AFM-XL/2",
   "architecture": "standard",
   "repeat": 1,
@@ -12,9 +28,7 @@
   "class_dropout_prob": 0.0,
   "input_size": 32,
   "num_classes": 1000,
-  "in_channels": 4,
   "depth": 28,
   "hidden_size": 1152,
-  "patch_size": 2,
   "num_heads": 16
 }

AFM-XL-2-2NFE-noguide/generator/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8a646088766ee3de076b7731984bad8fce5105d4a7166e84872dd0d41a6f5a1a
-size 2999303040
+oid sha256:d553e9ca93352c17ecbeb52136e74c721a4d6832368507e483a1bea30c5d75ff
+size 2999296584

AFM-XL-2-2NFE-noguide/generator/modeling_afm.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import copy
 import math
 from typing import Any, Dict, Optional
 
@@ -10,15 +9,10 @@ import torch
 import torch.nn.functional as F
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
 from torch import nn
 
 
-def get_dit_transformer_class():
-    from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
-
-    return DiTTransformer2DModel
-
-
 def build_dit_config(
     *,
     depth: int = 28,
@@ -54,110 +48,15 @@ def build_dit_config(
     }
 
 
-def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]:
-    converted = copy.deepcopy(state_dict)
-
-    converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous()
-    converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous()
-
-    if "t_embedder.mlp.0.weight" in converted:
-        timestep_weights = {
-            "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"),
-            "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"),
-            "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"),
-            "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"),
-        }
-    else:
-        timestep_weights = None
-    class_embedding = converted.pop("y_embedder.embedding_table.weight")
-    if class_embedding.shape[0] == 1000:
-        null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype)
-        class_embedding = torch.cat([class_embedding, null_class], dim=0)
-
-    for block_idx in range(depth):
-        if timestep_weights is not None:
-            for key, tensor in timestep_weights.items():
-                converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone()
-        converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = (
-            class_embedding.clone()
-        )
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.bias"
-        ]
-
-        q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0)
-        q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0)
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[
-            f"blocks.{block_idx}.attn.proj.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[
-            f"blocks.{block_idx}.attn.proj.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.bias"
-        ]
-
-        for suffix in (
-            "attn.qkv.weight",
-            "attn.qkv.bias",
-            "attn.proj.weight",
-            "attn.proj.bias",
-            "mlp.fc1.weight",
-            "mlp.fc1.bias",
-            "mlp.fc2.weight",
-            "mlp.fc2.bias",
-            "adaLN_modulation.1.weight",
-            "adaLN_modulation.1.bias",
-        ):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight")
-    converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias")
-    converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight")
-    converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias")
-
-    converted.pop("pos_embed", None)
-    for block_idx in range(depth):
-        for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()}
-
-
-def load_legacy_dit_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    depth: int,
-    patch_size: int,
-    hidden_size: int,
-    num_heads: int,
-) -> Dict[str, torch.Tensor]:
-    if "transformer_blocks.0.attn1.to_q.weight" in state_dict:
-        return state_dict
-    if any(key.startswith("blocks.") for key in state_dict):
-        return _convert_legacy_dit_state_dict(state_dict, depth)
-    return state_dict
+_PREFIX_KEYS = ("transformer.", "t_tgt_embedder.")
+
+
+def _wrap_dit_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    return {k if k.startswith(_PREFIX_KEYS) else f"transformer.{k}": v for k, v in state_dict.items()}
 
 
 def forward_dit_output(
-    transformer,
+    transformer: DiTTransformer2DModel,
     hidden_states: torch.Tensor,
     timestep: Optional[torch.Tensor],
     class_labels: torch.Tensor,
@@ -207,23 +106,6 @@ def forward_dit_output(
     )
 
 
-def remap_transformer_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    legacy_kwargs: Dict[str, int],
-    extra_prefixes: tuple[str, ...] = (),
-) -> Dict[str, torch.Tensor]:
-    state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs)
-    remapped: Dict[str, torch.Tensor] = {}
-    for key, value in state_dict.items():
-        if any(key.startswith(prefix) for prefix in extra_prefixes):
-            remapped[key] = value
-        elif key.startswith("transformer."):
-            remapped[key] = value
-        else:
-            remapped[f"transformer.{key}"] = value
-    return remapped
-
-
 def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor:
     if output.shape[1] == in_channels:
         return output
@@ -259,6 +141,14 @@ class TimestepEmbedder(nn.Module):
         return self.mlp(emb.to(dtype=weight_dtype))
 
 
+class _AFMCheckpointMixin(ModelMixin):
+    def _fix_state_dict_keys_on_load(self, state_dict):
+        wrapped = _wrap_dit_state_dict(state_dict)
+        state_dict.clear()
+        state_dict.update(wrapped)
+        super()._fix_state_dict_keys_on_load(state_dict)
+
+
 class Generator(nn.Module):
     def __init__(
         self,
@@ -279,14 +169,7 @@ class Generator(nn.Module):
         super().__init__()
         self.use_t_src = use_t_src
         self.use_t_tgt = use_t_tgt
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -306,20 +189,9 @@ class Generator(nn.Module):
             nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02)
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        extra = ("t_tgt_embedder.",)
-        has_extra = any(key.startswith(extra) for key in state_dict)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra)
-        if self.use_t_tgt and not has_extra:
-            strict = False
-        if not has_t_embedder:
+        if self.use_t_tgt and not any(key.startswith("t_tgt_embedder.") for key in state_dict):
             strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, t_src=None, t_tgt=None):
         if self.use_t_src and t_src is not None:
@@ -353,14 +225,7 @@ class GeneratorDeep(nn.Module):
     ):
         super().__init__()
         self.repeat = repeat
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -376,16 +241,7 @@ class GeneratorDeep(nn.Module):
         self.in_channels = in_channels
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        if not has_t_embedder:
-            strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, *args, **kwargs):
         hidden_states = self.transformer.pos_embed(x)
@@ -429,7 +285,7 @@ class GeneratorDeep(nn.Module):
         return split_output_channels(output, self.in_channels)
 
 
-class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
+class AFMGenerator2DModel(_AFMCheckpointMixin, Generator, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config
@@ -470,7 +326,7 @@ class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
         )
 
 
-class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin):
+class AFMGeneratorDeep2DModel(_AFMCheckpointMixin, GeneratorDeep, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config

AFM-XL-2-4NFE-guided/generator/config.json

@@ -1,6 +1,22 @@
 {
-  "_class_name": "AFMGenerator2DModel",
+  "_class_name": "DiTTransformer2DModel",
   "_diffusers_version": "0.36.0",
+  "sample_size": 32,
+  "num_layers": 28,
+  "num_attention_heads": 16,
+  "attention_head_dim": 72,
+  "in_channels": 4,
+  "out_channels": 4,
+  "patch_size": 2,
+  "attention_bias": true,
+  "activation_fn": "gelu-approximate",
+  "num_embeds_ada_norm": 1000,
+  "norm_type": "ada_norm_zero",
+  "norm_elementwise_affine": false,
+  "dropout": 0.0,
+  "norm_num_groups": 32,
+  "norm_eps": 1e-05,
+  "upcast_attention": false,
   "model_type": "AFM-XL/2",
   "architecture": "standard",
   "repeat": 1,
@@ -12,9 +28,7 @@
   "class_dropout_prob": 0.0,
   "input_size": 32,
   "num_classes": 1000,
-  "in_channels": 4,
   "depth": 28,
   "hidden_size": 1152,
-  "patch_size": 2,
   "num_heads": 16
 }

AFM-XL-2-4NFE-guided/generator/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0ba89ed56a0d1e898ec71dd953a7c28060fd04b440e85a41ccb00b385a93daa3
-size 2999303040
+oid sha256:4ebc0085e11680404c6745acd138f1f15e13a73b8e0d6c06a3bf6b1972d60206
+size 2999296584

AFM-XL-2-4NFE-guided/generator/modeling_afm.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import copy
 import math
 from typing import Any, Dict, Optional
 
@@ -10,15 +9,10 @@ import torch
 import torch.nn.functional as F
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
 from torch import nn
 
 
-def get_dit_transformer_class():
-    from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
-
-    return DiTTransformer2DModel
-
-
 def build_dit_config(
     *,
     depth: int = 28,
@@ -54,110 +48,15 @@ def build_dit_config(
     }
 
 
-def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]:
-    converted = copy.deepcopy(state_dict)
-
-    converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous()
-    converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous()
-
-    if "t_embedder.mlp.0.weight" in converted:
-        timestep_weights = {
-            "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"),
-            "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"),
-            "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"),
-            "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"),
-        }
-    else:
-        timestep_weights = None
-    class_embedding = converted.pop("y_embedder.embedding_table.weight")
-    if class_embedding.shape[0] == 1000:
-        null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype)
-        class_embedding = torch.cat([class_embedding, null_class], dim=0)
-
-    for block_idx in range(depth):
-        if timestep_weights is not None:
-            for key, tensor in timestep_weights.items():
-                converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone()
-        converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = (
-            class_embedding.clone()
-        )
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.bias"
-        ]
-
-        q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0)
-        q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0)
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[
-            f"blocks.{block_idx}.attn.proj.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[
-            f"blocks.{block_idx}.attn.proj.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.bias"
-        ]
-
-        for suffix in (
-            "attn.qkv.weight",
-            "attn.qkv.bias",
-            "attn.proj.weight",
-            "attn.proj.bias",
-            "mlp.fc1.weight",
-            "mlp.fc1.bias",
-            "mlp.fc2.weight",
-            "mlp.fc2.bias",
-            "adaLN_modulation.1.weight",
-            "adaLN_modulation.1.bias",
-        ):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight")
-    converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias")
-    converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight")
-    converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias")
-
-    converted.pop("pos_embed", None)
-    for block_idx in range(depth):
-        for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()}
-
-
-def load_legacy_dit_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    depth: int,
-    patch_size: int,
-    hidden_size: int,
-    num_heads: int,
-) -> Dict[str, torch.Tensor]:
-    if "transformer_blocks.0.attn1.to_q.weight" in state_dict:
-        return state_dict
-    if any(key.startswith("blocks.") for key in state_dict):
-        return _convert_legacy_dit_state_dict(state_dict, depth)
-    return state_dict
+_PREFIX_KEYS = ("transformer.", "t_tgt_embedder.")
+
+
+def _wrap_dit_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    return {k if k.startswith(_PREFIX_KEYS) else f"transformer.{k}": v for k, v in state_dict.items()}
 
 
 def forward_dit_output(
-    transformer,
+    transformer: DiTTransformer2DModel,
     hidden_states: torch.Tensor,
     timestep: Optional[torch.Tensor],
     class_labels: torch.Tensor,
@@ -207,23 +106,6 @@ def forward_dit_output(
     )
 
 
-def remap_transformer_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    legacy_kwargs: Dict[str, int],
-    extra_prefixes: tuple[str, ...] = (),
-) -> Dict[str, torch.Tensor]:
-    state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs)
-    remapped: Dict[str, torch.Tensor] = {}
-    for key, value in state_dict.items():
-        if any(key.startswith(prefix) for prefix in extra_prefixes):
-            remapped[key] = value
-        elif key.startswith("transformer."):
-            remapped[key] = value
-        else:
-            remapped[f"transformer.{key}"] = value
-    return remapped
-
-
 def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor:
     if output.shape[1] == in_channels:
         return output
@@ -259,6 +141,14 @@ class TimestepEmbedder(nn.Module):
         return self.mlp(emb.to(dtype=weight_dtype))
 
 
+class _AFMCheckpointMixin(ModelMixin):
+    def _fix_state_dict_keys_on_load(self, state_dict):
+        wrapped = _wrap_dit_state_dict(state_dict)
+        state_dict.clear()
+        state_dict.update(wrapped)
+        super()._fix_state_dict_keys_on_load(state_dict)
+
+
 class Generator(nn.Module):
     def __init__(
         self,
@@ -279,14 +169,7 @@ class Generator(nn.Module):
         super().__init__()
         self.use_t_src = use_t_src
         self.use_t_tgt = use_t_tgt
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -306,20 +189,9 @@ class Generator(nn.Module):
             nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02)
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        extra = ("t_tgt_embedder.",)
-        has_extra = any(key.startswith(extra) for key in state_dict)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra)
-        if self.use_t_tgt and not has_extra:
-            strict = False
-        if not has_t_embedder:
+        if self.use_t_tgt and not any(key.startswith("t_tgt_embedder.") for key in state_dict):
             strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, t_src=None, t_tgt=None):
         if self.use_t_src and t_src is not None:
@@ -353,14 +225,7 @@ class GeneratorDeep(nn.Module):
     ):
         super().__init__()
         self.repeat = repeat
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -376,16 +241,7 @@ class GeneratorDeep(nn.Module):
         self.in_channels = in_channels
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        if not has_t_embedder:
-            strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, *args, **kwargs):
         hidden_states = self.transformer.pos_embed(x)
@@ -429,7 +285,7 @@ class GeneratorDeep(nn.Module):
         return split_output_channels(output, self.in_channels)
 
 
-class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
+class AFMGenerator2DModel(_AFMCheckpointMixin, Generator, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config
@@ -470,7 +326,7 @@ class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
         )
 
 
-class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin):
+class AFMGeneratorDeep2DModel(_AFMCheckpointMixin, GeneratorDeep, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config

AFM-XL-2-56layer-1NFE-guided/generator/config.json

@@ -1,6 +1,22 @@
 {
-  "_class_name": "AFMGeneratorDeep2DModel",
+  "_class_name": "DiTTransformer2DModel",
   "_diffusers_version": "0.36.0",
+  "sample_size": 32,
+  "num_layers": 28,
+  "num_attention_heads": 16,
+  "attention_head_dim": 72,
+  "in_channels": 4,
+  "out_channels": 4,
+  "patch_size": 2,
+  "attention_bias": true,
+  "activation_fn": "gelu-approximate",
+  "num_embeds_ada_norm": 1000,
+  "norm_type": "ada_norm_zero",
+  "norm_elementwise_affine": false,
+  "dropout": 0.0,
+  "norm_num_groups": 32,
+  "norm_eps": 1e-05,
+  "upcast_attention": false,
   "model_type": "AFM-XL/2",
   "architecture": "deep",
   "repeat": 2,
@@ -12,9 +28,7 @@
   "class_dropout_prob": 0.0,
   "input_size": 32,
   "num_classes": 1000,
-  "in_channels": 4,
   "depth": 28,
   "hidden_size": 1152,
-  "patch_size": 2,
   "num_heads": 16
 }

AFM-XL-2-56layer-1NFE-guided/generator/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a35d5ec2bb5d92a4ff5bc98bd923ec5d34818d21213446c055010a3a7224a2ee
-size 2999303040
+oid sha256:39b4886846d53e4e664aa1eb70af1da6ad4a87c3f0f09d37f8e81b596ea319ce
+size 2999296584

AFM-XL-2-56layer-1NFE-guided/generator/modeling_afm.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import copy
 import math
 from typing import Any, Dict, Optional
 
@@ -10,15 +9,10 @@ import torch
 import torch.nn.functional as F
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
 from torch import nn
 
 
-def get_dit_transformer_class():
-    from diffusers.models.transformers.dit_transformer_2d import DiTTransformer2DModel
-
-    return DiTTransformer2DModel
-
-
 def build_dit_config(
     *,
     depth: int = 28,
@@ -54,110 +48,15 @@ def build_dit_config(
     }
 
 
-def _convert_legacy_dit_state_dict(state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]:
-    converted = copy.deepcopy(state_dict)
-
-    converted["pos_embed.proj.weight"] = converted.pop("x_embedder.proj.weight").clone().contiguous()
-    converted["pos_embed.proj.bias"] = converted.pop("x_embedder.proj.bias").clone().contiguous()
-
-    if "t_embedder.mlp.0.weight" in converted:
-        timestep_weights = {
-            "linear_1.weight": converted.pop("t_embedder.mlp.0.weight"),
-            "linear_1.bias": converted.pop("t_embedder.mlp.0.bias"),
-            "linear_2.weight": converted.pop("t_embedder.mlp.2.weight"),
-            "linear_2.bias": converted.pop("t_embedder.mlp.2.bias"),
-        }
-    else:
-        timestep_weights = None
-    class_embedding = converted.pop("y_embedder.embedding_table.weight")
-    if class_embedding.shape[0] == 1000:
-        null_class = torch.zeros(1, class_embedding.shape[1], dtype=class_embedding.dtype)
-        class_embedding = torch.cat([class_embedding, null_class], dim=0)
-
-    for block_idx in range(depth):
-        if timestep_weights is not None:
-            for key, tensor in timestep_weights.items():
-                converted[f"transformer_blocks.{block_idx}.norm1.emb.timestep_embedder.{key}"] = tensor.clone()
-        converted[f"transformer_blocks.{block_idx}.norm1.emb.class_embedder.embedding_table.weight"] = (
-            class_embedding.clone()
-        )
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.weight"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.norm1.linear.bias"] = converted[
-            f"blocks.{block_idx}.adaLN_modulation.1.bias"
-        ]
-
-        q, k, v = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.weight"], 3, dim=0)
-        q_bias, k_bias, v_bias = torch.chunk(converted[f"blocks.{block_idx}.attn.qkv.bias"], 3, dim=0)
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.weight"] = q
-        converted[f"transformer_blocks.{block_idx}.attn1.to_q.bias"] = q_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.weight"] = k
-        converted[f"transformer_blocks.{block_idx}.attn1.to_k.bias"] = k_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.weight"] = v
-        converted[f"transformer_blocks.{block_idx}.attn1.to_v.bias"] = v_bias
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.weight"] = converted[
-            f"blocks.{block_idx}.attn.proj.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.attn1.to_out.0.bias"] = converted[
-            f"blocks.{block_idx}.attn.proj.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.0.proj.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc1.bias"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.weight"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.weight"
-        ]
-        converted[f"transformer_blocks.{block_idx}.ff.net.2.bias"] = converted[
-            f"blocks.{block_idx}.mlp.fc2.bias"
-        ]
-
-        for suffix in (
-            "attn.qkv.weight",
-            "attn.qkv.bias",
-            "attn.proj.weight",
-            "attn.proj.bias",
-            "mlp.fc1.weight",
-            "mlp.fc1.bias",
-            "mlp.fc2.weight",
-            "mlp.fc2.bias",
-            "adaLN_modulation.1.weight",
-            "adaLN_modulation.1.bias",
-        ):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    converted["proj_out_1.weight"] = converted.pop("final_layer.adaLN_modulation.1.weight")
-    converted["proj_out_1.bias"] = converted.pop("final_layer.adaLN_modulation.1.bias")
-    converted["proj_out_2.weight"] = converted.pop("final_layer.linear.weight")
-    converted["proj_out_2.bias"] = converted.pop("final_layer.linear.bias")
-
-    converted.pop("pos_embed", None)
-    for block_idx in range(depth):
-        for suffix in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"):
-            converted.pop(f"blocks.{block_idx}.{suffix}", None)
-
-    return {key: tensor.detach().clone().contiguous() for key, tensor in converted.items()}
-
-
-def load_legacy_dit_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    depth: int,
-    patch_size: int,
-    hidden_size: int,
-    num_heads: int,
-) -> Dict[str, torch.Tensor]:
-    if "transformer_blocks.0.attn1.to_q.weight" in state_dict:
-        return state_dict
-    if any(key.startswith("blocks.") for key in state_dict):
-        return _convert_legacy_dit_state_dict(state_dict, depth)
-    return state_dict
+_PREFIX_KEYS = ("transformer.", "t_tgt_embedder.")
+
+
+def _wrap_dit_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    return {k if k.startswith(_PREFIX_KEYS) else f"transformer.{k}": v for k, v in state_dict.items()}
 
 
 def forward_dit_output(
-    transformer,
+    transformer: DiTTransformer2DModel,
     hidden_states: torch.Tensor,
     timestep: Optional[torch.Tensor],
     class_labels: torch.Tensor,
@@ -207,23 +106,6 @@ def forward_dit_output(
     )
 
 
-def remap_transformer_state_dict(
-    state_dict: Dict[str, torch.Tensor],
-    legacy_kwargs: Dict[str, int],
-    extra_prefixes: tuple[str, ...] = (),
-) -> Dict[str, torch.Tensor]:
-    state_dict = load_legacy_dit_state_dict(state_dict, **legacy_kwargs)
-    remapped: Dict[str, torch.Tensor] = {}
-    for key, value in state_dict.items():
-        if any(key.startswith(prefix) for prefix in extra_prefixes):
-            remapped[key] = value
-        elif key.startswith("transformer."):
-            remapped[key] = value
-        else:
-            remapped[f"transformer.{key}"] = value
-    return remapped
-
-
 def split_output_channels(output: torch.Tensor, in_channels: int) -> torch.Tensor:
     if output.shape[1] == in_channels:
         return output
@@ -259,6 +141,14 @@ class TimestepEmbedder(nn.Module):
         return self.mlp(emb.to(dtype=weight_dtype))
 
 
+class _AFMCheckpointMixin(ModelMixin):
+    def _fix_state_dict_keys_on_load(self, state_dict):
+        wrapped = _wrap_dit_state_dict(state_dict)
+        state_dict.clear()
+        state_dict.update(wrapped)
+        super()._fix_state_dict_keys_on_load(state_dict)
+
+
 class Generator(nn.Module):
     def __init__(
         self,
@@ -279,14 +169,7 @@ class Generator(nn.Module):
         super().__init__()
         self.use_t_src = use_t_src
         self.use_t_tgt = use_t_tgt
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -306,20 +189,9 @@ class Generator(nn.Module):
             nn.init.normal_(self.t_tgt_embedder.mlp[2].weight, std=0.02)
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        extra = ("t_tgt_embedder.",)
-        has_extra = any(key.startswith(extra) for key in state_dict)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs, extra_prefixes=extra)
-        if self.use_t_tgt and not has_extra:
-            strict = False
-        if not has_t_embedder:
+        if self.use_t_tgt and not any(key.startswith("t_tgt_embedder.") for key in state_dict):
             strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, t_src=None, t_tgt=None):
         if self.use_t_src and t_src is not None:
@@ -353,14 +225,7 @@ class GeneratorDeep(nn.Module):
     ):
         super().__init__()
         self.repeat = repeat
-        self._legacy_kwargs = {
-            "depth": depth,
-            "hidden_size": hidden_size,
-            "patch_size": patch_size,
-            "num_heads": num_heads,
-        }
-        dit_cls = get_dit_transformer_class()
-        self.transformer = dit_cls(
+        self.transformer = DiTTransformer2DModel(
             **build_dit_config(
                 depth=depth,
                 hidden_size=hidden_size,
@@ -376,16 +241,7 @@ class GeneratorDeep(nn.Module):
         self.in_channels = in_channels
 
     def load_state_dict(self, state_dict, strict: bool = True):
-        remapped = remap_transformer_state_dict(state_dict, self._legacy_kwargs)
-        has_t_embedder = any("t_embedder" in key for key in state_dict)
-        if not has_t_embedder:
-            strict = False
-        out = super().load_state_dict(remapped, strict=strict)
-        if not has_t_embedder:
-            for name, param in self.named_parameters():
-                if "timestep_embedder" in name:
-                    param.data.zero_()
-        return out
+        return super().load_state_dict(_wrap_dit_state_dict(state_dict), strict=strict)
 
     def forward(self, x, y, *args, **kwargs):
         hidden_states = self.transformer.pos_embed(x)
@@ -429,7 +285,7 @@ class GeneratorDeep(nn.Module):
         return split_output_channels(output, self.in_channels)
 
 
-class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
+class AFMGenerator2DModel(_AFMCheckpointMixin, Generator, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config
@@ -470,7 +326,7 @@ class AFMGenerator2DModel(Generator, ModelMixin, ConfigMixin):
         )
 
 
-class AFMGeneratorDeep2DModel(GeneratorDeep, ModelMixin, ConfigMixin):
+class AFMGeneratorDeep2DModel(_AFMCheckpointMixin, GeneratorDeep, ConfigMixin):
     config_name = "config.json"
 
     @register_to_config