nodes/silveroxides-ComfyUI-ModelUtils
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Match diffusion model layers across prefixes

Browse Source
This commit is contained in:
silveroxides
2026-09-23 02:36:18 +02:00
parent 57bd63e1eb
commit 1861d75aab
10 changed files with 354 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes/dora_extract_wd.py
+28 -7
View File
@@ -23,6 +23,7 @@ from .artifact_paths import canonical_model_artifact_path
from .adaptive_svd import ADAPTIVE_PARTIAL_MODES, adaptive_partial_svd
from .extraction_stream import dora_difference, paired_async_tensors, retry_cuda_oom_on_cpu
from .layer_parameters import parameter_input, resolve_layer_parameters
from .quantization_guard import DiffusionQuantization, diffusion_key_map
@@ -542,9 +543,21 @@ def extract_dora_from_files(
handler_b = MemoryEfficientSafeOpen(model_b_path, low_memory=lazy_load)
try:
keys_a = set(handler_a.keys())
keys_b = set(handler_b.keys())
weight_keys = sorted(k for k in keys_a if k.endswith(".weight"))
quantization_a = DiffusionQuantization(handler_a, model_a_path, "DoRA Extract model A")
quantization_b = DiffusionQuantization(handler_b, model_b_path, "DoRA Extract model B")
keys_a = set(quantization_a.data_keys)
keys_b = set(quantization_b.data_keys)
logical_a = diffusion_key_map(keys_a, "DoRA Extract model A")
logical_b = diffusion_key_map(keys_b, "DoRA Extract model B")
paired_keys_b = {
key_a: logical_b.get(logical) for logical, key_a in logical_a.items()
}
weight_keys = sorted(
key for key in keys_a
if key.endswith(".weight")
and key not in quantization_a.isolated_low_bit_keys
and paired_keys_b[key] not in quantization_b.isolated_low_bit_keys
)
if mode == "fixed":
profile = "extract:fixed"
defaults = {"linear_dim": linear_param, "conv_dim": conv_param,
@@ -573,15 +586,16 @@ def extract_dora_from_files(
layer_parameters, profile, weight_keys, defaults, node_name="DoRA Extract"
)
work_units = [
(key, key, key if key in keys_b else None)
(key, key, paired_keys_b[key])
for key in weight_keys
if (not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) if not include_mode else _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns))
and not (key not in keys_b and mismatch_mode == "skip")
and not (paired_keys_b[key] is None and mismatch_mode == "skip")
]
pbar = comfy.utils.ProgressBar(len(work_units))
stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
def _process_layer(key, cpu_a, cpu_b):
key_b = paired_keys_b[key]
values = resolved_parameters.get(key, defaults)
if mode == "fixed":
layer_linear_param = values["linear_dim"]
@@ -610,7 +624,7 @@ def extract_dora_from_files(
return "skipped", None
# Load tensors with pinned memory for CUDA
if key not in keys_b:
if key_b is None:
if mismatch_mode == "skip":
return "skipped", None
if mismatch_mode == "error":
@@ -709,9 +723,12 @@ def extract_dora_from_files(
writer_context = atomic_uel_writer(output_path)
writer = writer_context.__enter__()
stream = None
try:
stream = paired_async_tensors(
handler_a, handler_b, work_units, pin_memory=str(device).startswith("cuda")
handler_a, handler_b, work_units, pin_memory=str(device).startswith("cuda"),
quantization_a=quantization_a, quantization_b=quantization_b,
compute_dtype=torch.float32,
)
for key, cpu_a, cpu_b in tqdm(stream, total=len(work_units), desc="Extracting DoRA", unit="layers"):
status, layer_sd = _process_layer(key, cpu_a, cpu_b)
@@ -728,9 +745,13 @@ def extract_dora_from_files(
pbar.update(1)
stats["skipped"] += len(weight_keys) - len(work_units)
except BaseException as exc:
if stream is not None:
stream.close()
writer_context.__exit__(type(exc), exc, exc.__traceback__)
raise
else:
if stream is not None:
stream.close()
writer_context.__exit__(None, None, None)
print(f"[DoRA Extract] Done: {stats['extracted']} extracted, {stats['chunked']} chunked, "
nodes/dora_learned_wd.py
+28 -7
View File
@@ -20,6 +20,7 @@ from unifiedefficientloader import MemoryEfficientSafeOpen
from .uel_io import atomic_uel_writer
from .extraction_stream import dora_difference, paired_async_tensors, retry_cuda_oom_on_cpu
from .layer_parameters import parameter_input, resolve_layer_parameters
from .quantization_guard import DiffusionQuantization, diffusion_key_map
# Import the existing SV ranking and chunking logic so we don't have to duplicate it
from .dora_extract_wd import (
@@ -379,9 +380,21 @@ def extract_dora_learned_from_files(
handler_b = MemoryEfficientSafeOpen(model_b_path, low_memory=lazy_load)
try:
keys_a = set(handler_a.keys())
keys_b = set(handler_b.keys())
weight_keys = sorted(k for k in keys_a if k.endswith(".weight"))
quantization_a = DiffusionQuantization(handler_a, model_a_path, "Learned DoRA Extract model A")
quantization_b = DiffusionQuantization(handler_b, model_b_path, "Learned DoRA Extract model B")
keys_a = set(quantization_a.data_keys)
keys_b = set(quantization_b.data_keys)
logical_a = diffusion_key_map(keys_a, "Learned DoRA Extract model A")
logical_b = diffusion_key_map(keys_b, "Learned DoRA Extract model B")
paired_keys_b = {
key_a: logical_b.get(logical) for logical, key_a in logical_a.items()
}
weight_keys = sorted(
key for key in keys_a
if key.endswith(".weight")
and key not in quantization_a.isolated_low_bit_keys
and paired_keys_b[key] not in quantization_b.isolated_low_bit_keys
)
if mode == "fixed":
profile = "extract:fixed"
defaults = {"linear_dim": linear_param, "conv_dim": conv_param,
@@ -411,15 +424,16 @@ def extract_dora_learned_from_files(
node_name="Learned DoRA Extract",
)
work_units = [
(key, key, key if key in keys_b else None)
(key, key, paired_keys_b[key])
for key in weight_keys
if (not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) if not include_mode else _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns))
and not (key not in keys_b and mismatch_mode == "skip")
and not (paired_keys_b[key] is None and mismatch_mode == "skip")
]
pbar = comfy.utils.ProgressBar(len(work_units))
stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
def _process_layer(key, cpu_a, cpu_b):
key_b = paired_keys_b[key]
values = resolved_parameters.get(key, defaults)
if mode == "fixed":
layer_linear_param, layer_conv_param = values["linear_dim"], values["conv_dim"]
@@ -447,7 +461,7 @@ def extract_dora_learned_from_files(
if (_matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) if not include_mode else not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns)):
return "skipped", None
if key not in keys_b:
if key_b is None:
if mismatch_mode == "skip":
return "skipped", None
if mismatch_mode == "error":
@@ -564,9 +578,12 @@ def extract_dora_learned_from_files(
writer_context = atomic_uel_writer(output_path)
writer = writer_context.__enter__()
stream = None
try:
stream = paired_async_tensors(
handler_a, handler_b, work_units, pin_memory=str(device).startswith("cuda")
handler_a, handler_b, work_units, pin_memory=str(device).startswith("cuda"),
quantization_a=quantization_a, quantization_b=quantization_b,
compute_dtype=torch.float32,
)
for key, cpu_a, cpu_b in tqdm(stream, total=len(work_units), desc="Extracting Learned DoRA", unit="layers"):
status, layer_sd = _process_layer(key, cpu_a, cpu_b)
@@ -583,9 +600,13 @@ def extract_dora_learned_from_files(
pbar.update(1)
stats["skipped"] += len(weight_keys) - len(work_units)
except BaseException as exc:
if stream is not None:
stream.close()
writer_context.__exit__(type(exc), exc, exc.__traceback__)
raise
else:
if stream is not None:
stream.close()
writer_context.__exit__(None, None, None)
print(f"[Learned DoRA] Done: {stats['extracted']} extracted, {stats['chunked']} chunked, "
nodes/lora_extract_svd.py
+13 -7
View File
@@ -20,7 +20,7 @@ from .device_utils import (
from unifiedefficientloader import MemoryEfficientSafeOpen
from .uel_io import atomic_uel_writer
from .artifact_paths import canonical_model_artifact_path
from .quantization_guard import DiffusionQuantization
from .quantization_guard import DiffusionQuantization, diffusion_key_map
from .adaptive_svd import ADAPTIVE_PARTIAL_MODES, adaptive_partial_svd
from .extraction_stream import lora_difference, paired_async_tensors, retry_cuda_oom_on_cpu
from .layer_parameters import parameter_input, resolve_layer_parameters
@@ -556,6 +556,11 @@ def extract_lora_from_files(
low_bit_b = set(quantization_b.isolated_low_bit_keys)
keys_a = set(quantization_a.data_keys)
keys_b = set(quantization_b.data_keys)
logical_a = diffusion_key_map(keys_a, "LoRA Extract model A")
logical_b = diffusion_key_map(keys_b, "LoRA Extract model B")
paired_keys_b = {
key_a: logical_b.get(logical) for logical, key_a in logical_a.items()
}
direct_diff_suffixes = (".scale", ".lin")
parameter_layers = sorted(
k for k in keys_a
@@ -568,7 +573,7 @@ def extract_lora_from_files(
extraction_keys = [
k for k in keys_a
if k not in low_bit_a
and k not in low_bit_b
and paired_keys_b[k] not in low_bit_b
and (k.endswith(".weight")
or (
include_1d_diffs
@@ -629,15 +634,16 @@ def extract_lora_from_files(
node_name="LoRA Extract",
)
work_units = [
(key, key, key if key in keys_b else None)
(key, key, paired_keys_b[key])
for key in extraction_keys
if (not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) if not include_mode else _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns))
and not (key not in keys_b and mismatch_mode == "skip")
and not (paired_keys_b[key] is None and mismatch_mode == "skip")
]
pbar = comfy.utils.ProgressBar(len(work_units))
stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
def _process_layer(key, cpu_a, cpu_b):
key_b = paired_keys_b[key]
values = resolved_parameters.get(key, defaults)
if mode == "fixed":
layer_linear_param = values["linear_dim"]
@@ -674,14 +680,14 @@ def extract_lora_from_files(
return "skipped", None
source_dtypes = [quantization_a.logical_dtype(key, save_torch_dtype)]
if key in keys_b:
source_dtypes.append(quantization_b.logical_dtype(key, save_torch_dtype))
if key_b is not None:
source_dtypes.append(quantization_b.logical_dtype(key_b, save_torch_dtype))
layer_save_dtype = (
torch.float32 if torch.float32 in source_dtypes else save_torch_dtype
)
# Load tensors with pinned memory for CUDA
if key not in keys_b:
if key_b is None:
if mismatch_mode == "skip":
return "skipped", None
if mismatch_mode == "error":
nodes/merger.py
+30 -20
View File
@@ -13,7 +13,7 @@ from .device_utils import (
)
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
from .quantization_guard import DiffusionQuantization, inspect_low_bit_input, write_preserved_tensor
from .quantization_guard import DiffusionQuantization, diffusion_key_map, inspect_low_bit_input, write_preserved_tensor
from .lora_resize import canonical_lora_block_name, is_direct_diff_key, layer_tensor_keys, parse_lora_layers
from .layer_parameters import parameter_input, resolve_layer_parameters
from .uel_io import atomic_uel_writer, stream_work_units
@@ -109,17 +109,22 @@ class MergerLogic:
primary_handler = handlers[primary_model_name]
diffusion_quantizers = {}
diffusion_maps = {}
if model_type == "diffusion_models":
try:
diffusion_quantizers = {
name: DiffusionQuantization(handler, handler_paths[name], f"Model '{name}'")
for name, handler in handlers.items()
}
diffusion_maps = {
name: diffusion_key_map(quantizer.data_keys, f"Merger model {name!r}")
for name, quantizer in diffusion_quantizers.items()
}
except Exception:
for handler in handlers.values():
handler.__exit__(None, None, None)
raise
all_keys = diffusion_quantizers[primary_model_name].data_keys
all_keys = tuple(diffusion_maps[primary_model_name])
metadata = diffusion_quantizers[primary_model_name].output_metadata()
else:
all_keys = primary_handler.keys()
@@ -154,8 +159,7 @@ class MergerLogic:
for name, handler in handlers.items():
if name != primary_model_name:
secondary_keys = set(
diffusion_quantizers[name].data_keys
if name in diffusion_quantizers else handler.keys()
diffusion_maps[name] if name in diffusion_maps else handler.keys()
)
missing = primary_keys - secondary_keys
extra = secondary_keys - primary_keys
@@ -275,15 +279,14 @@ class MergerLogic:
)
for name in dict.fromkeys(requested_names):
if name in handlers and key in (
diffusion_quantizers[name].data_keys
if name in diffusion_quantizers else handlers[name].keys()
diffusion_maps[name] if name in diffusion_maps else handlers[name].keys()
):
source_key = diffusion_maps[name][key] if name in diffusion_maps else key
required = (
diffusion_quantizers[name].required_keys(key)
if name in diffusion_quantizers and key in diffusion_quantizers[name].quantized_keys
else [key]
diffusion_quantizers[name].required_keys(source_key)
if name in diffusion_quantizers else [source_key]
)
entries[name] = list(dict.fromkeys([key, *required]))
entries[name] = list(dict.fromkeys(required))
alpha_info = alpha_normalization.get(name, {}).get(key)
if alpha_info is not None:
entries[name].append(alpha_info[0])
@@ -331,33 +334,40 @@ class MergerLogic:
if diffusion_quantizers:
for name, quantizer in diffusion_quantizers.items():
if key in quantizer.quantized_keys and (name, key) in loaded:
source_key = diffusion_maps[name].get(key)
if source_key is None or (name, source_key) not in loaded:
continue
if source_key in quantizer.quantized_keys:
tensors = {
source_key: loaded[(name, source_key)]
for source_key in dict.fromkeys([key, *quantizer.required_keys(key)])
part: loaded[(name, part)]
for part in quantizer.required_keys(source_key)
}
loaded[(name, key)] = quantizer.decode(
key, tensors, process_dtype, device="cpu"
source_key, tensors, process_dtype, device="cpu"
)
else:
loaded[(name, key)] = loaded[(name, source_key)]
# Pre-load Model A's tensor with pinned memory for CUDA
cpu_tensor = loaded[(primary_model_name, key)]
# Determine original dtypes across all active models for this key
original_dtypes = [
diffusion_quantizers[primary_model_name].logical_dtype(key, process_dtype)
diffusion_quantizers[primary_model_name].logical_dtype(
diffusion_maps[primary_model_name][key], process_dtype
)
if primary_model_name in diffusion_quantizers
and key in diffusion_quantizers[primary_model_name].quantized_keys
else primary_handler.get_dtype(key)
]
for name, handler in handlers.items():
if name != primary_model_name and key in (
diffusion_quantizers[name].data_keys
if name in diffusion_quantizers else handler.keys()
diffusion_maps[name] if name in diffusion_maps else handler.keys()
):
original_dtypes.append(
diffusion_quantizers[name].logical_dtype(key, process_dtype)
if name in diffusion_quantizers and key in diffusion_quantizers[name].quantized_keys
diffusion_quantizers[name].logical_dtype(
diffusion_maps[name][key], process_dtype
)
if name in diffusion_quantizers
else handler.get_dtype(key)
)
nodes/model_analysis.py
+33 -8
View File
@@ -29,7 +29,7 @@ from .device_utils import (
)
from .merger import _compile_patterns, _matches_any_pattern, load_documentation_from_file
from .lora_resize import layer_tensor_keys, parse_lora_layers, validate_canonical_blocks
from .quantization_guard import DiffusionQuantization, inspect_low_bit_input
from .quantization_guard import DiffusionQuantization, diffusion_key_map, inspect_low_bit_input
FLOAT_DTYPES = {torch.float16, torch.bfloat16, torch.float32, torch.float64}
@@ -859,19 +859,36 @@ class ModelAnalysisLogic:
keys_b = set(
diffusion_quantizers[1].data_keys if diffusion_quantizers else handlers[1].keys()
)
diffusion_maps = None
if diffusion_quantizers:
diffusion_maps = (
diffusion_key_map(keys_a, "Model Analysis input A"),
diffusion_key_map(keys_b, "Model Analysis input B"),
)
keys_a = set(diffusion_maps[0])
keys_b = set(diffusion_maps[1])
glob_mode = bool(params.get("glob_patterns", False))
excluded_patterns = _compile_patterns(
params.get("exclude_patterns", ""),
glob_mode=glob_mode,
)
include_mode = bool(params.get("include_mode", False))
def excluded(key):
aliases = [key]
if diffusion_maps:
aliases.extend(
mapping[key] for mapping in diffusion_maps if key in mapping
)
matched = any(
_matches_any_pattern(alias, excluded_patterns, glob_mode=glob_mode)
for alias in aliases
)
return matched != include_mode
excluded_a = {
key for key in keys_a
if _matches_any_pattern(key, excluded_patterns, glob_mode=glob_mode) != include_mode
key for key in keys_a if excluded(key)
}
excluded_b = {
key for key in keys_b
if _matches_any_pattern(key, excluded_patterns, glob_mode=glob_mode) != include_mode
key for key in keys_b if excluded(key)
}
keys_a.difference_update(excluded_a)
keys_b.difference_update(excluded_b)
@@ -944,9 +961,17 @@ class ModelAnalysisLogic:
inventory["a_only"].extend(sorted(pass_a - pass_b))
inventory["b_only"].extend(sorted(pass_b - pass_a))
else:
units = [(key, key, key) for key in sorted(keys_a & keys_b)]
inventory["a_only"] = sorted(keys_a - keys_b)
inventory["b_only"] = sorted(keys_b - keys_a)
if diffusion_maps:
units = [
(key, diffusion_maps[0][key], diffusion_maps[1][key])
for key in sorted(keys_a & keys_b)
]
inventory["a_only"] = sorted(keys_a - keys_b)
inventory["b_only"] = sorted(keys_b - keys_a)
else:
units = [(key, key, key) for key in sorted(keys_a & keys_b)]
inventory["a_only"] = sorted(keys_a - keys_b)
inventory["b_only"] = sorted(keys_b - keys_a)
global_raw = RawStats()
layer_stats: dict[str, RawStats] = {}
block_stats: dict[str, RawStats] = {}
nodes/quantization_guard.py
+25
View File
@@ -39,6 +39,25 @@ _QUANT_SIDECARS = (
)
def diffusion_key_map(keys, label):
"""Map optional diffusion wrappers to each file's physical tensor keys."""
mapped = {}
for key in sorted(keys):
logical = key
for prefix in ("model.diffusion_model.", "diffusion_model."):
if key.startswith(prefix):
logical = key[len(prefix):]
break
previous = mapped.get(logical)
if previous is not None:
raise ValueError(
f"{label} has duplicate diffusion layer {logical!r}: "
f"{previous!r} and {key!r}"
)
mapped[logical] = key
return mapped
class DiffusionQuantization:
"""Index and decode Core quantized diffusion weights one layer at a time.
@@ -151,6 +170,12 @@ class DiffusionQuantization:
for name in self._keys
}
prefix = model_detection.unet_prefix_from_state_dict(state)
if self.quantized_keys and all(
name.startswith("model.diffusion_model.") for name in self.quantized_keys
):
prefix = "model.diffusion_model."
elif not any(name.startswith(prefix) for name in state):
prefix = ""
detected = {
name[len(prefix):]: value
for name, value in state.items() if name.startswith(prefix)
tests/test_adaptive_extraction_pipeline.py
+85
View File
@@ -1,4 +1,5 @@
import importlib
import json
import sys
import types
from collections import Counter
@@ -261,6 +262,90 @@ def test_lora_extraction_uses_only_async_uel_and_releases_sources(monkeypatch, t
]
@pytest.mark.parametrize("prefixed_input", ["a", "b"])
def test_lora_extraction_matches_prefixed_and_bare_models(monkeypatch, tmp_path, prefixed_input):
extraction = _load_module("nodes.lora_extract_svd")
paths = {name: tmp_path / f"{name}.safetensors" for name in ("a", "b")}
output = tmp_path / "mixed_prefix.safetensors"
for name in ("a", "b"):
key = f"{'model.diffusion_model.' if name == prefixed_input else ''}layer.weight"
value = torch.tensor([[2.0, 0.0], [0.0, 0.0]]) if name == "a" else torch.zeros((2, 2))
if name == prefixed_input:
stem = key[:-len(".weight")]
_write_uel(paths[name], {
key: value.to(torch.int8),
f"{stem}.weight_scale": torch.tensor(1.0),
f"{stem}.comfy_quant": torch.tensor(
list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8
),
})
else:
_write_uel(paths[name], {key: value})
monkeypatch.setattr(extraction, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(extraction, "cleanup_after_operation", lambda: None)
extraction.extract_lora_from_files(
str(paths["a"]), str(paths["b"]), "fixed", 1, 1, "cpu", "fp32", str(output),
linear_max_rank=1, conv_max_rank=1, force_clear_cache=False,
)
with MemoryEfficientSafeOpen(str(output), low_memory=True) as result:
keys = sorted(result.keys())
assert keys == [
"diffusion_model.layer.lora_A.weight",
"diffusion_model.layer.lora_B.weight",
]
stream = result.async_stream(keys, batch_size=1, prefetch_batches=1, pin_memory=False)
tensors = {}
try:
for key in keys:
batch = next(stream)
assert len(batch) == 1 and batch[0][0] == key
tensors[key] = batch[0][1]
finally:
for key in tensors:
result.mark_processed(key)
stream.close()
delta = tensors[keys[1]] @ tensors[keys[0]]
torch.testing.assert_close(delta, torch.tensor([[2.0, 0.0], [0.0, 0.0]]), atol=1e-4, rtol=0)
@pytest.mark.parametrize("module_name,function_name,extra", [
("nodes.dora_extract_wd", "extract_dora_from_files", {}),
("nodes.dora_learned_wd", "extract_dora_learned_from_files", {"optimize_iters": 0}),
])
@pytest.mark.parametrize("prefixed_input", ["a", "b"])
def test_dora_extraction_matches_prefixed_and_bare_models(
monkeypatch, tmp_path, module_name, function_name, extra, prefixed_input,
):
module = _load_module(module_name)
paths = {name: tmp_path / f"{name}.safetensors" for name in ("a", "b")}
output = tmp_path / "mixed_prefix_dora.safetensors"
for name in ("a", "b"):
key = f"{'model.diffusion_model.' if name == prefixed_input else ''}layer.weight"
value = torch.diag(torch.tensor([3.0, 1.0])) if name == "a" else torch.eye(2)
if name == prefixed_input:
stem = key[:-len(".weight")]
_write_uel(paths[name], {
key: value.to(torch.int8),
f"{stem}.weight_scale": torch.tensor(1.0),
f"{stem}.comfy_quant": torch.tensor(
list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8
),
})
else:
_write_uel(paths[name], {key: value})
monkeypatch.setattr(module, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(module, "cleanup_after_operation", lambda: None)
getattr(module, function_name)(
str(paths["a"]), str(paths["b"]), "fixed", 1, 1, "cpu", "fp32", str(output),
linear_max_rank=1, conv_max_rank=1, force_clear_cache=False, **extra,
)
with MemoryEfficientSafeOpen(str(output), low_memory=True) as result:
keys = set(result.keys())
assert "diffusion_model.layer.lora_down.weight" in keys
assert "diffusion_model.layer.lora_up.weight" in keys
assert not any("comfy_quant" in key or "model.diffusion_model." in key for key in keys)
@pytest.mark.parametrize("module_name, function_name, extra", [
("nodes.lora_extract_svd", "extract_lora_from_files", {"include_1d_diffs": False}),
("nodes.dora_extract_wd", "extract_dora_from_files", {}),
tests/test_lora_quantization_guard.py
+9 -6
View File
@@ -506,21 +506,24 @@ def test_diffusion_decoder_matches_core_quantized_tensor(tmp_path, modules):
handler.__exit__(None, None, None)
def test_padded_shape_comes_from_core_meta_architecture(monkeypatch, modules):
@pytest.mark.parametrize("prefix", ["model.", "model.diffusion_model."])
def test_padded_shape_comes_from_core_meta_architecture(monkeypatch, modules, prefix):
import comfy.model_detection
stem = f"{prefix}layer"
class HeaderOnly:
def keys(self):
return ["model.layer.weight", "model.layer.weight_scale", "model.layer.weight_scale_2"]
return [f"{stem}.weight", f"{stem}.weight_scale", f"{stem}.weight_scale_2"]
def metadata(self):
return {"_quantization_metadata": json.dumps({"layers": {"model.layer": {"format": "nvfp4"}}})}
return {"_quantization_metadata": json.dumps({"layers": {stem: {"format": "nvfp4"}}})}
def get_shape(self, key):
return {"model.layer.weight": (16, 8), "model.layer.weight_scale": (16, 1), "model.layer.weight_scale_2": ()}[key]
return {f"{stem}.weight": (16, 8), f"{stem}.weight_scale": (16, 1), f"{stem}.weight_scale_2": ()}[key]
def get_dtype(self, key):
return torch.uint8 if key == "model.layer.weight" else torch.float32
return torch.uint8 if key == f"{stem}.weight" else torch.float32
class Config:
def get_model(self, state, prefix, device):
@@ -534,7 +537,7 @@ def test_padded_shape_comes_from_core_meta_architecture(monkeypatch, modules):
monkeypatch.setattr(comfy.model_detection, "model_config_from_unet", lambda *args, **kwargs: Config())
decoder = modules["guard"].DiffusionQuantization(HeaderOnly(), "header-only", "Shape test")
assert decoder.logical_shape("model.layer.weight") == (3, 5)
assert decoder.logical_shape(f"{stem}.weight") == (3, 5)
def test_nvfp4_stream_decode_matches_core_full_dequantization(monkeypatch, tmp_path, modules):
tests/test_merger_filters.py
+56 -1
View File
@@ -168,7 +168,7 @@ def test_diffusion_merge_decodes_tensorwise_quantization_and_drops_sidecars(
output_path = tmp_path / "diffusion_models" / "quantized_merge.safetensors"
result = load_file(str(output_path))
key = "model.diffusion_model.quant.weight"
key = "quant.weight"
assert set(result) == {key, "dense.weight"}
assert result[key].tolist() == pytest.approx([4.0, 8.0])
assert result[key].dtype == torch.float32
@@ -176,3 +176,58 @@ def test_diffusion_merge_decodes_tensorwise_quantization_and_drops_sidecars(
with safe_open(str(output_path), framework="pt", device="cpu") as output:
metadata = output.metadata()
assert metadata in (None, {})
@pytest.mark.parametrize("prefixed_input", ["a", "b"])
def test_diffusion_merge_matches_prefixed_and_bare_models(monkeypatch, tmp_path, modules, prefixed_input):
merger, operations = modules
paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")}
for name in ("a", "b"):
key = f"{'model.diffusion_model.' if name == prefixed_input else ''}layer.weight"
if name == prefixed_input:
stem = key[:-len(".weight")]
save_file({
key: torch.tensor([1, 2], dtype=torch.int8),
f"{stem}.weight_scale": torch.tensor(2.0),
f"{stem}.comfy_quant": torch.tensor(
list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8
),
}, paths[name])
else:
save_file({key: torch.tensor([4.0, 8.0])}, paths[name])
monkeypatch.setattr(merger.folder_paths, "get_full_path", lambda _kind, name: paths.get(name))
monkeypatch.setattr(merger.folder_paths, "get_folder_paths", lambda _kind: [str(tmp_path)])
monkeypatch.setattr(merger.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(merger, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(merger, "cleanup_after_operation", lambda: None)
params = _params("mixed_prefix_merge")
params.update({"model_a": "a", "model_b": "b"})
merger.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b"}, "Weight-Sum",
operations.TWO_MODEL_MODES, params, "diffusion_models",
)
output = load_file(str(tmp_path / "diffusion_models" / "mixed_prefix_merge.safetensors"))
assert list(output) == ["layer.weight"]
torch.testing.assert_close(output["layer.weight"], torch.tensor([3.0, 6.0]))
def test_three_diffusion_models_match_across_prefixes(monkeypatch, tmp_path, modules):
merger, operations = modules
paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b", "c")}
save_file({"layer.weight": torch.tensor([1.0, 1.0])}, paths["a"])
save_file({"model.diffusion_model.layer.weight": torch.tensor([4.0, 4.0])}, paths["b"])
save_file({"layer.weight": torch.tensor([2.0, 2.0])}, paths["c"])
monkeypatch.setattr(merger.folder_paths, "get_full_path", lambda _kind, name: paths.get(name))
monkeypatch.setattr(merger.folder_paths, "get_folder_paths", lambda _kind: [str(tmp_path)])
monkeypatch.setattr(merger.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(merger, "prepare_for_large_operation", lambda *args: None)
monkeypatch.setattr(merger, "cleanup_after_operation", lambda: None)
params = _params("three_prefix_merge")
params.update({"model_a": "a", "model_b": "b", "model_c": "c"})
merger.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b", "model_c": "c"}, "Add-Difference",
operations.THREE_MODEL_MODES, params, "diffusion_models",
)
output = load_file(str(tmp_path / "diffusion_models" / "three_prefix_merge.safetensors"))
assert list(output) == ["layer.weight"]
torch.testing.assert_close(output["layer.weight"], torch.tensor([2.0, 2.0]))
tests/test_model_analysis.py
+47 -2
View File
@@ -366,8 +366,8 @@ def test_diffusion_analysis_decodes_int8_tensorwise_and_skips_sidecars(
)
rows = {row["key"]: row for row in csv.DictReader(StringIO(metrics_csv))}
assert set(rows) == {f"{prefix}.weight", "dense.weight"}
assert float(rows[f"{prefix}.weight"]["MAE"]) == pytest.approx(6.0)
assert set(rows) == {"layer.weight", "dense.weight"}
assert float(rows["layer.weight"]["MAE"]) == pytest.approx(6.0)
assert float(rows["dense.weight"]["MAE"]) == pytest.approx(3.0)
for path, keys in streamed:
if path.endswith("a.safetensors") or path.endswith("b.safetensors"):
@@ -377,6 +377,51 @@ def test_diffusion_analysis_decodes_int8_tensorwise_and_skips_sidecars(
assert sum(f"{prefix}.weight" in keys or "dense.weight" in keys for _, keys in streamed) == 4
@pytest.mark.parametrize("prefixed_input", ["a", "b"])
def test_diffusion_analysis_matches_prefixed_and_bare_layers(
monkeypatch, tmp_path, analysis, prefixed_input,
):
paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")}
prefix = "model.diffusion_model."
for name in ("a", "b"):
key = f"{prefix if name == prefixed_input else ''}layer.weight"
if name == prefixed_input:
stem = key[:-len(".weight")]
save_file({
key: torch.tensor([1, 2], dtype=torch.int8),
f"{stem}.weight_scale": torch.tensor(2.0),
f"{stem}.comfy_quant": torch.tensor(
list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8
),
}, paths[name])
else:
save_file({key: torch.tensor([4.0, 8.0])}, paths[name])
monkeypatch.setattr(analysis.folder_paths, "get_full_path", lambda _, name: paths[name])
monkeypatch.setattr(analysis, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(analysis, "cleanup_after_operation", lambda: None)
_, _, metrics_csv, _ = analysis.ModelAnalysisLogic.execute(
"a", "b", "diffusion_models", _params(),
)
rows = list(csv.DictReader(StringIO(metrics_csv)))
assert [row["key"] for row in rows] == ["layer.weight"]
assert float(rows[0]["MAE"]) == pytest.approx(3.0)
def test_diffusion_analysis_rejects_duplicate_prefix_variants(monkeypatch, tmp_path, analysis):
paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")}
save_file({
"layer.weight": torch.ones(2),
"model.diffusion_model.layer.weight": torch.ones(2),
}, paths["a"])
save_file({"layer.weight": torch.ones(2)}, paths["b"])
monkeypatch.setattr(analysis.folder_paths, "get_full_path", lambda _, name: paths[name])
monkeypatch.setattr(analysis, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(analysis, "cleanup_after_operation", lambda: None)
with pytest.raises(ValueError, match="duplicate diffusion layer 'layer.weight'"):
analysis.ModelAnalysisLogic.execute("a", "b", "diffusion_models", _params())
def test_embedding_similarity_alignment_is_explicit_opt_in(analysis):
a = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
b = torch.tensor([[0.0, 1.0], [1.0, 0.0]])