From ae2858bf6b5cedc8d49192c2839413fcbcf35f2d Mon Sep 17 00:00:00 2001 From: silveroxides Date: Sat, 3 Oct 2026 02:07:32 +0200 Subject: [PATCH] Add MiniMax H3 AdaLN fold node for diffusion models --- __init__.py | 4 +- nodes/minimax_h3_fold.py | 386 ++++++++++++++++++++++++++++++++++ tests/test_artifact_paths.py | 1 + tests/test_groups.toml | 5 + tests/test_minimax_h3_fold.py | 162 ++++++++++++++ tests/test_schema_tooltips.py | 1 + 6 files changed, 558 insertions(+), 1 deletion(-) create mode 100644 nodes/minimax_h3_fold.py create mode 100644 tests/test_minimax_h3_fold.py diff --git a/__init__.py b/__init__.py index 0e877bc..f2ba712 100644 --- a/__init__.py +++ b/__init__.py @@ -61,6 +61,7 @@ from .nodes.model_info_nodes import ( from .nodes.lora_rename import AnimaLoraRename from .nodes.minimax_h3_lora_convert import MiniMaxH3DiffusersLoRAConvert +from .nodes.minimax_h3_fold import MiniMaxH3FoldAdaLN from .nodes.dtype_conversion import DiffusionModelDtypeConversion from .nodes.layer_parameters import LayerParameterConfiguration from .nodes.lora_model_analysis import LoRAOnModelAnalysis @@ -120,8 +121,9 @@ class ModelUtilsExtension(ComfyExtension): # LoRA Utilities AnimaLoraRename, MiniMaxH3DiffusersLoRAConvert, - # Dtype Conversion + # Dtype Conversion & Folding DiffusionModelDtypeConversion, + MiniMaxH3FoldAdaLN, # Downloaders CheckpointInfoMetaDownloader, DiffusionModelInfoMetaDownloader, LoRAInfoMetaDownloader, EmbeddingInfoMetaDownloader, VAEInfoMetaDownloader, ControlNetInfoMetaDownloader, ManualPathInfoMetaDownloader, diff --git a/nodes/minimax_h3_fold.py b/nodes/minimax_h3_fold.py new file mode 100644 index 0000000..88f920f --- /dev/null +++ b/nodes/minimax_h3_fold.py @@ -0,0 +1,386 @@ +from __future__ import annotations + +import fnmatch +import json +import math +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 .device_utils import cleanup_after_operation +from .quantization_guard import inspect_low_bit_input +from .uel_io import atomic_uel_writer + + +def _compile_patterns(pattern_string: str, glob_mode: bool = False) -> list: + """Compile newline or whitespace separated patterns.""" + if not pattern_string or not pattern_string.strip(): + return [] + patterns = [] + for pattern in pattern_string.split(): + if not pattern: + continue + if glob_mode: + patterns.append(pattern) + else: + try: + patterns.append(re.compile(pattern)) + except re.error as e: + raise ValueError(f"Invalid regex pattern '{pattern}': {e}") from e + return patterns + + +def _matches_any_pattern(key: str, patterns: list, glob_mode: bool = False) -> bool: + if not patterns: + return False + if glob_mode: + return any(fnmatch.fnmatch(key, f"*{p}*") for p in patterns) + return any(p.search(key) is not None for p in patterns) + + +def time_embedding_curve( + w1: torch.Tensor, + b1: torch.Tensor, + w2: torch.Tensor, + b2: torch.Tensor, + grid_size: int, + device: torch.device, +) -> torch.Tensor: + """Sample silu(time_embedder(t)) on the curve grid.""" + t = torch.linspace(0.0, 1.0, grid_size, dtype=torch.float32, device=device) + half = w1.shape[1] // 2 + freqs = torch.exp( + -math.log(10000.0) + * torch.arange(half, dtype=torch.float32, device=device) + / half + ) + args = t[:, None] * freqs[None] + emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + w1_d, b1_d, w2_d, b2_d = (x.to(device, torch.float32) for x in (w1, b1, w2, b2)) + hidden = torch.nn.functional.silu(emb @ w1_d.T + b1_d) + return torch.nn.functional.silu(hidden @ w2_d.T + b2_d) + + +def _fold_curve_basis( + w1: torch.Tensor, + b1: torch.Tensor, + w2: torch.Tensor, + b2: torch.Tensor, + grid_size: int, + rank: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, float]: + """Calculate SVD curve basis and adaln_t_table.""" + compute_device = device + try: + curve = time_embedding_curve(w1, b1, w2, b2, grid_size, compute_device) + u, s, vh = torch.linalg.svd(curve, full_matrices=False) + except (torch.cuda.OutOfMemoryError, RuntimeError): + cleanup_after_operation() + compute_device = torch.device("cpu") + curve = time_embedding_curve(w1, b1, w2, b2, grid_size, compute_device) + u, s, vh = torch.linalg.svd(curve, full_matrices=False) + + target_rank = min(rank, vh.shape[0]) + table = u[:, :target_rank] * s[:target_rank] + basis = vh[:target_rank].T + error = ((table @ basis.T - curve).norm() / curve.norm()).item() + return table.float().cpu().contiguous(), basis.float().cpu().contiguous(), error + + +def _detect_time_embedder_keys(all_keys: list[str]) -> tuple[str, dict[str, str]]: + for prefix in ("", "diffusion_model.", "model.diffusion_model."): + candidates = { + "w1": f"{prefix}time_embedder.proj_in.weight", + "b1": f"{prefix}time_embedder.proj_in.bias", + "w2": f"{prefix}time_embedder.proj_out.weight", + "b2": f"{prefix}time_embedder.proj_out.bias", + } + if all(k in all_keys for k in candidates.values()): + return prefix, candidates + raise KeyError( + "Model is missing required time_embedder projection tensors (proj_in/proj_out)." + ) + + +def _is_adaln(key: str) -> bool: + return ".adaln_proj.linear." in key or key.startswith("adaln_proj.linear.") + + +def fold_minimax_h3_diffusion_model( + input_path: str, + output_path: str, + *, + exclude_patterns: str = "", + include_mode: bool = False, + discard_patterns: str = "", + glob_patterns: bool = False, + curve_rank: int = 8, + curve_grid: int = 1025, + process_device: str = "cuda", +) -> str: + """Stream and fold MiniMax H3 AdaLN linears in a diffusion model.""" + 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 model path.") + + device = torch.device( + process_device if torch.cuda.is_available() and process_device == "cuda" else "cpu" + ) + + compiled_excludes = _compile_patterns(exclude_patterns, glob_patterns) + compiled_discards = _compile_patterns(discard_patterns, glob_patterns) + + handler = MemoryEfficientSafeOpen(input_path, low_memory=True) + try: + inspect_low_bit_input(handler, f"Model ({input_path})", "MiniMax H3 Fold AdaLN") + all_keys = list(handler.keys()) + prefix, time_keys = _detect_time_embedder_keys(all_keys) + + # Step 1: Load time embedder and calculate basis + stream = handler.async_stream( + list(time_keys.values()), + batch_size=1, + prefetch_batches=1, + pin_memory=False, + ) + time_tensors = {} + try: + for batch in stream: + for k, t in batch: + time_tensors[k] = t + table, basis, error = _fold_curve_basis( + time_tensors[time_keys["w1"]], + time_tensors[time_keys["b1"]], + time_tensors[time_keys["w2"]], + time_tensors[time_keys["b2"]], + curve_grid, + curve_rank, + device, + ) + finally: + time_tensors.clear() + for k in time_keys.values(): + handler.mark_processed(k) + stream.close() + + # Determine which AdaLN layers are folded vs kept full width + any_adaln_kept_unfolded = False + for k in all_keys: + if _is_adaln(k) and not _matches_any_pattern(k, compiled_discards, glob_patterns): + matched = _matches_any_pattern(k, compiled_excludes, glob_patterns) + should_fold = matched if include_mode else not matched + if not should_fold: + any_adaln_kept_unfolded = True + break + + metadata = dict(handler.metadata() or {}) + if "config" in metadata: + try: + cfg = json.loads(metadata["config"]) + if isinstance(cfg, dict): + tf = cfg.setdefault("transformer", {}) + tf["adaln_curve_grid"] = table.shape[0] + tf["time_embed_dim"] = table.shape[1] + metadata["config"] = json.dumps(cfg, separators=(",", ":")) + except Exception: + pass + metadata["modelutils_adaln_fold"] = f"rank{table.shape[1]}_curve_grid_{table.shape[0]}" + + table_key = f"{prefix}adaln_t_table" + time_embedder_prefix = f"{prefix}time_embedder." + + counts = { + "folded": 0, + "preserved": 0, + "discarded": 0, + "time_embedder_dropped": 0, + } + + basis_device = basis.to(device) + + with atomic_uel_writer(output_path, metadata) as writer: + # Emit adaln_t_table + writer.write_batch([(table_key, table)]) + + # Filter remaining keys + remaining_keys = [k for k in all_keys if k != table_key] + progress = comfy.utils.ProgressBar(len(remaining_keys)) + + for key in remaining_keys: + if _matches_any_pattern(key, compiled_discards, glob_patterns): + counts["discarded"] += 1 + handler.mark_processed(key) + progress.update(1) + continue + + if key.startswith(time_embedder_prefix) and not any_adaln_kept_unfolded: + counts["time_embedder_dropped"] += 1 + handler.mark_processed(key) + progress.update(1) + continue + + # Stream this tensor + tstream = handler.async_stream( + [key], batch_size=1, prefetch_batches=1, pin_memory=False + ) + try: + batch = next(tstream) + tensor = batch[0][1] + + if _is_adaln(key): + matched = _matches_any_pattern(key, compiled_excludes, glob_patterns) + should_fold = matched if include_mode else not matched + if should_fold: + if key.endswith(".weight"): + if tensor.ndim != 2 or tensor.shape[1] != basis.shape[0]: + raise ValueError( + f"{key} has shape {tuple(tensor.shape)}, expected [*, {basis.shape[0]}]" + ) + folded_w = ( + tensor.to(device, torch.float32) @ basis_device + ).to(torch.float16).cpu().contiguous() + writer.write_batch([(key, folded_w)]) + else: + writer.write_batch([(key, tensor.to(torch.float16).contiguous())]) + counts["folded"] += 1 + else: + writer.write_batch([(key, tensor.contiguous())]) + counts["preserved"] += 1 + else: + writer.write_batch([(key, tensor.contiguous())]) + counts["preserved"] += 1 + finally: + handler.mark_processed(key) + tstream.close() + + progress.update(1) + finally: + handler.__exit__(None, None, None) + cleanup_after_operation() + + report_lines = [ + "MiniMax H3 AdaLN Fold complete", + f"Curve rank: {curve_rank}, grid points: {curve_grid}", + f"Reconstruction relative error: {error:.4e}", + f"AdaLN layers folded: {counts['folded']}", + f"Tensors preserved: {counts['preserved']}", + f"Tensors discarded: {counts['discarded']}", + f"Time embedder dropped: {counts['time_embedder_dropped']}", + f"Output size: {os.path.getsize(output_path)} bytes", + f"Output: {output_path}", + ] + return "\n".join(report_lines) + + +class MiniMaxH3FoldAdaLN(io.ComfyNode): + """Fold MiniMax H3 full-width AdaLN linears [*, 2688] onto rank-8 time-embedding curve basis.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="MiniMaxH3FoldAdaLN", + display_name="MiniMax H3 Fold AdaLN", + category="ModelUtils/DiffusionModels", + description="Folds MiniMax H3 full-width AdaLN linear layers onto a low-rank basis of the silu(time_embedder) curve, replacing the time embedder with adaln_t_table.", + inputs=[ + io.Combo.Input( + "model_name", + options=folder_paths.get_filename_list("diffusion_models"), + tooltip="MiniMax H3 diffusion model with full-width AdaLN linears [*, 2688] to fold.", + ), + io.String.Input( + "output_filename", + default="minimax_h3_folded", + tooltip="Output filename without extension, written under ComfyUI's diffusion_models directory.", + ), + io.String.Input( + "discard_patterns", + default="", + multiline=True, + tooltip="Newline-separated regex or glob patterns for tensors to omit entirely from the saved model (e.g. drop specific transformer blocks or refiner).", + ), + io.Boolean.Input( + "glob_patterns", + default=False, + tooltip="When True, filter patterns use shell globs (* matches any sequence, dots are literal). When False (default), patterns are Python regex matched as substrings.", + ), + io.Int.Input( + "curve_rank", + default=8, + min=1, + max=64, + step=1, + tooltip="Basis rank for the silu(time_embedder) curve (default 8, matching ComfyUI H3 checkpoints).", + ), + io.Int.Input( + "curve_grid", + default=1025, + min=65, + max=4097, + step=64, + tooltip="Number of interpolation grid points sampled for timestep t in [0, 1] (default 1025).", + ), + io.Combo.Input( + "process_device", + options=["cuda", "cpu"], + default="cuda", + tooltip="Device used for SVD curve decomposition and matrix projection; CUDA OOM automatically retries on CPU.", + ), + io.String.Input( + "exclude_patterns", + default="", + multiline=True, + tooltip="Newline-separated regex or glob patterns. Matching AdaLN layers remain unfolded at full width [*, 2688].", + ), + io.Boolean.Input( + "include_mode", + default=False, + tooltip="When True, use exclude_patterns as a whitelist: only matching AdaLN layers are folded; all others remain full width.", + ), + ], + outputs=[ + io.AnyType.Output(display_name="output_path"), + io.String.Output(display_name="report"), + ], + is_output_node=True, + ) + + @classmethod + def execute( + cls, + model_name: str, + output_filename: str, + discard_patterns: str = "", + glob_patterns: bool = False, + curve_rank: int = 8, + curve_grid: int = 1025, + process_device: str = "cuda", + exclude_patterns: str = "", + include_mode: bool = False, + ) -> io.NodeOutput: + input_path = folder_paths.get_full_path_or_raise( + "diffusion_models", model_name + ) + output_path, output_name = canonical_model_artifact_path( + "diffusion_models", output_filename + ) + report = fold_minimax_h3_diffusion_model( + input_path, + output_path, + exclude_patterns=exclude_patterns, + include_mode=include_mode, + discard_patterns=discard_patterns, + glob_patterns=glob_patterns, + curve_rank=curve_rank, + curve_grid=curve_grid, + process_device=process_device, + ) + return io.NodeOutput(output_name, report) diff --git a/tests/test_artifact_paths.py b/tests/test_artifact_paths.py index e17754d..00bf701 100644 --- a/tests/test_artifact_paths.py +++ b/tests/test_artifact_paths.py @@ -25,6 +25,7 @@ MIGRATED_SAVE_MODULES = ( "prunekeys.py", "lora_rename.py", "minimax_h3_lora_convert.py", + "minimax_h3_fold.py", ) diff --git a/tests/test_groups.toml b/tests/test_groups.toml index ba3198c..72f4fbb 100644 --- a/tests/test_groups.toml +++ b/tests/test_groups.toml @@ -12,6 +12,7 @@ paths = [ "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/minimax_h3_fold.py", "nodes/prunekeys.py", "nodes/renamekeys.py", "nodes/text_encoder_extract.py", ] python_tests = ["tests/test_artifact_paths.py"] @@ -72,3 +73,7 @@ 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"] + +[groups.minimax_h3_fold] +paths = ["nodes/minimax_h3_fold.py"] +python_tests = ["tests/test_minimax_h3_fold.py"] diff --git a/tests/test_minimax_h3_fold.py b/tests/test_minimax_h3_fold.py new file mode 100644 index 0000000..9f7da15 --- /dev/null +++ b/tests/test_minimax_h3_fold.py @@ -0,0 +1,162 @@ +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_fold_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 folder_module(): + return _load_module("nodes.minimax_h3_fold") + + +def _make_h3_checkpoint(prefix: str = "") -> dict[str, torch.Tensor]: + return { + f"{prefix}time_embedder.proj_in.weight": torch.eye(16, dtype=torch.float32), + f"{prefix}time_embedder.proj_in.bias": torch.zeros(16, dtype=torch.float32), + f"{prefix}time_embedder.proj_out.weight": torch.ones((2688, 16), dtype=torch.float32) * 0.01, + f"{prefix}time_embedder.proj_out.bias": torch.zeros(2688, dtype=torch.float32), + f"{prefix}blocks.0.adaln_proj.linear.weight": torch.ones((128, 2688), dtype=torch.float32), + f"{prefix}blocks.0.adaln_proj.linear.bias": torch.zeros(128, dtype=torch.float32), + f"{prefix}blocks.1.adaln_proj.linear.weight": torch.ones((128, 2688), dtype=torch.float32) * 2.0, + f"{prefix}blocks.1.adaln_proj.linear.bias": torch.zeros(128, dtype=torch.float32), + f"{prefix}blocks.0.attn.qkv_proj.weight": torch.eye(64, dtype=torch.float32), + } + + +def test_fold_h3_diffusion_model_exact_basis(folder_module, tmp_path): + source = tmp_path / "model_full.safetensors" + output = tmp_path / "model_folded.safetensors" + + original = _make_h3_checkpoint() + save_file(original, str(source)) + + report = folder_module.fold_minimax_h3_diffusion_model( + str(source), str(output), process_device="cpu" + ) + assert "MiniMax H3 AdaLN Fold complete" in report + assert "AdaLN layers folded: 4" in report + + converted = load_file(str(output)) + + assert "adaln_t_table" in converted + assert converted["adaln_t_table"].shape == (1025, 8) + assert converted["adaln_t_table"].dtype == torch.float32 + + # Folded AdaLN weight and bias + assert converted["blocks.0.adaln_proj.linear.weight"].shape == (128, 8) + assert converted["blocks.0.adaln_proj.linear.weight"].dtype == torch.float16 + assert converted["blocks.0.adaln_proj.linear.bias"].dtype == torch.float16 + + # time_embedder should be dropped + assert "time_embedder.proj_in.weight" not in converted + assert "time_embedder.proj_out.weight" not in converted + + # Unrelated layers preserved + assert "blocks.0.attn.qkv_proj.weight" in converted + assert torch.equal(converted["blocks.0.attn.qkv_proj.weight"], original["blocks.0.attn.qkv_proj.weight"]) + + +def test_fold_h3_discard_patterns(folder_module, tmp_path): + source = tmp_path / "model_discard.safetensors" + output = tmp_path / "model_discarded.safetensors" + + original = _make_h3_checkpoint() + save_file(original, str(source)) + + folder_module.fold_minimax_h3_diffusion_model( + str(source), + str(output), + discard_patterns="blocks.0.attn", + process_device="cpu", + ) + converted = load_file(str(output)) + + assert "blocks.0.attn.qkv_proj.weight" not in converted + assert "blocks.0.adaln_proj.linear.weight" in converted + + +def test_fold_h3_exclude_patterns_preserves_unfolded_and_retains_time_embedder(folder_module, tmp_path): + source = tmp_path / "model_partial.safetensors" + output = tmp_path / "model_partial_folded.safetensors" + + original = _make_h3_checkpoint() + save_file(original, str(source)) + + folder_module.fold_minimax_h3_diffusion_model( + str(source), + str(output), + exclude_patterns="blocks.1", + include_mode=False, + process_device="cpu", + ) + converted = load_file(str(output)) + + # Block 0 folded + assert converted["blocks.0.adaln_proj.linear.weight"].shape == (128, 8) + # Block 1 kept at full width + assert converted["blocks.1.adaln_proj.linear.weight"].shape == (128, 2688) + + # time_embedder MUST be retained because block 1 needs it + assert "time_embedder.proj_in.weight" in converted + assert "time_embedder.proj_out.weight" in converted + + +def test_node_schema_and_execution_contract(folder_module, monkeypatch, tmp_path): + schema = folder_module.MiniMaxH3FoldAdaLN.define_schema() + assert schema.node_id == "MiniMaxH3FoldAdaLN" + assert schema.is_output_node is True + + input_ids = [item.id for item in schema.inputs] + assert "model_name" in input_ids + assert "output_filename" in input_ids + assert "exclude_patterns" in input_ids + assert "include_mode" in input_ids + assert "discard_patterns" in input_ids + assert "glob_patterns" in input_ids + + # Every input must have a non-empty static tooltip + for item in schema.inputs: + assert hasattr(item, "tooltip") and item.tooltip.strip() + + output_displays = [item.display_name for item in schema.outputs] + assert output_displays == ["output_path", "report"] + + # Test execution + models_dir = tmp_path / "models" + diff_dir = models_dir / "diffusion_models" + diff_dir.mkdir(parents=True) + + source_file = diff_dir / "test_h3_model.safetensors" + save_file(_make_h3_checkpoint("diffusion_model."), str(source_file)) + + monkeypatch.setattr(folder_module.folder_paths, "models_dir", str(models_dir)) + monkeypatch.setattr( + folder_module.folder_paths, + "get_full_path_or_raise", + lambda folder, name: str(diff_dir / name), + ) + + result = folder_module.MiniMaxH3FoldAdaLN.execute( + model_name="test_h3_model.safetensors", + output_filename="sub/folded_h3", + process_device="cpu", + ) + + assert result.args[0] == "sub/folded_h3.safetensors" + assert "MiniMax H3 AdaLN Fold complete" in result.args[1] + assert (diff_dir / "sub" / "folded_h3.safetensors").is_file() diff --git a/tests/test_schema_tooltips.py b/tests/test_schema_tooltips.py index ab0f113..66017a5 100644 --- a/tests/test_schema_tooltips.py +++ b/tests/test_schema_tooltips.py @@ -54,6 +54,7 @@ FILTER_MODULES = [ "merger", "consensus_merger", "cwb_delta_lora_merger", "lodestone_merger", "lora_resize", "dtype_conversion", "model_analysis", "lora_extract_svd", "dora_extract_wd", "dora_learned_wd", "text_encoder_extract", + "minimax_h3_fold", ]