File size: 2,908 Bytes
7f15321 | 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 | """
patch_int8_add_format.py
Adds 'format': 'int8_tensorwise' to each per-layer comfy_quant config in an
INT8-Fast-saved safetensors so ComfyUI's stock Load Diffusion Model (PR #14636)
can dispatch it correctly. Writes a sibling file with '_fixed' appended to the
basename, preserving the original.
"""
import json
import os
import sys
import torch
from safetensors import safe_open
import safetensors.torch as sft
def patch_one(path_in: str) -> int:
if not os.path.isfile(path_in):
print(f" skip (not a file): {path_in}")
return 1
if not path_in.lower().endswith(".safetensors"):
print(f" skip (not .safetensors): {path_in}")
return 1
base, ext = os.path.splitext(path_in)
path_out = base + "_fixed" + ext
if os.path.exists(path_out):
print(f" output already exists, refusing to overwrite: {path_out}")
return 1
print(f" reading: {path_in}")
with safe_open(path_in, framework="pt") as f:
src_meta = dict(f.metadata() or {})
sd = {k: f.get_tensor(k) for k in f.keys()}
# Drop any malformed global header from previous patch attempts; per-layer
# comfy_quant byte tensors are the canonical source of truth.
if "_quantization_metadata" in src_meta:
print(" removing stale _quantization_metadata header")
src_meta.pop("_quantization_metadata", None)
patched = 0
already_ok = 0
for k in list(sd.keys()):
if not k.endswith(".comfy_quant"):
continue
try:
layer_conf = json.loads(sd[k].numpy().tobytes())
except Exception as e:
print(f" WARNING: could not decode {k}: {e}")
continue
if "format" in layer_conf:
already_ok += 1
continue
layer_conf["format"] = "int8_tensorwise"
sd[k] = torch.tensor(
list(json.dumps(layer_conf).encode("utf-8")),
dtype=torch.uint8,
)
patched += 1
if patched == 0 and already_ok == 0:
print(" ERROR: no per-layer comfy_quant byte tensors found.")
print(" This file doesn't look like it was saved by INT8-Fast's INT8ModelSave.")
return 1
if patched == 0:
print(f" nothing to do: all {already_ok} layers already have 'format'")
return 0
print(f" writing: {path_out}")
sft.save_file(sd, path_out, metadata=src_meta)
print(f" patched {patched} layer(s); {already_ok} already had 'format'")
return 0
def main(argv):
if len(argv) < 2:
print("Usage: patch_int8_add_format.py <file.safetensors> [more files...]")
return 1
rc = 0
for p in argv[1:]:
print(f"\n--- {os.path.basename(p)} ---")
rc |= patch_one(p)
return rc
if __name__ == "__main__":
sys.exit(main(sys.argv)) |