Add comments, remove qkv fusion

Author: Warlord-K

src/diffusers/loaders/__init__.py

@@ -89,10 +89,10 @@ def text_encoder_attn_modules(text_encoder):
                 AmusedLoraLoaderMixin,
                 CogVideoXLoraLoaderMixin,
                 FluxLoraLoaderMixin,
+                AuraFlowLoraLoaderMixin,
                 LoraLoaderMixin,
                 Mochi1LoraLoaderMixin,
                 SD3LoraLoaderMixin,
-                AuraFlowLoraLoaderMixin,
                 StableDiffusionLoraLoaderMixin,
                 StableDiffusionXLLoraLoaderMixin,
             )

src/diffusers/loaders/lora_pipeline.py

@@ -1649,9 +1649,9 @@ class AuraFlowLoraLoaderMixin(LoraBaseMixin):
 
     _lora_loadable_modules = ["transformer"]
     transformer_name = TRANSFORMER_NAME
-    text_encoder_name = TEXT_ENCODER_NAME
 
     @classmethod
+    # Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.lora_state_dict
     @validate_hf_hub_args
     def lora_state_dict(
         cls,
@@ -1742,6 +1742,7 @@ def lora_state_dict(
 
         return state_dict
 
+    # Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.load_lora_weights
     def load_lora_weights(
         self, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], adapter_name=None, **kwargs
     ):
@@ -1788,6 +1789,7 @@ def load_lora_weights(
 
 
     @classmethod
+    # Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.load_lora_into_transformer
     def load_lora_into_transformer(cls, state_dict, transformer, adapter_name=None, _pipeline=None):
         """
         This will load the LoRA layers specified in `state_dict` into `transformer`.
@@ -1866,6 +1868,7 @@ def load_lora_into_transformer(cls, state_dict, transformer, adapter_name=None,
             # Unsafe code />
 
     @classmethod
+    # Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.save_lora_weights
     def save_lora_weights(
         cls,
         save_directory: Union[str, os.PathLike],
@@ -1913,6 +1916,7 @@ def save_lora_weights(
             safe_serialization=safe_serialization,
         )
 
+    # Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.fuse_lora
     def fuse_lora(
         self,
         components: List[str] = ["transformer"],
@@ -1956,6 +1960,7 @@ def fuse_lora(
             components=components, lora_scale=lora_scale, safe_fusing=safe_fusing, adapter_names=adapter_names
         )
 
+    # Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.unfuse_lora
     def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
         r"""
         Reverses the effect of

src/diffusers/models/transformers/auraflow_transformer_2d.py

@@ -20,7 +20,8 @@
 import torch.nn.functional as F
 
 from ...configuration_utils import ConfigMixin, register_to_config
-from ...utils import is_torch_version, logging
+from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
+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 ..attention_processor import (
     Attention,
@@ -32,8 +33,6 @@
 from ..modeling_outputs import Transformer2DModelOutput
 from ..modeling_utils import ModelMixin
 from ..normalization import AdaLayerNormZero, FP32LayerNorm
-from ...utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
-from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
 
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
@@ -344,8 +343,8 @@ def __init__(
 
         self.gradient_checkpointing = False
 
-    @property
     # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
+    @property
     def attn_processors(self) -> Dict[str, AttentionProcessor]:
         r"""
         Returns:
@@ -453,7 +452,7 @@ def forward(
         hidden_states: torch.FloatTensor,
         encoder_hidden_states: torch.FloatTensor = None,
         timestep: torch.LongTensor = None,
-        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
         return_dict: bool = True,
     ) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
         height, width = hidden_states.shape[-2:]
@@ -466,18 +465,18 @@ def forward(
         encoder_hidden_states = torch.cat(
             [self.register_tokens.repeat(encoder_hidden_states.size(0), 1, 1), encoder_hidden_states], dim=1
         )
-        if joint_attention_kwargs is not None:
-            joint_attention_kwargs = joint_attention_kwargs.copy()
-            lora_scale = joint_attention_kwargs.pop("scale", 1.0)
+        if attention_kwargs is not None:
+            attention_kwargs = attention_kwargs.copy()
+            lora_scale = attention_kwargs.pop("scale", 1.0)
         else:
             lora_scale = 1.0
         if USE_PEFT_BACKEND:
             # weight the lora layers by setting `lora_scale` for each PEFT layer
             scale_lora_layers(self, lora_scale)
         else:
-            if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
+            if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
                 logger.warning(
-                    "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
+                    "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
                 )
         # MMDiT blocks.
         for index_block, block in enumerate(self.joint_transformer_blocks):

src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py

@@ -18,13 +18,13 @@
 from transformers import T5Tokenizer, UMT5EncoderModel
 
 from ...image_processor import VaeImageProcessor
+from ...loaders import AuraFlowLoraLoaderMixin
 from ...models import AuraFlowTransformer2DModel, AutoencoderKL
 from ...models.attention_processor import AttnProcessor2_0, FusedAttnProcessor2_0, XFormersAttnProcessor
 from ...schedulers import FlowMatchEulerDiscreteScheduler
 from ...utils import logging, replace_example_docstring
 from ...utils.torch_utils import randn_tensor
 from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput
-from ...loaders import AuraFlowLoraLoaderMixin
 
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name