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 d9a82aea08a9..4c1b5d4b49c0 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 + +
+ MiniMax-H3 generates video and its soundtrack together. A single transformer denoises one packed sequence containing the text conditioning, conditioning media, and target video and audio latents. There is no separate vocoder and no audio post-hoc pass: video and audio come out of the same denoising loop. You can find the original MiniMax-H3 checkpoints under the [MiniMaxAI](https://huggingface.co/MiniMaxAI) organization. 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_base.py b/src/diffusers/loaders/lora_base.py index 1c8a66976b81..0009f04df9e2 100644 --- a/src/diffusers/loaders/lora_base.py +++ b/src/diffusers/loaders/lora_base.py @@ -46,7 +46,7 @@ set_weights_and_activate_adapters, ) from ..utils.peft_utils import _create_lora_config -from ..utils.state_dict_utils import _load_sft_state_dict_metadata +from ..utils.state_dict_utils import _load_sft_file_metadata, _load_sft_state_dict_metadata if is_transformers_available(): @@ -208,7 +208,15 @@ def _fetch_state_dict( user_agent, allow_pickle, metadata=None, + return_file_metadata=False, ): + """ + `metadata` is diffusers' own LoRA adapter metadata, parsed out of the file's `__metadata__`. With + `return_file_metadata`, that `__metadata__` is returned in full as a third element. It is `None` whenever the + weights did not come from a safetensors file — a state dict passed in memory, or a pickled checkpoint — since there + is nowhere else for a header to live. + """ + file_metadata = None model_file = None if not isinstance(pretrained_model_name_or_path_or_dict, dict): # Let's first try to load .safetensors weights @@ -239,6 +247,8 @@ def _fetch_state_dict( ) state_dict = safetensors.torch.load_file(model_file, device="cpu") metadata = _load_sft_state_dict_metadata(model_file) + if return_file_metadata: + file_metadata = _load_sft_file_metadata(model_file) except (IOError, safetensors.SafetensorError) as e: if not allow_pickle: @@ -246,6 +256,7 @@ def _fetch_state_dict( # try loading non-safetensors weights model_file = None metadata = None + file_metadata = None pass if model_file is None: @@ -270,6 +281,8 @@ def _fetch_state_dict( else: state_dict = pretrained_model_name_or_path_or_dict + if return_file_metadata: + return state_dict, metadata, file_metadata return state_dict, metadata diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index 00ccc0c42078..d2f903b7a873 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -3117,3 +3117,167 @@ 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. + + Every known producer trains 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, musubi-tuner under a flattened + `lora_unet_` one — 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 + - musubi-tuner's flattened `lora_unet_blocks_0_attn_qkv_proj` -> `blocks.0.attn.qkv_proj` + - `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()} + + # musubi-tuner (kohya sd-scripts) writes one flat module name per key under a `lora_unet_` prefix, with every `.` + # collapsed to `_`. H3's own module names contain underscores (`qkv_proj`, `token_refiner`, `final_layer`, + # `video_out`), so the dot path is recovered by matching the whole flattened name against the module vocabulary + # rather than by splitting on `_`. + if any(k.startswith("lora_unet_") for k in state_dict): + flattened_modules = [ + (r"blocks_(\d+)_attn_(qkv|out)_proj", r"blocks.\1.attn.\2_proj"), + (r"blocks_(\d+)_mlp_fc([12])", r"blocks.\1.mlp.fc\2"), + (r"blocks_(\d+)_adaln_proj_linear", r"blocks.\1.adaln_proj.linear"), + (r"token_refiner_blocks_(\d+)_attn_(qkv|out)_proj", r"token_refiner.blocks.\1.attn.\2_proj"), + (r"token_refiner_blocks_(\d+)_mlp_fc([12])", r"token_refiner.blocks.\1.mlp.fc\2"), + (r"(video|audio)_patch_proj", r"\1_patch_proj"), + (r"condition_proj", "condition_proj"), + (r"time_embedder_proj_(in|out)", r"time_embedder.proj_\1"), + (r"final_layer_adaln_proj_linear", "final_layer.adaln_proj.linear"), + (r"final_layer_(video|audio)_out", r"final_layer.\1_out"), + ] + unflattened = {} + for key, value in state_dict.items(): + module, _, suffix = key.removeprefix("lora_unet_").partition(".") + dotted = None + for pattern, replacement in flattened_modules: + if re.fullmatch(pattern, module): + dotted = re.sub(pattern, replacement, module) + break + if dotted is None: + raise ValueError( + f"`{key}` does not name a MiniMax-H3 module in musubi-tuner's flattened `lora_unet_` layout." + ) + unflattened[f"{dotted}.{suffix}"] = value + state_dict = unflattened + + 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 9ccddf0b40d6..aee68486c68d 100644 --- a/src/diffusers/loaders/lora_pipeline.py +++ b/src/diffusers/loaders/lora_pipeline.py @@ -20,6 +20,7 @@ from ..utils import ( deprecate, + get_peft_kwargs, get_submodule_by_name, is_bitsandbytes_available, is_gguf_available, @@ -54,6 +55,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, @@ -73,6 +75,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"} @@ -6736,6 +6740,381 @@ 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 ships two independent DiT partitions with identical module names (a LoRA for one loads into the other + and degrades output), so routing is explicit: converted state dicts target `transformer.`; reach `transformer_ref` + with a `transformer_ref.`-prefixed file or `load_into_transformer_ref=True`. + + LoRAs trained against a pruned checkpoint (the `*_pruned_*` files in + [Comfy-Org/MiniMax-H3](https://huggingface.co/Comfy-Org/MiniMax-H3); + [joyfox/MiniMax-H3-Turbo](https://huggingface.co/joyfox/MiniMax-H3-Turbo) is one) fail with a size mismatch. + Alpha-less files load at `alpha == rank` (scale 1.0, the convention stated by e.g. + [larryvrh/MiniMax-H3-Turbo-Lora](https://huggingface.co/larryvrh/MiniMax-H3-Turbo-Lora)); a `__metadata__` `alpha` + entry (e.g. [lightx2v/Minimax-h3-Turbo](https://huggingface.co/lightx2v/Minimax-h3-Turbo)'s 8-step file), when + present, is honored instead. + """ + + _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"} + + # `return_file_metadata=True` because some H3 LoRAs record their training alpha in the file's `__metadata__`. + state_dict, metadata, file_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, + return_file_metadata=True, + ) + + # Read before the conversion below, which folds per-module alphas into the weights and drops the scalars. + has_alpha_tensors = any(k.endswith(".alpha") for k in state_dict) + + 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} + + # A peft dump: diffusers module names carrying peft's `.default.` infix and no component prefix. The missing + # prefix is what keeps this from shadowing a file diffusers itself wrote. + is_unprefixed_diffusers_format = any(".default.weight" in k for k in state_dict) and not any( + k.startswith((f"{cls.transformer_name}.", f"{cls.transformer_ref_name}.")) for k in state_dict + ) + if is_unprefixed_diffusers_format: + state_dict = {f"{cls.transformer_name}.{k.replace('.default.', '.')}": v for k, v in state_dict.items()} + + is_non_diffusers_format = any( + k.startswith(("diffusion_model.", "blocks.", "token_refiner.", "final_layer.", "lora_unet_")) + for k in state_dict + ) + if is_non_diffusers_format: + state_dict = _convert_non_diffusers_minimax_h3_lora_to_diffusers(state_dict) + + # Published H3 LoRAs are alpha-less and mixed-rank (64 on attention and FFN, 16 on the AdaLN projections), and + # `get_peft_kwargs` would scale one of the two rank groups by `alpha / r`, so `alpha == rank` is pinned below. + # + # An alpha-less file may instead record its training alpha in the `__metadata__`, honored as a uniform network + # alpha (`alpha / rank` per module). Per-module scalars win, the fold above having already applied them. + network_alpha = None + if file_metadata is not None and "alpha" in file_metadata and not has_alpha_tensors: + try: + network_alpha = float(file_metadata["alpha"]) + except ValueError: + logger.warning( + f"This LoRA file records `alpha` as {file_metadata['alpha']!r} in its `__metadata__`. MiniMax-H3" + " reads that entry as the network alpha every module was trained with, and it has to be a number," + " so it is being ignored — the adapter is loaded with `alpha == rank` instead." + ) + else: + logger.info( + f"Using the network alpha {network_alpha} this LoRA file records in its `__metadata__`; every " + "module is scaled by `alpha / rank`." + ) + + # Without a `__metadata__` alpha, no metadata is built: `get_peft_kwargs` already derives `alpha == rank` for + # alpha-less files, which is the fold's contract too. + if metadata is None and network_alpha is not None: + 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 + ) + lora_config_kwargs["lora_alpha"] = network_alpha + lora_config_kwargs["alpha_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. + """ + 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 not transformer_state_dict and not transformer_ref_state_dict: + logger.warning( + f"No LoRA keys associated to {type(self).__name__} found with the prefix " + f"`{self.transformer_name}.` or `{self.transformer_ref_name}.`, so loading this state dict would " + f"target no module at all and nothing was loaded. This is an unrecognized MiniMax-H3 LoRA layout; " + f"its first keys are {sorted(state_dict)[:4]}. Open an issue if you think it's unexpected: " + "https://github.com/huggingface/diffusers/issues/new" + ) + return + + if transformer_ref_state_dict and 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"`.' + ) + + 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. + self.load_lora_into_transformer_ref( + transformer_state_dict, + transformer_ref=transformer_ref, + prefix=self.transformer_name, + 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: + self.load_lora_into_transformer_ref( + transformer_ref_state_dict, + transformer_ref=transformer_ref, + prefix=self.transformer_ref_name, + 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. + """ + # 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 load_lora_into_transformer_ref( + cls, + state_dict, + transformer_ref, + prefix, + adapter_name=None, + _pipeline=None, + low_cpu_mem_usage=False, + hotswap: bool = False, + metadata=None, + ): + """ + Load LoRA layers into the `transformer_ref` partition. `prefix` is the component name the keys carry, which is + `transformer_ref` for a file that names the partition and `transformer` for one routed here by + `load_into_transformer_ref=True`. See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more + details. + """ + logger.info(f"Loading {prefix}.") + transformer_ref.load_lora_adapter( + state_dict, + prefix=prefix, + 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/src/diffusers/utils/state_dict_utils.py b/src/diffusers/utils/state_dict_utils.py index cdfdb7b1df04..00832171dcde 100644 --- a/src/diffusers/utils/state_dict_utils.py +++ b/src/diffusers/utils/state_dict_utils.py @@ -529,15 +529,21 @@ def state_dict_all_zero(state_dict, filter_str=None): return all(torch.all(param == 0).item() for param in state_dict.values()) -def _load_sft_state_dict_metadata(model_file: str): +def _load_sft_file_metadata(model_file: str): + """A safetensors file's own `__metadata__`, minus safetensors' `format` entry.""" import safetensors.torch - from ..loaders.lora_base import LORA_ADAPTER_METADATA_KEY - with safetensors.torch.safe_open(model_file, framework="pt", device="cpu") as f: metadata = f.metadata() or {} metadata.pop("format", None) + return metadata + + +def _load_sft_state_dict_metadata(model_file: str): + from ..loaders.lora_base import LORA_ADAPTER_METADATA_KEY + + metadata = _load_sft_file_metadata(model_file) if metadata: raw = metadata.get(LORA_ADAPTER_METADATA_KEY) return json.loads(raw) if raw else None 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..997cfe91c1dd --- /dev/null +++ b/tests/lora/test_lora_layers_minimax_h3.py @@ -0,0 +1,203 @@ +# 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 os +import tempfile + +import pytest +import safetensors.torch +import torch + +from diffusers.modular_pipelines import MiniMaxH3Blocks, MiniMaxH3ModularPipeline +from diffusers.utils import is_peft_available, logging + +from ..testing_utils import CaptureLogger, 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: loading the layouts + that circulate, the alpha handling that gets each of them to its trained scale, 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 test_load_lora_weights_warns_when_nothing_is_targeted(self): + r""" + An unrecognized layout that keeps the substring `lora` in every key passes the format check and then filters to + nothing in both partitions. Neither partition branch would run, so without this warning the load is a silent + no-op. The message follows `PeftAdapterMixin.load_lora_adapter`'s wording, which is what every single-denoiser + model emits in the same situation. + """ + pipe = self.get_pipeline() + state_dict = { + "some_other_model.layers.0.lora_A.weight": torch.randn(4, 24), + "some_other_model.layers.0.lora_B.weight": torch.randn(24, 4), + } + + logger = logging.get_logger("diffusers.loaders.lora_pipeline") + logger.setLevel(logging.WARNING) + with CaptureLogger(logger) as cap_logger: + pipe.load_lora_weights(state_dict, adapter_name="dummy") + + assert cap_logger.out.startswith("No LoRA keys associated to MiniMaxH3ModularPipeline") + assert "some_other_model.layers.0.lora_A.weight" in cap_logger.out + # Nothing was loaded into either partition. + for component in [pipe.transformer, pipe.transformer_ref]: + assert "dummy" not in getattr(component, "peft_config", {}) + assert not [module for module in component.modules() if "dummy" in getattr(module, "scaling", {})] + + def save_with_file_metadata(self, state_dict, tmpdir, alpha="8"): + r""" + One producer records the alpha it trained with in the safetensors `__metadata__` instead of in per-module + scalars, so the value exists only on disk — a state dict handed over in memory cannot carry it. + """ + weight_name = "pytorch_lora_weights.safetensors" + safetensors.torch.save_file( + state_dict, os.path.join(tmpdir, weight_name), metadata={"floating_dtype": "bfloat16", "alpha": alpha} + ) + return weight_name + + def test_load_lora_weights_honors_the_metadata_alpha(self): + r""" + A file with one uniform rank, no `.alpha` scalars and `alpha` "8" in its own `__metadata__`: rank 128 + against alpha 8 is a trained scale of 0.0625; synthesizing `alpha == rank` would apply the adapter 16x too + strongly. + """ + state_dict = self.get_dummy_diffusers_lora_state_dict(prefix="transformer", rank=128, adaln_rank=128) + + with tempfile.TemporaryDirectory() as tmpdir: + weight_name = self.save_with_file_metadata(state_dict, tmpdir) + + pipe = self.get_pipeline() + pipe.load_lora_weights(tmpdir, weight_name=weight_name, adapter_name="dummy") + + injected = [module for module in pipe.transformer.modules() if "dummy" in getattr(module, "scaling", {})] + assert len(injected) == 6 + assert {module.scaling["dummy"] for module in injected} == {0.0625} + + 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_load_lora_weights_into_transformer_ref(self): + pipe = self.get_pipeline() + + pipe.load_lora_weights( + self.get_dummy_diffusers_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_diffusers_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_diffusers_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_diffusers_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