wangzhang

/

gemma-4-E4B-it-abliterix

like 14

Hi Matthew, thanks a lot for doing this — that sounds very useful.

I’d prefer to start with the script and the exact conversion steps first, so I can reproduce the repack locally and verify that it only changes the packaging format, not the actual weights.

If everything checks out, I’m happy for you to share the single-safetensors repack as a separate Hugging Face repo or PR, as long as the original model attribution and Gemma license terms are preserved. I can also link to it from the model card for people who want to use it with ComfyUI or similar tools.

Could you share:

• the script
• the exact command you used
• the expected output file size / sha256 hash
• which programs you tested it with
• whether anything besides repacking was changed

Thanks again — really appreciate the contribution.

Absolutely, I'm packaging everything up now. Do you have an alternate means of communication or an email? I can post it all here but was unsure if you want this to be private while you verify everything. I can also try to find a secure file sender that only you can retrieve, but I'd still need some sort of contact info.

Thanks — I appreciate it.

To keep everything transparent and safe, I’d prefer not to use private file senders or email for the initial review.

You can post the script, exact commands, and conversion notes here in the discussion. If the repacked safetensors file is too large, the cleanest option would be to create a separate Hugging Face repo under your account and share the link here.

Please include:

• the script
• the exact command used
• expected output file size
• sha256 hash
• what you tested it with
• confirmation that this is only a packaging/repack change and does not modify the weights

Once I can verify it, I’m happy to link your repack from the model card for users who want a single-safetensors version.

Thanks again for putting this together.

Full disclaimer, Claude helped me build this script…I’m not a coder so my depth of knowledge is limited to knowing basic script formats and only surface-level understanding.

I can't post the python file or a .txt file in this discussion but I put the full text of the script at the bottom of this post.

For the process, I cloned your repo into a local folder on my machine. The script needs three file paths set at the top before running. The paths contain XXXXXXXX to show where they need to be updated before running.

• ABLITERATED_REPO_DIR – points to the location of your cloned repo.
• COMFY_REFERENCE – since I use comfyui, this points to the official Comfy-Org/Gemma4 file gemma4_e4b_it_bf16.safetensors file that I had saved in my ComfyUI text encoders folder.
• OUTPUT_FILE – where you want the new file saved

Once those are set, I ran the script in python:
python repack_gemma4.py

The file size is 14.9 GB
SHA256: BA756436A278E5C683535A432B12425278D365274CFB303220F6E7F86D7D92AF
It should contain 719 tensors total to match the Comfy-Org reference file

I tested it and use it in ComfyUI, V0.21.1 and the built in TextGenerate node that supports Gemma models. No errors are produced when running it.

Pulled this from AI (again, sorry, I’m not an expert here):

It only repackages your model — no weight changes, no quantization, no precision conversion. Specifically:

1. Combines your sharded files into one single safetensors file
2. Renames the key prefix from "model.language_model." to "model." (which is what ComfyUI expects)
3. Adds back the K and V projection tensors for layers 24-41 by copying them from layer 23 (these are shared in your HF format but need to exist explicitly for ComfyUI)
4. For the k_norm tensors on layers 24-41, sources them from the Comfy reference file because those layers use a different shape than layer 23. This is structural only — k_norm doesn't carry the refusal-direction information that abliteration modifies (that lives in q/k/v/o_proj and mlp.down_proj), so your abliteration work is fully preserved
5. Adds in the non-language-model parts from the Comfy reference file (vision model, tokenizer, etc.) so it's a complete drop-in replacement
6. Uses a manual file writer to work around a Python overflow error that hits on one of the large tensors

The file stays in bf16 precision throughout. All your abliteration changes from the original weights are preserved exactly.

The script is below this line and should be saved as a .py file to run in python.

import sys
import struct
import json
from pathlib import Path
from safetensors import safe_open
import torch

ABLITERATED_REPO_DIR = r"C:\XXXXXXXX\models\text_encoders\gemma-4-E4B-it-abliterix"
COMFY_REFERENCE = r"C:\XXXXXXXX\models\text_encoders\gemma4_e4b_it_bf16.safetensors"
OUTPUT_FILE = r"C:\XXXXXXXX\XXXXXXXX\XXXXXXXX\gemma4_e4b_abliterated_bf16.safetensors"

