Repository: huggingface/diffusers
File: tests/models/test_modeling_common.py (Python)

Code before:

            # load the first adapter
            torch.manual_seed(0)
            init_dict, _ = self.prepare_init_args_and_inputs_for_common()
            model = self.model_class(**init_dict).to(torch_device)

            if do_compile or (rank0 != rank1):
                # no need to prepare if the model is not compiled or if the ranks are identical
                model.enable_lora_hotswap(target_rank=max_rank)

            file_name0 = os.path.join(os.path.join(tmp_dirname, "0"), "pytorch_lora_weights.safetensors")
            file_name1 = os.path.join(os.path.join(tmp_dirname, "1"), "pytorch_lora_weights.safetensors")
            model.load_lora_adapter(file_name0, safe_serialization=True, adapter_name="adapter0", prefix=None)

            if do_compile:
                model = torch.compile(model, mode="reduce-overhead", dynamic=different_resolutions is not None)

            with torch.inference_mode():
                # additionally check if dynamic compilation works.
                if different_resolutions is not None:
                    for height, width in self.different_shapes_for_compilation:
                        new_inputs_dict = self.prepare_dummy_input(height=height, width=width)
                        _ = model(**new_inputs_dict)
                else:
                    output0_after = model(**inputs_dict)["sample"]
                    assert torch.allclose(output0_before, output0_after, atol=tol, rtol=tol)

            # hotswap the 2nd adapter
            model.load_lora_adapter(file_name1, adapter_name="adapter0", hotswap=True, prefix=None)

            # we need to call forward to potentially trigger recompilation
            with torch.inference_mode():
                if different_resolutions is not None:
                    for height, width in self.different_shapes_for_compilation:
                        new_inputs_dict = self.prepare_dummy_input(height=height, width=width)
                        _ = model(**new_inputs_dict)
                else:
                    output1_after = model(**inputs_dict)["sample"]
                    assert torch.allclose(output1_before, output1_after, atol=tol, rtol=tol)

            # check error when not passing valid adapter name
            name = "does-not-exist"
            msg = f"Trying to hotswap LoRA adapter '{name}' but there is no existing adapter by that name"
            with self.assertRaisesRegex(ValueError, msg):
                model.load_lora_adapter(file_name1, adapter_name=name, hotswap=True, prefix=None)

    @parameterized.expand([(11, 11), (7, 13), (13, 7)])  # important to test small to large and vice versa
    def test_hotswapping_model(self, rank0, rank1):
        self.check_model_hotswap(
            do_compile=False, rank0=rank0, rank1=rank1, target_modules0=["to_q", "to_k", "to_v", "to_out.0"]
        )

Review comment (suggestion):
We don't need assertions when testing for different shapes as that would make the test unnecessarily complicated. We just check for recompilations.

Code after:
            # load the first adapter
            torch.manual_seed(0)
            init_dict, _ = self.prepare_init_args_and_inputs_for_common()
            model = self.model_class(**init_dict).to(torch_device)

            if do_compile or (rank0 != rank1):
                # no need to prepare if the model is not compiled or if the ranks are identical
                model.enable_lora_hotswap(target_rank=max_rank)

            file_name0 = os.path.join(os.path.join(tmp_dirname, "0"), "pytorch_lora_weights.safetensors")
            file_name1 = os.path.join(os.path.join(tmp_dirname, "1"), "pytorch_lora_weights.safetensors")
            model.load_lora_adapter(file_name0, safe_serialization=True, adapter_name="adapter0", prefix=None)

            if do_compile:
                model = torch.compile(model, mode="reduce-overhead", dynamic=different_shapes is not None)

            with torch.inference_mode():
                # additionally check if dynamic compilation works.
                if different_shapes is not None:
                    for height, width in different_shapes:
                        new_inputs_dict = self.prepare_dummy_input(height=height, width=width)
                        _ = model(**new_inputs_dict)
                else:
                    output0_after = model(**inputs_dict)["sample"]
                    assert torch.allclose(output0_before, output0_after, atol=tol, rtol=tol)

            # hotswap the 2nd adapter
            model.load_lora_adapter(file_name1, adapter_name="adapter0", hotswap=True, prefix=None)

            # we need to call forward to potentially trigger recompilation
            with torch.inference_mode():
                if different_shapes is not None:
                    for height, width in different_shapes:
                        new_inputs_dict = self.prepare_dummy_input(height=height, width=width)
                        _ = model(**new_inputs_dict)
                else:
                    output1_after = model(**inputs_dict)["sample"]
                    assert torch.allclose(output1_before, output1_after, atol=tol, rtol=tol)

            # check error when not passing valid adapter name
            name = "does-not-exist"
            msg = f"Trying to hotswap LoRA adapter '{name}' but there is no existing adapter by that name"
            with self.assertRaisesRegex(ValueError, msg):
                model.load_lora_adapter(file_name1, adapter_name=name, hotswap=True, prefix=None)

    @parameterized.expand([(11, 11), (7, 13), (13, 7)])  # important to test small to large and vice versa
    def test_hotswapping_model(self, rank0, rank1):
        self.check_model_hotswap(
            do_compile=False, rank0=rank0, rank1=rank1, target_modules0=["to_q", "to_k", "to_v", "to_out.0"]
        )