mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 01:04:32 +02:00
feat(model): load comfyui comfy_quant int8 checkpoints via sdnq
Detect comfy_quant markers in native single-file transformer loads and adopt the pre-quantized tensors as SDNQ int8 layers instead of rejecting them. ComfyUI int8_tensorwise is a strict subset of SDNQ symmetric int8, so weights and scales are taken bit-exact with no dequantize-requantize round trip; quantized matmul and scale-dtype settings apply as usual. - detect_comfy_quant/remap_comfy_quant helpers plus a prequantized builder in native_transformer; file markers dictate the layer set - unsupported formats and wrong-arch markers fall back to the base repo - offline unit tests and an opt-in real-file krea2 test
This commit is contained in:
CalamitousFelicitousness
@@ -41,6 +41,7 @@ loading the adapter from the base repo.
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Callable
|
||||
@@ -69,6 +70,8 @@ DEFAULT_IGNORED_PREFIXES: tuple[str, ...] = (
|
||||
"text_encoders.",
|
||||
"vae.",
|
||||
)
|
||||
COMFY_QUANT_MARKER = ".comfy_quant"
|
||||
COMFY_SUPPORTED_FORMATS: tuple[str, ...] = ("int8_tensorwise",)
|
||||
|
||||
|
||||
class OverrideArchMismatch(Exception):
|
||||
@@ -415,6 +418,66 @@ def check_forbidden_markers(
|
||||
)
|
||||
|
||||
|
||||
def detect_comfy_quant(state_dict: dict, type_name: str) -> tuple[set[str], str] | None:
|
||||
"""Detect ComfyUI ``comfy_quant`` pre-quantized layers in a state dict.
|
||||
|
||||
ComfyUI's quantized checkpoints mark each quantized layer with a
|
||||
``<name>.comfy_quant`` uint8 tensor whose bytes are a JSON object naming
|
||||
the storage format, alongside a ``<name>.weight_scale`` tensor. Returns
|
||||
the set of marked module names and the format string, or ``None`` when no
|
||||
markers are present. Raises :class:`OverrideArchMismatch` for malformed
|
||||
markers or unsupported formats so the caller's base-repo fallback engages.
|
||||
"""
|
||||
marker_keys = [key for key in state_dict if key.endswith(COMFY_QUANT_MARKER)]
|
||||
if not marker_keys:
|
||||
return None
|
||||
marked: set[str] = set()
|
||||
formats: set[str] = set()
|
||||
for key in marker_keys:
|
||||
name = key[: -len(COMFY_QUANT_MARKER)]
|
||||
try:
|
||||
meta = json.loads(state_dict[key].cpu().numpy().tobytes())
|
||||
fmt = meta["format"]
|
||||
except Exception as e:
|
||||
raise OverrideArchMismatch(
|
||||
f"Load model: type={type_name} native_transformer comfy_quant marker "
|
||||
f"for {name!r} is malformed ({type(e).__name__}: {e})"
|
||||
) from e
|
||||
marked.add(name)
|
||||
formats.add(fmt)
|
||||
unsupported = sorted(formats - set(COMFY_SUPPORTED_FORMATS))
|
||||
if unsupported:
|
||||
log.error(
|
||||
f'Load model: type={type_name} quant=comfy format={",".join(unsupported)} not supported '
|
||||
f'(supported: {",".join(COMFY_SUPPORTED_FORMATS)})'
|
||||
)
|
||||
raise OverrideArchMismatch(
|
||||
f"Load model: type={type_name} native_transformer comfy_quant format "
|
||||
f"{', '.join(unsupported)} not supported"
|
||||
)
|
||||
return marked, next(iter(formats))
|
||||
|
||||
|
||||
def remap_comfy_quant(state_dict: dict, marked_names: set[str]) -> dict:
|
||||
"""Translate comfy_quant tensor naming to SDNQ naming.
|
||||
|
||||
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.
|
||||
"""
|
||||
scale_suffix = ".weight_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(scale_suffix) and key[: -len(scale_suffix)] in marked_names:
|
||||
remapped[f"{key[: -len(scale_suffix)]}.scale"] = value.reshape(-1, 1)
|
||||
continue
|
||||
remapped[key] = value
|
||||
return remapped
|
||||
|
||||
|
||||
def partition_siblings(
|
||||
state_dict: dict,
|
||||
siblings: dict[str, SiblingSpec],
|
||||
@@ -552,6 +615,159 @@ def build_component_quantized(
|
||||
return component
|
||||
|
||||
|
||||
def build_component_prequantized(
|
||||
*,
|
||||
component_name: str,
|
||||
state_dict: dict,
|
||||
config: dict,
|
||||
cls: type,
|
||||
marked_names: set[str],
|
||||
dtype,
|
||||
acceptable_missing: tuple[str, ...],
|
||||
zero_init_missing: tuple[str, ...] = (),
|
||||
**kwargs,
|
||||
) -> object:
|
||||
"""Build a component from a comfy_quant pre-quantized state dict, mapping
|
||||
the marked layers onto SDNQ int8 layers without dequantizing.
|
||||
|
||||
ComfyUI ``int8_tensorwise`` is a strict subset of SDNQ symmetric int8:
|
||||
same storage layout (unpacked int8 ``[out, in]``), same dequant math
|
||||
(``weight * scale``, no zero point), so the file's tensors are adopted
|
||||
bit-exact. The model is built under ``init_empty_weights`` and each marked
|
||||
Linear is swapped for an SDNQ wrapper with per-tensor dequant geometry
|
||||
(``group_size=-1``, scalar scale); the file dictates which layers are
|
||||
quantized, independent of the user's quantization settings. Weights load
|
||||
through ``SDNQQuantizer``'s pre-quantized path, which preserves the int8
|
||||
weight dtype and the fp32 scale (unlike :func:`build_component_quantized`,
|
||||
floating-point SDNQ params are deliberately not cast to the target dtype).
|
||||
|
||||
Layers are assembled in canonical dequant layout first;
|
||||
``apply_sdnq_options_to_model`` then enables quantized matmul per the
|
||||
user's settings (transposing eligible layers), matching the order used by
|
||||
``modules.sdnq.loader.load_sdnq_model``.
|
||||
|
||||
Caller is responsible for prefix stripping; the state dict arrives with
|
||||
its comfy marker keys intact and is remapped here.
|
||||
"""
|
||||
import rich.progress as rp
|
||||
from accelerate import init_empty_weights
|
||||
from accelerate.utils import set_module_tensor_to_device
|
||||
from diffusers.utils import get_module_from_name
|
||||
from modules.sdnq.quantizer import SDNQConfig, SDNQQuantizer
|
||||
from modules.sdnq.dequantizer import SDNQDequantizer
|
||||
from modules.sdnq.layers import get_sdnq_wrapper_class
|
||||
from modules.sdnq.forward import get_forward_func
|
||||
from modules.sdnq.loader import apply_sdnq_options_to_model
|
||||
|
||||
target_dtype = dtype if dtype is not None else devices.dtype
|
||||
sd = remap_comfy_quant(state_dict, marked_names)
|
||||
|
||||
quantization_config = SDNQConfig(
|
||||
weights_dtype="int8",
|
||||
quantized_matmul_dtype="int8",
|
||||
group_size=-1,
|
||||
use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul,
|
||||
dequantize_fp32=shared.opts.sdnq_dequantize_fp32,
|
||||
add_skip_keys=False,
|
||||
modules_to_not_convert=[],
|
||||
)
|
||||
quantizer = SDNQQuantizer(quantization_config, pre_quantized=True)
|
||||
quantizer.torch_dtype = target_dtype
|
||||
|
||||
with init_empty_weights(include_buffers=False):
|
||||
component = cls.from_config(config, **kwargs)
|
||||
|
||||
dequant_forward = get_forward_func("Linear", "int8", False)
|
||||
for name in sorted(marked_names):
|
||||
try:
|
||||
parent, child = get_module_from_name(component, name)
|
||||
linear = getattr(parent, child)
|
||||
except (AttributeError, ValueError) as e:
|
||||
raise OverrideArchMismatch(
|
||||
f"Load model: transformer=native {component_name} comfy_quant marked "
|
||||
f"module {name!r} not found in {cls.__name__}"
|
||||
) from e
|
||||
if not isinstance(linear, torch.nn.Linear):
|
||||
raise OverrideArchMismatch(
|
||||
f"Load model: transformer=native {component_name} comfy_quant marked "
|
||||
f"module {name!r} is {linear.__class__.__name__}, expected Linear"
|
||||
)
|
||||
linear.sdnq_dequantizer = SDNQDequantizer(
|
||||
result_dtype=target_dtype,
|
||||
result_shape=None,
|
||||
original_shape=torch.Size((linear.out_features, linear.in_features)),
|
||||
original_stride=(linear.in_features, 1),
|
||||
quantized_weight_shape=torch.Size((linear.out_features, linear.in_features)),
|
||||
weights_dtype="int8",
|
||||
quantized_matmul_dtype="int8",
|
||||
hadamard_group_size=256,
|
||||
group_size=-1,
|
||||
svd_rank=32,
|
||||
svd_steps=8,
|
||||
use_quantized_matmul=False,
|
||||
re_quantize_for_matmul=False,
|
||||
use_stochastic_rounding=False,
|
||||
use_hadamard=False,
|
||||
layer_class_name="Linear",
|
||||
)
|
||||
wrapped = get_sdnq_wrapper_class(linear, dequant_forward)
|
||||
wrapped.scale = torch.nn.Parameter(torch.empty((1, 1), dtype=torch.float32, device="meta"), requires_grad=False)
|
||||
wrapped.zero_point = None
|
||||
wrapped.svd_up = None
|
||||
wrapped.svd_down = None
|
||||
setattr(parent, child, wrapped)
|
||||
|
||||
target_device = (
|
||||
devices.cpu if shared.opts.diffusers_offload_mode != "none"
|
||||
else devices.device
|
||||
)
|
||||
|
||||
expected_keys = set(component.state_dict().keys())
|
||||
loaded_keys: set[str] = set()
|
||||
unexpected: list[str] = []
|
||||
total = len(sd)
|
||||
|
||||
pbar = rp.Progress(
|
||||
rp.TextColumn(f'[cyan]Load {component_name}:'),
|
||||
rp.BarColumn(),
|
||||
rp.MofNCompleteColumn(),
|
||||
rp.TaskProgressColumn(),
|
||||
rp.TimeRemainingColumn(),
|
||||
rp.TimeElapsedColumn(),
|
||||
rp.TextColumn('[cyan]{task.description}'),
|
||||
console=console,
|
||||
)
|
||||
with pbar:
|
||||
task = pbar.add_task(total=total, description=cls.__name__)
|
||||
for name, value in sd.items():
|
||||
if name in expected_keys:
|
||||
if quantizer.check_if_quantized_param(component, value, name):
|
||||
quantizer.create_quantized_param(component, value, name, target_device, dtype=target_dtype)
|
||||
else:
|
||||
if torch.is_floating_point(value):
|
||||
value = value.to(target_dtype)
|
||||
set_module_tensor_to_device(component, name, target_device, value=value, dtype=target_dtype)
|
||||
loaded_keys.add(name)
|
||||
else:
|
||||
unexpected.append(name)
|
||||
pbar.update(task, advance=1)
|
||||
|
||||
missing = sorted(expected_keys - loaded_keys)
|
||||
missing = materialize_zero_init(
|
||||
component, missing, zero_init_missing, device=target_device, dtype=target_dtype
|
||||
)
|
||||
validate_state_dict_load(component_name, missing, unexpected, acceptable_missing)
|
||||
|
||||
component = quantizer._process_model_after_weight_loading(component) # pylint: disable=protected-access
|
||||
component = apply_sdnq_options_to_model(
|
||||
component,
|
||||
dtype=target_dtype,
|
||||
dequantize_fp32=shared.opts.sdnq_dequantize_fp32,
|
||||
use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul,
|
||||
)
|
||||
return component
|
||||
|
||||
|
||||
def build_component(
|
||||
*,
|
||||
component_name: str,
|
||||
@@ -571,11 +787,16 @@ def build_component(
|
||||
"""Convert (if needed), instantiate, load weights, dtype-cast, quantize,
|
||||
and offload-place a single component. Raises on any hard failure.
|
||||
|
||||
For the transformer component under SDNQ, the per-tensor pre-mode path
|
||||
in :func:`build_component_quantized` is used so quantization is applied
|
||||
in flight (one layer's worth of bf16 in memory at a time). All other
|
||||
cases (siblings, non-quantized loads, NVIDIAModelOptConfig, layerwise
|
||||
quant) go through the standard load_state_dict + post-quantize path.
|
||||
Transformer state dicts carrying ComfyUI ``comfy_quant`` markers are
|
||||
dispatched to :func:`build_component_prequantized`, which adopts the
|
||||
file's int8 tensors as SDNQ layers regardless of quantization settings
|
||||
(the converter and ``quant_args`` are bypassed: the file is already
|
||||
quantized). For the transformer component under SDNQ, the per-tensor
|
||||
pre-mode path in :func:`build_component_quantized` is used so
|
||||
quantization is applied in flight (one layer's worth of bf16 in memory
|
||||
at a time). All other cases (siblings, non-quantized loads,
|
||||
NVIDIAModelOptConfig, layerwise quant) go through the standard
|
||||
load_state_dict + post-quantize path.
|
||||
|
||||
``dtype`` overrides ``devices.dtype`` when supplied; otherwise the global
|
||||
default is used. ``modules_to_not_convert`` and ``modules_dtype_dict``
|
||||
@@ -584,6 +805,27 @@ def build_component(
|
||||
reach ``cls.from_config`` for both construction paths.
|
||||
"""
|
||||
try:
|
||||
comfy_quant = detect_comfy_quant(state_dict, cls.__name__) if component_name == "transformer" else None
|
||||
if comfy_quant is not None:
|
||||
marked_names, comfy_format = comfy_quant
|
||||
log.info(
|
||||
f'Load model: transformer=native {component_name} quant=comfy '
|
||||
f'format={comfy_format} layers={len(marked_names)} keys={len(state_dict)} cls={cls.__name__}'
|
||||
)
|
||||
component = build_component_prequantized(
|
||||
component_name=component_name,
|
||||
state_dict=state_dict,
|
||||
config=config,
|
||||
cls=cls,
|
||||
marked_names=marked_names,
|
||||
dtype=dtype,
|
||||
acceptable_missing=acceptable_missing,
|
||||
zero_init_missing=zero_init_missing,
|
||||
**kwargs,
|
||||
)
|
||||
devices.torch_gc()
|
||||
return component
|
||||
|
||||
if converter is not None:
|
||||
log.debug(f'Load model: transformer=native {component_name} converter={converter.__name__} keys={len(state_dict)}')
|
||||
try:
|
||||
|
||||
@@ -8,8 +8,12 @@
|
||||
- Zero-init regression: a checkpoint that omits the dormant last.up/last.down residual branch,
|
||||
after materialize_zero_init, produces the same output as the base whose up is zeroed. Port
|
||||
only, so it runs without the reference.
|
||||
- comfy_quant real-file (opt-in): loads an actual ComfyUI int8_tensorwise Krea2 single file
|
||||
through the native loader and verifies SDNQ adoption plus a finite tiny forward. Enabled by
|
||||
setting $KREA2_COMFY_FILE to the .safetensors path; needs network for the base repo config
|
||||
($KREA2_COMFY_REPO, default CalamitousFelicitousness/Krea-2-Base-Diffusers).
|
||||
|
||||
No server, no checkpoint.
|
||||
No server, no checkpoint (except the opt-in comfy_quant test).
|
||||
"""
|
||||
|
||||
import importlib.util
|
||||
@@ -89,8 +93,8 @@ def run_parity(mmdit, port):
|
||||
print("PARITY OK")
|
||||
|
||||
|
||||
def load_materialize_zero_init():
|
||||
"""Import the real loader helper. native_transformer pulls in modules.shared, which needs
|
||||
def bootstrap_repo():
|
||||
"""Make repo modules importable. native_transformer pulls in modules.shared, which needs
|
||||
cmd_args parsed first, so bootstrap it the same way the native-transformer suite does."""
|
||||
repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
if repo not in sys.path:
|
||||
@@ -106,6 +110,10 @@ def load_materialize_zero_init():
|
||||
sys.argv = orig_argv
|
||||
installer.add_args(modules.cmd_args.parser)
|
||||
modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([])
|
||||
|
||||
|
||||
def load_materialize_zero_init():
|
||||
bootstrap_repo()
|
||||
from pipelines.native_transformer import materialize_zero_init
|
||||
return materialize_zero_init
|
||||
|
||||
@@ -159,12 +167,59 @@ def run_zero_init_regression(port):
|
||||
print("ZERO-INIT OK")
|
||||
|
||||
|
||||
def run_comfy_quant_real_file():
|
||||
"""Opt-in end-to-end check against a real ComfyUI int8_tensorwise Krea2 file: the native
|
||||
loader must adopt every marked linear as an SDNQ int8 layer and produce a finite output on
|
||||
a tiny forward. $KREA2_COMFY_LAYERS overrides the expected layer count (default 224)."""
|
||||
path = os.environ.get("KREA2_COMFY_FILE")
|
||||
if not path:
|
||||
print("COMFY REAL-FILE SKIPPED (set KREA2_COMFY_FILE to enable)")
|
||||
return
|
||||
assert os.path.exists(path), f"KREA2_COMFY_FILE not found: {path}"
|
||||
expected_layers = int(os.environ.get("KREA2_COMFY_LAYERS", "224"))
|
||||
repo_id = os.environ.get("KREA2_COMFY_REPO", "CalamitousFelicitousness/Krea-2-Base-Diffusers")
|
||||
|
||||
bootstrap_repo()
|
||||
from pipelines import native_transformer as nt
|
||||
from pipelines.krea2 import KREA2_SPEC
|
||||
|
||||
transformer, siblings = nt.load(local_file=path, repo_id=repo_id, spec=KREA2_SPEC, diffusers_cfg={})
|
||||
assert siblings == {}
|
||||
|
||||
sdnq_layers = [m for m in transformer.modules() if m.__class__.__name__ == "SDNQLinear"]
|
||||
print(f"comfy_quant real file: {len(sdnq_layers)} SDNQ layers")
|
||||
assert len(sdnq_layers) == expected_layers, f"expected {expected_layers} SDNQ layers, got {len(sdnq_layers)}"
|
||||
assert all(m.weight.dtype == torch.int8 for m in sdnq_layers), "all adopted weights must stay int8"
|
||||
assert transformer.blocks[0].attn.wq.__class__.__name__ == "SDNQLinear"
|
||||
assert getattr(transformer, "quantization_config", None) is not None
|
||||
|
||||
cfg = transformer.config
|
||||
param = next(p for p in transformer.parameters() if p.is_floating_point())
|
||||
device, dtype = param.device, param.dtype
|
||||
batch, txtlen, imglen = 1, 3, 4
|
||||
seq = txtlen + imglen
|
||||
gen = torch.Generator().manual_seed(1)
|
||||
img = torch.randn(batch, imglen, cfg.channels * cfg.patch ** 2, generator=gen).to(device=device, dtype=dtype)
|
||||
context = torch.randn(batch, txtlen, cfg.txtlayers, cfg.txtdim, generator=gen).to(device=device, dtype=dtype)
|
||||
timestep = torch.rand(batch, generator=gen).to(device=device, dtype=dtype)
|
||||
pos = torch.randint(0, 16, (batch, seq, 3), generator=gen).float().to(device=device)
|
||||
mask = torch.ones(batch, seq, dtype=torch.bool, device=device)
|
||||
with torch.no_grad():
|
||||
out = transformer(
|
||||
hidden_states=img, encoder_hidden_states=context, timestep=timestep,
|
||||
position_ids=pos, attention_mask=mask, return_dict=False,
|
||||
)[0]
|
||||
assert torch.isfinite(out).all(), "forward produced non-finite values"
|
||||
print("COMFY REAL-FILE OK")
|
||||
|
||||
|
||||
def main():
|
||||
port = load_port()
|
||||
# Parity runs first in pristine torch state; the regression imports modules afterwards.
|
||||
mmdit = load_reference()
|
||||
run_parity(mmdit, port)
|
||||
run_zero_init_regression(port)
|
||||
run_comfy_quant_real_file()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -8,6 +8,8 @@ Covers the pure helpers that own per-arch knob handling:
|
||||
- ``strip_prefix`` for single/multi prefix detection and mixed-prefix rejection
|
||||
- ``partition_siblings`` for inline-sibling key partitioning
|
||||
- ``check_forbidden_markers`` for structural-mismatch rejection
|
||||
- ``detect_comfy_quant`` for ComfyUI comfy_quant marker detection and format gating
|
||||
- ``remap_comfy_quant`` for comfy_quant -> SDNQ key translation
|
||||
- ``is_noop_converter`` for diffusers no-op lambda detection
|
||||
- ``validate_state_dict_load`` for unexpected / missing key handling
|
||||
- ``make_default_spec`` default-spec synthesis with diffusers converter pickup
|
||||
@@ -294,6 +296,111 @@ def test_forbidden_markers_empty_tuple_no_op():
|
||||
nt.check_forbidden_markers(sd, (), 'Test', '/tmp/x.safetensors')
|
||||
|
||||
|
||||
# ============================================================
|
||||
# detect_comfy_quant / remap_comfy_quant
|
||||
# ============================================================
|
||||
|
||||
def comfy_marker(fmt: str) -> torch.Tensor:
|
||||
import json
|
||||
return torch.tensor(list(json.dumps({'format': fmt}).encode()), dtype=torch.uint8)
|
||||
|
||||
|
||||
def test_detect_comfy_no_markers_returns_none():
|
||||
sd = {'blocks.0.attn.wq.weight': torch.zeros(2), 'blocks.0.attn.wq.bias': torch.zeros(2)}
|
||||
assert nt.detect_comfy_quant(sd, 'Test') is None
|
||||
|
||||
|
||||
def test_detect_comfy_valid_markers():
|
||||
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.comfy_quant': comfy_marker('int8_tensorwise'),
|
||||
'blocks.0.mlp.up.weight': torch.zeros((4, 4), dtype=torch.int8),
|
||||
'blocks.0.mlp.up.weight_scale': torch.tensor(0.25),
|
||||
'blocks.0.mlp.up.comfy_quant': comfy_marker('int8_tensorwise'),
|
||||
'norm.weight': torch.zeros(4),
|
||||
}
|
||||
detected = nt.detect_comfy_quant(sd, 'Test')
|
||||
assert detected is not None
|
||||
marked, fmt = detected
|
||||
assert marked == {'blocks.0.attn.wq', 'blocks.0.mlp.up'}
|
||||
assert fmt == 'int8_tensorwise'
|
||||
|
||||
|
||||
def test_detect_comfy_unsupported_format_raises():
|
||||
sd = {
|
||||
'blocks.0.attn.wq.weight': torch.zeros((4, 4)),
|
||||
'blocks.0.attn.wq.comfy_quant': comfy_marker('float8_e4m3fn_scaled'),
|
||||
}
|
||||
try:
|
||||
nt.detect_comfy_quant(sd, 'Test')
|
||||
raise AssertionError('expected OverrideArchMismatch')
|
||||
except nt.OverrideArchMismatch as e:
|
||||
assert 'float8_e4m3fn_scaled' in str(e)
|
||||
|
||||
|
||||
def test_detect_comfy_malformed_marker_raises():
|
||||
sd = {
|
||||
'blocks.0.attn.wq.weight': torch.zeros((4, 4)),
|
||||
'blocks.0.attn.wq.comfy_quant': torch.tensor(list(b'not json'), dtype=torch.uint8),
|
||||
}
|
||||
try:
|
||||
nt.detect_comfy_quant(sd, 'Test')
|
||||
raise AssertionError('expected OverrideArchMismatch')
|
||||
except nt.OverrideArchMismatch as e:
|
||||
assert 'malformed' in str(e)
|
||||
|
||||
|
||||
def test_detect_comfy_marker_missing_format_field_raises():
|
||||
import json
|
||||
sd = {
|
||||
'blocks.0.attn.wq.comfy_quant': torch.tensor(list(json.dumps({'fmt': 'x'}).encode()), dtype=torch.uint8),
|
||||
}
|
||||
try:
|
||||
nt.detect_comfy_quant(sd, 'Test')
|
||||
raise AssertionError('expected OverrideArchMismatch')
|
||||
except nt.OverrideArchMismatch as e:
|
||||
assert 'malformed' in str(e)
|
||||
|
||||
|
||||
def test_remap_comfy_renames_and_reshapes_scale():
|
||||
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.comfy_quant': comfy_marker('int8_tensorwise'),
|
||||
}
|
||||
out = nt.remap_comfy_quant(sd, {'blocks.0.attn.wq'})
|
||||
assert set(out.keys()) == {'blocks.0.attn.wq.weight', 'blocks.0.attn.wq.scale'}
|
||||
assert out['blocks.0.attn.wq.scale'].shape == (1, 1)
|
||||
assert out['blocks.0.attn.wq.scale'].item() == 0.5
|
||||
assert out['blocks.0.attn.wq.weight'].dtype == torch.int8
|
||||
|
||||
|
||||
def test_remap_comfy_passes_unmarked_keys_verbatim():
|
||||
unrelated_scale = torch.tensor(2.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.comfy_quant': comfy_marker('int8_tensorwise'),
|
||||
'norm.weight': torch.ones(4),
|
||||
'other.weight_scale': unrelated_scale,
|
||||
}
|
||||
out = nt.remap_comfy_quant(sd, {'blocks.0.attn.wq'})
|
||||
assert 'norm.weight' in out
|
||||
assert out['other.weight_scale'] is unrelated_scale, 'weight_scale outside marked names must pass through untouched'
|
||||
|
||||
|
||||
def test_remap_comfy_does_not_mutate_input():
|
||||
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.comfy_quant': comfy_marker('int8_tensorwise'),
|
||||
}
|
||||
keys_before = set(sd.keys())
|
||||
nt.remap_comfy_quant(sd, {'blocks.0.attn.wq'})
|
||||
assert set(sd.keys()) == keys_before, 'input dict must not be mutated (read_state_dict caches it)'
|
||||
|
||||
|
||||
# ============================================================
|
||||
# is_noop_converter
|
||||
# ============================================================
|
||||
@@ -831,6 +938,202 @@ def test_load_converter_crash_raises_mismatch():
|
||||
os.unlink(path)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Integration: comfy_quant pre-quantized load
|
||||
# ============================================================
|
||||
# Same patching strategy as the plain load tests, plus pinned SDNQ opts so the
|
||||
# builder's settings reads are deterministic. dtype=float32 makes the dequant
|
||||
# math bit-exact against a manual int8 * scale reference.
|
||||
|
||||
def comfy_fixture(dim: int) -> dict:
|
||||
"""comfy_quant int8_tensorwise export shape: in_proj carries int8 weight +
|
||||
scalar fp32 scale + marker; out_proj and biases stay plain f16."""
|
||||
return {
|
||||
'model.diffusion_model.in_proj.weight': torch.randint(-128, 127, (dim, dim), dtype=torch.int8),
|
||||
'model.diffusion_model.in_proj.weight_scale': torch.tensor(0.03125, dtype=torch.float32),
|
||||
'model.diffusion_model.in_proj.comfy_quant': comfy_marker('int8_tensorwise'),
|
||||
'model.diffusion_model.in_proj.bias': torch.randn(dim, dtype=torch.float16),
|
||||
'model.diffusion_model.out_proj.weight': torch.randn(dim, dim, dtype=torch.float16),
|
||||
'model.diffusion_model.out_proj.bias': torch.zeros(dim, dtype=torch.float16),
|
||||
}
|
||||
|
||||
|
||||
class ComfyTestEnv:
|
||||
"""Patches quant helpers + config fetch and pins the SDNQ opts the
|
||||
prequantized builder reads, restoring everything on exit."""
|
||||
|
||||
def __init__(self, dim: int):
|
||||
self.dim = dim
|
||||
|
||||
def __enter__(self):
|
||||
from modules import model_quant, shared
|
||||
self.shared = shared
|
||||
self.orig_fetch = nt.fetch_component_config
|
||||
self.orig_get_dit = model_quant.get_dit_args
|
||||
self.orig_get_qtype = model_quant.get_quant_type
|
||||
self.orig_do_post = model_quant.do_post_load_quant
|
||||
self.model_quant = model_quant
|
||||
nt.fetch_component_config = lambda repo, sub: {'dim': self.dim}
|
||||
model_quant.get_dit_args = lambda *a, **k: ({}, {})
|
||||
model_quant.get_quant_type = lambda *a, **k: None
|
||||
model_quant.do_post_load_quant = lambda *a, **k: None
|
||||
self.orig_opts = {
|
||||
'sdnq_use_quantized_matmul': shared.opts.sdnq_use_quantized_matmul,
|
||||
'sdnq_dequantize_fp32': shared.opts.sdnq_dequantize_fp32,
|
||||
'diffusers_offload_mode': shared.opts.diffusers_offload_mode,
|
||||
}
|
||||
shared.opts.data['sdnq_use_quantized_matmul'] = False
|
||||
shared.opts.data['sdnq_dequantize_fp32'] = True
|
||||
shared.opts.data['diffusers_offload_mode'] = 'model' # force CPU placement
|
||||
return self
|
||||
|
||||
def __exit__(self, *exc):
|
||||
nt.fetch_component_config = self.orig_fetch
|
||||
self.model_quant.get_dit_args = self.orig_get_dit
|
||||
self.model_quant.get_quant_type = self.orig_get_qtype
|
||||
self.model_quant.do_post_load_quant = self.orig_do_post
|
||||
for key, value in self.orig_opts.items():
|
||||
self.shared.opts.data[key] = value
|
||||
return False
|
||||
|
||||
|
||||
def test_load_comfy_int8_end_to_end():
|
||||
"""Full pipeline on a comfy_quant fixture: the marked linear becomes an
|
||||
SDNQ int8 layer holding the file's exact tensors, the unmarked linear
|
||||
loads plain, and the forward pass matches a dequantized reference."""
|
||||
fd, path = tempfile.mkstemp(suffix='.safetensors')
|
||||
try:
|
||||
dim = 8
|
||||
raw = comfy_fixture(dim)
|
||||
write_fixture(raw, fd, path)
|
||||
|
||||
with ComfyTestEnv(dim):
|
||||
spec = nt.TransformerSpec(cls=MockMiniTransformer)
|
||||
transformer, siblings = nt.load(
|
||||
local_file=path,
|
||||
repo_id='fake/repo',
|
||||
spec=spec,
|
||||
diffusers_cfg={},
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
assert siblings == {}
|
||||
in_proj = transformer.in_proj
|
||||
assert in_proj.__class__.__name__ == 'SDNQLinear', f'marked layer is {in_proj.__class__.__name__}'
|
||||
dq = in_proj.sdnq_dequantizer
|
||||
assert dq.weights_dtype == 'int8'
|
||||
assert dq.group_size == -1
|
||||
assert dq.use_quantized_matmul is False
|
||||
assert dq.original_shape == (dim, dim)
|
||||
|
||||
# File tensors adopted bit-exact: int8 codes and fp32 scalar scale.
|
||||
assert in_proj.weight.dtype == torch.int8
|
||||
assert torch.equal(in_proj.weight.detach().cpu(), raw['model.diffusion_model.in_proj.weight'])
|
||||
assert in_proj.scale.dtype == torch.float32
|
||||
assert tuple(in_proj.scale.shape) == (1, 1)
|
||||
assert in_proj.scale.item() == raw['model.diffusion_model.in_proj.weight_scale'].item()
|
||||
assert in_proj.zero_point is None
|
||||
|
||||
# Dequant matches comfy semantics exactly: fp = int8.float() * scale.
|
||||
expected = raw['model.diffusion_model.in_proj.weight'].float() * raw['model.diffusion_model.in_proj.weight_scale']
|
||||
dequantized = dq(in_proj.weight, in_proj.scale, zero_point=None, svd_up=None, svd_down=None)
|
||||
assert torch.equal(dequantized.detach().cpu(), expected)
|
||||
|
||||
# Forward parity against a reference built from the dequantized weight.
|
||||
x = torch.randn(2, dim)
|
||||
out = in_proj(x)
|
||||
ref = torch.nn.functional.linear(x, expected, in_proj.bias.detach().cpu())
|
||||
assert torch.allclose(out.detach().cpu(), ref, atol=1e-6)
|
||||
|
||||
# Unmarked layer loads as a plain Linear cast to the target dtype.
|
||||
assert transformer.out_proj.__class__ is torch.nn.Linear
|
||||
assert transformer.out_proj.weight.dtype == torch.float32
|
||||
assert torch.allclose(
|
||||
transformer.out_proj.weight.detach().cpu(),
|
||||
raw['model.diffusion_model.out_proj.weight'].float(),
|
||||
)
|
||||
|
||||
# Marked as SDNQ-quantized so downstream never re-quantizes.
|
||||
assert getattr(transformer, 'quantization_config', None) is not None
|
||||
finally:
|
||||
if os.path.exists(path):
|
||||
os.unlink(path)
|
||||
|
||||
|
||||
def test_load_comfy_unsupported_format_raises_mismatch():
|
||||
"""A comfy_quant file in a format sdnext cannot adopt must surface as
|
||||
OverrideArchMismatch so load_transformer falls back to the base repo."""
|
||||
fd, path = tempfile.mkstemp(suffix='.safetensors')
|
||||
try:
|
||||
dim = 8
|
||||
raw = comfy_fixture(dim)
|
||||
raw['model.diffusion_model.in_proj.comfy_quant'] = comfy_marker('float8_e4m3fn_scaled')
|
||||
write_fixture(raw, fd, path)
|
||||
|
||||
with ComfyTestEnv(dim):
|
||||
spec = nt.TransformerSpec(cls=MockMiniTransformer)
|
||||
raised = False
|
||||
try:
|
||||
nt.load(local_file=path, repo_id='fake/repo', spec=spec, diffusers_cfg={})
|
||||
except nt.OverrideArchMismatch as e:
|
||||
raised = True
|
||||
assert 'float8_e4m3fn_scaled' in str(e)
|
||||
assert raised, 'expected OverrideArchMismatch'
|
||||
finally:
|
||||
if os.path.exists(path):
|
||||
os.unlink(path)
|
||||
|
||||
|
||||
def test_load_comfy_marker_for_unknown_module_raises_mismatch():
|
||||
"""A marker naming a module the target class does not have means the file
|
||||
belongs to a different arch; must fall back, not crash."""
|
||||
fd, path = tempfile.mkstemp(suffix='.safetensors')
|
||||
try:
|
||||
dim = 8
|
||||
raw = comfy_fixture(dim)
|
||||
raw['model.diffusion_model.ghost.weight'] = torch.zeros((dim, dim), dtype=torch.int8)
|
||||
raw['model.diffusion_model.ghost.weight_scale'] = torch.tensor(1.0)
|
||||
raw['model.diffusion_model.ghost.comfy_quant'] = comfy_marker('int8_tensorwise')
|
||||
write_fixture(raw, fd, path)
|
||||
|
||||
with ComfyTestEnv(dim):
|
||||
spec = nt.TransformerSpec(cls=MockMiniTransformer)
|
||||
raised = False
|
||||
try:
|
||||
nt.load(local_file=path, repo_id='fake/repo', spec=spec, diffusers_cfg={})
|
||||
except nt.OverrideArchMismatch as e:
|
||||
raised = True
|
||||
assert 'ghost' in str(e)
|
||||
assert raised, 'expected OverrideArchMismatch'
|
||||
finally:
|
||||
if os.path.exists(path):
|
||||
os.unlink(path)
|
||||
|
||||
|
||||
def test_build_component_comfy_preempts_sdnq_fresh_quant():
|
||||
"""When SDNQ on-load quant settings are active (quant_type=SDNQConfig), a
|
||||
comfy_quant file must still take the pre-quantized path: fresh quant of
|
||||
int8 data would corrupt it. quant_args={} proves the comfy branch never
|
||||
reaches build_component_quantized (which requires a quantization_config)."""
|
||||
dim = 8
|
||||
raw = comfy_fixture(dim)
|
||||
sd = {k[len('model.diffusion_model.'):]: v for k, v in raw.items()}
|
||||
with ComfyTestEnv(dim):
|
||||
component = nt.build_component(
|
||||
component_name='transformer',
|
||||
state_dict=sd,
|
||||
config={'dim': dim},
|
||||
cls=MockMiniTransformer,
|
||||
converter=None,
|
||||
acceptable_missing=('rope.',),
|
||||
quant_args={},
|
||||
quant_type='SDNQConfig',
|
||||
dtype=torch.float32,
|
||||
)
|
||||
assert component.in_proj.__class__.__name__ == 'SDNQLinear'
|
||||
assert component.in_proj.weight.dtype == torch.int8
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Run
|
||||
# ============================================================
|
||||
@@ -879,6 +1182,20 @@ def run_all():
|
||||
]:
|
||||
run_test(cat, fn)
|
||||
|
||||
log.warning('=== comfy_quant detection / remap ===')
|
||||
cat = category('comfy')
|
||||
for fn in [
|
||||
test_detect_comfy_no_markers_returns_none,
|
||||
test_detect_comfy_valid_markers,
|
||||
test_detect_comfy_unsupported_format_raises,
|
||||
test_detect_comfy_malformed_marker_raises,
|
||||
test_detect_comfy_marker_missing_format_field_raises,
|
||||
test_remap_comfy_renames_and_reshapes_scale,
|
||||
test_remap_comfy_passes_unmarked_keys_verbatim,
|
||||
test_remap_comfy_does_not_mutate_input,
|
||||
]:
|
||||
run_test(cat, fn)
|
||||
|
||||
log.warning('=== noop_converter detection ===')
|
||||
cat = category('noop')
|
||||
for fn in [
|
||||
@@ -939,6 +1256,16 @@ def run_all():
|
||||
]:
|
||||
run_test(cat, fn)
|
||||
|
||||
log.warning('=== comfy_quant load ===')
|
||||
cat = category('comfy_load')
|
||||
for fn in [
|
||||
test_load_comfy_int8_end_to_end,
|
||||
test_load_comfy_unsupported_format_raises_mismatch,
|
||||
test_load_comfy_marker_for_unknown_module_raises_mismatch,
|
||||
test_build_component_comfy_preempts_sdnq_fresh_quant,
|
||||
]:
|
||||
run_test(cat, fn)
|
||||
|
||||
log.warning('=== Results ===')
|
||||
total_passed = 0
|
||||
total_failed = 0
|
||||
|
||||
Reference in New Issue
Block a user