KV_OWNER_LAYER = 23
KV_SHARED_LAYERS = list(range(24, 42))
KV_PROJ_SUFFIXES = [
"self_attn.k_proj.weight",
"self_attn.v_proj.weight",
]
K_NORM_SUFFIX = "self_attn.k_norm.weight"

_TORCH_TO_ST = {
torch.float32: "F32", torch.float64: "F64", torch.float16: "F16",
torch.bfloat16: "BF16", torch.int64: "I64", torch.int32: "I32",
torch.int16: "I16", torch.int8: "I8", torch.uint8: "U8", torch.bool: "BOOL",
}
_ST_SIZE = {
"F32": 4, "F64": 8, "F16": 2, "BF16": 2,
"I64": 8, "I32": 4, "I16": 2, "I8": 1, "U8": 1, "BOOL": 1,
}

def tensor_to_bytes(t):
t = t.contiguous().cpu()
if t.dtype == torch.bfloat16:
return t.view(torch.int16).numpy().tobytes()
return t.numpy().tobytes()

def save_safetensors_large(tensors, filepath):
header = {}
offset = 0
for name, t in tensors.items():
if t.dtype not in _TORCH_TO_ST:
raise ValueError(f"Unsupported dtype {t.dtype} for tensor {name}")
st_dtype = _TORCH_TO_ST[t.dtype]
shape = list(t.shape)
nbytes = t.numel() * _ST_SIZE[st_dtype]
header[name] = {
"dtype": st_dtype,
"shape": shape,
"data_offsets": [offset, offset + nbytes],
}
offset += nbytes

header_bytes = json.dumps(header, separators=(",", ":")).encode("utf-8")
pad = (8 - (len(header_bytes) % 8)) % 8
header_bytes = header_bytes + b" " * pad

print(f"      Header: {len(header_bytes)} bytes, total payload: {offset / (1024**3):.2f} GB")
print(f"      Writing to disk (streaming)...")

n = len(tensors)
with open(filepath, "wb") as fp:
    fp.write(struct.pack("<Q", len(header_bytes)))
    fp.write(header_bytes)
    for i, (name, t) in enumerate(tensors.items(), 1):
        chunk = tensor_to_bytes(t)
        fp.write(chunk)
        if i % 200 == 0 or i == n:
            print(f"        wrote {i}/{n} tensors")

def main():
print("=" * 70)
print("Gemma 4 E4B Abliterated Repack (v3: corrected k_norm handling)")
print("=" * 70)
print(f"Abliterated source: {ABLITERATED_REPO_DIR}")
print(f"Reference file: {COMFY_REFERENCE}")
print(f"Output file: {OUTPUT_FILE}")
print()

if not Path(ABLITERATED_REPO_DIR).is_dir():
    print(f"ERROR: Abliterated repo folder not found: {ABLITERATED_REPO_DIR}")
    sys.exit(1)
if not Path(COMFY_REFERENCE).is_file():
    print(f"ERROR: Reference file not found: {COMFY_REFERENCE}")
    sys.exit(1)

shards = sorted(Path(ABLITERATED_REPO_DIR).glob("*.safetensors"))
if not shards:
    print(f"ERROR: No .safetensors files found in {ABLITERATED_REPO_DIR}")
    sys.exit(1)
print(f"[1/5] Found {len(shards)} shard(s) in abliterated repo:")
for s in shards:
    print(f"      {s.name}  ({s.stat().st_size / (1024**3):.2f} GB)")

print(f"\n[2/5] Loading and renaming abliterated language model weights...")
abliterated_lm = {}
for shard in shards:
    with safe_open(shard, framework="pt") as f:
        for hf_key in f.keys():
            if hf_key.startswith("model.language_model."):
                comfy_key = hf_key.replace("model.language_model.", "model.", 1)
                abliterated_lm[comfy_key] = f.get_tensor(hf_key)
    print(f"      Read {shard.name}")
print(f"      Loaded {len(abliterated_lm)} LM tensors (expected 665)")

print(f"\n[3/5] Loading reference Comfy file (for non-LM + k_norm fills)...")
ref_full = {}
with safe_open(COMFY_REFERENCE, framework="pt") as f:
    for key in f.keys():
        ref_full[key] = f.get_tensor(key)
non_lm_components = {k: v for k, v in ref_full.items() if not k.startswith("model.")}
print(f"      Loaded {len(ref_full)} total tensors from reference")
print(f"      Non-LM tensors: {len(non_lm_components)}")
component_prefixes = sorted(set(k.split('.')[0] for k in non_lm_components.keys()))
print(f"      Non-LM components: {', '.join(component_prefixes)}")

