Fix test_lora_fuse_nan.

Author: hameerabbasi

tests/lora/utils.py

@@ -1563,7 +1563,6 @@ def test_simple_inference_with_text_denoiser_multi_adapter_weighted(self):
                 "output with no lora and output with lora disabled should give same results",
             )
 
-    @skip_mps
     @pytest.mark.xfail(
         condition=torch.device(torch_device).type == "cpu" and is_torch_version(">=", "2.5"),
         reason="Test currently fails on CPU and PyTorch 2.5.1 but not on PyTorch 2.4.1.",
@@ -1595,17 +1594,26 @@ def test_lora_fuse_nan(self):
                     ].weight += float("inf")
                 else:
                     named_modules = [name for name, _ in pipe.transformer.named_modules()]
-                    tower_name = (
-                        "transformer_blocks"
-                        if any(name == "transformer_blocks" for name in named_modules)
-                        else "blocks"
-                    )
-                    transformer_tower = getattr(pipe.transformer, tower_name)
-                    has_attn1 = any("attn1" in name for name in named_modules)
-                    if has_attn1:
-                        transformer_tower[0].attn1.to_q.lora_A["adapter-1"].weight += float("inf")
-                    else:
-                        transformer_tower[0].attn.to_q.lora_A["adapter-1"].weight += float("inf")
+                    possible_tower_names = [
+                        "transformer_blocks",
+                        "blocks",
+                        "joint_transformer_blocks",
+                        "single_transformer_blocks",
+                    ]
+                    filtered_tower_names = [
+                        tower_name for tower_name in possible_tower_names if hasattr(pipe.transformer, tower_name)
+                    ]
+                    if len(filtered_tower_names) == 0:
+                        pytest.xfail(
+                            reason=f"`pipe.transformer` didn't have any of the following attributes: {possible_tower_names}."
+                        )
+                    for tower_name in filtered_tower_names:
+                        transformer_tower = getattr(pipe.transformer, tower_name)
+                        has_attn1 = any("attn1" in name for name in named_modules)
+                        if has_attn1:
+                            transformer_tower[0].attn1.to_q.lora_A["adapter-1"].weight += float("inf")
+                        else:
+                            transformer_tower[0].attn.to_q.lora_A["adapter-1"].weight += float("inf")
 
             # with `safe_fusing=True` we should see an Error
             with self.assertRaises(ValueError):