tests: add lora tester mixin for modular and migrate minimax h3. - #14605
tests: add lora tester mixin for modular and migrate minimax h3.#14605sayakpaul wants to merge 4 commits into
Conversation
| # A modular pipeline takes its inputs from its blocks rather than from a `__call__` signature, and the denoiser | ||
| # block of every LoRA-capable one declares `InputParam.template("attention_kwargs")`. | ||
| if issubclass(pipeline_class, ModularPipeline): | ||
| return "attention_kwargs" |
There was a problem hiding this comment.
I don't like the pipeline-specific condition here but I think it's a fair compromise.
| pass | ||
|
|
||
|
|
||
| class TestMiniMaxH3ModularPipelineLoRA(MiniMaxH3ModularPipelineTesterConfig, ModularLoraTesterMixin): |
There was a problem hiding this comment.
The extra methods are specific to the transformer and transformer_ref related voodoo going on in the pipeline.
| class ModularLoraMemoryTesterMixin(BaseModularPipelineOutputMixin, LoraMemoryTesterMixin): | ||
| """LoRA x offloading tests for modular pipelines: group offloading composed with `load_lora_weights`.""" | ||
|
|
||
| @pytest.mark.skip( |
There was a problem hiding this comment.
We can add testing for this later.
| "transformer_ref": ["to_q", "to_k", "to_v", "to_out.0"], | ||
| } | ||
|
|
||
| def get_dummy_diffusers_lora_state_dict(self, transformer, prefix="transformer", rank=8, adaln_rank=2): |
There was a problem hiding this comment.
In the process, I also optimized the runtime of these tests. Previously, we were building the entire pipeline just to get the transformer from it. We already have all the components available when we're calling this method. So, this method now simply reuses them unlike:
diffusers/tests/lora/test_lora_layers_minimax_h3.py
Lines 104 to 113 in ee9f497
With more and more pipelines being released with Modular and supporting LoRA, it's important for us to ensure we test them as good as we do for standard pipelines that support LoRAs.
This PR, hence, adds a LoRA tester mixin class to our testing suite for Modular Pipelines. The mixin overrides
tests/pipelines/testing_utils/lora.pyand migrates the existing MiniMax H3 LoRA testing suite with it (as the existing one already makes use of Modular).We can add a LoRA tester class for the remaining Modular Pipelines that support it in a follow-up PR.