diff --git a/__init__.py b/__init__.py index f634ba9..0e877bc 100644 --- a/__init__.py +++ b/__init__.py @@ -60,6 +60,7 @@ from .nodes.model_info_nodes import ( from .nodes.lora_rename import AnimaLoraRename +from .nodes.minimax_h3_lora_convert import MiniMaxH3DiffusersLoRAConvert from .nodes.dtype_conversion import DiffusionModelDtypeConversion from .nodes.layer_parameters import LayerParameterConfiguration from .nodes.lora_model_analysis import LoRAOnModelAnalysis @@ -118,6 +119,7 @@ class ModelUtilsExtension(ComfyExtension): LoRAMergeToModel, # LoRA Utilities AnimaLoraRename, + MiniMaxH3DiffusersLoRAConvert, # Dtype Conversion DiffusionModelDtypeConversion, # Downloaders diff --git a/nodes/minimax_h3_lora_convert.py b/nodes/minimax_h3_lora_convert.py new file mode 100644 index 0000000..4930cf5 --- /dev/null +++ b/nodes/minimax_h3_lora_convert.py @@ -0,0 +1,333 @@ +from __future__ import annotations + +from dataclasses import dataclass +import os +import re + +import comfy.utils +from comfy_api.latest import io +import folder_paths +import torch +from unifiedefficientloader import MemoryEfficientSafeOpen + +from .artifact_paths import canonical_model_artifact_path +from .quantization_guard import inspect_low_bit_input +from .uel_io import atomic_uel_writer + +_FACTOR_RE = re.compile( + r"^(?P.+)\.lora_(?P[AB])(?:\.default)?\.weight$" +) +_BLOCK_RE = re.compile(r"^transformer_blocks\.(?P\d+)\.(?P.+)$") +_REFINER_RE = re.compile(r"^token_refiner\.refiner_blocks\.(?P\d+)\.(?P.+)$") +_QKV_RE = re.compile(r"^attn\.to_(?P[qkv])$") + +_DIRECT_TAILS = { + "attn.to_out.0": "attn.out_proj", + "ff.net.0.proj": "mlp.fc1", + "ff.net.2": "mlp.fc2", +} + +_EXTRA_PREFIX_MAP = { + "context_embedder": "condition_proj", + "proj_in": "video_patch_proj", + "proj_out": "final_layer.video_out", + "audio_proj_in": "audio_patch_proj", + "audio_proj_out": "final_layer.audio_out", + "norm_out.norm": "final_layer.norm", + "norm_out.linear": "final_layer.adaln_proj.linear", + "time_embedder.linear_1": "time_embedder.proj_in", + "time_embedder.linear_2": "time_embedder.proj_out", +} + + +@dataclass(frozen=True) +class _ConversionUnit: + target: str + sources: tuple[str, ...] + qkv: bool + is_swiglu: bool = False + + +def _target_module(module: str) -> tuple[str, str | None, bool]: + """Map Diffusers module name to ComfyUI native target module. + + Returns: + (target_module, qkv_part, is_swiglu_fc1) + """ + clean_mod = module.removeprefix("transformer.") + match = _BLOCK_RE.fullmatch(clean_mod) + if match: + prefix = f"diffusion_model.blocks.{match.group('index')}" + tail = match.group("tail") + else: + match = _REFINER_RE.fullmatch(clean_mod) + if match: + prefix = f"diffusion_model.token_refiner.blocks.{match.group('index')}" + tail = match.group("tail") + else: + for old_pfx, new_pfx in _EXTRA_PREFIX_MAP.items(): + if clean_mod.startswith(old_pfx): + target = f"diffusion_model.{new_pfx}{clean_mod[len(old_pfx):]}" + return target, None, False + raise ValueError(f"Unsupported MiniMax H3 Diffusers module: {module}") + + qkv = _QKV_RE.fullmatch(tail) + if qkv: + return f"{prefix}.attn.qkv_proj", qkv.group("part"), False + + if tail in _DIRECT_TAILS: + target = f"{prefix}.{_DIRECT_TAILS[tail]}" + is_swiglu = (tail == "ff.net.0.proj") + return target, None, is_swiglu + + raise ValueError(f"Unsupported MiniMax H3 Diffusers module tail: {module}") + + +def _build_plan(keys: list[str]) -> list[_ConversionUnit]: + """Build conversion work units from input keys.""" + factors: dict[str, dict[str, str]] = {} + for key in keys: + match = _FACTOR_RE.fullmatch(key) + if not match: + raise ValueError( + f"Input is not a compatible MiniMax H3 Diffusers PEFT LoRA; unsupported key: {key}" + ) + mod = match.group("module") + factor = match.group("factor") + pair = factors.setdefault(mod, {}) + if factor in pair: + raise ValueError(f"Duplicate LoRA factor for: {mod}") + pair[factor] = key + + missing = [mod for mod, pair in factors.items() if set(pair.keys()) != {"A", "B"}] + if missing: + raise ValueError(f"Incomplete LoRA A/B pair for: {missing[0]}") + + direct_units: list[_ConversionUnit] = [] + qkv_groups: dict[str, dict[str, tuple[str, str]]] = {} + + for mod, pair in factors.items(): + target, qkv_part, is_swiglu = _target_module(mod) + source_pair = (pair["A"], pair["B"]) + if qkv_part is None: + direct_units.append( + _ConversionUnit( + target=target, + sources=source_pair, + qkv=False, + is_swiglu=is_swiglu, + ) + ) + else: + qkv_groups.setdefault(target, {})[qkv_part] = source_pair + + qkv_units: list[_ConversionUnit] = [] + for target, parts in qkv_groups.items(): + if set(parts.keys()) != {"q", "k", "v"}: + raise ValueError(f"Incomplete Q/K/V LoRA group for: {target}") + sources = tuple(k for part in ("q", "k", "v") for k in parts[part]) + qkv_units.append(_ConversionUnit(target=target, sources=sources, qkv=True)) + + units = sorted(direct_units + qkv_units, key=lambda u: u.target) + targets = [u.target for u in units] + if len(targets) != len(set(targets)): + raise ValueError("Conversion would produce duplicate ComfyUI module keys") + + return units + + +def _validate_pair(down: torch.Tensor, up: torch.Tensor, label: str) -> None: + if down.ndim != 2 or up.ndim != 2: + raise ValueError(f"LoRA factors for '{label}' must be 2D matrices") + if down.shape[0] != up.shape[1]: + raise ValueError( + f"LoRA factor rank mismatch for '{label}': down={list(down.shape)}, up={list(up.shape)}" + ) + + +def _convert_qkv( + tensors: dict[str, torch.Tensor], + sources: tuple[str, ...], +) -> tuple[torch.Tensor, torch.Tensor]: + """Fuse separate Q, K, V LoRA factors into unified block-diagonal qkv_proj factors.""" + pairs = [ + (tensors[sources[i]], tensors[sources[i + 1]]) + for i in range(0, 6, 2) + ] + for part, (down, up) in zip(("Q", "K", "V"), pairs): + _validate_pair(down, up, part) + + in_dims = {down.shape[1] for down, _ in pairs} + if len(in_dims) != 1: + raise ValueError("Q/K/V LoRA input dimensions do not match") + + dtypes = {t.dtype for pair in pairs for t in pair} + if len(dtypes) != 1: + raise ValueError("Q/K/V LoRA dtypes do not match") + + source_dtype = pairs[0][0].dtype + ranks = [down.shape[0] for down, _ in pairs] + output_dims = [up.shape[0] for _, up in pairs] + total_rank = sum(ranks) + total_out = sum(output_dims) + + # Concatenate down matrices vertically: [total_rank, in_dim] + exact_down = torch.cat([down for down, _ in pairs], dim=0).contiguous() + + # Place up matrices block-diagonally: [total_out, total_rank] + exact_up = torch.zeros( + (total_out, total_rank), + dtype=source_dtype, + device=pairs[0][1].device, + ) + row = 0 + col = 0 + for (_, part_up), out_dim, part_rank in zip(pairs, output_dims, ranks): + exact_up[row : row + out_dim, col : col + part_rank].copy_(part_up) + row += out_dim + col += part_rank + + return exact_down, exact_up.contiguous() + + +def _convert_swiglu( + down: torch.Tensor, up: torch.Tensor, label: str +) -> tuple[torch.Tensor, torch.Tensor]: + """Swap SwiGLU output halves from Diffusers [value; gate] to ComfyUI [gate; value].""" + _validate_pair(down, up, label) + if up.shape[0] % 2 != 0: + raise ValueError( + f"SwiGLU LoRA B factor for '{label}' has odd output dimension {up.shape[0]}" + ) + val_b, gate_b = up.chunk(2, dim=0) + swapped_up = torch.cat([gate_b, val_b], dim=0).contiguous() + return down.contiguous(), swapped_up + + +def convert_minimax_h3_diffusers_lora( + input_path: str, + output_path: str, +) -> str: + """Stream and convert Diffusers MiniMax H3 LoRA into ComfyUI native format.""" + if os.path.normcase(os.path.abspath(input_path)) == os.path.normcase( + os.path.abspath(output_path) + ): + raise ValueError("Output path must differ from the source LoRA path") + + handler = MemoryEfficientSafeOpen(input_path, low_memory=True) + try: + inspect_low_bit_input(handler, f"LoRA ({input_path})", "MiniMax H3 LoRA Convert") + keys = list(handler.keys()) + source_count = len(keys) + units = _build_plan(keys) + + metadata = {str(k): str(v) for k, v in (handler.metadata() or {}).items()} + metadata.update( + { + "modelutils_conversion": "diffusers_minimax_h3_to_comfyui", + "modelutils_qkv_fusion": "lossless_block_diagonal", + "modelutils_mlp_swiglu": "swapped_gate_value", + } + ) + + progress = comfy.utils.ProgressBar(len(units)) + output_count = 0 + fused_qkv_count = 0 + + with atomic_uel_writer(output_path, metadata) as writer: + for unit in units: + stream = handler.async_stream( + list(unit.sources), + batch_size=1, + prefetch_batches=1, + pin_memory=False, + ) + tensors = {} + try: + for batch in stream: + for k, t in batch: + tensors[k] = t + + if unit.qkv: + down, up = _convert_qkv(tensors, unit.sources) + fused_qkv_count += 1 + elif unit.is_swiglu: + down_key, up_key = unit.sources + down, up = _convert_swiglu( + tensors[down_key], tensors[up_key], unit.target + ) + else: + down_key, up_key = unit.sources + down = tensors[down_key].contiguous() + up = tensors[up_key].contiguous() + _validate_pair(down, up, unit.target) + + writer.write_batch( + [ + (f"{unit.target}.lora_A.weight", down), + (f"{unit.target}.lora_B.weight", up), + ] + ) + output_count += 2 + finally: + tensors.clear() + for key in unit.sources: + handler.mark_processed(key) + stream.close() + + progress.update(1) + finally: + handler.__exit__(None, None, None) + + lines = [ + "MiniMax H3 Diffusers LoRA conversion complete", + f"Source tensors: {source_count}", + f"Output tensors: {output_count}", + f"Fused QKV modules: {fused_qkv_count}", + f"Output size: {os.path.getsize(output_path)} bytes", + f"Output: {output_path}", + ] + return "\n".join(lines) + + +class MiniMaxH3DiffusersLoRAConvert(io.ComfyNode): + """Convert Diffusers MiniMax H3 LoRAs to ComfyUI's native layout.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="MiniMaxH3DiffusersLoRAConvert", + display_name="MiniMax H3 Diffusers LoRA Convert", + category="ModelUtils/LoRA", + description="Converts Diffusers/PEFT MiniMax H3 LoRAs to ComfyUI's native fused layout (lossless block-diagonal QKV and SwiGLU-aligned MLP).", + inputs=[ + io.Combo.Input( + "lora_name", + options=folder_paths.get_filename_list("loras"), + tooltip="Diffusers-format MiniMax H3 LoRA to convert.", + ), + io.String.Input( + "output_filename", + default="minimax_h3_comfyui", + tooltip="Output filename without extension, written under ComfyUI's LoRA directory.", + ), + ], + outputs=[ + io.AnyType.Output(display_name="output_path"), + io.String.Output(display_name="conversion_report"), + ], + is_output_node=True, + ) + + @classmethod + def execute( + cls, + lora_name: str, + output_filename: str, + ) -> io.NodeOutput: + input_path = folder_paths.get_full_path_or_raise("loras", lora_name) + output_path, output_name = canonical_model_artifact_path( + "loras", output_filename + ) + report = convert_minimax_h3_diffusers_lora(input_path, output_path) + return io.NodeOutput(output_name, report) diff --git a/tests/test_artifact_paths.py b/tests/test_artifact_paths.py index 26dcb92..e17754d 100644 --- a/tests/test_artifact_paths.py +++ b/tests/test_artifact_paths.py @@ -24,6 +24,7 @@ MIGRATED_SAVE_MODULES = ( "renamekeys.py", "prunekeys.py", "lora_rename.py", + "minimax_h3_lora_convert.py", ) diff --git a/tests/test_groups.toml b/tests/test_groups.toml index f533741..ba3198c 100644 --- a/tests/test_groups.toml +++ b/tests/test_groups.toml @@ -11,6 +11,7 @@ paths = [ "nodes/dora_learned_wd.py", "nodes/dtype_conversion.py", "nodes/cwb_delta_lora_merger.py", "nodes/lodestone_merger.py", "nodes/lora_extract_svd.py", "nodes/lora_merger.py", "nodes/lora_rename.py", "nodes/lora_resize.py", "nodes/merger.py", + "nodes/minimax_h3_lora_convert.py", "nodes/prunekeys.py", "nodes/renamekeys.py", "nodes/text_encoder_extract.py", ] python_tests = ["tests/test_artifact_paths.py"] @@ -67,3 +68,7 @@ python_tests = ["tests/test_layer_parameters.py"] [groups.dtype_conversion] paths = ["nodes/dtype_conversion.py"] python_tests = ["tests/test_dtype_conversion.py"] + +[groups.minimax_h3_lora_convert] +paths = ["nodes/minimax_h3_lora_convert.py"] +python_tests = ["tests/test_minimax_h3_lora_convert.py"] diff --git a/tests/test_minimax_h3_lora_convert.py b/tests/test_minimax_h3_lora_convert.py new file mode 100644 index 0000000..22655ab --- /dev/null +++ b/tests/test_minimax_h3_lora_convert.py @@ -0,0 +1,178 @@ +import importlib +from pathlib import Path +import sys +import types + +import pytest +from safetensors.torch import load_file, save_file +import torch + +REPO_ROOT = Path(__file__).resolve().parents[1] +PACKAGE_NAME = "modelutils_minimax_convert_tests" + + +def _load_module(name: str): + if PACKAGE_NAME not in sys.modules: + package = types.ModuleType(PACKAGE_NAME) + package.__path__ = [str(REPO_ROOT)] + sys.modules[PACKAGE_NAME] = package + return importlib.import_module(f"{PACKAGE_NAME}.{name}") + + +@pytest.fixture(scope="module") +def converter(): + return _load_module("nodes.minimax_h3_lora_convert") + + +def _make_source_tensors(prefix: str = "transformer_blocks.0") -> dict[str, torch.Tensor]: + tensors = {} + for offset, part in enumerate(("q", "k", "v"), start=1): + # Down: rank 2, in_dim 3 + tensors[f"{prefix}.attn.to_{part}.lora_A.default.weight"] = torch.tensor( + [[1.0, float(offset), 0.0], [0.0, 1.0, float(offset)]], + dtype=torch.float32, + ) + # Up: out_dim 2, rank 2 + tensors[f"{prefix}.attn.to_{part}.lora_B.default.weight"] = torch.tensor( + [[float(offset), 0.0], [0.0, float(offset) + 1.0]], + dtype=torch.float32, + ) + + # out_proj: rank 2, in_dim 3, out_dim 2 + tensors[f"{prefix}.attn.to_out.0.lora_A.default.weight"] = torch.arange( + 6, dtype=torch.float32 + ).reshape(2, 3) + tensors[f"{prefix}.attn.to_out.0.lora_B.default.weight"] = torch.arange( + 4, dtype=torch.float32 + ).reshape(2, 2) + + # mlp.fc1 (SwiGLU): rank 2, in_dim 3, out_dim 4 (halves of size 2) + tensors[f"{prefix}.ff.net.0.proj.lora_A.default.weight"] = torch.ones( + (2, 3), dtype=torch.float32 + ) + # Val half: rows 0..1 (all 1s), Gate half: rows 2..3 (all 2s) + tensors[f"{prefix}.ff.net.0.proj.lora_B.default.weight"] = torch.cat( + [torch.full((2, 2), 1.0), torch.full((2, 2), 2.0)], dim=0 + ) + + # mlp.fc2: rank 2, in_dim 4, out_dim 3 + tensors[f"{prefix}.ff.net.2.lora_A.default.weight"] = torch.ones( + (2, 4), dtype=torch.float32 + ) + tensors[f"{prefix}.ff.net.2.lora_B.default.weight"] = torch.ones( + (3, 2), dtype=torch.float32 + ) + + return tensors + + +def test_full_rank_qkv_conversion_is_exact(converter, tmp_path): + source = tmp_path / "source_exact.safetensors" + output = tmp_path / "output_exact.safetensors" + + original = _make_source_tensors("token_refiner.refiner_blocks.0") + save_file(original, str(source)) + + report = converter.convert_minimax_h3_diffusers_lora(str(source), str(output)) + assert "MiniMax H3 Diffusers LoRA conversion complete" in report + assert "Fused QKV modules: 1" in report + + converted = load_file(str(output)) + + down = converted["diffusion_model.token_refiner.blocks.0.attn.qkv_proj.lora_A.weight"] + up = converted["diffusion_model.token_refiner.blocks.0.attn.qkv_proj.lora_B.weight"] + + # Ranks: 2 + 2 + 2 = 6, in_dim: 3, out_dim: 2 + 2 + 2 = 6 + assert down.shape == (6, 3) + assert up.shape == (6, 6) + + # Mathematical delta equality: up @ down must equal stacked (up_q @ down_q, up_k @ down_k, up_v @ down_v) + fused_delta = up @ down + expected_delta = torch.cat( + [ + original[f"token_refiner.refiner_blocks.0.attn.to_{p}.lora_B.default.weight"] + @ original[f"token_refiner.refiner_blocks.0.attn.to_{p}.lora_A.default.weight"] + for p in ("q", "k", "v") + ], + dim=0, + ) + assert torch.allclose(fused_delta, expected_delta, atol=1e-5, rtol=1e-5) + + +def test_swiglu_mlp_fc1_halves_swapped(converter, tmp_path): + source = tmp_path / "source_swiglu.safetensors" + output = tmp_path / "output_swiglu.safetensors" + + original = _make_source_tensors("transformer_blocks.0") + save_file(original, str(source)) + + converter.convert_minimax_h3_diffusers_lora(str(source), str(output)) + converted = load_file(str(output)) + + fc1_b = converted["diffusion_model.blocks.0.mlp.fc1.lora_B.weight"] + orig_b = original["transformer_blocks.0.ff.net.0.proj.lora_B.default.weight"] + + # In source: top half (rows 0..1) is 1.0, bottom half (rows 2..3) is 2.0 + # In converted: top half must be 2.0 (gate), bottom half must be 1.0 (value) + assert torch.equal(fc1_b[:2], orig_b[2:]) + assert torch.equal(fc1_b[2:], orig_b[:2]) + + +def test_rejects_incomplete_qkv_group(converter, tmp_path): + source = _make_source_tensors() + del source["transformer_blocks.0.attn.to_v.lora_B.default.weight"] + + path = tmp_path / "incomplete.safetensors" + save_file(source, str(path)) + + with pytest.raises(ValueError, match="Incomplete LoRA A/B pair"): + converter.convert_minimax_h3_diffusers_lora( + str(path), str(tmp_path / "unused.safetensors") + ) + + +def test_rejects_unsupported_module(converter, tmp_path): + source = {"unknown_layer.lora_A.weight": torch.ones(2, 2), "unknown_layer.lora_B.weight": torch.ones(2, 2)} + path = tmp_path / "unknown.safetensors" + save_file(source, str(path)) + + with pytest.raises(ValueError, match="Unsupported MiniMax H3 Diffusers module"): + converter.convert_minimax_h3_diffusers_lora( + str(path), str(tmp_path / "unused.safetensors") + ) + + +def test_node_schema_and_execution_contract(converter, monkeypatch, tmp_path): + schema = converter.MiniMaxH3DiffusersLoRAConvert.define_schema() + assert schema.node_id == "MiniMaxH3DiffusersLoRAConvert" + assert schema.is_output_node is True + + input_ids = [item.id for item in schema.inputs] + assert input_ids == ["lora_name", "output_filename"] + + output_displays = [item.display_name for item in schema.outputs] + assert output_displays == ["output_path", "conversion_report"] + + # Test execution through node + models_dir = tmp_path / "models" + loras_dir = models_dir / "loras" + loras_dir.mkdir(parents=True) + + source_file = loras_dir / "test_diffusers_lora.safetensors" + save_file(_make_source_tensors(), str(source_file)) + + monkeypatch.setattr(converter.folder_paths, "models_dir", str(models_dir)) + monkeypatch.setattr( + converter.folder_paths, + "get_full_path_or_raise", + lambda folder, name: str(loras_dir / name), + ) + + result = converter.MiniMaxH3DiffusersLoRAConvert.execute( + lora_name="test_diffusers_lora.safetensors", + output_filename="sub/converted_h3", + ) + + assert result.args[0] == "sub/converted_h3.safetensors" + assert "MiniMax H3 Diffusers LoRA conversion complete" in result.args[1] + assert (loras_dir / "sub" / "converted_h3.safetensors").is_file()