More docs and small refactors

guiyrt · guiyrt · commit 43d2e77de148 · 2024-12-17T00:43:31.000Z
diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml
@@ -234,6 +234,8 @@
       title: Textual Inversion
     - local: api/loaders/unet
       title: UNet
+    - local: api/loaders/transformers_sd3
+      title: SD3Transformer2D
     - local: api/loaders/peft
       title: PEFT
     title: Loaders
diff --git a/docs/source/en/api/attnprocessor.md b/docs/source/en/api/attnprocessor.md
@@ -53,5 +53,5 @@ An attention processor is a class for applying different types of attention mech
 ## AttnProcessorNPU
 [[autodoc]] models.attention_processor.AttnProcessorNPU
 
-## IPAdapterJointAttnProcessor2_0
-[[autodoc]] models.attention_processor.IPAdapterJointAttnProcessor2_0
+## SD3IPAdapterJointAttnProcessor2_0
+[[autodoc]] models.attention_processor.SD3IPAdapterJointAttnProcessor2_0
diff --git a/docs/source/en/api/loaders/transformers_sd3.md b/docs/source/en/api/loaders/transformers_sd3.md
@@ -0,0 +1,29 @@
+<!--Copyright 2024 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+-->
+
+# SD3Transformer2D
+
+This class is useful when *only* loading weights into a [`SD3Transformer2DModel`]. If you need to load weights into the text encoder or a text encoder and SD3Transformer2DModel, check [`SD3LoraLoaderMixin`](lora#diffusers.loaders.SD3LoraLoaderMixin) class instead.
+
+The [`SD3Transformer2DLoadersMixin`] class currently only loads IP-Adapter weights, but will be used in the future to save weights and load LoRAs.
+
+<Tip>
+
+To learn more about how to load LoRA weights, see the [LoRA](../../using-diffusers/loading_adapters#lora) loading guide.
+
+</Tip>
+
+## SD3Transformer2DLoadersMixin
+
+[[autodoc]] loaders.transformers_sd3.SD3Transformer2DLoadersMixin
+    - all
+    - _load_ip_adapter_weights
diff --git a/src/diffusers/loaders/__init__.py b/src/diffusers/loaders/__init__.py
@@ -56,6 +56,7 @@ def text_encoder_attn_modules(text_encoder):
 if is_torch_available():
     _import_structure["single_file_model"] = ["FromOriginalModelMixin"]
 
+    _import_structure["transformers_sd3"] = ["SD3Transformer2DLoadersMixin"]
     _import_structure["unet"] = ["UNet2DConditionLoadersMixin"]
     _import_structure["utils"] = ["AttnProcsLayers"]
     if is_transformers_available():
@@ -82,6 +83,7 @@ def text_encoder_attn_modules(text_encoder):
 if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
     if is_torch_available():
         from .single_file_model import FromOriginalModelMixin
+        from .transformers_sd3 import SD3Transformer2DLoadersMixin
         from .unet import UNet2DConditionLoadersMixin
         from .utils import AttnProcsLayers
 
diff --git a/src/diffusers/loaders/ip_adapter.py b/src/diffusers/loaders/ip_adapter.py
@@ -40,9 +40,9 @@
     AttnProcessor2_0,
     IPAdapterAttnProcessor,
     IPAdapterAttnProcessor2_0,
-    IPAdapterJointAttnProcessor2_0,
     IPAdapterXFormersAttnProcessor,
     JointAttnProcessor2_0,
+    SD3IPAdapterJointAttnProcessor2_0,
 )
 
 
@@ -369,7 +369,7 @@ def is_ip_adapter_active(self) -> bool:
         scales = [
             attn_proc.scale
             for attn_proc in self.transformer.attn_processors.values()
-            if isinstance(attn_proc, IPAdapterJointAttnProcessor2_0)
+            if isinstance(attn_proc, SD3IPAdapterJointAttnProcessor2_0)
         ]
 
         return len(scales) > 0 and any(scale > 0 for scale in scales)
@@ -379,7 +379,7 @@ def load_ip_adapter(
         self,
         pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
         subfolder: str,
-        weight_name: str,
+        weight_name: str = "ip-adapter.safetensors",
         image_encoder_folder: Optional[str] = "image_encoder",
         **kwargs,
     ) -> None:
@@ -396,7 +396,7 @@ def load_ip_adapter(
             subfolder (`str`):
                 The subfolder location of a model file within a larger model repository on the Hub or locally. If a
                 list is passed, it should have the same length as `weight_name`.
-            weight_name (`str`):
+            weight_name (`str`, defaults to "ip-adapter.safetensors"):
                 The name of the weight file to load. If a list is passed, it should have the same length as
                 `subfolder`.
             image_encoder_folder (`str`, *optional*, defaults to `image_encoder`):
@@ -547,7 +547,7 @@ def set_ip_adapter_scale(self, scale: float) -> None:
 
         """
         for attn_processor in self.transformer.attn_processors.values():
-            if isinstance(attn_processor, IPAdapterJointAttnProcessor2_0):
+            if isinstance(attn_processor, SD3IPAdapterJointAttnProcessor2_0):
                 attn_processor.scale = scale
 
     def unload_ip_adapter(self) -> None:
@@ -577,7 +577,9 @@ def unload_ip_adapter(self) -> None:
 
         # Restore original attention processors layers
         attn_procs = {
-            name: (JointAttnProcessor2_0() if isinstance(value, IPAdapterJointAttnProcessor2_0) else value.__class__())
+            name: (
+                JointAttnProcessor2_0() if isinstance(value, SD3IPAdapterJointAttnProcessor2_0) else value.__class__()
+            )
             for name, value in self.transformer.attn_processors.items()
         }
         self.transformer.set_attn_processor(attn_procs)
diff --git a/src/diffusers/loaders/transformers_sd3.py b/src/diffusers/loaders/transformers_sd3.py
@@ -0,0 +1,89 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Dict
+
+from ..models.attention_processor import SD3IPAdapterJointAttnProcessor2_0
+from ..models.embeddings import IPAdapterTimeImageProjection
+from ..models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT, load_model_dict_into_meta
+
+
+class SD3Transformer2DLoadersMixin:
+    """Load IP-Adapters and LoRA layers into a `[SD3Transformer2DModel]`."""
+
+    def _load_ip_adapter_weights(self, state_dict: Dict, low_cpu_mem_usage: bool = _LOW_CPU_MEM_USAGE_DEFAULT) -> None:
+        """Sets IP-Adapter attention processors, image projection, and loads state_dict.
+
+        Args:
+            state_dict (`Dict`):
+                State dict with keys "ip_adapter", which contains parameters for attention processors, and
+                "image_proj", which contains parameters for image projection net.
+            low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`):
+                Speed up model loading only loading the pretrained weights and not initializing the weights. This also
+                tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model.
+                Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this
+                argument to `True` will raise an error.
+        """
+        # IP-Adapter cross attention parameters
+        hidden_size = self.config.attention_head_dim * self.config.num_attention_heads
+        ip_hidden_states_dim = self.config.attention_head_dim * self.config.num_attention_heads
+        timesteps_emb_dim = state_dict["ip_adapter"]["0.norm_ip.linear.weight"].shape[1]
+
+        # Dict where key is transformer layer index, value is attention processor's state dict
+        # ip_adapter state dict keys example: "0.norm_ip.linear.weight"
+        layer_state_dict = {idx: {} for idx in range(len(self.attn_processors))}
+        for key, weights in state_dict["ip_adapter"].items():
+            idx, name = key.split(".", maxsplit=1)
+            layer_state_dict[int(idx)][name] = weights
+
+        # Create IP-Adapter attention processor
+        attn_procs = {}
+        for idx, name in enumerate(self.attn_processors.keys()):
+            attn_procs[name] = SD3IPAdapterJointAttnProcessor2_0(
+                hidden_size=hidden_size,
+                ip_hidden_states_dim=ip_hidden_states_dim,
+                head_dim=self.config.attention_head_dim,
+                timesteps_emb_dim=timesteps_emb_dim,
+            ).to(self.device, dtype=self.dtype)
+
+            if not low_cpu_mem_usage:
+                attn_procs[name].load_state_dict(layer_state_dict[idx], strict=True)
+            else:
+                load_model_dict_into_meta(
+                    attn_procs[name], layer_state_dict[idx], device=self.device, dtype=self.dtype
+                )
+
+        self.set_attn_processor(attn_procs)
+
+        # Image projetion parameters
+        embed_dim = state_dict["image_proj"]["proj_in.weight"].shape[1]
+        output_dim = state_dict["image_proj"]["proj_out.weight"].shape[0]
+        hidden_dim = state_dict["image_proj"]["proj_in.weight"].shape[0]
+        heads = state_dict["image_proj"]["layers.0.attn.to_q.weight"].shape[0] // 64
+        num_queries = state_dict["image_proj"]["latents"].shape[1]
+        timestep_in_dim = state_dict["image_proj"]["time_embedding.linear_1.weight"].shape[1]
+
+        # Image projection
+        self.image_proj = IPAdapterTimeImageProjection(
+            embed_dim=embed_dim,
+            output_dim=output_dim,
+            hidden_dim=hidden_dim,
+            heads=heads,
+            num_queries=num_queries,
+            timestep_in_dim=timestep_in_dim,
+        ).to(device=self.device, dtype=self.dtype)
+
+        if not low_cpu_mem_usage:
+            self.image_proj.load_state_dict(state_dict, strict=True)
+        else:
+            load_model_dict_into_meta(self.image_proj, state_dict, device=self.device, dtype=self.dtype)
diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py
@@ -5160,7 +5160,7 @@ def __call__(
         return hidden_states
 
 
-class IPAdapterJointAttnProcessor2_0(torch.nn.Module):
+class SD3IPAdapterJointAttnProcessor2_0(torch.nn.Module):
     """
     Attention processor for IP-Adapter used typically in processing the SD3-like self-attention projections, with
     additional image-based information and timestep embeddings.
@@ -5844,7 +5844,7 @@ def __call__(
     IPAdapterAttnProcessor,
     IPAdapterAttnProcessor2_0,
     IPAdapterXFormersAttnProcessor,
-    IPAdapterJointAttnProcessor2_0,
+    SD3IPAdapterJointAttnProcessor2_0,
     PAGIdentitySelfAttnProcessor2_0,
     PAGCFGIdentitySelfAttnProcessor2_0,
     LoRAAttnProcessor,
diff --git a/src/diffusers/models/transformers/transformer_sd3.py b/src/diffusers/models/transformers/transformer_sd3.py
@@ -19,19 +19,19 @@
 
 from ...configuration_utils import ConfigMixin, register_to_config
 from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
+from ...loaders.transformers_sd3 import SD3Transformer2DLoadersMixin
 from ...models.attention import FeedForward, JointTransformerBlock
 from ...models.attention_processor import (
     Attention,
     AttentionProcessor,
     FusedJointAttnProcessor2_0,
-    IPAdapterJointAttnProcessor2_0,
     JointAttnProcessor2_0,
 )
-from ...models.modeling_utils import ModelMixin, load_model_dict_into_meta
+from ...models.modeling_utils import ModelMixin
 from ...models.normalization import AdaLayerNormContinuous, AdaLayerNormZero
 from ...utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
 from ...utils.torch_utils import maybe_allow_in_graph
-from ..embeddings import CombinedTimestepTextProjEmbeddings, IPAdapterTimeImageProjection, PatchEmbed
+from ..embeddings import CombinedTimestepTextProjEmbeddings, PatchEmbed
 from ..modeling_outputs import Transformer2DModelOutput
 
 
@@ -104,7 +104,9 @@ def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor):
         return hidden_states
 
 
-class SD3Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
+class SD3Transformer2DModel(
+    ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, SD3Transformer2DLoadersMixin
+):
     """
     The Transformer model introduced in Stable Diffusion 3.
 
@@ -331,89 +333,6 @@ def _set_gradient_checkpointing(self, module, value=False):
         if hasattr(module, "gradient_checkpointing"):
             module.gradient_checkpointing = value
 
-    def _load_ip_adapter_weights(self, state_dict: Dict, low_cpu_mem_usage: bool) -> None:
-        """Sets IP-Adapter attention processors, image projection, and loads state_dict.
-
-        Args:
-            state_dict (`Dict`):
-                PyTorch state dict with keys "ip_adapter", which contains parameters for attention processors, and
-                "image_proj", which contains parameters for image projection net.
-            low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`):
-                Speed up model loading only loading the pretrained weights and not initializing the weights. This also
-                tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model.
-                Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this
-                argument to `True` will raise an error.
-        """
-        # IP-Adapter cross attention parameters
-        hidden_size = self.config.attention_head_dim * self.config.num_attention_heads
-        ip_hidden_states_dim = self.config.attention_head_dim * self.config.num_attention_heads
-        timesteps_emb_dim = state_dict["ip_adapter"]["0.norm_ip.linear.weight"].shape[1]
-
-        # Dict where key is transformer layer index, value is attention processor's state dict
-        # ip_adapter state dict keys example: "0.norm_ip.linear.weight"
-        layer_state_dict = {idx: {} for idx in range(len(self.attn_processors))}
-        for key, weights in state_dict["ip_adapter"].items():
-            idx, name = key.split(".", maxsplit=1)
-            layer_state_dict[int(idx)][name] = weights
-
-        # Create IP-Adapter attention processor
-        attn_procs = {}
-        for idx, name in enumerate(self.attn_processors.keys()):
-            attn_procs[name] = IPAdapterJointAttnProcessor2_0(
-                hidden_size=hidden_size,
-                ip_hidden_states_dim=ip_hidden_states_dim,
-                head_dim=self.config.attention_head_dim,
-                timesteps_emb_dim=timesteps_emb_dim,
-            ).to(self.device, dtype=self.dtype)
-
-            if not low_cpu_mem_usage:
-                attn_procs[name].load_state_dict(layer_state_dict[idx], strict=True)
-            else:
-                load_model_dict_into_meta(
-                    attn_procs[name], layer_state_dict[idx], device=self.device, dtype=self.dtype
-                )
-
-        self.set_attn_processor(attn_procs)
-
-        # Convert image_proj state dict to diffusers
-        image_proj_state_dict = {}
-        for key, value in state_dict["image_proj"].items():
-            if key.startswith("layers."):
-                idx = key.split(".")[1]
-                key = key.replace(f"layers.{idx}.0.norm1", f"layers.{idx}.ln0")
-                key = key.replace(f"layers.{idx}.0.norm2", f"layers.{idx}.ln1")
-                key = key.replace(f"layers.{idx}.0.to_q", f"layers.{idx}.attn.to_q")
-                key = key.replace(f"layers.{idx}.0.to_kv", f"layers.{idx}.attn.to_kv")
-                key = key.replace(f"layers.{idx}.0.to_out", f"layers.{idx}.attn.to_out.0")
-                key = key.replace(f"layers.{idx}.1.0", f"layers.{idx}.adaln_norm")
-                key = key.replace(f"layers.{idx}.1.1", f"layers.{idx}.ff.net.0.proj")
-                key = key.replace(f"layers.{idx}.1.3", f"layers.{idx}.ff.net.2")
-                key = key.replace(f"layers.{idx}.2.1", f"layers.{idx}.adaln_proj")
-            image_proj_state_dict[key] = value
-
-        # Image projetion parameters
-        embed_dim = image_proj_state_dict["proj_in.weight"].shape[1]
-        output_dim = image_proj_state_dict["proj_out.weight"].shape[0]
-        hidden_dim = image_proj_state_dict["proj_in.weight"].shape[0]
-        heads = image_proj_state_dict["layers.0.attn.to_q.weight"].shape[0] // 64
-        num_queries = image_proj_state_dict["latents"].shape[1]
-        timestep_in_dim = image_proj_state_dict["time_embedding.linear_1.weight"].shape[1]
-
-        # Image projection
-        self.image_proj = IPAdapterTimeImageProjection(
-            embed_dim=embed_dim,
-            output_dim=output_dim,
-            hidden_dim=hidden_dim,
-            heads=heads,
-            num_queries=num_queries,
-            timestep_in_dim=timestep_in_dim,
-        ).to(device=self.device, dtype=self.dtype)
-
-        if not low_cpu_mem_usage:
-            self.image_proj.load_state_dict(image_proj_state_dict, strict=True)
-        else:
-            load_model_dict_into_meta(self.image_proj, image_proj_state_dict, device=self.device, dtype=self.dtype)
-
     def forward(
         self,
         hidden_states: torch.FloatTensor,
