Fix AuraFlow LoRA tests by applying to the right denoiser layers.

Co-authored-by: AstraliteHeart <81396681+AstraliteHeart@users.noreply.github.com>

Author: hameerabbasi

tests/lora/test_lora_layers_auraflow.py

@@ -74,6 +74,7 @@ class AuraFlowLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests):
     tokenizer_cls, tokenizer_id = AutoTokenizer, "hf-internal-testing/tiny-random-t5"
     text_encoder_cls, text_encoder_id = UMT5EncoderModel, "hf-internal-testing/tiny-random-umt5"
     text_encoder_target_modules = ["q", "k", "v", "o"]
+    denoiser_target_modules = ["to_q", "to_k", "to_v", "to_out.0", "linear_1"]
 
     @property
     def output_shape(self):

tests/lora/utils.py

@@ -104,6 +104,7 @@ class PeftLoraLoaderMixinTests:
     vae_kwargs = None
 
     text_encoder_target_modules = ["q_proj", "k_proj", "v_proj", "out_proj"]
+    denoiser_target_modules = ["to_q", "to_k", "to_v", "to_out.0"]
 
     def get_dummy_components(self, scheduler_cls=None, use_dora=False):
         if self.unet_kwargs and self.transformer_kwargs:
@@ -157,7 +158,7 @@ def get_dummy_components(self, scheduler_cls=None, use_dora=False):
         denoiser_lora_config = LoraConfig(
             r=rank,
             lora_alpha=rank,
-            target_modules=["to_q", "to_k", "to_v", "to_out.0"],
+            target_modules=self.denoiser_target_modules,
             init_lora_weights=False,
             use_dora=use_dora,
         )
@@ -2040,7 +2041,7 @@ def test_lora_B_bias(self):
         bias_values = {}
         denoiser = pipe.unet if self.unet_kwargs is not None else pipe.transformer
         for name, module in denoiser.named_modules():
-            if any(k in name for k in ["to_q", "to_k", "to_v", "to_out.0"]):
+            if any(k in name for k in self.denoiser_target_modules):
                 if module.bias is not None:
                     bias_values[name] = module.bias.data.clone()