tools: merge script for rebuilding transformer from fullgreed_bf16

tools/fullgreed_edit_merge_bf16.py

@@ -0,0 +1,122 @@
+"""Rebuild Fullgreed-Edit from fullgreed_bf16 (uploaded 2026-07-10).
+
+merged = Z-Image-Edit (Anjoe/z-image-edit, fp32) + (fullgreed_bf16 - Z-Image base bf16)
+
+Key mapping comfy(453) -> diffusers-edit(521): unfuse attention.qkv -> to_q/to_k/to_v
+(order q,k,v), out -> to_out.0, q_norm/k_norm -> norm_q/norm_k,
+final_layer.* -> all_final_layer.2-1.*, x_embedder -> all_x_embedder.2-1 with the
+delta applied to cols 0-63 only (cols 64-127 = ref-latent branch, kept from Edit).
+Output bf16, streamed to disk so peak RAM stays ~1 GB. Verifies 5 random tensors
+by recompute, then uploads to John2386/fullgreed-edit/transformer/.
+
+Run on a Colab CPU runtime. Needs HF_TOKEN in the environment (from Colab Secrets).
+"""
+import hashlib, json, os, random, struct
+import torch
+from huggingface_hub import hf_hub_download, HfApi
+from safetensors import safe_open
+
+DIM = 3840
+OUT = "/content/diffusion_pytorch_model.safetensors"
+REPO_OUT = "John2386/fullgreed-edit"
+PATH_IN_REPO = "transformer/diffusion_pytorch_model.safetensors"
+TOKEN = os.environ["HF_TOKEN"]  # fail fast if missing
+
+
+def edit_to_comfy(k):
+    if k == "all_x_embedder.2-1.weight":
+        return "x_embedder.weight", ("padcols", 128)
+    if k == "all_x_embedder.2-1.bias":
+        return "x_embedder.bias", None
+    if k.startswith("all_final_layer.2-1."):
+        return "final_layer." + k[len("all_final_layer.2-1."):], None
+    if ".attention.norm_q." in k:
+        return k.replace(".attention.norm_q.", ".attention.q_norm."), None
+    if ".attention.norm_k." in k:
+        return k.replace(".attention.norm_k.", ".attention.k_norm."), None
+    if ".attention.to_out.0." in k:
+        return k.replace(".attention.to_out.0.", ".attention.out."), None
+    for i, proj in enumerate(("to_q", "to_k", "to_v")):
+        tag = f".attention.{proj}.weight"
+        if k.endswith(tag):
+            return k[: -len(tag)] + ".attention.qkv.weight", ("rows", i * DIM, (i + 1) * DIM)
+    return k, None
+
+
+def delta_for(fg, fb, key):
+    ck, spec = edit_to_comfy(key)
+    d = fg.get_tensor(ck).to(torch.float32) - fb.get_tensor(ck).to(torch.float32)
+    if spec is None:
+        return d
+    if spec[0] == "rows":
+        return d[spec[1]:spec[2]]
+    if spec[0] == "padcols":
+        return torch.nn.functional.pad(d, (0, spec[1] - d.shape[1]))
+    raise ValueError(spec)
+
+
+def merged_tensor(fe, fg, fb, key):
+    return (fe.get_tensor(key).to(torch.float32) + delta_for(fg, fb, key)).to(torch.bfloat16)
+
+
+print("downloading inputs...", flush=True)
+EDIT = hf_hub_download("Anjoe/z-image-edit", "transformer/diffusion_pytorch_model.safetensors")
+BASE = hf_hub_download("Comfy-Org/z_image", "split_files/diffusion_models/z_image_bf16.safetensors")
+GREED = hf_hub_download("John2386/fullgreed", "fullgreed_bf16.safetensors")
+
+fe = safe_open(EDIT, framework="pt")
+fb = safe_open(BASE, framework="pt")
+fg = safe_open(GREED, framework="pt")
+keys = sorted(fe.keys())
+assert len(keys) == 521, len(keys)
+
+shapes = {k: list(fe.get_slice(k).get_shape()) for k in keys}
+header, off = {"__metadata__": {"format": "pt",
+                                "merge": "Anjoe/z-image-edit + (John2386/fullgreed fullgreed_bf16 - Comfy-Org/z_image bf16)",
+                                "merge_scale": "1.0", "built": "2026-07-11"}}, 0
+for k in keys:
+    n = 2  # bf16 bytes
+    for s in shapes[k]:
+        n *= s
+    header[k] = {"dtype": "BF16", "shape": shapes[k], "data_offsets": [off, off + n]}
+    off += n
+hb = json.dumps(header, separators=(",", ":")).encode()
+
+print(f"writing {off/1e9:.2f} GB to {OUT}", flush=True)
+with open(OUT, "wb") as f:
+    f.write(struct.pack("<Q", len(hb)))
+    f.write(hb)
+    for i, k in enumerate(keys):
+        t = merged_tensor(fe, fg, fb, k)
+        assert list(t.shape) == shapes[k], k
+        f.write(t.contiguous().view(torch.uint8).numpy().tobytes())
+        if i % 50 == 0:
+            print(f"  {i}/521 {k}", flush=True)
+
+print("verifying...", flush=True)
+fo = safe_open(OUT, framework="pt")
+okeys = list(fo.keys())
+assert len(okeys) == 521, len(okeys)
+random.seed(0)
+for k in random.sample(okeys, 5) + ["all_x_embedder.2-1.weight"]:
+    want = merged_tensor(fe, fg, fb, k)
+    got = fo.get_tensor(k)
+    assert torch.equal(want, got), k
+    print(f"  ok {k} {list(got.shape)}", flush=True)
+# ref-latent columns must be pure Edit weights (delta zero there)
+xe = fo.get_tensor("all_x_embedder.2-1.weight")
+assert torch.equal(xe[:, 64:], fe.get_tensor("all_x_embedder.2-1.weight").to(torch.bfloat16)[:, 64:])
+print("  ok ref-latent cols untouched", flush=True)
+
+sha = hashlib.sha256()
+with open(OUT, "rb") as f:
+    for chunk in iter(lambda: f.read(1 << 24), b""):
+        sha.update(chunk)
+print("sha256:", sha.hexdigest())
+print("size:", os.path.getsize(OUT))
+
+print("uploading...", flush=True)
+api = HfApi(token=TOKEN)
+api.upload_file(path_or_fileobj=OUT, path_in_repo=PATH_IN_REPO, repo_id=REPO_OUT,
+                commit_message="transformer: rebuild edit merge from fullgreed_bf16 (2026-07-10 weights)")
+print("DONE — uploaded to", REPO_OUT + "/" + PATH_IN_REPO)