#!/usr/bin/env python3
"""Make a ComfyUI-GGUF-compatible copy of an AtomicChat Qwen-Image-2.1(-Turbo) GGUF.

Problem: the AtomicChat GGUF stores the 1-D RMSNorm weights
(transformer_blocks.N.attn.norm_q.weight / norm_k.weight, shape [128])
quantized as Q8_0. stable-diffusion.cpp dequantizes them, but ComfyUI-GGUF has
no quantized RMSNorm op, so the fused rms_rope in ComfyUI receives the raw Q8_0
bytes (128 values -> 4 blocks x 34 bytes = 136) and fails with
  "Expected weight ... shape [136] and normalized_shape = [128]".

Fix: rewrite the file with every quantized 1-D tensor dequantized to F32 (the
convention llama.cpp and city96 convert.py use for norms). All other tensors are
copied byte-for-byte; the header keeps its original (empty) KV table so
stable-diffusion.cpp still loads it the same way. The original is never touched.

usage: python fix_atomicchat_gguf.py IN.gguf OUT.gguf   (needs: pip install gguf numpy)
"""
import struct, sys, numpy as np, gguf
from gguf.quants import dequantize
from gguf.constants import GGMLQuantizationType as Q

src, dst = sys.argv[1], sys.argv[2]
assert src != dst, "refusing to overwrite the original"
r = gguf.GGUFReader(src)
ALIGN = 32
NOQ = {Q.F32, Q.F16, Q.BF16, Q.F64}

def s(b): return struct.pack("<Q", len(b)) + b

items, fixed = [], []
for t in r.tensors:
    ne = [int(x) for x in t.shape]  # ggml order (fastest dim first)
    if len(ne) == 1 and t.tensor_type not in NOQ:
        data = dequantize(t.data, t.tensor_type).astype(np.float32).reshape(-1).tobytes()
        items.append((t.name, ne, Q.F32, data)); fixed.append((t.name, t.tensor_type.name, ne))
    else:
        items.append((t.name, ne, t.tensor_type, t.data.tobytes()))

# copy any KV pairs verbatim (raw bytes from the source header)
kv_raw, kv_n = b"", 0
for f in r.fields.values():
    if f.name.startswith("GGUF."): continue
    kv_raw += b"".join(p.tobytes() for p in f.parts); kv_n += 1

hdr = b"GGUF" + struct.pack("<I", 3) + struct.pack("<Q", len(items)) + struct.pack("<Q", kv_n) + kv_raw
off = 0; offs = []
for name, ne, tt, data in items:
    offs.append(off); off += len(data); off += (-off) % ALIGN
for (name, ne, tt, data), o in zip(items, offs):
    hdr += s(name.encode()) + struct.pack("<I", len(ne)) + b"".join(struct.pack("<Q", d) for d in ne)
    hdr += struct.pack("<I", int(tt)) + struct.pack("<Q", o)
hdr += b"\0" * ((-len(hdr)) % ALIGN)
with open(dst, "wb") as f:
    f.write(hdr)
    for name, ne, tt, data in items:
        f.write(data); f.write(b"\0" * ((-len(data)) % ALIGN))
print(f"dequantized {len(fixed)} 1-D tensors to F32:")
for n, q, ne in fixed[:4]: print("  ", n, q, ne)
if len(fixed) > 4: print("   ...")
# verify
r2 = gguf.GGUFReader(dst)
assert len(r2.tensors) == len(r.tensors)
for a, b in zip(r.tensors, r2.tensors):
    assert a.name == b.name and list(a.shape) == list(b.shape)
    if a.tensor_type == b.tensor_type: assert np.array_equal(a.data, b.data), a.name
    else: assert np.array_equal(dequantize(a.data, a.tensor_type).reshape(-1), np.asarray(b.data).reshape(-1)), a.name
print("verified: every tensor identical; fixed ones equal their dequantized values")