| """
|
| 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()}
|
|
|
|
|
|
|
| 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)) |
