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

Code before:

    def test_gradient_checkpointing_is_applied(self):
        expected_set = {"FluxTransformer2DModel"}
        super().test_gradient_checkpointing_is_applied(expected_set=expected_set)

    # The test exists for cases like
    # https://github.com/huggingface/diffusers/issues/11874
    def test_lora_exclude_modules(self):
        from peft import LoraConfig, get_peft_model_state_dict, inject_adapter_in_model, set_peft_model_state_dict

        lora_rank = 4
        target_module = "single_transformer_blocks.0.proj_out"
        adapter_name = "foo"
        init_dict, _ = self.prepare_init_args_and_inputs_for_common()
        model = self.model_class(**init_dict).to(torch_device)

        state_dict = model.state_dict()
        target_mod_shape = state_dict[f"{target_module}.weight"].shape
        lora_state_dict = {
            f"{target_module}.lora_A.weight": torch.ones(lora_rank, target_mod_shape[1]) * 22,
            f"{target_module}.lora_B.weight": torch.ones(target_mod_shape[0], lora_rank) * 33,
        }
        config = LoraConfig(
            r=lora_rank, target_modules=["single_transformer_blocks.0.proj_out"], exclude_modules=["proj_out"]
        )
        inject_adapter_in_model(config, model, adapter_name=adapter_name, state_dict=lora_state_dict)
        set_peft_model_state_dict(model, lora_state_dict, adapter_name)
        retrieved_lora_state_dict = get_peft_model_state_dict(model, adapter_name=adapter_name)
        assert len(retrieved_lora_state_dict) == len(lora_state_dict)
        assert (retrieved_lora_state_dict["single_transformer_blocks.0.proj_out.lora_A.weight"] == 22).all()
        assert (retrieved_lora_state_dict["single_transformer_blocks.0.proj_out.lora_B.weight"] == 33).all()


class FluxTransformerCompileTests(TorchCompileTesterMixin, unittest.TestCase):
    model_class = FluxTransformer2DModel
    different_shapes_for_compilation = [(4, 4), (4, 8), (8, 8)]

    def prepare_init_args_and_inputs_for_common(self):
        return FluxTransformerTests().prepare_init_args_and_inputs_for_common()

    def prepare_dummy_input(self, height, width):
        return FluxTransformerTests().prepare_dummy_input(height=height, width=width)


class FluxTransformerLoRAHotSwapTests(LoraHotSwappingForModelTesterMixin, unittest.TestCase):
    model_class = FluxTransformer2DModel
    different_shapes_for_compilation = [(4, 4), (4, 8), (8, 8)]

    def prepare_init_args_and_inputs_for_common(self):
        return FluxTransformerTests().prepare_init_args_and_inputs_for_common()


Review comment (suggestion):
Passing `exclude_modules` should no longer be necessary (or even passing `target_modules`, for that matter).

Code after:

    def test_gradient_checkpointing_is_applied(self):
        expected_set = {"FluxTransformer2DModel"}
        super().test_gradient_checkpointing_is_applied(expected_set=expected_set)

    # The test exists for cases like
    # https://github.com/huggingface/diffusers/issues/11874
    @unittest.skipIf(not is_peft_available(), "Only with PEFT")
    def test_lora_exclude_modules(self):
        from peft import LoraConfig, get_peft_model_state_dict, inject_adapter_in_model, set_peft_model_state_dict

        lora_rank = 4
        target_module = "single_transformer_blocks.0.proj_out"
        adapter_name = "foo"
        init_dict, _ = self.prepare_init_args_and_inputs_for_common()
        model = self.model_class(**init_dict).to(torch_device)

        state_dict = model.state_dict()
        target_mod_shape = state_dict[f"{target_module}.weight"].shape
        lora_state_dict = {
            f"{target_module}.lora_A.weight": torch.ones(lora_rank, target_mod_shape[1]) * 22,
            f"{target_module}.lora_B.weight": torch.ones(target_mod_shape[0], lora_rank) * 33,
        }
        # Passing exclude_modules should no longer be necessary (or even passing target_modules, for that matter).
        config = LoraConfig(
            r=lora_rank, target_modules=["single_transformer_blocks.0.proj_out"], exclude_modules=["proj_out"]
        )
        inject_adapter_in_model(config, model, adapter_name=adapter_name, state_dict=lora_state_dict)
        set_peft_model_state_dict(model, lora_state_dict, adapter_name)
        retrieved_lora_state_dict = get_peft_model_state_dict(model, adapter_name=adapter_name)
        assert len(retrieved_lora_state_dict) == len(lora_state_dict)
        assert (retrieved_lora_state_dict["single_transformer_blocks.0.proj_out.lora_A.weight"] == 22).all()
        assert (retrieved_lora_state_dict["single_transformer_blocks.0.proj_out.lora_B.weight"] == 33).all()


class FluxTransformerCompileTests(TorchCompileTesterMixin, unittest.TestCase):
    model_class = FluxTransformer2DModel
    different_shapes_for_compilation = [(4, 4), (4, 8), (8, 8)]

    def prepare_init_args_and_inputs_for_common(self):
        return FluxTransformerTests().prepare_init_args_and_inputs_for_common()

    def prepare_dummy_input(self, height, width):
        return FluxTransformerTests().prepare_dummy_input(height=height, width=width)


class FluxTransformerLoRAHotSwapTests(LoraHotSwappingForModelTesterMixin, unittest.TestCase):
    model_class = FluxTransformer2DModel
    different_shapes_for_compilation = [(4, 4), (4, 8), (8, 8)]

    def prepare_init_args_and_inputs_for_common(self):