use _from_deprecated_attn_block check re: @patrickvonplaten

Author: williamberman

src/diffusers/models/attention_processor.py

@@ -68,6 +68,7 @@ def __init__(
         eps: float = 1e-5,
         rescale_output_factor: float = 1.0,
         residual_connection: bool = False,
+        _from_deprecated_attn_block=False,
         processor: Optional["AttnProcessor"] = None,
     ):
         super().__init__()
@@ -78,6 +79,10 @@ def __init__(
         self.rescale_output_factor = rescale_output_factor
         self.residual_connection = residual_connection
 
+        # we make use of this private variable to know whether this class is loaded
+        # with an deprecated state dict so that we can convert it on the fly
+        self._from_deprecated_attn_block = _from_deprecated_attn_block
+
         self.scale_qk = scale_qk
         self.scale = dim_head**-0.5 if self.scale_qk else 1.0
 

src/diffusers/models/modeling_utils.py

@@ -583,7 +583,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
                 if device_map is None:
                     param_device = "cpu"
                     state_dict = load_state_dict(model_file, variant=variant)
-                    cls.convert_deprecated_attention_blocks(state_dict)
+                    model._convert_deprecated_attention_blocks(state_dict)
                     # move the params from meta device to cpu
                     missing_keys = set(model.state_dict().keys()) - set(state_dict.keys())
                     if len(missing_keys) > 0:
@@ -626,7 +626,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
                 model = cls.from_config(config, **unused_kwargs)
 
                 state_dict = load_state_dict(model_file, variant=variant)
-                cls.convert_deprecated_attention_blocks(state_dict)
+                model._convert_deprecated_attention_blocks(state_dict)
 
                 model, missing_keys, unexpected_keys, mismatched_keys, error_msgs = cls._load_pretrained_model(
                     model,
@@ -763,42 +763,6 @@ def _find_mismatched_keys(
 
         return model, missing_keys, unexpected_keys, mismatched_keys, error_msgs
 
-    @classmethod
-    def convert_deprecated_attention_blocks(cls, state_dict):
-        # We check for the deprecated attention block via the `proj_attn` layer in the state dict
-        # The only other class with a layer called `proj_attn` is in `SelfAttention1d` which is used
-        # by only the top level model, `UNet1DModel`. Since, `UNet1DModel` wont have any of the deprecated
-        # attention blocks, we can just early return.
-        if cls.__name__ == "UNet1DModel":
-            return
-
-        deprecated_attention_block_paths = []
-
-        for k in state_dict.keys():
-            if "proj_attn.weight" in k:
-                index = k.index("proj_attn.weight")
-                path = k[: index - 1]
-                deprecated_attention_block_paths.append(path)
-
-        for path in deprecated_attention_block_paths:
-            # group_norm path stays the same
-
-            # query -> to_q
-            state_dict[f"{path}.to_q.weight"] = state_dict.pop(f"{path}.query.weight")
-            state_dict[f"{path}.to_q.bias"] = state_dict.pop(f"{path}.query.bias")
-
-            # key -> to_k
-            state_dict[f"{path}.to_k.weight"] = state_dict.pop(f"{path}.key.weight")
-            state_dict[f"{path}.to_k.bias"] = state_dict.pop(f"{path}.key.bias")
-
-            # value -> to_v
-            state_dict[f"{path}.to_v.weight"] = state_dict.pop(f"{path}.value.weight")
-            state_dict[f"{path}.to_v.bias"] = state_dict.pop(f"{path}.value.bias")
-
-            # proj_attn -> to_out.0
-            state_dict[f"{path}.to_out.0.weight"] = state_dict.pop(f"{path}.proj_attn.weight")
-            state_dict[f"{path}.to_out.0.bias"] = state_dict.pop(f"{path}.proj_attn.bias")
-
     @property
     def device(self) -> device:
         """
@@ -841,3 +805,34 @@ def num_parameters(self, only_trainable: bool = False, exclude_embeddings: bool
             return sum(p.numel() for p in non_embedding_parameters if p.requires_grad or not only_trainable)
         else:
             return sum(p.numel() for p in self.parameters() if p.requires_grad or not only_trainable)
+
+    def _convert_deprecated_attention_blocks(self, state_dict):
+        deprecated_attention_block_paths = []
+
+        def recursive_find_attn_block(name, module):
+            if hasattr(module, "_from_deprecated_attn_block") and module._from_deprecated_attn_block:
+                deprecated_attention_block_paths.append(name)
+            for sub_name, sub_module in module.named_children():
+                sub_name = sub_name if name == "" else f"{name}.{sub_name}"
+                recursive_find_attn_block(sub_name, sub_module)
+
+        recursive_find_attn_block("", self)
+
+        for path in deprecated_attention_block_paths:
+            # group_norm path stays the same
+
+            # query -> to_q
+            state_dict[f"{path}.to_q.weight"] = state_dict.pop(f"{path}.query.weight")
+            state_dict[f"{path}.to_q.bias"] = state_dict.pop(f"{path}.query.bias")
+
+            # key -> to_k
+            state_dict[f"{path}.to_k.weight"] = state_dict.pop(f"{path}.key.weight")
+            state_dict[f"{path}.to_k.bias"] = state_dict.pop(f"{path}.key.bias")
+
+            # value -> to_v
+            state_dict[f"{path}.to_v.weight"] = state_dict.pop(f"{path}.value.weight")
+            state_dict[f"{path}.to_v.bias"] = state_dict.pop(f"{path}.value.bias")
+
+            # proj_attn -> to_out.0
+            state_dict[f"{path}.to_out.0.weight"] = state_dict.pop(f"{path}.proj_attn.weight")
+            state_dict[f"{path}.to_out.0.bias"] = state_dict.pop(f"{path}.proj_attn.bias")

src/diffusers/models/unet_2d_blocks.py

@@ -437,6 +437,7 @@ def __init__(
                         residual_connection=True,
                         bias=True,
                         upcast_softmax=True,
+                        _from_deprecated_attn_block=True,
                     )
                 )
             else:
@@ -725,6 +726,7 @@ def __init__(
                     residual_connection=True,
                     bias=True,
                     upcast_softmax=True,
+                    _from_deprecated_attn_block=True,
                 )
             )
 
@@ -1078,6 +1080,7 @@ def __init__(
                     residual_connection=True,
                     bias=True,
                     upcast_softmax=True,
+                    _from_deprecated_attn_block=True,
                 )
             )
 
@@ -1156,6 +1159,7 @@ def __init__(
                     residual_connection=True,
                     bias=True,
                     upcast_softmax=True,
+                    _from_deprecated_attn_block=True,
                 )
             )
 
@@ -1730,6 +1734,7 @@ def __init__(
                     residual_connection=True,
                     bias=True,
                     upcast_softmax=True,
+                    _from_deprecated_attn_block=True,
                 )
             )
 
@@ -2068,6 +2073,7 @@ def __init__(
                     residual_connection=True,
                     bias=True,
                     upcast_softmax=True,
+                    _from_deprecated_attn_block=True,
                 )
             )
 
@@ -2144,6 +2150,7 @@ def __init__(
                 residual_connection=True,
                 bias=True,
                 upcast_softmax=True,
+                _from_deprecated_attn_block=True,
             )
         )