QrusherZA

/

H3_Turbo_ComfyUI

like 36

Hey, so this LoRA work is not done by me, i only uploaded my fixes i have done for comfyui for the original LoRA, the original poster has detailed this on his repository

https://huggingface.co/larryvrh/MiniMax-H3-Turbo-Lora

Yes, fully understand that. I mean, you've got code that is modifying Larry's tensors to produce ones usable in ComfyUI. Is that something You'd share?

yeah sure

from __future__ import annotations

import argparse
from pathlib import Path

import torch
from safetensors import safe_open
from safetensors.torch import save_file

LORA_A_SUFFIX = ".lora_A.weight"
LORA_B_SUFFIX = ".lora_B.weight"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Convert a raw MiniMax H3 Diffusers/PEFT LoRA into the "
            "generic key format expected by ComfyUI."
        )
    )

    parser.add_argument(
        "--input",
        required=True,
        type=Path,
        help="Input .safetensors LoRA",
    )

    parser.add_argument(
        "--output",
        required=True,
        type=Path,
        help="Output ComfyUI-compatible .safetensors LoRA",
    )

    parser.add_argument(
        "--alpha",
        type=float,
        default=None,
        help=(
            "Optional LoRA alpha. Use the lora_alpha value from "
            "adapter_config.json when available."
        ),
    )

    parser.add_argument(
        "--drop-adaln",
        action="store_true",
        help=(
            "Remove AdaLN tensors. This may be needed when applying a "
            "full-model LoRA to certain curve-form or quantized H3 checkpoints."
        ),
    )

    return parser.parse_args()


def convert_key(key: str) -> str:
    # Already compatible.
    if key.startswith("diffusion_model."):
        return key

    # The uploaded log shows raw keys beginning with blocks.*, final_layer.*,
    # and token_refiner.*, so they only require this wrapper prefix.
    return f"diffusion_model.{key}"


def load_safetensors(path: Path) -> tuple[dict[str, torch.Tensor], dict[str, str]]:
    tensors: dict[str, torch.Tensor] = {}

    with safe_open(str(path), framework="pt", device="cpu") as file:
        metadata = file.metadata() or {}

        for key in file.keys():
            tensors[key] = file.get_tensor(key)

    metadata = {str(key): str(value) for key, value in metadata.items()}
    return tensors, metadata


def validate_pairs(tensors: dict[str, torch.Tensor]) -> list[str]:
    problems: list[str] = []

    for key, tensor_a in tensors.items():
        if not key.endswith(LORA_A_SUFFIX):
            continue

        base = key[: -len(LORA_A_SUFFIX)]
        b_key = f"{base}{LORA_B_SUFFIX}"

        if b_key not in tensors:
            problems.append(f"Missing B tensor for: {key}")
            continue

        tensor_b = tensors[b_key]

        if tensor_a.ndim < 2 or tensor_b.ndim < 2:
            problems.append(
                f"Unexpected dimensions: {key}={tuple(tensor_a.shape)}, "
                f"{b_key}={tuple(tensor_b.shape)}"
            )
            continue

        # Diffusers convention:
        # A/down = [rank, input]
        # B/up   = [output, rank]
        if tensor_a.shape[0] != tensor_b.shape[1]:
            problems.append(
                f"Rank mismatch: {key}={tuple(tensor_a.shape)}, "
                f"{b_key}={tuple(tensor_b.shape)}"
            )

    return problems


