nodes/silveroxides-ComfyUI-ModelUtils
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add MiniMax H3 AdaLN unfold node

Browse Source
This commit is contained in:
silveroxides
2026-10-05 14:54:00 +02:00
parent f383a501a5
commit 2ee3f30c32
3 changed files with 378 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+2 -1
View File
@@ -61,7 +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.minimax_h3_fold import MiniMaxH3FoldAdaLN, MiniMaxH3UnfoldAdaLN
from .nodes.dtype_conversion import DiffusionModelDtypeConversion
from .nodes.layer_parameters import LayerParameterConfiguration
from .nodes.lora_model_analysis import LoRAOnModelAnalysis
@@ -124,6 +124,7 @@ class ModelUtilsExtension(ComfyExtension):
# Dtype Conversion & Folding
DiffusionModelDtypeConversion,
MiniMaxH3FoldAdaLN,
MiniMaxH3UnfoldAdaLN,
# Downloaders
CheckpointInfoMetaDownloader, DiffusionModelInfoMetaDownloader, LoRAInfoMetaDownloader, EmbeddingInfoMetaDownloader,
VAEInfoMetaDownloader, ControlNetInfoMetaDownloader, ManualPathInfoMetaDownloader,
nodes/minimax_h3_fold.py
+292
View File
@@ -393,3 +393,295 @@ class MiniMaxH3FoldAdaLN(io.ComfyNode):
process_device=process_device,
)
return io.NodeOutput(output_name, report)
def unfold_minimax_h3_diffusion_model(
model_a_path: str,
model_b_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 reconstruct unfolded MiniMax H3 AdaLN linears using a reference base model."""
if os.path.normcase(os.path.abspath(model_a_path)) == os.path.normcase(
os.path.abspath(output_path)
):
raise ValueError("Output path must differ from the folded model (Model A) path.")
if os.path.normcase(os.path.abspath(model_b_path)) == os.path.normcase(
os.path.abspath(output_path)
):
raise ValueError("Output path must differ from the base model (Model B) 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_a = MemoryEfficientSafeOpen(model_a_path, low_memory=True)
handler_b = MemoryEfficientSafeOpen(model_b_path, low_memory=True)
try:
inspect_low_bit_input(handler_a, f"Model A ({model_a_path})", "MiniMax H3 Unfold AdaLN")
inspect_low_bit_input(handler_b, f"Model B ({model_b_path})", "MiniMax H3 Unfold AdaLN")
all_keys_a = list(handler_a.keys())
all_keys_b = list(handler_b.keys())
prefix_b, time_keys_b = _detect_time_embedder_keys(all_keys_b)
prefix_a = "diffusion_model." if any(k.startswith("diffusion_model.") for k in all_keys_a) else ""
# Step 1: Compute basis and table from base model (Model B)
tstream_b = handler_b.async_stream(
list(time_keys_b.values()),
batch_size=1,
prefetch_batches=1,
pin_memory=False,
)
time_tensors_b = {}
try:
for batch in tstream_b:
for k, t in batch:
time_tensors_b[k] = t
_, basis, error = _fold_curve_basis(
time_tensors_b[time_keys_b["w1"]],
time_tensors_b[time_keys_b["b1"]],
time_tensors_b[time_keys_b["w2"]],
time_tensors_b[time_keys_b["b2"]],
curve_grid,
curve_rank,
device,
)
finally:
for k in time_keys_b.values():
handler_b.mark_processed(k)
tstream_b.close()
basis_device = basis.to(device)
metadata = dict(handler_a.metadata() or {})
if "config" in metadata:
try:
cfg = json.loads(metadata["config"])
if isinstance(cfg, dict):
tf = cfg.setdefault("transformer", {})
tf.pop("adaln_curve_grid", None)
tf["time_embed_dim"] = basis.shape[0]
metadata["config"] = json.dumps(cfg, separators=(",", ":"))
except Exception:
pass
metadata["modelutils_adaln_unfold"] = f"reconstructed_rank{basis.shape[1]}_with_base"
table_key_a = f"{prefix_a}adaln_t_table"
counts = {
"unfolded": 0,
"preserved": 0,
"discarded": 0,
"time_embedder_restored": 0,
}
with atomic_uel_writer(output_path, metadata) as writer:
# Emit the 4 time_embedder tensors restored from Model B
for _, base_k in time_keys_b.items():
target_k = f"{prefix_a}{base_k.removeprefix(prefix_b)}"
if not _matches_any_pattern(target_k, compiled_discards, glob_patterns):
writer.write_batch([(target_k, time_tensors_b[base_k])])
counts["time_embedder_restored"] += 1
time_tensors_b.clear()
remaining_keys_a = [k for k in all_keys_a if k != table_key_a]
progress = comfy.utils.ProgressBar(len(remaining_keys_a))
for key in remaining_keys_a:
if _matches_any_pattern(key, compiled_discards, glob_patterns):
counts["discarded"] += 1
handler_a.mark_processed(key)
progress.update(1)
continue
astream = handler_a.async_stream(
[key], batch_size=1, prefetch_batches=1, pin_memory=False
)
try:
tensor_a = next(astream)[0][1]
if _is_adaln(key):
matched = _matches_any_pattern(key, compiled_excludes, glob_patterns)
should_unfold = matched if include_mode else not matched
if should_unfold and key.endswith(".weight"):
base_key = f"{prefix_b}{key.removeprefix(prefix_a)}"
if base_key not in all_keys_b:
raise KeyError(
f"Matching base layer '{base_key}' not found in Model B."
)
bstream = handler_b.async_stream(
[base_key], batch_size=1, prefetch_batches=1, pin_memory=False
)
try:
base_w = next(bstream)[0][1]
if base_w.ndim != 2 or base_w.shape[1] != basis.shape[0]:
raise ValueError(
f"{base_key} in Model B has shape {tuple(base_w.shape)}, expected [*, {basis.shape[0]}]"
)
bw_f32 = base_w.to(device, torch.float32)
w_folded = tensor_a.to(device, torch.float32)
w_diff = w_folded - bw_f32 @ basis_device
w_unfolded = (bw_f32 + w_diff @ basis_device.T).to(base_w.dtype).cpu().contiguous()
writer.write_batch([(key, w_unfolded)])
counts["unfolded"] += 1
finally:
handler_b.mark_processed(base_key)
bstream.close()
else:
writer.write_batch([(key, tensor_a.contiguous())])
counts["preserved"] += 1
else:
writer.write_batch([(key, tensor_a.contiguous())])
counts["preserved"] += 1
finally:
handler_a.mark_processed(key)
astream.close()
progress.update(1)
finally:
handler_a.__exit__(None, None, None)
handler_b.__exit__(None, None, None)
cleanup_after_operation()
report_lines = [
"MiniMax H3 AdaLN Unfold complete",
f"Curve rank: {curve_rank}, grid points: {curve_grid}",
f"Base basis reconstruction error: {error:.4e}",
f"AdaLN layers unfolded: {counts['unfolded']}",
f"Time embedder tensors restored: {counts['time_embedder_restored']}",
f"Tensors preserved: {counts['preserved']}",
f"Tensors discarded: {counts['discarded']}",
f"Output size: {os.path.getsize(output_path)} bytes",
f"Output: {output_path}",
]
return "\n".join(report_lines)
class MiniMaxH3UnfoldAdaLN(io.ComfyNode):
"""Reconstruct full-width AdaLN linears [*, 2688] from a folded model and reference base model."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="MiniMaxH3UnfoldAdaLN",
display_name="MiniMax H3 Unfold AdaLN",
category="ModelUtils/DiffusionModels",
description="Reconstructs full-width MiniMax H3 AdaLN linears [*, 2688] from a folded model (Model A) using a reference non-folded base model (Model B) to supply the time embedder and orthogonal dimensions.",
inputs=[
io.Combo.Input(
"model_a",
options=folder_paths.get_filename_list("diffusion_models"),
tooltip="Select the folded MiniMax H3 model (Model A) to unfold.",
),
io.Combo.Input(
"model_b",
options=folder_paths.get_filename_list("diffusion_models"),
tooltip="Select the reference original/unfolded MiniMax H3 base model (Model B) to supply the time embedder and missing dimensions.",
),
io.String.Input(
"output_filename",
default="minimax_h3_unfolded",
tooltip="Name for the reconstructed unfolded model file, saved in your diffusion_models folder.",
),
io.String.Input(
"discard_patterns",
default="",
multiline=True,
tooltip="Layers to completely remove from the final file (one name or pattern per line). Useful to strip parts you do not need.",
),
io.Boolean.Input(
"glob_patterns",
default=False,
tooltip="Enable simple wildcard matching (* to match anything). If unchecked, uses regular expressions.",
),
io.Int.Input(
"curve_rank",
default=8,
min=1,
max=64,
step=1,
tooltip="Basis rank used when the model was folded (default 8, matching standard ComfyUI H3 checkpoints).",
),
io.Int.Input(
"curve_grid",
default=1025,
min=65,
max=4097,
step=64,
tooltip="Sampling resolution used when the model was folded (default 1025, matching standard ComfyUI H3 checkpoints).",
),
io.Combo.Input(
"process_device",
options=["cuda", "cpu"],
default="cuda",
tooltip="Hardware device used to perform the math. Use cuda (GPU) for speed, or cpu if you run low on VRAM.",
),
io.String.Input(
"exclude_patterns",
default="",
multiline=True,
tooltip="Specific layers to skip unfolding, keeping them at their folded size (one name or pattern per line).",
),
io.Boolean.Input(
"include_mode",
default=False,
tooltip="Invert the exclude filter: only unfold layers matching the pattern, and keep everything else untouched.",
),
],
outputs=[
io.AnyType.Output(display_name="output_path"),
io.String.Output(display_name="report"),
],
is_output_node=True,
)
@classmethod
def execute(
cls,
model_a: str,
model_b: 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:
model_a_path = folder_paths.get_full_path_or_raise(
"diffusion_models", model_a
)
model_b_path = folder_paths.get_full_path_or_raise(
"diffusion_models", model_b
)
output_path, output_name = canonical_model_artifact_path(
"diffusion_models", output_filename
)
report = unfold_minimax_h3_diffusion_model(
model_a_path,
model_b_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)
tests/test_minimax_h3_fold.py
+84
View File
@@ -160,3 +160,87 @@ def test_node_schema_and_execution_contract(folder_module, monkeypatch, tmp_path
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()
def test_unfold_h3_diffusion_model_reconstructs_with_base(folder_module, tmp_path):
base_file = tmp_path / "model_base.safetensors"
folded_file = tmp_path / "model_folded.safetensors"
unfolded_file = tmp_path / "model_unfolded.safetensors"
original_base = _make_h3_checkpoint()
save_file(original_base, str(base_file))
# Fold base model to create folded model A
folder_module.fold_minimax_h3_diffusion_model(
str(base_file), str(folded_file), process_device="cpu"
)
# Unfold model A using base model B
report = folder_module.unfold_minimax_h3_diffusion_model(
str(folded_file), str(base_file), str(unfolded_file), process_device="cpu"
)
assert "MiniMax H3 AdaLN Unfold complete" in report
assert "AdaLN layers unfolded: 2" in report
reconstructed = load_file(str(unfolded_file))
# adaln_t_table should be dropped in unfolded output
assert "adaln_t_table" not in reconstructed
# time_embedder restored from base model B
assert "time_embedder.proj_in.weight" in reconstructed
assert "time_embedder.proj_out.weight" in reconstructed
assert torch.equal(reconstructed["time_embedder.proj_in.weight"], original_base["time_embedder.proj_in.weight"])
assert torch.equal(reconstructed["time_embedder.proj_out.weight"], original_base["time_embedder.proj_out.weight"])
# AdaLN linear weights reconstructed to full width [128, 2688]
assert reconstructed["blocks.0.adaln_proj.linear.weight"].shape == (128, 2688)
assert reconstructed["blocks.1.adaln_proj.linear.weight"].shape == (128, 2688)
# Other tensors preserved
assert torch.equal(reconstructed["blocks.0.attn.qkv_proj.weight"], original_base["blocks.0.attn.qkv_proj.weight"])
def test_unfold_h3_node_schema_and_execution(folder_module, monkeypatch, tmp_path):
schema = folder_module.MiniMaxH3UnfoldAdaLN.define_schema()
assert schema.node_id == "MiniMaxH3UnfoldAdaLN"
assert schema.is_output_node is True
input_ids = [item.id for item in schema.inputs]
assert input_ids == [
"model_a", "model_b", "output_filename", "discard_patterns",
"glob_patterns", "curve_rank", "curve_grid", "process_device",
"exclude_patterns", "include_mode",
]
assert input_ids[-1] == "include_mode"
for item in schema.inputs:
assert hasattr(item, "tooltip") and item.tooltip.strip()
models_dir = tmp_path / "models"
diff_dir = models_dir / "diffusion_models"
diff_dir.mkdir(parents=True)
base_file = diff_dir / "base.safetensors"
folded_file = diff_dir / "folded.safetensors"
save_file(_make_h3_checkpoint(), str(base_file))
folder_module.fold_minimax_h3_diffusion_model(str(base_file), str(folded_file), process_device="cpu")
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.MiniMaxH3UnfoldAdaLN.execute(
model_a="folded.safetensors",
model_b="base.safetensors",
output_filename="sub/unfolded_h3",
process_device="cpu",
)
assert result.args[0] == "sub/unfolded_h3.safetensors"
assert "MiniMax H3 AdaLN Unfold complete" in result.args[1]
assert (diff_dir / "sub" / "unfolded_h3.safetensors").is_file()