Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-17 08:19:11 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat(model): support comfy_quant header metadata containers

Browse Source 
Newer quantized checkpoints record per-layer formats in the safetensors
header _quantization_metadata instead of marker tensors. The native loader
now reads the header map, re-keys it through the same prefix strip as the
tensors, and transcodes it into marker tensors so both container forms
share one detection path; header entries win over markers.

- drop optional input_scale sidecars for marked layers
- map full_precision_matrix_mult onto the sdnq per-layer matmul exclusion list
- log the detection source (markers, header, both)
This commit is contained in:
CalamitousFelicitousness
2026-07-11 08:48:52 +01:00
parent 9bb94dcfff
commit d005ac5aa1
2 changed files with 303 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pipelines/native_transformer.py
+83 -4
View File
@@ -274,8 +274,12 @@ def load(
quant_type = model_quant.get_quant_type(quant_args)
state_dict = sd_models.read_state_dict(local_file, what="transformer")
metadata_layers = read_quantization_metadata(local_file)
state_dict = drop_companion_keys(state_dict, spec.ignored_prefixes, spec.cls.__name__)
state_dict, detected_prefix = strip_prefix(state_dict, spec.prefixes, spec.cls.__name__)
if metadata_layers and detected_prefix:
# header metadata names mirror the file's tensor naming, so they carry the same prefix
metadata_layers = {name[len(detected_prefix):] if name.startswith(detected_prefix) else name: meta for name, meta in metadata_layers.items()}
check_forbidden_markers(state_dict, spec.forbidden_markers, spec.cls.__name__, local_file)
transformer_sd, sibling_sds = partition_siblings(state_dict, spec.siblings)
del state_dict
@@ -304,6 +308,7 @@ def load(
modules_to_not_convert=modules_to_not_convert,
modules_dtype_dict=modules_dtype_dict,
converter_handles_quant=spec.converter_handles_quant,
metadata_layers=metadata_layers,
**kwargs,
)
del transformer_sd
@@ -427,6 +432,64 @@ def check_forbidden_markers(
)
def read_quantization_metadata(local_file: str) -> dict[str, dict] | None:
"""Read per-layer quantization info from the safetensors header metadata.
Newer comfy_quant checkpoints record layer formats in the header
``__metadata__`` under ``_quantization_metadata`` (a JSON string with a
``layers`` map keyed by tensor naming) instead of per-layer marker
tensors. Returns the layers map, or ``None`` when the header carries no
such entry. Raises :class:`OverrideArchMismatch` when the entry is
present but malformed, so the caller's base-repo fallback engages.
"""
from modules.model_probe import read_safetensors_header
try:
header = read_safetensors_header(local_file)
except Exception as e:
log.debug(f'Load model: file="{local_file}" header metadata unreadable ({e})')
return None
meta = (header.get("__metadata__") or {}).get("_quantization_metadata")
if meta is None:
return None
try:
parsed = json.loads(meta) if isinstance(meta, str) else meta
layers = parsed["layers"]
except Exception as e:
raise OverrideArchMismatch(
f'Load model: file="{local_file}" native_transformer _quantization_metadata '
f"is malformed ({type(e).__name__}: {e})"
) from e
if not isinstance(layers, dict) or not all(isinstance(entry, dict) for entry in layers.values()):
raise OverrideArchMismatch(
f'Load model: file="{local_file}" native_transformer _quantization_metadata '
f"layers map is malformed"
)
return layers
def transcode_quant_metadata(state_dict: dict, metadata_layers: dict[str, dict]) -> tuple[dict, str]:
"""Synthesize per-layer marker tensors from header quantization metadata.
Header metadata and marker tensors describe the same per-layer format
dicts; converging on markers lets detection, converters, and remapping
handle both forms identically. Entries without a matching ``.weight``
(companion components filtered earlier) are skipped; existing markers
are overwritten by their header entry. Returns a new dict plus the
detection source (``header`` or ``both``) for logging.
"""
sd = dict(state_dict)
had_markers = any(key.endswith(COMFY_QUANT_MARKER) for key in sd)
count = 0
for name, meta in metadata_layers.items():
if f"{name}.weight" not in sd:
continue
sd[f"{name}{COMFY_QUANT_MARKER}"] = torch.tensor(list(json.dumps(meta).encode("utf-8")), dtype=torch.uint8)
count += 1
if count == 0:
return state_dict, "markers"
return sd, "both" if had_markers else "header"
def detect_comfy_quant(state_dict: dict, type_name: str) -> tuple[dict[str, dict], str] | None:
"""Detect ``comfy_quant`` pre-quantized layers in a state dict.
@@ -490,14 +553,19 @@ def remap_comfy_quant(state_dict: dict, marked_names: set[str]) -> dict:
Renames ``<name>.weight_scale`` to ``<name>.scale`` (reshaped to 2-D,
``[]`` becomes ``[1, 1]``, since SDNQ transposes scales in place) and
drops the ``<name>.comfy_quant`` markers. Returns a new dict; the input
(which may be the cached state dict) is not mutated.
drops the ``<name>.comfy_quant`` markers plus optional
``<name>.input_scale`` activation-calibration sidecars (SDNQ re-derives
activation scales dynamically). Returns a new dict; the input (which may
be the cached state dict) is not mutated.
"""
scale_suffix = ".weight_scale"
input_scale_suffix = ".input_scale"
remapped: dict = {}
for key, value in state_dict.items():
if key.endswith(COMFY_QUANT_MARKER) and key[: -len(COMFY_QUANT_MARKER)] in marked_names:
continue
if key.endswith(input_scale_suffix) and key[: -len(input_scale_suffix)] in marked_names:
continue
if key.endswith(scale_suffix) and key[: -len(scale_suffix)] in marked_names:
remapped[f"{key[: -len(scale_suffix)]}.scale"] = value.reshape(-1, 1)
continue
@@ -699,6 +767,9 @@ def build_component_prequantized(
f"{comfy_format} expects {storage_dtype} weights but {name!r} has {found}"
)
# layers flagged full_precision_matrix_mult stay on the dequant path even
# when the user enables quantized matmul (exact-name match on the config list)
full_precision_mm = [f"{name}.weight" for name in sorted(marked_names) if marker_meta[name].get("full_precision_matrix_mult")]
quantization_config = SDNQConfig(
weights_dtype=weights_dtype,
quantized_matmul_dtype=matmul_dtype,
@@ -707,6 +778,7 @@ def build_component_prequantized(
dequantize_fp32=shared.opts.sdnq_dequantize_fp32,
add_skip_keys=False,
modules_to_not_convert=[],
modules_to_not_use_matmul=full_precision_mm,
)
quantizer = SDNQQuantizer(quantization_config, pre_quantized=True)
quantizer.torch_dtype = target_dtype
@@ -844,6 +916,7 @@ def build_component(
modules_to_not_convert: list | None = None,
modules_dtype_dict: dict | None = None,
converter_handles_quant: bool = False,
metadata_layers: dict[str, dict] | None = None,
**kwargs,
) -> object:
"""Convert (if needed), instantiate, load weights, dtype-cast, quantize,
@@ -851,7 +924,9 @@ def build_component(
Transformer state dicts carrying ``comfy_quant`` markers dispatch
to :func:`build_component_prequantized` (``quant_args`` are bypassed: the
file is already quantized); a ``converter_handles_quant`` converter runs
file is already quantized). ``metadata_layers`` (header-metadata quant
info) is transcoded into markers first, so both container forms share
one path; a ``converter_handles_quant`` converter runs
before that detection, float-oriented converters after it. Under SDNQ the
transformer uses the per-tensor pre-mode path in
:func:`build_component_quantized` so quantization is applied in flight.
@@ -866,6 +941,10 @@ def build_component(
reach ``cls.from_config`` for both construction paths.
"""
try:
quant_source = "markers"
if component_name == "transformer" and metadata_layers:
state_dict, quant_source = transcode_quant_metadata(state_dict, metadata_layers)
if converter is not None and converter_handles_quant and component_name == "transformer":
state_dict = apply_converter(converter, state_dict, cls, component_name)
converter = None # consumed; must not run again on the non-comfy path below
@@ -876,7 +955,7 @@ def build_component(
convrot_count = sum(1 for meta in marker_meta.values() if meta.get("convrot"))
log.info(
f'Load model: transformer=native {component_name} quant=comfy '
f'format={comfy_format} layers={len(marker_meta)} convrot={convrot_count} keys={len(state_dict)} cls={cls.__name__}'
f'format={comfy_format} layers={len(marker_meta)} convrot={convrot_count} source={quant_source} keys={len(state_dict)} cls={cls.__name__}'
)
component = build_component_prequantized(
component_name=component_name,
test/test-native-transformer.py
+220 -2
View File
@@ -435,6 +435,121 @@ def test_remap_comfy_does_not_mutate_input():
assert set(sd.keys()) == keys_before, 'input dict must not be mutated (read_state_dict caches it)'
def test_remap_comfy_drops_input_scale():
"""Optional activation-calibration sidecars are dropped for marked layers
(SDNQ derives activation scales dynamically); unmarked ones pass through."""
unrelated = torch.tensor(1.0)
sd = {
'blocks.0.attn.wq.weight': torch.zeros((4, 4), dtype=torch.int8),
'blocks.0.attn.wq.weight_scale': torch.tensor(0.5),
'blocks.0.attn.wq.input_scale': torch.tensor(0.125),
'blocks.0.attn.wq.comfy_quant': comfy_marker('int8_tensorwise'),
'other.input_scale': unrelated,
}
out = nt.remap_comfy_quant(sd, {'blocks.0.attn.wq'})
assert 'blocks.0.attn.wq.input_scale' not in out
assert out['other.input_scale'] is unrelated
# ============================================================
# header _quantization_metadata
# ============================================================
def quant_metadata(layers: dict) -> dict:
import json
return {'_quantization_metadata': json.dumps({'format_version': '1.0', 'layers': layers})}
def test_read_quant_metadata_absent_returns_none():
fd, path = tempfile.mkstemp(suffix='.safetensors')
try:
write_fixture({'a.weight': torch.zeros(2)}, fd, path)
assert nt.read_quantization_metadata(path) is None
finally:
os.unlink(path)
def test_read_quant_metadata_parses_layers():
fd, path = tempfile.mkstemp(suffix='.safetensors')
try:
layers = {'blocks.0.attn.wq': {'format': 'int8_tensorwise'}}
write_fixture({'blocks.0.attn.wq.weight': torch.zeros(2)}, fd, path, metadata=quant_metadata(layers))
assert nt.read_quantization_metadata(path) == layers
finally:
os.unlink(path)
def test_read_quant_metadata_malformed_raises():
fd, path = tempfile.mkstemp(suffix='.safetensors')
try:
write_fixture({'a.weight': torch.zeros(2)}, fd, path, metadata={'_quantization_metadata': 'not json'})
try:
nt.read_quantization_metadata(path)
raise AssertionError('expected OverrideArchMismatch')
except nt.OverrideArchMismatch as e:
assert 'malformed' in str(e)
finally:
os.unlink(path)
def test_read_quant_metadata_non_dict_layers_raises():
import json
fd, path = tempfile.mkstemp(suffix='.safetensors')
try:
payload = json.dumps({'layers': {'blocks.0.attn.wq': 'int8_tensorwise'}})
write_fixture({'a.weight': torch.zeros(2)}, fd, path, metadata={'_quantization_metadata': payload})
try:
nt.read_quantization_metadata(path)
raise AssertionError('expected OverrideArchMismatch')
except nt.OverrideArchMismatch as e:
assert 'malformed' in str(e)
finally:
os.unlink(path)
def test_transcode_quant_metadata_synthesizes_markers():
"""Header entries become marker tensors for layers whose weight is present;
entries without a matching weight (companion components) are skipped."""
sd = {
'blocks.0.attn.wq.weight': torch.zeros((4, 4), dtype=torch.int8),
'blocks.0.attn.wq.weight_scale': torch.tensor(0.5),
}
layers = {
'blocks.0.attn.wq': {'format': 'int8_tensorwise'},
'text_encoder.mlp.up': {'format': 'int8_tensorwise'},
}
out, source = nt.transcode_quant_metadata(sd, layers)
assert source == 'header'
assert 'blocks.0.attn.wq.comfy_quant' in out
assert 'text_encoder.mlp.up.comfy_quant' not in out
assert 'blocks.0.attn.wq.comfy_quant' not in sd, 'input dict must not be mutated'
marked, fmt = nt.detect_comfy_quant(out, 'Test')
assert fmt == 'int8_tensorwise'
assert marked['blocks.0.attn.wq'] == {'format': 'int8_tensorwise'}
def test_transcode_quant_metadata_header_wins():
"""When both forms are present the header entry overwrites the marker,
matching the reference loader's precedence."""
sd = {
'blocks.0.attn.wq.weight': torch.zeros((4, 4), dtype=torch.int8),
'blocks.0.attn.wq.comfy_quant': comfy_marker('float8_e4m3fn'),
}
layers = {'blocks.0.attn.wq': {'format': 'int8_tensorwise', 'full_precision_matrix_mult': True}}
out, source = nt.transcode_quant_metadata(sd, layers)
assert source == 'both'
marked, fmt = nt.detect_comfy_quant(out, 'Test')
assert fmt == 'int8_tensorwise'
assert marked['blocks.0.attn.wq'].get('full_precision_matrix_mult') is True
def test_transcode_quant_metadata_no_matches_is_noop():
sd = {'blocks.0.attn.wq.weight': torch.zeros((4, 4), dtype=torch.int8)}
out, source = nt.transcode_quant_metadata(sd, {'unrelated.layer': {'format': 'int8_tensorwise'}})
assert source == 'markers'
assert out is sd
# ============================================================
# is_noop_converter
# ============================================================
@@ -639,9 +754,9 @@ class MockKwargsTransformer(MockMiniTransformer):
return cls(dim=config['dim'])
def write_fixture(state_dict_keys: dict, fd: int, path: str) -> str:
def write_fixture(state_dict_keys: dict, fd: int, path: str, metadata: dict | None = None) -> str:
os.close(fd)
safetensors.torch.save_file(state_dict_keys, path)
safetensors.torch.save_file(state_dict_keys, path, metadata=metadata)
return path
@@ -1229,6 +1344,98 @@ def test_load_comfy_fp8_end_to_end():
os.unlink(path)
def test_load_comfy_metadata_end_to_end():
"""Marker-less container: quant info only in the header
_quantization_metadata, layer names carrying the file's tensor prefix.
The load path must re-key through the prefix strip and land on the same
prequantized build as the marker form."""
fd, path = tempfile.mkstemp(suffix='.safetensors')
try:
dim = 8
raw = comfy_fixture(dim)
del raw['model.diffusion_model.in_proj.comfy_quant']
layers = {'model.diffusion_model.in_proj': {'format': 'int8_tensorwise'}}
write_fixture(raw, fd, path, metadata=quant_metadata(layers))
with ComfyTestEnv(dim):
spec = nt.TransformerSpec(cls=MockMiniTransformer)
transformer, _ = nt.load(
local_file=path,
repo_id='fake/repo',
spec=spec,
diffusers_cfg={},
dtype=torch.float32,
)
in_proj = transformer.in_proj
assert in_proj.__class__.__name__ == 'SDNQLinear'
assert in_proj.weight.dtype == torch.int8
assert torch.equal(in_proj.weight.detach().cpu(), raw['model.diffusion_model.in_proj.weight'])
assert transformer.out_proj.__class__ is torch.nn.Linear
finally:
if os.path.exists(path):
os.unlink(path)
def test_load_comfy_metadata_bare_names_end_to_end():
"""Marker-less container without a tensor prefix (official Comfy-Org
exports): metadata names match the bare keys and need no re-keying."""
fd, path = tempfile.mkstemp(suffix='.safetensors')
try:
dim = 8
raw = {key[len('model.diffusion_model.'):]: value for key, value in comfy_fixture(dim).items()}
del raw['in_proj.comfy_quant']
raw['in_proj.input_scale'] = torch.tensor(0.125, dtype=torch.float32)
layers = {'in_proj': {'format': 'int8_tensorwise'}}
write_fixture(raw, fd, path, metadata=quant_metadata(layers))
with ComfyTestEnv(dim):
spec = nt.TransformerSpec(cls=MockMiniTransformer)
transformer, _ = nt.load(
local_file=path,
repo_id='fake/repo',
spec=spec,
diffusers_cfg={},
dtype=torch.float32,
)
assert transformer.in_proj.__class__.__name__ == 'SDNQLinear'
assert torch.equal(transformer.in_proj.weight.detach().cpu(), raw['in_proj.weight'])
finally:
if os.path.exists(path):
os.unlink(path)
def test_load_comfy_full_precision_mm_excluded_from_matmul():
"""Layers flagged full_precision_matrix_mult land on the config's
per-layer matmul exclusion list, which apply_sdnq_options_to_model
honors when the user enables quantized matmul."""
fd, path = tempfile.mkstemp(suffix='.safetensors')
try:
dim = 8
import json
raw = comfy_fixture(dim)
payload = {'format': 'int8_tensorwise', 'full_precision_matrix_mult': True}
raw['model.diffusion_model.in_proj.comfy_quant'] = torch.tensor(list(json.dumps(payload).encode()), dtype=torch.uint8)
write_fixture(raw, fd, path)
with ComfyTestEnv(dim):
spec = nt.TransformerSpec(cls=MockMiniTransformer)
transformer, _ = nt.load(
local_file=path,
repo_id='fake/repo',
spec=spec,
diffusers_cfg={},
dtype=torch.float32,
)
assert 'in_proj.weight' in transformer.quantization_config.modules_to_not_use_matmul
assert transformer.in_proj.sdnq_dequantizer.use_quantized_matmul is False
finally:
if os.path.exists(path):
os.unlink(path)
def test_detect_comfy_mixed_formats_raises():
"""One file mixing int8 and fp8 layers has no single SDNQ mapping; reject."""
sd = {
@@ -1735,6 +1942,14 @@ def run_all():
test_remap_comfy_renames_and_reshapes_scale,
test_remap_comfy_passes_unmarked_keys_verbatim,
test_remap_comfy_does_not_mutate_input,
test_remap_comfy_drops_input_scale,
test_read_quant_metadata_absent_returns_none,
test_read_quant_metadata_parses_layers,
test_read_quant_metadata_malformed_raises,
test_read_quant_metadata_non_dict_layers_raises,
test_transcode_quant_metadata_synthesizes_markers,
test_transcode_quant_metadata_header_wins,
test_transcode_quant_metadata_no_matches_is_noop,
]:
run_test(cat, fn)
@@ -1803,6 +2018,9 @@ def run_all():
for fn in [
test_load_comfy_int8_end_to_end,
test_load_comfy_fp8_end_to_end,
test_load_comfy_metadata_end_to_end,
test_load_comfy_metadata_bare_names_end_to_end,
test_load_comfy_full_precision_mm_excluded_from_matmul,
test_detect_comfy_mixed_formats_raises,
test_load_comfy_marker_dtype_mismatch_raises,
test_load_comfy_unsupported_format_raises_mismatch,