diff --git a/docs/source/en/api/loaders/lora.md b/docs/source/en/api/loaders/lora.md index 03592a2cd5cf..4722d082822c 100644 --- a/docs/source/en/api/loaders/lora.md +++ b/docs/source/en/api/loaders/lora.md @@ -38,6 +38,7 @@ LoRA is a fast and lightweight training method that inserts and trains a signifi - [`Flux2LoraLoaderMixin`] provides similar functions for [Flux2](https://huggingface.co/docs/diffusers/main/en/api/pipelines/flux2). - [`ErnieImageLoraLoaderMixin`] provides similar functions for [Ernie-Image](https://huggingface.co/docs/diffusers/main/en/api/pipelines/ernie_image). - [`LTX2LoraLoaderMixin`] provides similar functions for [Flux2](https://huggingface.co/docs/diffusers/main/en/api/pipelines/ltx2). +- [`MiniMaxH3LoraLoaderMixin`] provides similar functions for [MiniMax-H3](https://huggingface.co/docs/diffusers/main/en/api/pipelines/minimax_h3). - [`LoraBaseMixin`] provides a base class with several utility methods to fuse, unfuse, unload, LoRAs and more. > [!TIP] @@ -157,6 +158,10 @@ LoRA is a fast and lightweight training method that inserts and trains a signifi [[autodoc]] loaders.lora_pipeline.Krea2LoraLoaderMixin +## MiniMaxH3LoraLoaderMixin + +[[autodoc]] loaders.lora_pipeline.MiniMaxH3LoraLoaderMixin + ## LoraBaseMixin [[autodoc]] loaders.lora_base.LoraBaseMixin diff --git a/docs/source/en/api/pipelines/minimax_h3.md b/docs/source/en/api/pipelines/minimax_h3.md index 5e027004b825..e6e001baf8b4 100644 --- a/docs/source/en/api/pipelines/minimax_h3.md +++ b/docs/source/en/api/pipelines/minimax_h3.md @@ -11,6 +11,12 @@ specific language governing permissions and limitations under the License. --> # MiniMax-H3 +
+ + LoRA + +
+ > [!TIP] > MiniMax-H3 is not part of a diffusers release yet. Install diffusers from the pull request to use it: diff --git a/src/diffusers/loaders/__init__.py b/src/diffusers/loaders/__init__.py index 1c6693bd0c08..828744386453 100644 --- a/src/diffusers/loaders/__init__.py +++ b/src/diffusers/loaders/__init__.py @@ -91,6 +91,7 @@ def text_encoder_attn_modules(text_encoder): "Ideogram4LoraLoaderMixin", "ErnieImageLoraLoaderMixin", "CosmosLoraLoaderMixin", + "MiniMaxH3LoraLoaderMixin", ] _import_structure["textual_inversion"] = ["TextualInversionLoaderMixin"] _import_structure["ip_adapter"] = [ @@ -139,6 +140,7 @@ def text_encoder_attn_modules(text_encoder): LTX2LoraLoaderMixin, LTXVideoLoraLoaderMixin, Lumina2LoraLoaderMixin, + MiniMaxH3LoraLoaderMixin, Mochi1LoraLoaderMixin, QwenImageLoraLoaderMixin, SanaLoraLoaderMixin, diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index 07e3351685e8..20396686a63b 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -3122,3 +3122,134 @@ def _convert_non_diffusers_ace_step_lora_to_diffusers(state_dict): converted_state_dict[new_key] = state_dict.pop(key) return converted_state_dict + + +def _convert_non_diffusers_minimax_h3_lora_to_diffusers(state_dict): + """Convert a non-diffusers MiniMax-H3 LoRA state dict onto `MiniMaxH3Transformer3DModel`'s module names. + + Both known producers train against the original checkpoint's module names — ai-toolkit under a `diffusion_model.` + prefix, the reference `generate.py` / ComfyUI checkpoints under no prefix at all — so the prefix is optional and + the module names are what identifies the format. Handles: + + - `diffusion_model.` prefix removal, and bare `blocks.` / `token_refiner.` / `final_layer.` keys + - `lora_down`/`lora_up` (kohya) -> `lora_A`/`lora_B`, with `.alpha` folded into the weights + - fused `attn.qkv_proj` -> split `to_q`/`to_k`/`to_v`; `attn.out_proj` -> `to_out.0` + - `mlp.fc1` -> `ff.net.0.proj` with its two output halves swapped, `mlp.fc2` -> `ff.net.2` + - `blocks.` -> `transformer_blocks.`, `token_refiner.blocks.` -> `token_refiner.refiner_blocks.`, and the + `final_layer.` / patch / condition / timestep projections onto their diffusers names + + The result is prefixed with `transformer.`, the partition every published H3 LoRA is trained against; + `MiniMaxH3LoraLoaderMixin.load_lora_weights` is what redirects it to `transformer_ref` when asked. + """ + state_dict = {k.removeprefix("diffusion_model."): v for k, v in state_dict.items()} + + is_kohya = any(".lora_down.weight" in k for k in state_dict) + down_suffix = ".lora_down.weight" if is_kohya else ".lora_A.weight" + up_suffix = ".lora_up.weight" if is_kohya else ".lora_B.weight" + + def pull(base): + """Pop the (lora_A, lora_B) pair for a module path with any `.alpha` folded in, or None if absent.""" + down_key = base + down_suffix + if down_key not in state_dict: + return None + down = state_dict.pop(down_key) + up = state_dict.pop(base + up_suffix) + alpha = state_dict.pop(base + ".alpha", None) + if alpha is not None: + # LoRA is scaled by `alpha / rank` in the forward pass; split the factor between down and up. + scale_down, scale_up = alpha.item() / down.shape[0], 1.0 + while scale_down * 2 < scale_up: + scale_down *= 2 + scale_up /= 2 + down, up = down * scale_down, up * scale_up + return down, up + + converted_state_dict = {} + + # The projections outside the block stack. `final_layer.norm` and the `norm1`/`norm2`/`q_norm`/`k_norm` RMSNorms + # carry no LoRA-able Linear, so they have no entry. + standalone_renames = { + "video_patch_proj": "proj_in", + "audio_patch_proj": "audio_proj_in", + "condition_proj": "context_embedder", + "time_embedder.proj_in": "time_embedder.linear_1", + "time_embedder.proj_out": "time_embedder.linear_2", + "final_layer.adaln_proj.linear": "norm_out.linear", + "final_layer.video_out": "proj_out", + "final_layer.audio_out": "audio_proj_out", + } + for source, target in standalone_renames.items(): + pair = pull(source) + if pair is not None: + down, up = pair + converted_state_dict[f"{target}.lora_A.weight"] = down + converted_state_dict[f"{target}.lora_B.weight"] = up + + # The main stack and the text token refiner hold the same block layout, except that a refiner block has no AdaLN + # projection. + block_specs = [ + (r"blocks\.(\d+)\.", "blocks", "transformer_blocks"), + (r"token_refiner\.blocks\.(\d+)\.", "token_refiner.blocks", "token_refiner.refiner_blocks"), + ] + for pattern, source_prefix, target_prefix in block_specs: + num_layers = 0 + for key in state_dict: + match = re.match(pattern, key) + if match: + num_layers = max(num_layers, int(match.group(1)) + 1) + + for i in range(num_layers): + source = f"{source_prefix}.{i}" + target = f"{target_prefix}.{i}" + + # Fused qkv -> split to_q / to_k / to_v (shared down/lora_A, chunk up/lora_B in thirds). Both producers + # consume the fused rows as `[q_all; k_all; v_all]`, so no per-head de-interleave is involved. + qkv = pull(f"{source}.attn.qkv_proj") + if qkv is not None: + down, up = qkv + if up.shape[0] % 3 != 0: + raise ValueError( + f"`{source}.attn.qkv_proj` has {up.shape[0]} output rows, which is not divisible by 3. " + "This is not a fused MiniMax-H3 QKV projection." + ) + up_q, up_k, up_v = torch.chunk(up, 3, dim=0) + for proj, up_proj in (("to_q", up_q), ("to_k", up_k), ("to_v", up_v)): + converted_state_dict[f"{target}.attn.{proj}.lora_A.weight"] = down.clone() + converted_state_dict[f"{target}.attn.{proj}.lora_B.weight"] = up_proj.contiguous() + + # `fc1` stays fused, as diffusers' `SwiGLU` also fuses its two projections, but the reference computes + # `fc2(silu(gate) * value)` from a fused `[gate; value]` while `SwiGLU` computes `value * silu(gate)` from + # a fused `[value; gate]`, so the two halves swap places. `lora_A` is untouched: the swap is a permutation + # of output rows, so it applies to `lora_B` alone. + fc1 = pull(f"{source}.mlp.fc1") + if fc1 is not None: + down, up = fc1 + if up.shape[0] % 2 != 0: + raise ValueError( + f"`{source}.mlp.fc1` has {up.shape[0]} output rows, which is not even. This is not a fused " + "MiniMax-H3 SwiGLU projection." + ) + up_gate, up_value = up.chunk(2, dim=0) + converted_state_dict[f"{target}.ff.net.0.proj.lora_A.weight"] = down + converted_state_dict[f"{target}.ff.net.0.proj.lora_B.weight"] = torch.cat( + [up_value, up_gate], dim=0 + ).contiguous() + + for source_module, target_module in ( + ("attn.out_proj", "attn.to_out.0"), + ("mlp.fc2", "ff.net.2"), + ("adaln_proj.linear", "adaln_proj.linear"), + ): + pair = pull(f"{source}.{source_module}") + if pair is not None: + down, up = pair + converted_state_dict[f"{target}.{target_module}.lora_A.weight"] = down + converted_state_dict[f"{target}.{target_module}.lora_B.weight"] = up + + if len(state_dict) > 0: + raise ValueError( + f"`state_dict` should be empty at this point but has {sorted(state_dict.keys())}. " + "This may be an unsupported MiniMax-H3 LoRA layout." + ) + + return {f"transformer.{k}": v for k, v in converted_state_dict.items()} diff --git a/src/diffusers/loaders/lora_pipeline.py b/src/diffusers/loaders/lora_pipeline.py index 8de23d81528c..b092a5c18f1e 100644 --- a/src/diffusers/loaders/lora_pipeline.py +++ b/src/diffusers/loaders/lora_pipeline.py @@ -21,6 +21,7 @@ from ..utils import ( USE_PEFT_BACKEND, deprecate, + get_peft_kwargs, get_submodule_by_name, is_bitsandbytes_available, is_gguf_available, @@ -56,6 +57,7 @@ _convert_non_diffusers_ltx2_lora_to_diffusers, _convert_non_diffusers_ltxv_lora_to_diffusers, _convert_non_diffusers_lumina2_lora_to_diffusers, + _convert_non_diffusers_minimax_h3_lora_to_diffusers, _convert_non_diffusers_qwen_lora_to_diffusers, _convert_non_diffusers_wan_lora_to_diffusers, _convert_non_diffusers_z_image_lora_to_diffusers, @@ -81,6 +83,8 @@ UNET_NAME = "unet" TRANSFORMER_NAME = "transformer" LTX2_CONNECTOR_NAME = "connectors" +# MiniMax-H3 ships two independently trained DiT partitions in one repository, under two component names. +MINIMAX_H3_TRANSFORMER_REF_NAME = "transformer_ref" _MODULE_NAME_TO_ATTRIBUTE_MAP_FLUX = {"x_embedder": "in_channels"} @@ -7041,6 +7045,329 @@ def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs): super().unfuse_lora(components=components, **kwargs) +class MiniMaxH3LoraLoaderMixin(LoraBaseMixin): + r""" + Load LoRA layers into [`MiniMaxH3Transformer3DModel`]. Specific to [`MiniMaxH3ModularPipeline`]. + + MiniMax-H3 holds two independently trained DiT partitions in one repository — `transformer/` for the `t2va` and + `fl2va` workflows, `transformer_ref/` for `ref2va` — and a workflow loads only its own. The two are separate + checkpoints with nothing tied between them, and their module names are identical, so a LoRA trained against one + loads without error into the other and silently produces garbage. Nothing in a published H3 LoRA records which + partition it was trained against, so the routing is explicit: a converted state dict targets `transformer.`, and + `transformer_ref` is reached either by a `transformer_ref.`-prefixed file (what `save_lora_weights` writes) or by + passing `load_into_transformer_ref=True`. + + Two things to know about third-party H3 LoRAs. LoRAs trained against a *pruned* checkpoint do not load: pruned + releases replace the timestep MLP with a small interpolation table, so their AdaLN projections take an 8-wide + input instead of `time_embed_dim`, and the update cannot be mapped onto the released checkpoint — loading fails + with a size mismatch naming the module. And published H3 LoRAs carry no alpha information while applying as + `W + lora_B @ lora_A`, so mixed-rank files are loaded with `alpha == rank` per module (effective scale exactly + 1.0); their updates are also small enough relative to the base weights that [`~MiniMaxH3LoraLoaderMixin.fuse_lora`] + into bfloat16 discards most of the update — prefer the default unfused path. + """ + + _lora_loadable_modules = ["transformer", "transformer_ref"] + transformer_name = TRANSFORMER_NAME + transformer_ref_name = MINIMAX_H3_TRANSFORMER_REF_NAME + + @classmethod + @validate_hf_hub_args + def lora_state_dict( + cls, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + **kwargs, + ): + r""" + See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details. + """ + cache_dir = kwargs.pop("cache_dir", None) + force_download = kwargs.pop("force_download", False) + proxies = kwargs.pop("proxies", None) + local_files_only = kwargs.pop("local_files_only", None) + token = kwargs.pop("token", None) + revision = kwargs.pop("revision", None) + subfolder = kwargs.pop("subfolder", None) + weight_name = kwargs.pop("weight_name", None) + use_safetensors = kwargs.pop("use_safetensors", None) + return_lora_metadata = kwargs.pop("return_lora_metadata", False) + + allow_pickle = False + if use_safetensors is None: + use_safetensors = True + allow_pickle = True + + user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"} + + state_dict, metadata = _fetch_state_dict( + pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict, + weight_name=weight_name, + use_safetensors=use_safetensors, + local_files_only=local_files_only, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + token=token, + revision=revision, + subfolder=subfolder, + user_agent=user_agent, + allow_pickle=allow_pickle, + ) + + is_dora_scale_present = any("dora_scale" in k for k in state_dict) + if is_dora_scale_present: + warn_msg = "It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new." + logger.warning(warn_msg) + state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k} + + # ai-toolkit writes the original checkpoint's module names under a `diffusion_model.` prefix, while the + # reference `generate.py` / ComfyUI checkpoints carry no prefix at all, so the module names are what + # identifies a non-diffusers file. + is_non_diffusers_format = any( + k.startswith(("diffusion_model.", "blocks.", "token_refiner.", "final_layer.")) for k in state_dict + ) + if is_non_diffusers_format: + state_dict = _convert_non_diffusers_minimax_h3_lora_to_diffusers(state_dict) + + # Every published MiniMax-H3 LoRA is alpha-less, MIXED-RANK (rank 64 on attention and FFN modules, rank 16 + # on the AdaLN projections) and applies as `W + lora_B @ lora_A`, i.e. at an effective scale of 1.0 *per + # module*. `get_peft_kwargs` reads `lora_alpha` off whichever rank it happens to see first and never + # re-derives it, so one of the two rank groups would be silently scaled by `alpha / r`. The `LoraConfig` is + # therefore built here with `alpha == rank` everywhere and passed on as metadata, which `load_lora_adapter` + # uses in place of its own `get_peft_kwargs` inference — that inference recovers everything else (ranks, + # target modules), just not this alpha correction. Keying the synthesis off the absence of alpha information + # rather than off the conversion is deliberate: the same file also circulates pre-converted to diffusers + # keys, and that copy needs the same treatment. + if metadata is None and not any(k.endswith(".alpha") for k in state_dict): + metadata = {} + for prefix in (cls.transformer_name, cls.transformer_ref_name): + component_state_dict = { + k.removeprefix(f"{prefix}."): v for k, v in state_dict.items() if k.startswith(f"{prefix}.") + } + # `^` anchors each pattern to a full module name, as `load_lora_adapter` does for the ranks it derives. + rank = {f"^{k}": v.shape[1] for k, v in component_state_dict.items() if "lora_B" in k and v.ndim > 1} + if not rank: + continue + lora_config_kwargs = get_peft_kwargs( + rank, network_alpha_dict=None, peft_state_dict=component_state_dict, is_unet=False + ) + # The same fix-up `PeftAdapterMixin.load_lora_adapter` applies to SAI control LoRAs. + lora_config_kwargs["lora_alpha"] = lora_config_kwargs["r"] + lora_config_kwargs["alpha_pattern"] = lora_config_kwargs["rank_pattern"] + metadata.update(_pack_dict_with_prefix(lora_config_kwargs, prefix)) + metadata = metadata or None + + out = (state_dict, metadata) if return_lora_metadata else state_dict + return out + + def load_lora_weights( + self, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + adapter_name: str | None = None, + hotswap: bool = False, + load_into_transformer_ref: bool = False, + **kwargs, + ): + """ + Load LoRA layers into `transformer` or, with `load_into_transformer_ref=True`, into `transformer_ref`. See + [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for more details. + + Args: + load_into_transformer_ref (`bool`, defaults to `False`): + Load the `transformer.`-prefixed layers into the `transformer_ref` partition — the one the `ref2va` + workflow denoises with — instead of `transformer`. Only needed when both partitions are loaded: a + pipeline that holds `transformer_ref` alone routes there on its own. + """ + if not USE_PEFT_BACKEND: + raise ValueError("PEFT backend is required for this method.") + + low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA) + + # if a dict is passed, copy it instead of modifying it inplace + if isinstance(pretrained_model_name_or_path_or_dict, dict): + pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy() + + kwargs["return_lora_metadata"] = True + state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) + + is_correct_format = all("lora" in key for key in state_dict.keys()) + if not is_correct_format: + raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + + # A workflow loads only its own partition, so `getattr(..., None)` — never the `hasattr` ternary the + # single-denoiser mixins use — is what tells the two apart. + transformer = getattr(self, self.transformer_name, None) + transformer_ref = getattr(self, self.transformer_ref_name, None) + + transformer_state_dict = {k: v for k, v in state_dict.items() if k.startswith(f"{self.transformer_name}.")} + transformer_ref_state_dict = { + k: v for k, v in state_dict.items() if k.startswith(f"{self.transformer_ref_name}.") + } + + if transformer is None and transformer_ref is None: + logger.warning( + f"No denoiser to load the LoRA into: this pipeline holds neither `{self.transformer_name}` nor " + f"`{self.transformer_ref_name}`. Skipping." + ) + return + + if transformer_state_dict: + # `transformer.`-prefixed layers go to `transformer_ref` when the caller asks for it, and also when + # `transformer_ref` is the only partition present — which is what `workflow="ref2va"` loads. + into_ref = load_into_transformer_ref or transformer is None + if into_ref and transformer_ref is None: + raise ValueError( + f"`load_into_transformer_ref=True` needs a `{self.transformer_ref_name}` component, which this " + 'pipeline does not have. Load it with `workflow="ref2va"`, or drop the argument to load into ' + f"`{self.transformer_name}`." + ) + if not into_ref and transformer_ref is not None and not transformer_ref_state_dict: + logger.warning( + f"Both MiniMax-H3 partitions are loaded and this LoRA does not say which one it was trained " + f"against, so it is going into `{self.transformer_name}` — the partition every published H3 LoRA " + f"so far targets. Pass `load_into_transformer_ref=True` for the `{self.transformer_ref_name}` " + "partition instead." + ) + if into_ref: + # `transformer.`-prefixed layers into the other partition, so the prefix and the target differ. + transformer_ref.load_lora_adapter( + transformer_state_dict, + prefix=self.transformer_name, + network_alphas=None, + adapter_name=adapter_name, + metadata=metadata, + _pipeline=self, + low_cpu_mem_usage=low_cpu_mem_usage, + hotswap=hotswap, + ) + else: + self.load_lora_into_transformer( + transformer_state_dict, + transformer=transformer, + adapter_name=adapter_name, + metadata=metadata, + _pipeline=self, + low_cpu_mem_usage=low_cpu_mem_usage, + hotswap=hotswap, + ) + + if transformer_ref_state_dict: + if transformer_ref is None: + raise ValueError( + f"This LoRA has `{self.transformer_ref_name}.`-prefixed layers but the pipeline does not hold a " + f'`{self.transformer_ref_name}` component. Load it with `workflow="ref2va"`.' + ) + transformer_ref.load_lora_adapter( + transformer_ref_state_dict, + prefix=self.transformer_ref_name, + network_alphas=None, + adapter_name=adapter_name, + metadata=metadata, + _pipeline=self, + low_cpu_mem_usage=low_cpu_mem_usage, + hotswap=hotswap, + ) + + @classmethod + # Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.load_lora_into_transformer with SD3Transformer2DModel->MiniMaxH3Transformer3DModel + def load_lora_into_transformer( + cls, + state_dict, + transformer, + adapter_name=None, + _pipeline=None, + low_cpu_mem_usage=False, + hotswap: bool = False, + metadata=None, + ): + """ + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details. + """ + if low_cpu_mem_usage and is_peft_version("<", "0.13.0"): + raise ValueError( + "`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`." + ) + + # Load the layers corresponding to transformer. + logger.info(f"Loading {cls.transformer_name}.") + transformer.load_lora_adapter( + state_dict, + network_alphas=None, + adapter_name=adapter_name, + metadata=metadata, + _pipeline=_pipeline, + low_cpu_mem_usage=low_cpu_mem_usage, + hotswap=hotswap, + ) + + @classmethod + def save_lora_weights( + cls, + save_directory: str | os.PathLike, + transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, + transformer_ref_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, + is_main_process: bool = True, + weight_name: str = None, + save_function: Callable = None, + safe_serialization: bool = True, + transformer_lora_adapter_metadata: dict | None = None, + transformer_ref_lora_adapter_metadata: dict | None = None, + ): + r""" + Save the LoRA layers of one or both MiniMax-H3 partitions. Which partition a LoRA belongs to is not recoverable + from its keys, so this is the only way to publish an H3 LoRA that records it. See + [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. + """ + lora_layers = {} + lora_metadata = {} + + if transformer_lora_layers: + lora_layers[cls.transformer_name] = transformer_lora_layers + lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata + if transformer_ref_lora_layers: + lora_layers[cls.transformer_ref_name] = transformer_ref_lora_layers + lora_metadata[cls.transformer_ref_name] = transformer_ref_lora_adapter_metadata + + if not lora_layers: + raise ValueError( + "You must pass at least one of `transformer_lora_layers` or `transformer_ref_lora_layers`." + ) + + cls._save_lora_weights( + save_directory=save_directory, + lora_layers=lora_layers, + lora_metadata=lora_metadata, + is_main_process=is_main_process, + weight_name=weight_name, + save_function=save_function, + safe_serialization=safe_serialization, + ) + + def fuse_lora( + self, + components: list[str] = ["transformer", "transformer_ref"], + lora_scale: float = 1.0, + safe_fusing: bool = False, + adapter_names: list[str] | None = None, + **kwargs, + ): + r""" + See [`~loaders.StableDiffusionLoraLoaderMixin.fuse_lora`] for more details. + """ + super().fuse_lora( + components=components, + lora_scale=lora_scale, + safe_fusing=safe_fusing, + adapter_names=adapter_names, + **kwargs, + ) + + def unfuse_lora(self, components: list[str] = ["transformer", "transformer_ref"], **kwargs): + r""" + See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. + """ + super().unfuse_lora(components=components, **kwargs) + + class LoraLoaderMixin(StableDiffusionLoraLoaderMixin): def __init__(self, *args, **kwargs): deprecation_message = "LoraLoaderMixin is deprecated and this will be removed in a future version. Please use `StableDiffusionLoraLoaderMixin`, instead." diff --git a/src/diffusers/modular_pipelines/minimax_h3/modular_pipeline.py b/src/diffusers/modular_pipelines/minimax_h3/modular_pipeline.py index d39c84b7e3b3..b0f75c64af1d 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/minimax_h3/modular_pipeline.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from ...loaders import MiniMaxH3LoraLoaderMixin from ...utils import logging from ..modular_pipeline import ModularPipeline @@ -146,7 +147,7 @@ def audio_latent_num_frames( return int(round(num_frames / fps * latents_per_second)) -class MiniMaxH3ModularPipeline(ModularPipeline): +class MiniMaxH3ModularPipeline(ModularPipeline, MiniMaxH3LoraLoaderMixin): """ A ModularPipeline for joint video + audio generation with MiniMax-H3: the `t2va` (text only) and `fl2va` (first and/or last keyframe) workflows against the `transformer/` checkpoint partition, and the `ref2va` (omni-reference) diff --git a/tests/lora/test_lora_layers_minimax_h3.py b/tests/lora/test_lora_layers_minimax_h3.py new file mode 100644 index 000000000000..d5fae696ff2d --- /dev/null +++ b/tests/lora/test_lora_layers_minimax_h3.py @@ -0,0 +1,284 @@ +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import tempfile + +import pytest +import torch + +from diffusers.modular_pipelines import MiniMaxH3Blocks, MiniMaxH3ModularPipeline +from diffusers.utils import is_peft_available + +from ..testing_utils import require_peft_backend + + +if is_peft_available(): + from peft.utils import get_peft_model_state_dict + + +@require_peft_backend +class TestMiniMaxH3LoraLayers: + """ + The MiniMax-H3 LoRA surface that is specific to this model and its two checkpoint partitions: the conversion of + the two circulating non-diffusers formats (fused projections, original module names, no alpha keys), the + `alpha == rank` metadata synthesis for alpha-less files, and the routing between the `transformer` and + `transformer_ref` partitions. Generic LoRA behavior is not tested here. + """ + + pipeline_class = MiniMaxH3ModularPipeline + pipeline_blocks_class = MiniMaxH3Blocks + pretrained_model_name_or_path = "hf-internal-testing/tiny-minimax-h3-modular-pipe" + + def get_pipeline(self): + pipeline = self.pipeline_blocks_class().init_pipeline(self.pretrained_model_name_or_path) + pipeline.load_components(dtype=torch.float32) + pipeline.set_progress_bar_config(disable=None) + return pipeline + + def get_dummy_lora_state_dict(self, prefix="diffusion_model.", rank=4, adaln_rank=None): + r""" + A LoRA in the layout both real-world producers emit: the *original* checkpoint's module names, fused + `attn.qkv_proj` and `mlp.fc1`, and no `.alpha`. ai-toolkit prefixes them with `diffusion_model.`; the one + public H3 LoRA carries no prefix at all, which `prefix=""` reproduces. `adaln_rank` makes the file mixed-rank, + as that LoRA is. + """ + transformer = self.get_pipeline().transformer + config = transformer.config + hidden = config.hidden_size + inner = config.num_attention_heads * config.attention_head_dim + adaln_rank = adaln_rank or rank + + state_dict = {} + for source, in_features, out_features, module_rank in [ + ("blocks.0.attn.qkv_proj", hidden, 3 * inner, rank), + ("blocks.0.attn.out_proj", inner, hidden, rank), + ("blocks.0.mlp.fc1", hidden, 2 * config.ffn_dim, rank), + ("blocks.0.mlp.fc2", config.ffn_dim, hidden, rank), + ("blocks.0.adaln_proj.linear", config.time_embed_dim, 6 * 3 * hidden, adaln_rank), + ("token_refiner.blocks.0.attn.qkv_proj", hidden, 3 * inner, rank), + ("final_layer.adaln_proj.linear", config.time_embed_dim, 2 * hidden, adaln_rank), + ]: + state_dict[f"{prefix}{source}.lora_A.weight"] = torch.randn(module_rank, in_features) + state_dict[f"{prefix}{source}.lora_B.weight"] = torch.randn(out_features, module_rank) + return state_dict + + def test_lora_state_dict_conversion(self): + r"""The original module names map onto the diffusers ones, fused projections split, `mlp.fc1` halves swap.""" + state_dict = self.get_dummy_lora_state_dict() + rank = state_dict["diffusion_model.blocks.0.attn.qkv_proj.lora_A.weight"].shape[0] + fused_up = state_dict["diffusion_model.blocks.0.mlp.fc1.lora_B.weight"] + qkv_up = state_dict["diffusion_model.blocks.0.attn.qkv_proj.lora_B.weight"] + + converted = self.pipeline_class.lora_state_dict(state_dict) + + assert "transformer.transformer_blocks.0.attn.to_q.lora_A.weight" in converted + assert "transformer.transformer_blocks.0.attn.to_out.0.lora_B.weight" in converted + assert "transformer.transformer_blocks.0.ff.net.0.proj.lora_A.weight" in converted + assert "transformer.transformer_blocks.0.ff.net.2.lora_B.weight" in converted + assert "transformer.transformer_blocks.0.adaln_proj.linear.lora_A.weight" in converted + assert "transformer.token_refiner.refiner_blocks.0.attn.to_v.lora_B.weight" in converted + assert "transformer.norm_out.linear.lora_B.weight" in converted + assert not any("qkv_proj" in key or "fc1" in key or "final_layer" in key for key in converted) + + # The fused QKV splits into three row blocks that share `lora_A`. + inner = qkv_up.shape[0] // 3 + for index, projection in enumerate(["to_q", "to_k", "to_v"]): + prefix = f"transformer.transformer_blocks.0.attn.{projection}" + assert torch.equal(converted[f"{prefix}.lora_B.weight"], qkv_up[index * inner : (index + 1) * inner]) + assert torch.equal( + converted[f"{prefix}.lora_A.weight"], + converted["transformer.transformer_blocks.0.attn.to_q.lora_A.weight"], + ) + + # `mlp.fc1` is `[gate; value]` and `SwiGLU.proj` is `[value; gate]`, so `lora_B`'s halves swap and `lora_A` + # is untouched. A key-name-only assertion would pass with the swap missing, which is a silent quality bug. + ffn_dim = fused_up.shape[0] // 2 + swapped = converted["transformer.transformer_blocks.0.ff.net.0.proj.lora_B.weight"] + assert torch.equal(swapped, torch.cat([fused_up[ffn_dim:], fused_up[:ffn_dim]])) + assert torch.equal( + converted["transformer.transformer_blocks.0.ff.net.0.proj.lora_A.weight"], + state_dict["diffusion_model.blocks.0.mlp.fc1.lora_A.weight"], + ) + assert all(value.shape[1] == rank or value.shape[0] == rank for value in converted.values()) + + def test_lora_state_dict_conversion_without_a_prefix(self): + r"""The one public H3 LoRA has no prefix at all, so the module names are what identifies the format.""" + converted = self.pipeline_class.lora_state_dict(self.get_dummy_lora_state_dict(prefix="")) + + assert "transformer.transformer_blocks.0.attn.to_k.lora_B.weight" in converted + assert all(key.startswith("transformer.") for key in converted) + + def test_lora_state_dict_conversion_raises_on_an_unknown_module(self): + state_dict = self.get_dummy_lora_state_dict() + state_dict["diffusion_model.blocks.0.not_a_module.lora_A.weight"] = torch.randn(4, 8) + state_dict["diffusion_model.blocks.0.not_a_module.lora_B.weight"] = torch.randn(8, 4) + + with pytest.raises(ValueError, match="not_a_module"): + self.pipeline_class.lora_state_dict(state_dict) + + def test_lora_state_dict_synthesizes_unit_scale_metadata(self): + r""" + A non-diffusers H3 LoRA has no alpha information and applies as `W + lora_B @ lora_A`. `get_peft_kwargs` reads + `lora_alpha` off the first rank it sees and never re-derives it, so a mixed-rank file — which the public turbo + LoRA is — would have its majority-rank modules scaled by `alpha / r`. The converted metadata pins + `alpha == rank` for every module instead. + """ + state_dict = self.get_dummy_lora_state_dict(rank=8, adaln_rank=2) + + _, metadata = self.pipeline_class.lora_state_dict(state_dict, return_lora_metadata=True) + + assert metadata["transformer.r"] == 8 + assert metadata["transformer.lora_alpha"] == 8 + assert metadata["transformer.alpha_pattern"] == metadata["transformer.rank_pattern"] + assert set(metadata["transformer.rank_pattern"].values()) == {2} + assert "^norm_out.linear" in metadata["transformer.rank_pattern"] + + def get_dummy_diffusers_lora_state_dict(self, prefix="transformer", rank=8, adaln_rank=2): + r""" + The same adapter already converted to diffusers keys and republished — which is how the public turbo LoRA also + circulates. Mixed-rank, still alpha-less, so it needs the same treatment as the original layout even though no + conversion runs. + """ + transformer = self.get_pipeline().transformer + config = transformer.config + hidden = config.hidden_size + inner = config.num_attention_heads * config.attention_head_dim + + state_dict = {} + for module, in_features, out_features, module_rank in [ + ("transformer_blocks.0.attn.to_q", hidden, inner, rank), + ("transformer_blocks.0.attn.to_out.0", inner, hidden, rank), + ("transformer_blocks.0.ff.net.0.proj", hidden, 2 * config.ffn_dim, rank), + ("transformer_blocks.0.ff.net.2", config.ffn_dim, hidden, rank), + ("transformer_blocks.0.adaln_proj.linear", config.time_embed_dim, 6 * 3 * hidden, adaln_rank), + ("norm_out.linear", config.time_embed_dim, 2 * hidden, adaln_rank), + ]: + state_dict[f"{prefix}.{module}.lora_A.weight"] = torch.randn(module_rank, in_features) + state_dict[f"{prefix}.{module}.lora_B.weight"] = torch.randn(out_features, module_rank) + return state_dict + + @pytest.mark.parametrize("prefix", ["transformer", "transformer_ref"], ids=["transformer", "transformer_ref"]) + def test_load_lora_weights_diffusers_format_mixed_rank(self, prefix): + r""" + A mixed-rank, alpha-less adapter already in diffusers keys bypasses the converter, so the alpha handling cannot + live there: without it `get_peft_kwargs` takes `lora_alpha` from whichever rank it sees first and one of the + two rank groups is applied at `alpha / r`. + """ + pipe = self.get_pipeline() + + pipe.load_lora_weights(self.get_dummy_diffusers_lora_state_dict(prefix=prefix), adapter_name="dummy") + + component = getattr(pipe, prefix) + injected = [module for module in component.modules() if "dummy" in getattr(module, "scaling", {})] + assert len(injected) == 6 + assert {module.scaling["dummy"] for module in injected} == {1.0} + for module in injected: + assert module.lora_alpha["dummy"] == module.r["dummy"] + assert component.transformer_blocks[0].attn.to_q.r["dummy"] == 8 + assert component.transformer_blocks[0].adaln_proj.linear.r["dummy"] == 2 + assert component.norm_out.linear.r["dummy"] == 2 + + def test_lora_state_dict_respects_existing_metadata(self): + r"""A file that carries diffusers' own `lora_adapter_metadata` must not have it overwritten.""" + pipe = self.get_pipeline() + pipe.load_lora_weights(self.get_dummy_diffusers_lora_state_dict(), adapter_name="dummy") + layers = get_peft_model_state_dict(pipe.transformer, adapter_name="dummy") + saved_metadata = {"r": 8, "lora_alpha": 8, "rank_pattern": {}, "alpha_pattern": {}, "target_modules": ["x"]} + + with tempfile.TemporaryDirectory() as tmpdir: + self.pipeline_class.save_lora_weights( + tmpdir, transformer_lora_layers=layers, transformer_lora_adapter_metadata=saved_metadata + ) + _, metadata = self.pipeline_class.lora_state_dict(tmpdir, return_lora_metadata=True) + + assert metadata["transformer.target_modules"] == ["x"] + + @pytest.mark.parametrize("prefix", ["diffusion_model.", ""], ids=["ai_toolkit", "unprefixed"]) + def test_load_lora_weights(self, prefix): + r""" + A mixed-rank, alpha-less file — the public turbo LoRA's shape — has to reach every module at its own rank and + at an effective scale of exactly 1.0. + """ + pipe = self.get_pipeline() + + pipe.load_lora_weights( + self.get_dummy_lora_state_dict(prefix=prefix, rank=8, adaln_rank=2), adapter_name="dummy" + ) + + assert "dummy" in pipe.transformer.peft_config + # Both partitions are loaded here and the file does not say which one it targets, so only `transformer` gets it. + assert "dummy" not in getattr(pipe.transformer_ref, "peft_config", {}) + + injected = [module for module in pipe.transformer.modules() if "dummy" in getattr(module, "scaling", {})] + # 3 split qkv + to_out.0 + the two ff Linears + adaln, the refiner's 3 split qkv, and norm_out + assert len(injected) == 11 + assert {module.scaling["dummy"] for module in injected} == {1.0} + for module in injected: + assert module.lora_A["dummy"].weight.shape[0] == module.r["dummy"] + assert module.lora_alpha["dummy"] == module.r["dummy"] + assert pipe.transformer.transformer_blocks[0].attn.to_q.r["dummy"] == 8 + assert pipe.transformer.transformer_blocks[0].adaln_proj.linear.r["dummy"] == 2 + assert pipe.transformer.norm_out.linear.r["dummy"] == 2 + + def test_load_lora_weights_into_transformer_ref(self): + pipe = self.get_pipeline() + + pipe.load_lora_weights(self.get_dummy_lora_state_dict(), adapter_name="dummy", load_into_transformer_ref=True) + + assert "dummy" in pipe.transformer_ref.peft_config + assert "dummy" not in getattr(pipe.transformer, "peft_config", {}) + + def test_save_load_lora_weights_round_trip(self): + r""" + `save_lora_weights` is the only mechanism that records which partition a LoRA belongs to, so the round trip + has to preserve it. + """ + pipe = self.get_pipeline() + pipe.load_lora_weights(self.get_dummy_lora_state_dict(), adapter_name="dummy", load_into_transformer_ref=True) + layers = get_peft_model_state_dict(pipe.transformer_ref, adapter_name="dummy") + + with tempfile.TemporaryDirectory() as tmpdir: + self.pipeline_class.save_lora_weights(tmpdir, transformer_ref_lora_layers=layers) + reloaded = self.pipeline_class.lora_state_dict(tmpdir) + + assert reloaded + assert all(key.startswith("transformer_ref.") for key in reloaded) + + fresh = self.get_pipeline() + fresh.load_lora_weights(reloaded, adapter_name="dummy") + assert "dummy" in fresh.transformer_ref.peft_config + assert "dummy" not in getattr(fresh.transformer, "peft_config", {}) + + def test_load_lora_weights_routes_to_the_only_partition(self): + r""" + `workflow="ref2va"` loads `transformer_ref` and no `transformer`, and nothing in a published H3 LoRA says which + partition it targets, so the one partition that is present is the unambiguous destination. + """ + pipe = self.pipeline_blocks_class().get_workflow("ref2va").init_pipeline(self.pretrained_model_name_or_path) + pipe.load_components(dtype=torch.float32) + assert getattr(pipe, "transformer", None) is None + + state_dict = self.get_dummy_lora_state_dict() + pipe.load_lora_weights(state_dict, adapter_name="dummy") + + assert "dummy" in pipe.transformer_ref.peft_config + + def test_load_lora_weights_raises_without_the_requested_partition(self): + pipe = self.pipeline_blocks_class().get_workflow("t2va").init_pipeline(self.pretrained_model_name_or_path) + pipe.load_components(dtype=torch.float32) + assert getattr(pipe, "transformer_ref", None) is None + + state_dict = self.get_dummy_lora_state_dict() + with pytest.raises(ValueError, match="load_into_transformer_ref"): + pipe.load_lora_weights(state_dict, load_into_transformer_ref=True) diff --git a/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py b/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py index 366cb366b220..82b3cb5bd174 100644 --- a/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py +++ b/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. + import numpy as np import pytest import torch