def main() -> None:
    args = parse_args()

    if not args.input.is_file():
        raise FileNotFoundError(f"Input file does not exist: {args.input}")

    if args.input.resolve() == args.output.resolve():
        raise ValueError("Input and output paths must be different.")

    source, metadata = load_safetensors(args.input)

    converted: dict[str, torch.Tensor] = {}
    renamed = 0
    removed_adaln = 0

    for old_key, tensor in source.items():
        new_key = convert_key(old_key)

        if args.drop_adaln and ".adaln_proj." in new_key:
            removed_adaln += 1
            continue

        if new_key in converted:
            raise RuntimeError(f"Conversion produced a duplicate key: {new_key}")

        converted[new_key] = tensor

        if new_key != old_key:
            renamed += 1

    problems = validate_pairs(converted)

    if problems:
        print("\nValidation warnings:")
        for problem in problems[:30]:
            print(f"  - {problem}")

        if len(problems) > 30:
            print(f"  - ...and {len(problems) - 30} more")

    alpha_added = 0

    if args.alpha is not None:
        alpha_value = torch.tensor(args.alpha, dtype=torch.float32)

        for key in list(converted):
            if not key.endswith(LORA_A_SUFFIX):
                continue

            base = key[: -len(LORA_A_SUFFIX)]
            alpha_key = f"{base}.alpha"

            if alpha_key not in converted:
                converted[alpha_key] = alpha_value.clone()
                alpha_added += 1

    metadata["comfyui_conversion"] = "MiniMax H3 diffusion_model prefix added"
    metadata["comfyui_source_file"] = args.input.name

    args.output.parent.mkdir(parents=True, exist_ok=True)
    save_file(converted, str(args.output), metadata=metadata)

    print("\nConversion complete")
    print(f"Input:          {args.input}")
    print(f"Output:         {args.output}")
    print(f"Input tensors:  {len(source)}")
    print(f"Output tensors: {len(converted)}")
    print(f"Renamed:        {renamed}")
    print(f"AdaLN removed:  {removed_adaln}")
    print(f"Alpha added:    {alpha_added}")
    print(f"Pair warnings:  {len(problems)}")

    print("\nExample converted keys:")
    for key in list(converted)[:12]:
        print(f"  {key}")


if __name__ == "__main__":
    main()

but then you will also need a script to strip those adaln layers if you not using bf16

from pathlib import Path

from safetensors import safe_open
from safetensors.torch import save_file

INPUT_LORA = Path(
    r"C:\ComfyUI\ComfyUI_windows_portable\ComfyUI\models\loras"
    r"\minimax_h3_turbo_4step_ema_ckpt850_comfyui.safetensors"
)

OUTPUT_LORA = Path(
    r"C:\ComfyUI\ComfyUI_windows_portable\ComfyUI\models\loras"
    r"\minimax_h3_turbo_4step_ema_ckpt850_comfyui_fixed.safetensors"
)

# These belong to the full-width timestep/AdaLN path.
INCOMPATIBLE_PARTS = (
    ".adaln_proj.",
    ".time_embedder.",
    ".adaln_t_table",
)


def should_remove(key: str) -> bool:
    return any(part in key for part in INCOMPATIBLE_PARTS)


def main() -> None:
    if not INPUT_LORA.exists():
        raise FileNotFoundError(INPUT_LORA)

    kept = {}
    removed = []

    with safe_open(
        INPUT_LORA,
        framework="pt",
        device="cpu",
    ) as source:
        metadata = source.metadata() or {}

        for key in source.keys():
            if should_remove(key):
                removed.append(key)
                continue

            kept[key] = source.get_tensor(key)

    output_metadata = dict(metadata)
    output_metadata.update(
        {
            "h3_pruned_compatible": "true",
            "removed_full_width_adaln_entries": str(len(removed)),
            "source_lora": INPUT_LORA.name,
        }
    )

    save_file(
        kept,
        OUTPUT_LORA,
        metadata=output_metadata,
    )

    print(f"Input:   {INPUT_LORA}")
    print(f"Output:  {OUTPUT_LORA}")
    print(f"Kept:    {len(kept)} tensors")
    print(f"Removed: {len(removed)} tensors")

    categories = {
        "AdaLN": sum(".adaln_proj." in key for key in removed),
        "time_embedder": sum(".time_embedder." in key for key in removed),
        "table": sum(".adaln_t_table" in key for key in removed),
    }

    print("\nRemoved categories:")
    for name, count in categories.items():
        print(f"  {name}: {count}")


if __name__ == "__main__":
    main()