fixes

Author: sayakpaul

tests/lora/utils.py

@@ -1784,18 +1784,15 @@ def test_missing_keys_warning(self):
         missing_key = [k for k in state_dict if "lora_A" in k][0]
         del state_dict[missing_key]
 
-        logger = (
-            logging.get_logger("diffusers.loaders.unet")
-            if self.unet_kwargs is not None
-            else logging.get_logger("diffusers.loaders.peft")
-        )
+        logger = logging.get_logger("diffusers.loaders.peft")
         logger.setLevel(30)
         with CaptureLogger(logger) as cap_logger:
             pipe.load_lora_weights(state_dict)
 
         # Since the missing key won't contain the adapter name ("default_0").
         # Also strip out the component prefix (such as "unet." from `missing_key`).
         component = list({k.split(".")[0] for k in state_dict})[0]
+        print(f"{cap_logger.out=}")
         self.assertTrue(missing_key.replace(f"{component}.", "") in cap_logger.out.replace("default_0.", ""))
 
     def test_unexpected_keys_warning(self):
@@ -1823,11 +1820,7 @@ def test_unexpected_keys_warning(self):
         unexpected_key = [k for k in state_dict if "lora_A" in k][0] + ".diffusers_cat"
         state_dict[unexpected_key] = torch.tensor(1.0, device=torch_device)
 
-        logger = (
-            logging.get_logger("diffusers.loaders.unet")
-            if self.unet_kwargs is not None
-            else logging.get_logger("diffusers.loaders.peft")
-        )
+        logger = logging.get_logger("diffusers.loaders.peft")
         logger.setLevel(30)
         with CaptureLogger(logger) as cap_logger:
             pipe.load_lora_weights(state_dict)