| |
| """Export a MOSSGramSolve edit (D = Delta G_f (G_f + mu G_g + eps I)^-1 per linear) as a ComfyUI full-difference LoRA (.diff keys). |
| Solved per linear on the GPU in float32 (the ridge keeps it well conditioned for rho >= ~0.01); only one linear's tensors are resident. |
| |
| comfyui_venv/bin/python workflows/moss/export_gramsolve.py <out_name> --mu 0 --rho 1.0 --layers 21-41 |
| Writes models/loras/<out_name>.safetensors (bf16) and prints per-linear |D|/|Delta|. |
| """ |
| import argparse, os, sys |
| import torch |
| from safetensors import safe_open |
| from safetensors.torch import save_file |
|
|
| ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..") |
| sys.path.insert(0, os.path.join(ROOT, "custom_nodes", "moss_block_lab")) |
| ap = argparse.ArgumentParser() |
| ap.add_argument("out_name") |
| ap.add_argument("--mu", type=float, default=0.0) |
| ap.add_argument("--rho", type=float, default=1.0) |
| ap.add_argument("--layers", default="21-41") |
| ap.add_argument("--flat", default="flat") |
| ap.add_argument("--gen", default="gen") |
| ap.add_argument("--device", default="cuda") |
| a = ap.parse_args() |
|
|
|
|
| def parse_layers(spec): |
| out = set() |
| for part in spec.split(","): |
| lo, _, hi = part.partition("-") |
| out.update(range(int(lo), int(hi or lo) + 1)) |
| return sorted(out) |
|
|
|
|
| G = os.path.join(ROOT, "output", "moss", "grams") |
| gf = safe_open(os.path.join(G, a.flat + ".safetensors"), "pt") |
| gg = safe_open(os.path.join(G, a.gen + ".safetensors"), "pt") if a.mu > 0 else None |
| fs = safe_open(os.path.join(ROOT, "models", "diffusion_models", "moss_image_sft_bf16.safetensors"), "pt") |
| fb = safe_open(os.path.join(ROOT, "models", "diffusion_models", "moss_image_base_bf16.safetensors"), "pt") |
| out = {} |
| for i in parse_layers(a.layers): |
| for name in ("qkv", "o", "ffn_up_gate", "ffn_down"): |
| k = f"layers.{i}.proj.{name}" |
| dev = a.device |
| delta = fb.get_tensor(k + ".weight").to(dev, torch.float32) - fs.get_tensor(k + ".weight").to(dev, torch.float32) |
| A = gf.get_tensor(k + ".G").to(dev) |
| rhs = A.clone() |
| if gg is not None: |
| A = A + a.mu * gg.get_tensor(k + ".G").to(dev) |
| A.diagonal().add_(a.rho * A.diagonal().sum() / A.shape[0]) |
| D = torch.cholesky_solve(rhs @ delta.T, torch.linalg.cholesky(A)).T |
| out[f"diffusion_model.{k}.diff"] = D.to(torch.bfloat16).cpu().contiguous() |
| print(f"{k}: |D|/|Delta| = {float(D.norm() / delta.norm()):.3f}", flush=True) |
| path = os.path.join(ROOT, "models", "loras", a.out_name + ".safetensors") |
| save_file(out, path, metadata={"mu": str(a.mu), "rho": str(a.rho), "layers": a.layers, "grams": f"{a.flat},{a.gen}", |
| "recipe": "D = (W_base - W_sft) G_f (G_f + mu G_g + rho tr(.)/d I)^-1"}) |
| print("wrote", path, round(os.path.getsize(path) / 1e9, 2), "GB") |
|
|
