File size: 5,031 Bytes
be0f394 16364ff be0f394 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | """Quantize MOSS-Image components into ComfyUI-native ConvRot formats (comfy-kitchen layouts).
dit int4 : convrot_w4a4 (group 256, kernel group 64), header _quantization_metadata (= toxicdog/MOSS-Image-SFT-INT4-ConvRot-ComfyUI)
dit int8 : int8_tensorwise + convrot (per-row scale), per-layer .comfy_quant tensors (= Abiray/MOSS-Image-SFT-INT8-convrot)
te int8 : same as Abiray text encoder (embed_tokens + LM linears except in_proj_a/b + vision block linears)
"""
import argparse, json, os, re, sys, time
import torch
from safetensors import safe_open
from safetensors.torch import save_file
from comfy_kitchen.tensor import TensorWiseINT8Layout, TensorCoreConvRotW4A4Layout
DIT_RE = re.compile(r"^layers\.\d+\.proj\.(q|k|v|o|ffn_up_gate|ffn_down)\.weight$")
TE_RE = re.compile(r"^(language_model\.embed_tokens|language_model\.layers\.\d+\.(linear_attn\.(in_proj_qkv|in_proj_z|out_proj)|mlp\.(gate|up|down)_proj|self_attn\.(q|k|v|o)_proj)|visual\.blocks\.\d+\.(attn\.(qkv|proj)|mlp\.linear_fc[12]))\.weight$")
def te_group(k):
for g in (256, 64, 16, 4):
if k % g == 0:
return g
return None
def load_index(d, index):
m = json.load(open(os.path.join(d, index)))["weight_map"]
by = {}
for k, f in m.items():
by.setdefault(f, []).append(k)
return by
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--component", choices=["dit", "te"], required=True)
ap.add_argument("--mode", choices=["int8", "int4"], required=True)
ap.add_argument("--src", required=True, help="MOSS-Image pipeline dir")
ap.add_argument("--out", required=True)
ap.add_argument("--device", default="cpu")
ap.add_argument("--threads", type=int, default=32)
a = ap.parse_args()
torch.set_num_threads(a.threads)
if os.path.exists(a.out):
sys.exit(f"exists: {a.out}")
sub, index = ("transformer", "diffusion_pytorch_model.safetensors.index.json") if a.component == "dit" else ("text_encoder", "model.safetensors.index.json")
d = os.path.join(a.src, sub)
rx = DIT_RE if a.component == "dit" else TE_RE
out, layers, stats = {}, {}, []
t0 = time.time()
for fn, keys in sorted(load_index(d, index).items()):
with safe_open(os.path.join(d, fn), "pt") as f:
for k in sorted(keys):
w = f.get_tensor(k)
if not rx.match(k):
out[k] = w.contiguous()
continue
name = k[: -len(".weight")]
wf = w.to(a.device, torch.float32)
if a.mode == "int4":
q, p = TensorCoreConvRotW4A4Layout.quantize(wf, convrot_groupsize=256)
deq = TensorCoreConvRotW4A4Layout.dequantize(q, p)
out[name + ".weight"] = q.cpu().contiguous()
out[name + ".weight_scale"] = p.scale.float().reshape(-1).cpu().contiguous()
layers[name] = {"format": "convrot_w4a4", "convrot_groupsize": 256, "quant_group_size": 64}
else:
g = 256 if a.component == "dit" else te_group(wf.shape[1])
q, p = TensorWiseINT8Layout.quantize(wf, is_weight=True, per_channel=True, convrot=True, convrot_groupsize=g)
deq = TensorWiseINT8Layout.dequantize(q, p)
out[name + ".weight"] = q.cpu().contiguous()
out[name + ".weight_scale"] = p.scale.float().reshape(-1, 1).cpu().contiguous()
conf = {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": g}
out[name + ".comfy_quant"] = torch.frombuffer(bytearray(json.dumps(conf).encode()), dtype=torch.uint8).clone()
layers[name] = conf
deq = deq.float()
rel = ((deq - wf).norm() / wf.norm()).item()
cos = torch.nn.functional.cosine_similarity(deq.flatten(), wf.flatten(), dim=0).item()
stats.append((name, list(wf.shape), rel, cos))
print(f"{name} {list(wf.shape)} rel={rel:.5f} cos={cos:.6f}", flush=True)
if a.component == "te":
# HF / ComfyUI-core naming (model.language_model.*, model.visual.*); also accepted by the MOSS-Image plugin
out = {("model." + k if k.startswith(("language_model.", "visual.")) else k): v for k, v in out.items()}
meta = {"format": "pt", "model": "moss-image"}
if a.mode == "int4":
meta["_quantization_metadata"] = json.dumps({"format_version": "1.0", "layers": layers})
tmp = a.out + ".partial"
os.makedirs(os.path.dirname(os.path.abspath(a.out)), exist_ok=True)
save_file(out, tmp, metadata=meta)
os.rename(tmp, a.out)
json.dump(stats, open(a.out + ".stats.json", "w"))
rels = [s[2] for s in stats]
print(f"DONE {a.out} layers={len(stats)} tensors={len(out)} mean_rel={sum(rels)/len(rels):.5f} max_rel={max(rels):.5f} min_cos={min(s[3] for s in stats):.6f} size={os.path.getsize(a.out)} t={time.time()-t0:.0f}s")
if __name__ == "__main__":
main()
|