diff --git a/nodes/dora_extract_wd.py b/nodes/dora_extract_wd.py index 98bc19f..46789ad 100644 --- a/nodes/dora_extract_wd.py +++ b/nodes/dora_extract_wd.py @@ -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, " diff --git a/nodes/dora_learned_wd.py b/nodes/dora_learned_wd.py index 270117a..5014112 100644 --- a/nodes/dora_learned_wd.py +++ b/nodes/dora_learned_wd.py @@ -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, " diff --git a/nodes/lora_extract_svd.py b/nodes/lora_extract_svd.py index b1c6f7f..2372711 100644 --- a/nodes/lora_extract_svd.py +++ b/nodes/lora_extract_svd.py @@ -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": diff --git a/nodes/merger.py b/nodes/merger.py index e5d1a58..b72fb15 100644 --- a/nodes/merger.py +++ b/nodes/merger.py @@ -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) ) diff --git a/nodes/model_analysis.py b/nodes/model_analysis.py index 85110dd..cde897b 100644 --- a/nodes/model_analysis.py +++ b/nodes/model_analysis.py @@ -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] = {} diff --git a/nodes/quantization_guard.py b/nodes/quantization_guard.py index ca2e399..1e0d208 100644 --- a/nodes/quantization_guard.py +++ b/nodes/quantization_guard.py @@ -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) diff --git a/tests/test_adaptive_extraction_pipeline.py b/tests/test_adaptive_extraction_pipeline.py index f4f4eba..2f33437 100644 --- a/tests/test_adaptive_extraction_pipeline.py +++ b/tests/test_adaptive_extraction_pipeline.py @@ -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", {}), diff --git a/tests/test_lora_quantization_guard.py b/tests/test_lora_quantization_guard.py index 424e927..8cbc056 100644 --- a/tests/test_lora_quantization_guard.py +++ b/tests/test_lora_quantization_guard.py @@ -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): diff --git a/tests/test_merger_filters.py b/tests/test_merger_filters.py index 7af35da..b6f7665 100644 --- a/tests/test_merger_filters.py +++ b/tests/test_merger_filters.py @@ -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])) diff --git a/tests/test_model_analysis.py b/tests/test_model_analysis.py index fd4c88b..f739de3 100644 --- a/tests/test_model_analysis.py +++ b/tests/test_model_analysis.py @@ -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]])