print(f"\n[4/5] Reconstructing shared K/V projections and k_norm tensors...")
print(f"      K/V projections: clone from layer {KV_OWNER_LAYER} to layers {KV_SHARED_LAYERS[0]}-{KV_SHARED_LAYERS[-1]}")
print(f"      k_norm: source from reference (correct local-attention shape)")
proj_added = 0
knorm_from_abl = 0
knorm_from_ref = 0

for layer_idx in KV_SHARED_LAYERS:
    # k_proj / v_proj: cloned from layer 23 (shared in HF)
    for suffix in KV_PROJ_SUFFIXES:
        owner_key = f"model.layers.{KV_OWNER_LAYER}.{suffix}"
        target_key = f"model.layers.{layer_idx}.{suffix}"
        if target_key in abliterated_lm:
            continue
        if owner_key not in abliterated_lm:
            print(f"      ERROR: Owner key missing: {owner_key}")
            sys.exit(1)
        abliterated_lm[target_key] = abliterated_lm[owner_key].clone()
        proj_added += 1

    # k_norm: per-layer, smaller shape; source from reference if absent
    knorm_key = f"model.layers.{layer_idx}.{K_NORM_SUFFIX}"
    if knorm_key in abliterated_lm:
        knorm_from_abl += 1
    else:
        if knorm_key not in ref_full:
            print(f"      ERROR: k_norm missing in both abliterix and reference: {knorm_key}")
            sys.exit(1)
        abliterated_lm[knorm_key] = ref_full[knorm_key].clone()
        knorm_from_ref += 1

print(f"      Added {proj_added} K/V projection tensors (cloned from layer {KV_OWNER_LAYER})")
print(f"      k_norm sources: {knorm_from_abl} kept from abliterix, {knorm_from_ref} pulled from reference")
print(f"      Total LM tensors now: {len(abliterated_lm)} (expected 719)")

print(f"\n[5/5] Validating against reference and saving...")
combined_dict = {**abliterated_lm, **non_lm_components}
ref_key_order = list(ref_full.keys())

combined_keys = set(combined_dict.keys())
ref_keys = set(ref_key_order)
missing = ref_keys - combined_keys
extra = combined_keys - ref_keys
if missing:
    print(f"      WARNING: {len(missing)} keys in reference but NOT in output:")
    for k in sorted(missing)[:10]:
        print(f"        {k}")
if extra:
    print(f"      WARNING: {len(extra)} keys in output but NOT in reference:")
    for k in sorted(extra)[:10]:
        print(f"        {k}")
if not missing and not extra:
    print(f"      Key set matches reference exactly ({len(combined_dict)} tensors)")

shape_mismatches = []
for k in combined_dict:
    if k in ref_full:
        ref_shape = tuple(ref_full[k].shape)
        our_shape = tuple(combined_dict[k].shape)
        if ref_shape != our_shape:
            shape_mismatches.append((k, our_shape, ref_shape))
if shape_mismatches:
    print(f"      WARNING: {len(shape_mismatches)} shape mismatch(es) vs reference:")
    for k, ours, ref in shape_mismatches[:10]:
        print(f"        {k}: ours={ours}, reference={ref}")
    if len(shape_mismatches) > 10:
        print(f"        ... and {len(shape_mismatches)-10} more")
else:
    print(f"      All tensor shapes match reference.")

combined_ordered = {k: combined_dict[k] for k in ref_key_order if k in combined_dict}

Path(OUTPUT_FILE).parent.mkdir(parents=True, exist_ok=True)
save_safetensors_large(combined_ordered, OUTPUT_FILE)

output_size_gb = Path(OUTPUT_FILE).stat().st_size / (1024**3)
print(f"      Saved {output_size_gb:.2f} GB")

print()
print("=" * 70)
print(f"DONE. New file:")
print(f"  {OUTPUT_FILE}")
print()
print("Next steps:")
print("  1. Fully quit ComfyUI Desktop (system tray -> Quit)")
print("  2. Relaunch ComfyUI Desktop")
print("  3. In your Load CLIP node, select the new file")
print("  4. Run the workflow")
print("=" * 70)

if name == "main":
try:
main()
except KeyboardInterrupt:
print("\nInterrupted.")
sys.exit(1)
except Exception as e:
print(f"\nERROR: {e}")
import traceback
traceback.print_exc()
sys.exit(1)