tighten test

Author: sayakpaul

src/diffusers/loaders/lora_pipeline.py

@@ -2284,7 +2284,6 @@ def unload_lora_weights(self):
             transformer._transformer_norm_layers = None
 
         if getattr(transformer, "_overwritten_params", None) is not None:
-            print(f"{transformer._overwritten_params.keys()=}")
             overwritten_params = transformer._overwritten_params
             module_names = set()
 

tests/lora/test_lora_layers_flux.py

@@ -479,6 +479,10 @@ def test_lora_unload_with_parameter_expanded_shapes(self):
         self.assertTrue(cap_logger.out.startswith("Expanding the nn.Linear input/output features for module"))
 
         control_pipe.unload_lora_weights()
+        self.assertTrue(
+            control_pipe.transformer.config.in_channels == num_channels_without_control,
+            f"Expected {num_channels_without_control} channels in the modified transformer but has {control_pipe.transformer.config.in_channels=}",
+        )
         loaded_pipe = FluxPipeline.from_pipe(control_pipe)
         self.assertTrue(
             loaded_pipe.transformer.config.in_channels == num_channels_without_control,