CornLogic

/

10EROS_1.4_Int8_ConvRot

like 16

Hi is it possible to share your script you used to merge the DMD lora onto the diff model. I am trying to write one but it's not working out.

this is what i have so far, and it's breaking silently, want to pull my hair out 😞

import torch
from safetensors import safe_open
from safetensors.torch import load_file, save_file
from tqdm import tqdm
import time

BASE_MODEL = r"C:\ComfyUI\ComfyUI_windows_portable\ComfyUI\models\checkpoints\sulphur_dev_bf16.safetensors"
LORA = r"C:\ComfyUI\ComfyUI_windows_portable\ComfyUI\models\loras\LTX2.3_DMD_reshaped_r256.safetensors"
OUTPUT = r"C:\ComfyUI\ComfyUI_windows_portable\ComfyUI\models\checkpoints\sulphur_dev_bf16_DMD.safetensors"
STRENGTH = 1.0

# Preserve base metadata (LTX-2.3 carries audio VAE config here — dropping it
# breaks audio on the merged checkpoint).
print("Reading base metadata...")
with safe_open(BASE_MODEL, framework="pt", device="cpu") as f:
    meta = f.metadata() or {}

print("Loading base model...")
base = load_file(BASE_MODEL)

print("Loading LoRA...")
lora = load_file(LORA)

merged = dict(base)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print(f"Using device: {device}")

# Find every lora_down tensor
down_keys = [
    k
    for k in lora.keys()
    if (k.endswith(".lora_A.weight") or k.endswith(".lora_down.weight"))
]

applied = 0
skipped_missing = 0
skipped_shape = 0
skipped_nd = 0

with torch.no_grad():
    print("Building lookup ...")
    lookup = {}

    for k in merged:
        lookup[k.split("model.")[-1]] = k

    for down_key in tqdm(down_keys):
        try:
            up_key = down_key.replace(".lora_A.weight", ".lora_B.weight")
            alpha_key = down_key.replace(".lora_A.weight", ".alpha")

            suffix = down_key.replace(".lora_A.weight", ".weight")

            suffix = suffix.split("diffusion_model.")[-1]

            base_key = lookup.get(suffix)

            if base_key not in merged:
                skipped_missing += 1
                continue

            if up_key not in lora:
                print(f"Missing lora_up for {down_key}")
                skipped_missing += 1
                continue

            W = merged[base_key]

            # Only 2D Linear weights are LoRA-targetable. Anything else (modulation
            # tables, norms) must be left alone.
            if W.dim() != 2:
                print(f"Non-2D, skipping {base_key} {tuple(W.shape)}")
                skipped_nd += 1
                continue

            A = lora[down_key]
            B = lora[up_key]
            rank = A.shape[0]

            if alpha_key in lora:
                alpha = float(lora[alpha_key].item())
            else:
                alpha = rank

            scale = (alpha / rank) * STRENGTH

            # Verify matrix multiplication dimensions BEFORE allocating delta
            if B.shape[1] != A.shape[0]:
                print(
                    f"Rank mismatch: {base_key} " f"A{tuple(A.shape)} B{tuple(B.shape)}"
                )
                skipped_shape += 1
                continue

            # Verify resulting matrix will match the checkpoint weight
            expected = (B.shape[0], A.shape[1])

            if expected != W.shape:
                print(
                    f"Shape mismatch: {base_key} "
                    f"W{tuple(W.shape)} Expected{expected}"
                )
                skipped_shape += 1
                continue

            # Safe to compute
            t = time.perf_counter()

            delta = torch.matmul(B.float(), A.float())

            elapsed = time.perf_counter() - t

            delta.mul_(scale)

            Wf = W.float()

            Wf.add_(delta)

            merged[base_key] = Wf.to(W.dtype)

            del Wf

            applied += 1

            del A, B, delta

            if applied % 50 == 0:
                print(
                    f"{applied}/{len(down_keys)} "
                    f"{base_key} "
                    f"{time.perf_counter()-t:.3f}s",
                    flush=True,
                )

            if applied % 100 == 0:
                print(
                    f"{applied}/{len(down_keys)} " f"elapsed={elapsed:.3f}s", flush=True
                )
        except Exception as e:
            print("\nFAILED")
            print("Layer:", down_key)
            print("Base :", base_key if "base_key" in locals() else "<unknown>")
            print("Error:", repr(e))
            raise


print(f"\nApplied:          {applied}")
print(f"Skipped (no base):{skipped_missing}")
print(f"Skipped (shape):  {skipped_shape}")
print(f"Skipped (non-2D): {skipped_nd}")

if applied == 0:
    print(
        "\nWARNING: 0 keys applied — nothing was merged. Check that the LoRA "
        "key names match the base (naming convention mismatch?)."
    )

print("\nSaving...")
save_file(merged, OUTPUT, metadata=meta)
print("Done!")