From 57bd63e1eb2bbceba457907a06bc77dcef1e34db Mon Sep 17 00:00:00 2001 From: silveroxides Date: Wed, 23 Sep 2026 01:53:25 +0200 Subject: [PATCH] Support quantized diffusion model inputs --- nodes/extraction_stream.py | 69 ++++- nodes/lora_extract_svd.py | 37 ++- nodes/lora_model_analysis.py | 27 +- nodes/lora_resize.py | 26 +- nodes/merger.py | 68 ++++- nodes/model_analysis.py | 139 +++++++-- nodes/quantization_guard.py | 321 +++++++++++++++++++++ tests/test_adaptive_extraction_pipeline.py | 42 +++ tests/test_lora_model_analysis.py | 16 + tests/test_lora_quantization_guard.py | 149 ++++++++++ tests/test_merger_filters.py | 46 +++ tests/test_model_analysis.py | 46 +++ 12 files changed, 917 insertions(+), 69 deletions(-) diff --git a/nodes/extraction_stream.py b/nodes/extraction_stream.py index 4c6eca0..b4fbea9 100644 --- a/nodes/extraction_stream.py +++ b/nodes/extraction_stream.py @@ -109,24 +109,69 @@ def dora_difference(cpu_a, cpu_b, device, *, keep_sources=False): return (*compute(cpu_a.float(), cpu_b_float), "cpu") -def paired_async_tensors(handler_a, handler_b, work_units, *, pin_memory: bool): - """Yield ordered A/B source tensors and release them after each work unit.""" - keys_a = [unit[1] for unit in work_units if unit[1] is not None] - keys_b = [unit[2] for unit in work_units if unit[2] is not None] +def paired_async_tensors( + handler_a, + handler_b, + work_units, + *, + pin_memory: bool, + quantization_a=None, + quantization_b=None, + compute_dtype=torch.float32, +): + """Yield logical A/B tensors and release their source keys per work unit.""" + def source_keys(quantization, logical_key): + if logical_key is None: + return () + return tuple( + quantization.required_keys(logical_key) + if quantization is not None + else (logical_key,) + ) + + required_a = [source_keys(quantization_a, unit[1]) for unit in work_units] + required_b = [source_keys(quantization_b, unit[2]) for unit in work_units] + keys_a = [key for keys in required_a for key in keys] + keys_b = [key for keys in required_b for key in keys] cursor_a = AsyncTensorCursor(handler_a, keys_a, pin_memory=pin_memory) cursor_b = AsyncTensorCursor(handler_b, keys_b, pin_memory=pin_memory) try: - for logical_key, key_a, key_b in work_units: - tensor_a = cursor_a.take(key_a) if key_a is not None else None - tensor_b = cursor_b.take(key_b) if key_b is not None else None + for (logical_key, key_a, key_b), keys_for_a, keys_for_b in zip( + work_units, required_a, required_b + ): + tensors_a = {} + tensors_b = {} + consumed_a = [] + consumed_b = [] + tensor_a = tensor_b = None try: + for source_key in keys_for_a: + tensors_a[source_key] = cursor_a.take(source_key) + consumed_a.append(source_key) + for source_key in keys_for_b: + tensors_b[source_key] = cursor_b.take(source_key) + consumed_b.append(source_key) + tensor_a = ( + quantization_a.decode( + key_a, tensors_a, compute_dtype, device="cpu" + ) + if quantization_a is not None and key_a in quantization_a.quantized_keys + else tensors_a.get(key_a) + ) + tensor_b = ( + quantization_b.decode( + key_b, tensors_b, compute_dtype, device="cpu" + ) + if quantization_b is not None and key_b in quantization_b.quantized_keys + else tensors_b.get(key_b) + ) yield logical_key, tensor_a, tensor_b finally: - del tensor_a, tensor_b - if key_a is not None: - cursor_a.release(key_a) - if key_b is not None: - cursor_b.release(key_b) + del tensor_a, tensor_b, tensors_a, tensors_b + for source_key in consumed_a: + cursor_a.release(source_key) + for source_key in consumed_b: + cursor_b.release(source_key) cursor_a.finish() cursor_b.finish() finally: diff --git a/nodes/lora_extract_svd.py b/nodes/lora_extract_svd.py index c11c737..b1c6f7f 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 inspect_low_bit_input +from .quantization_guard import DiffusionQuantization 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 @@ -546,17 +546,23 @@ def extract_lora_from_files( handler_b = MemoryEfficientSafeOpen(model_b_path, low_memory=lazy_load) try: - low_bit_a = inspect_low_bit_input(handler_a, f"Model A ({model_a_path})", "LoRA Extract") - low_bit_b = inspect_low_bit_input(handler_b, f"Model B ({model_b_path})", "LoRA Extract") - keys_a = set(handler_a.keys()) - keys_b = set(handler_b.keys()) + quantization_a = DiffusionQuantization( + handler_a, model_a_path, f"Model A ({model_a_path})" + ) + quantization_b = DiffusionQuantization( + handler_b, model_b_path, f"Model B ({model_b_path})" + ) + low_bit_a = set(quantization_a.isolated_low_bit_keys) + low_bit_b = set(quantization_b.isolated_low_bit_keys) + keys_a = set(quantization_a.data_keys) + keys_b = set(quantization_b.data_keys) direct_diff_suffixes = (".scale", ".lin") parameter_layers = sorted( k for k in keys_a if k.endswith(".weight") or ( include_1d_diffs - and (len(handler_a.get_shape(k)) == 1 or k.endswith(direct_diff_suffixes)) + and (len(quantization_a.logical_shape(k)) == 1 or k.endswith(direct_diff_suffixes)) ) ) extraction_keys = [ @@ -566,7 +572,7 @@ def extract_lora_from_files( and (k.endswith(".weight") or ( include_1d_diffs - and (len(handler_a.get_shape(k)) == 1 or k.endswith(direct_diff_suffixes)) + and (len(quantization_a.logical_shape(k)) == 1 or k.endswith(direct_diff_suffixes)) )) ] extraction_keys.sort() @@ -667,9 +673,9 @@ def extract_lora_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 - source_dtypes = [handler_a.get_dtype(key)] + source_dtypes = [quantization_a.logical_dtype(key, save_torch_dtype)] if key in keys_b: - source_dtypes.append(handler_b.get_dtype(key)) + source_dtypes.append(quantization_b.logical_dtype(key, save_torch_dtype)) layer_save_dtype = ( torch.float32 if torch.float32 in source_dtypes else save_torch_dtype ) @@ -775,9 +781,16 @@ def extract_lora_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 LoRA", unit="layers"): status, layer_sd = _process_layer(key, cpu_a, cpu_b) @@ -794,9 +807,13 @@ def extract_lora_from_files( pbar.update(1) stats["skipped"] += len(extraction_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"[LoRA Extract] Done: {stats['extracted']} extracted, {stats['chunked']} chunked, " diff --git a/nodes/lora_model_analysis.py b/nodes/lora_model_analysis.py index 53b3ce2..20a3399 100644 --- a/nodes/lora_model_analysis.py +++ b/nodes/lora_model_analysis.py @@ -25,15 +25,15 @@ from .model_analysis import ( _compute_pair_on_device, _fmt, _is_cuda_oom, _is_float_dtype, _layer_name, _metrics, _release_cuda_oom, _render_comparison, _render_cwb, ) -from .quantization_guard import inspect_low_bit_input +from .quantization_guard import DiffusionQuantization, inspect_low_bit_input from .uel_io import stream_work_units -def _plan_application(base, lora, base_path, params): +def _plan_application(base, lora, base_path, params, base_quant): pairs, passthrough = parse_lora_layers(lora.keys()) validate_canonical_blocks(pairs, "LoRA-on-model analysis") layer_map = build_lora_layer_map([{"pairs": pairs}], base_path) - targets = {canonical_lora_block_name(key): key for key in base.keys()} + targets = {canonical_lora_block_name(key): key for key in base_quant.data_keys} patches = {} mapped = set() unmapped_targets = [] @@ -59,10 +59,10 @@ def _plan_application(base, lora, base_path, params): inventory = {name: [] for name in ("a_only", "b_only", "shape", "nonfloat", "low_bit", "excluded_a", "excluded_b")} unmatched = sorted(set(pairs) - mapped) + unmapped_targets patterns = _compile_patterns(params.get("exclude_patterns", ""), glob_mode=params.get("glob_patterns", False)) - low_base = inspect_low_bit_input(base, "Base diffusion model", "LoRA-on-model analysis") + low_base = base_quant.isolated_low_bit_keys low_lora = inspect_low_bit_input(lora, "LoRA", "LoRA-on-model analysis") units = [] - for key in sorted(base.keys()): + for key in sorted(base_quant.data_keys): if _matches_any_pattern(key, patterns, glob_mode=params.get("glob_patterns", False)) != params.get("include_mode", False): inventory["excluded_a"].append(key) inventory["excluded_b"].append(key) @@ -76,10 +76,10 @@ def _plan_application(base, lora, base_path, params): if key in low_base or any(source in low_lora and source not in alpha_keys for source in source_keys): inventory["low_bit"].append(key) continue - if not _is_float_dtype(base.get_dtype(key)): + if not _is_float_dtype(base_quant.logical_dtype(key, torch.float32)): inventory["nonfloat"].append(key) continue - entries = {0: [key]} + entries = {0: base_quant.required_keys(key)} if source_keys: entries[1] = source_keys units.append((key, entries)) @@ -170,15 +170,22 @@ def analyze_lora_on_model(model_a, model_b, params): layer_stats, block_stats, block_cwb = {}, {}, {} try: with MemoryEfficientSafeOpen(base_path, low_memory=True) as base, MemoryEfficientSafeOpen(lora_path, low_memory=True) as lora: - units, patches, inventory, unmatched, passthrough = _plan_application(base, lora, base_path, params) + base_quant = DiffusionQuantization(base, base_path, "Base diffusion model") + units, patches, inventory, unmatched, passthrough = _plan_application(base, lora, base_path, params, base_quant) device = params["process_device"] top_limit = params["top_weight_differences"] with closing(stream_work_units({0: base, 1: lora}, units, pin_memory=str(device).startswith("cuda"))) as stream: for key, loaded in stream: try: + if key in base_quant.quantized_keys: + loaded[(0, key)] = base_quant.decode( + key, + {source_key: loaded[(0, source_key)] for source_key in base_quant.required_keys(key)}, + torch.float32, + ) raw, cwb, top, fallback = _analyze_unit(key, loaded, patches.get(key, []), strength, device, top_limit) - shape = tuple(base.get_shape(key)) - records.append(TensorRecord(key, shape, str(base.get_dtype(key)), str(torch.float32), raw, cwb, fallback)) + shape = base_quant.logical_shape(key) + records.append(TensorRecord(key, shape, str(base_quant.logical_dtype(key, torch.float32)), str(torch.float32), raw, cwb, fallback)) global_raw.add(raw) layer_stats.setdefault(_layer_name(key), RawStats()).add(raw) block_stats.setdefault(_block_name(key), RawStats()).add(raw) diff --git a/nodes/lora_resize.py b/nodes/lora_resize.py index 533f639..975dcaa 100644 --- a/nodes/lora_resize.py +++ b/nodes/lora_resize.py @@ -23,7 +23,7 @@ from .uel_io import atomic_uel_writer, stream_work_units from .artifact_paths import canonical_model_artifact_path from .lora_alpha import lora_alpha_scale, normalize_lora_pair from .layer_parameters import parameter_input, resolve_layer_parameters -from .quantization_guard import inspect_low_bit_input, layer_has_low_bit, write_preserved_tensor +from .quantization_guard import DiffusionQuantization, inspect_low_bit_input, layer_has_low_bit, write_preserved_tensor from typing import Optional, Dict, Tuple, List @@ -1528,9 +1528,10 @@ def merge_loras_to_model( try: - base_low_bit_keys = inspect_low_bit_input( - base_handler, f"Base model ({base_model_path})", "LoRA Merge To Model" + base_quant = DiffusionQuantization( + base_handler, base_model_path, "LoRA Merge To Model base model" ) + base_low_bit_keys = base_quant.isolated_low_bit_keys lora_low_bit_keys = [ inspect_low_bit_input( handler, f"LoRA {i + 1} ({lora_paths[i]})", "LoRA Merge To Model" @@ -1597,7 +1598,7 @@ def merge_loras_to_model( return result.replace(".", "_") - base_keys = list(base_handler.keys()) + base_keys = list(base_quant.data_keys) base_aliases = set() for key in base_keys: normalized = key @@ -1668,7 +1669,7 @@ def merge_loras_to_model( skip_patterns = _compile_patterns(skip_patterns_str) # Preserve metadata from base model - base_metadata = base_handler.metadata().copy() if base_handler.metadata() else {} + base_metadata = base_quant.output_metadata() base_metadata["merge_comment"] = f"Merged {len(lora_paths)} LoRAs with weights: {lora_weights}" # Build output path before loop @@ -1746,7 +1747,7 @@ def merge_loras_to_model( skipped = any(pattern.search(planned_base_key) for pattern in skip_patterns) != include_mode entries = {} if planned_base_key not in base_low_bit_keys and not skipped: - entries[-1] = [planned_base_key] + entries[-1] = base_quant.required_keys(planned_base_key) for info, _, _, direct_key, block_keys in planned_direct: if not layer_has_low_bit(block_keys, info["low_bit_keys"]): entries.setdefault(info["index"], []).append(direct_key) @@ -1845,7 +1846,14 @@ def merge_loras_to_model( continue # Load base weight only after guarded and skipped keys are classified. - cpu_base = loaded[(-1, base_key)] + cpu_base = ( + base_quant.decode( + base_key, + {source_key: loaded[(-1, source_key)] for source_key in base_quant.required_keys(base_key)}, + torch.float32, + ) + if base_key in base_quant.quantized_keys else loaded[(-1, base_key)] + ) if direct_contributions or low_rank_contributions: if device == 'cuda': @@ -1854,7 +1862,7 @@ def merge_loras_to_model( base_weight = cpu_base.to(device=device, dtype=torch.float32) del cpu_base - source_dtypes = [base_handler.get_dtype(base_key)] + source_dtypes = [base_quant.logical_dtype(base_key, save_dtype)] applied_1d_diff = False applied = False @@ -2069,7 +2077,7 @@ def merge_loras_to_model( ) else: target_dtype = select_output_dtype( - [base_handler.get_dtype(base_key)], + [base_quant.logical_dtype(base_key, save_dtype)], save_dtype, ) writer.write(base_key, cpu_base.to(target_dtype).contiguous()) diff --git a/nodes/merger.py b/nodes/merger.py index 23a643d..e5d1a58 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 inspect_low_bit_input, write_preserved_tensor +from .quantization_guard import DiffusionQuantization, 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 @@ -98,16 +98,32 @@ class MergerLogic: prepare_for_large_operation(total_size_gb * 1.2, torch.device(process_device)) handlers = {} + handler_paths = {} for name in model_names.values(): if name and name != "None": path = folder_paths.get_full_path(model_type, name) if not path: raise FileNotFoundError(f"Model '{name}' not found.") handlers[name] = MemoryEfficientSafeOpen(path, low_memory=True) + handler_paths[name] = path primary_handler = handlers[primary_model_name] - all_keys = primary_handler.keys() - metadata = primary_handler.metadata() + diffusion_quantizers = {} + if model_type == "diffusion_models": + try: + diffusion_quantizers = { + name: DiffusionQuantization(handler, handler_paths[name], f"Model '{name}'") + for name, handler in handlers.items() + } + except Exception: + for handler in handlers.values(): + handler.__exit__(None, None, None) + raise + all_keys = diffusion_quantizers[primary_model_name].data_keys + metadata = diffusion_quantizers[primary_model_name].output_metadata() + else: + all_keys = primary_handler.keys() + metadata = primary_handler.metadata() low_bit_keys_by_name = {} if model_type == "loras": @@ -137,7 +153,10 @@ class MergerLogic: primary_keys = set(all_keys) for name, handler in handlers.items(): if name != primary_model_name: - secondary_keys = set(handler.keys()) + secondary_keys = set( + diffusion_quantizers[name].data_keys + if name in diffusion_quantizers else handler.keys() + ) missing = primary_keys - secondary_keys extra = secondary_keys - primary_keys if missing: @@ -255,8 +274,16 @@ class MergerLogic: if label != "A" ) for name in dict.fromkeys(requested_names): - if name in handlers and key in handlers[name].keys(): - entries[name] = [key] + if name in handlers and key in ( + diffusion_quantizers[name].data_keys + if name in diffusion_quantizers else handlers[name].keys() + ): + required = ( + diffusion_quantizers[name].required_keys(key) + if name in diffusion_quantizers and key in diffusion_quantizers[name].quantized_keys + else [key] + ) + entries[name] = list(dict.fromkeys([key, *required])) alpha_info = alpha_normalization.get(name, {}).get(key) if alpha_info is not None: entries[name].append(alpha_info[0]) @@ -302,14 +329,37 @@ class MergerLogic: if scale != 1.0: loaded[(name, key)] = loaded[(name, key)] * scale + if diffusion_quantizers: + for name, quantizer in diffusion_quantizers.items(): + if key in quantizer.quantized_keys and (name, key) in loaded: + tensors = { + source_key: loaded[(name, source_key)] + for source_key in dict.fromkeys([key, *quantizer.required_keys(key)]) + } + loaded[(name, key)] = quantizer.decode( + key, tensors, process_dtype, device="cpu" + ) + # 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 = [primary_handler.get_dtype(key)] + original_dtypes = [ + diffusion_quantizers[primary_model_name].logical_dtype(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 handler.keys(): - original_dtypes.append(handler.get_dtype(key)) + if name != primary_model_name and key in ( + diffusion_quantizers[name].data_keys + if name in diffusion_quantizers 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 + else handler.get_dtype(key) + ) target_dtype = determine_target_dtype(original_dtypes, save_torch_dtype, override_dtype) preserve_model_a_1d = model_type == "loras" and cpu_tensor.ndim == 1 and not include_1d_diffs diff --git a/nodes/model_analysis.py b/nodes/model_analysis.py index 4e3f460..85110dd 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 inspect_low_bit_input +from .quantization_guard import DiffusionQuantization, inspect_low_bit_input FLOAT_DTYPES = {torch.float16, torch.bfloat16, torch.float32, torch.float64} @@ -573,6 +573,47 @@ def _load_paired_unit(handlers, units, *, pin_memory: bool): return loaded +def _load_diffusion_pair(handlers, keys_a, keys_b, *, pin_memory: bool): + """Load one diffusion layer and its quantization sidecars from each input.""" + streams = [] + tensors = [{}, {}] + + def collect(stream, expected_keys, result): + for expected_key in expected_keys: + batch = next(stream) + for key, tensor in batch: + result[key] = tensor + if len(batch) != 1: + raise RuntimeError("Asynchronous diffusion analysis expected one tensor per batch.") + key, _ = batch[0] + if key != expected_key: + raise RuntimeError("Asynchronous diffusion analysis returned tensors out of order.") + + try: + for handler, keys in zip(handlers, (keys_a, keys_b)): + streams.append(handler.async_stream( + keys, + batch_size=1, + prefetch_batches=1, + pin_memory=pin_memory, + )) + with ThreadPoolExecutor(max_workers=2) as executor: + future_a = executor.submit(collect, streams[0], keys_a, tensors[0]) + future_b = executor.submit(collect, streams[1], keys_b, tensors[1]) + future_a.result() + future_b.result() + return tensors + except BaseException: + for handler, loaded in zip(handlers, tensors): + for key in loaded: + handler.mark_processed(key) + loaded.clear() + raise + finally: + for stream in streams: + stream.close() + + def _average_layer_metrics(layer_stats: dict[str, RawStats]) -> dict[str, float | None]: metrics = [_metrics(raw) for raw in layer_stats.values()] result = {} @@ -798,13 +839,26 @@ class ModelAnalysisLogic: handlers = [MemoryEfficientSafeOpen(path, low_memory=True) for path in paths] try: - low_bit = [ - inspect_low_bit_input(handler, f"Input {index + 1} ({path})", "Model Analysis") - for index, (handler, path) in enumerate(zip(handlers, paths)) - ] + diffusion_quantizers = [] + if model_type == "diffusion_models": + diffusion_quantizers = [ + DiffusionQuantization(handler, path, f"Input {index + 1} ({path})") + for index, (handler, path) in enumerate(zip(handlers, paths)) + ] + if diffusion_quantizers: + low_bit = [quantizer.isolated_low_bit_keys for quantizer in diffusion_quantizers] + else: + low_bit = [ + inspect_low_bit_input(handler, f"Input {index + 1} ({path})", "Model Analysis") + for index, (handler, path) in enumerate(zip(handlers, paths)) + ] similarity_alignment = bool(params.get("cwb_similarity_alignment", False)) - keys_a = set(handlers[0].keys()) - keys_b = set(handlers[1].keys()) + keys_a = set( + diffusion_quantizers[0].data_keys if diffusion_quantizers else handlers[0].keys() + ) + keys_b = set( + diffusion_quantizers[1].data_keys if diffusion_quantizers else handlers[1].keys() + ) glob_mode = bool(params.get("glob_patterns", False)) excluded_patterns = _compile_patterns( params.get("exclude_patterns", ""), @@ -940,15 +994,34 @@ class ModelAnalysisLogic: inventory["low_bit"].append(f"{key}: A={key_a}, B={key_b}") skipped_units += 1 continue - shape_a = tuple(handlers[0].get_shape(key_a)) - shape_b = tuple(handlers[1].get_shape(key_b)) + shape_a = tuple( + diffusion_quantizers[0].logical_shape(key_a) + if diffusion_quantizers and key_a in diffusion_quantizers[0].quantized_keys + else handlers[0].get_shape(key_a) + ) + shape_b = tuple( + diffusion_quantizers[1].logical_shape(key_b) + if diffusion_quantizers and key_b in diffusion_quantizers[1].quantized_keys + else handlers[1].get_shape(key_b) + ) if shape_a != shape_b: inventory["shape"].append(f"{key}: A={shape_a}, B={shape_b}") skipped_units += 1 continue ready_units.append(( key, key_a, key_b, shape_a, - handlers[0].get_dtype(key_a), handlers[1].get_dtype(key_b), + diffusion_quantizers[0].logical_dtype(key_a, torch.float32) + if diffusion_quantizers and key_a in diffusion_quantizers[0].quantized_keys + else handlers[0].get_dtype(key_a), + diffusion_quantizers[1].logical_dtype(key_b, torch.float32) + if diffusion_quantizers and key_b in diffusion_quantizers[1].quantized_keys + else handlers[1].get_dtype(key_b), + list(dict.fromkeys([key_a, *(diffusion_quantizers[0].required_keys(key_a) + if diffusion_quantizers and key_a in diffusion_quantizers[0].quantized_keys + else [])])), + list(dict.fromkeys([key_b, *(diffusion_quantizers[1].required_keys(key_b) + if diffusion_quantizers and key_b in diffusion_quantizers[1].quantized_keys + else [])])), )) progress = comfy.utils.ProgressBar(len(units) + len(lora_pairs)) if skipped_units: @@ -956,13 +1029,33 @@ class ModelAnalysisLogic: with torch.no_grad(): iterator = tqdm(ready_units, desc="Analyzing model tensors", unit="tensors") - for key, key_a, key_b, shape_a, dtype_a, dtype_b in iterator: - loaded = _load_paired_unit( - handlers, [(key, key_a, key_b)], - pin_memory=str(params["process_device"]).startswith("cuda"), - ) - _, _, _, tensor_a, tensor_b = loaded[0] + for key, key_a, key_b, shape_a, dtype_a, dtype_b, required_a, required_b in iterator: + loaded = None + paired = None + tensor_a = tensor_b = None + if diffusion_quantizers: + loaded = _load_diffusion_pair( + handlers, required_a, required_b, + pin_memory=str(params["process_device"]).startswith("cuda"), + ) + else: + paired = _load_paired_unit( + handlers, [(key, key_a, key_b)], + pin_memory=str(params["process_device"]).startswith("cuda"), + ) try: + if diffusion_quantizers: + if key_a in diffusion_quantizers[0].quantized_keys: + loaded[0][key_a] = diffusion_quantizers[0].decode( + key_a, loaded[0], torch.float32, device="cpu" + ) + if key_b in diffusion_quantizers[1].quantized_keys: + loaded[1][key_b] = diffusion_quantizers[1].decode( + key_b, loaded[1], torch.float32, device="cpu" + ) + tensor_a, tensor_b = loaded[0][key_a], loaded[1][key_b] + else: + _, _, _, tensor_a, tensor_b = paired[0] if not _is_float_dtype(dtype_a) or not _is_float_dtype(dtype_b): equal = torch.equal(tensor_a, tensor_b) inventory["nonfloat"].append( @@ -980,9 +1073,17 @@ class ModelAnalysisLogic: fallback_keys.append(key) finally: del tensor_a, tensor_b - loaded.clear() - handlers[0].mark_processed(key_a) - handlers[1].mark_processed(key_b) + if loaded is not None: + loaded[0].clear() + loaded[1].clear() + for source_key in required_a: + handlers[0].mark_processed(source_key) + for source_key in required_b: + handlers[1].mark_processed(source_key) + else: + paired.clear() + handlers[0].mark_processed(key_a) + handlers[1].mark_processed(key_b) clear_after_unit() progress.update(1) diff --git a/nodes/quantization_guard.py b/nodes/quantization_guard.py index 347bbdc..ca2e399 100644 --- a/nodes/quantization_guard.py +++ b/nodes/quantization_guard.py @@ -32,6 +32,327 @@ UPSTREAM_WRITER_STORAGE_CODES = { } +_QUANT_SIDECARS = ( + "comfy_quant", "weight_scale", "weight_scale_2", "input_scale", + "pre_quant_scale", "weight_s_rel", "weight_s_channel", + "weight_correction", "weight_codebook", "scale_weight", "scale_input", +) + + +class DiffusionQuantization: + """Index and decode Core quantized diffusion weights one layer at a time. + + The handler is a UEL low-memory handler. Header inspection uses only its + keys, shapes, dtypes, and metadata; tensor values are supplied to ``decode`` + by the caller's bounded UEL work unit. + """ + + def __init__(self, handler, model_path: str, label: str): + self.handler = handler + self.model_path = model_path + self.label = label + self._keys = tuple(handler.keys()) + self._key_set = set(self._keys) + self._metadata = dict(handler.metadata() or {}) + self._layer_configs = {} + self._architecture_shapes = None + self._legacy_scaled_fp8 = any(key.endswith("scaled_fp8") for key in self._keys) + + encoded = self._metadata.get("_quantization_metadata") + if encoded: + try: + layers = json.loads(encoded).get("layers", {}) + except (TypeError, ValueError) as exc: + raise ValueError( + f"[{label}] Invalid _quantization_metadata in {model_path}: {exc}" + ) from exc + if not isinstance(layers, dict): + raise ValueError(f"[{label}] Invalid quantization layers metadata in {model_path}") + self._layer_configs.update(layers) + + quantized = set() + for sidecar in self._keys: + if not sidecar.endswith(".comfy_quant"): + continue + stem = sidecar[:-len(".comfy_quant")] + weight_key = f"{stem}.weight" + if weight_key not in self._key_set: + raise ValueError(f"[{label}] Quantization sidecar has no weight: {sidecar}") + quantized.add(weight_key) + + for stem in self._layer_configs: + weight_key = f"{stem}.weight" + if weight_key in self._key_set: + quantized.add(weight_key) + + if self._legacy_scaled_fp8: + for sidecar in self._keys: + if sidecar.endswith(".scale_weight"): + stem = sidecar[:-len(".scale_weight")] + weight_key = f"{stem}.weight" + if weight_key in self._key_set: + quantized.add(weight_key) + self._layer_configs.setdefault(stem, {"format": "float8_e4m3fn"}) + + self.quantized_keys = frozenset(quantized) + sidecars = set() + for key in self.quantized_keys: + stem = key[:-len(".weight")] + for name in _QUANT_SIDECARS: + candidate = f"{stem}.{name}" + if candidate in self._key_set: + sidecars.add(candidate) + for key in self._keys: + if key.endswith(".comfy_quant") and any( + key == f"{weight[:-len('.weight')]}.comfy_quant" for weight in self.quantized_keys + ): + sidecars.add(key) + if self._legacy_scaled_fp8: + sidecars.update(key for key in self._keys if key.endswith("scaled_fp8")) + self._sidecars = frozenset(sidecars) + self.data_keys = tuple(key for key in self._keys if key not in self._sidecars) + + self.isolated_low_bit_keys = { + key for key in self.data_keys + if key not in self.quantized_keys and _is_low_bit(handler, key) + } + if len(self.isolated_low_bit_keys) > MAX_ISOLATED_LOW_BIT_TENSORS: + raise ValueError( + f"[{label}] {model_path} contains {len(self.isolated_low_bit_keys)} low-bit tensors; " + "casted or quantized inputs are unsupported" + ) + if self.isolated_low_bit_keys: + logging.warning( + "[%s] %s contains %d isolated low-bit tensor(s); preserving/excluding unchanged: %s", + label, model_path, len(self.isolated_low_bit_keys), + ", ".join(sorted(self.isolated_low_bit_keys)), + ) + + def required_keys(self, key: str) -> list[str]: + if key not in self.quantized_keys: + return [key] + stem = key[:-len(".weight")] + return [key] + [ + f"{stem}.{name}" for name in _QUANT_SIDECARS + if f"{stem}.{name}" in self._key_set + ] + + def _architecture_shape(self, key: str) -> tuple[int, ...]: + if self._architecture_shapes is None: + from comfy import model_detection + + # Core detection examines keys, shapes and dtypes. Meta tensors supply + # that inventory without materializing any checkpoint tensor data. + state = { + name: torch.empty( + tuple(self.handler.get_shape(name)), + dtype=self.handler.get_dtype(name), device="meta", + ) + for name in self._keys + } + prefix = model_detection.unet_prefix_from_state_dict(state) + detected = { + name[len(prefix):]: value + for name, value in state.items() if name.startswith(prefix) + } + try: + config = model_detection.model_config_from_unet( + detected, "", metadata=self._metadata, + ) + if config is None: + raise ValueError("Core could not detect this diffusion architecture") + with torch.device("meta"): + model = config.get_model(detected, "", device=torch.device("meta")) + modules = dict(getattr(model, "diffusion_model", model).named_modules()) + shapes = {} + for weight_key in self.quantized_keys: + module_name = weight_key[len(prefix):].removesuffix(".weight") + module = modules.get(module_name) + if module is None: + continue + shape = getattr(module, "_orig_shape", None) + if shape is None: + weight = getattr(module, "weight", None) + shape = getattr(weight, "shape", None) + if shape is not None: + shapes[weight_key] = tuple(shape) + self._architecture_shapes = shapes + except Exception as exc: + raise ValueError( + f"[{self.label}] Cannot determine exact quantized layer shapes from " + f"Core architecture for {self.model_path}: {exc}" + ) from exc + finally: + state.clear() + detected.clear() + if "model" in locals(): + del model + + shape = self._architecture_shapes.get(key) + if shape is None: + raise ValueError( + f"[{self.label}] Core architecture has no logical shape for quantized layer {key}" + ) + stored = tuple(self.handler.get_shape(key)) + if len(shape) != 2 or len(stored) != 2: + raise ValueError(f"[{self.label}] Invalid padded quantized shape for {key}: {shape}, {stored}") + stem = key[:-len(".weight")] + config = self._layer_configs.get(stem) + nvfp4 = (isinstance(config, dict) and config.get("format") == "nvfp4") or ( + f"{stem}.weight_scale_2" in self._key_set + ) + if nvfp4: + expected = ((shape[0] + 15) // 16 * 16, (shape[1] + 15) // 16 * 8) + else: + expected = ((shape[0] + 31) // 32 * 32, (shape[1] + 31) // 32 * 32) + if stored != expected: + raise ValueError( + f"[{self.label}] Core architecture shape {shape} disagrees with packed storage {stored} for {key}" + ) + return shape + + def _format(self, key: str, tensors=None): + stem = key[:-len(".weight")] + config = self._layer_configs.get(stem) + sidecar = f"{stem}.comfy_quant" + if tensors is None and config is None and sidecar in self._key_set: + stream = self.handler.async_stream( + [sidecar], batch_size=1, prefetch_batches=1, pin_memory=False, + ) + yielded_keys = [] + try: + batch = next(stream) + yielded_keys = [name for name, _ in batch] + if len(batch) != 1 or batch[0][0] != sidecar: + raise RuntimeError(f"[{self.label}] Unexpected quantization sidecar for {key}") + tensors = {sidecar: batch[0][1]} + finally: + try: + for yielded_key in yielded_keys: + self.handler.mark_processed(yielded_key) + finally: + stream.close() + if tensors is not None and sidecar in tensors: + raw = tensors[sidecar].detach().to(device="cpu").contiguous().reshape(-1).tolist() + try: + config = json.loads(bytes(raw).decode("utf-8")) + except (UnicodeDecodeError, TypeError, ValueError) as exc: + raise ValueError(f"[{self.label}] Invalid quantization sidecar {sidecar}: {exc}") from exc + self._layer_configs[stem] = config + if config is None: + raise ValueError(f"[{self.label}] Missing quantization config for {key}") + if not isinstance(config, dict): + raise ValueError(f"[{self.label}] Invalid quantization config for {key}: expected object") + fmt = config.get("format") + try: + from comfy.quant_ops import QUANT_ALGOS + except ImportError as exc: + raise RuntimeError("ComfyUI quantization support is unavailable") from exc + if fmt not in QUANT_ALGOS: + raise ValueError(f"[{self.label}] Unsupported ComfyUI quantization format {fmt!r} for {key}") + return fmt, config + + def logical_shape(self, key: str) -> tuple[int, ...]: + if key not in self.quantized_keys: + return tuple(self.handler.get_shape(key)) + fmt, _ = self._format(key) + if fmt == "nvfp4" or fmt == "mxfp8": + return self._architecture_shape(key) + shape = tuple(self.handler.get_shape(key)) + if fmt in {"convrot_w4a4", "asym_w4a8_int8"}: + if len(shape) != 2: + raise ValueError(f"[{self.label}] Expected packed matrix weight for {key}, got {shape}") + return shape[0], shape[1] * 2 + return shape + + def logical_dtype(self, key: str, compute_dtype: torch.dtype) -> torch.dtype: + if key not in self.quantized_keys: + return self.handler.get_dtype(key) + return compute_dtype + + def decode(self, key: str, tensors, compute_dtype: torch.dtype, device="cpu") -> torch.Tensor: + if key not in self.quantized_keys: + raise ValueError(f"[{self.label}] {key} is not a recognized quantized diffusion weight") + fmt, config = self._format(key, tensors) + try: + from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor, get_layout_class + except ImportError as exc: + raise RuntimeError("ComfyUI quantization support is unavailable") from exc + + stem = key[:-len(".weight")] + device = torch.device(device) + shape = self.logical_shape(key) + + def tensor(name, *, dtype=None, required=False): + tensor_key = f"{stem}.{name}" + if tensor_key not in tensors and name == "weight_scale": + legacy_key = f"{stem}.scale_weight" + if legacy_key in tensors: + tensor_key = legacy_key + value = tensors.get(tensor_key) + if value is None: + if required: + raise ValueError(f"[{self.label}] Missing quantization tensor {tensor_key} for {key}") + return None + value = value.to(device=device) + if dtype is not None: + value = value.view(dtype=dtype) + return value + + params_conf = config.get("params", {}) + if not isinstance(params_conf, dict): + params_conf = {} + if fmt in ("float8_e4m3fn", "float8_e5m2"): + values = {"scale": tensor("weight_scale", required=True)} + elif fmt == "mxfp8": + values = {"scale": tensor("weight_scale", dtype=torch.float8_e8m0fnu, required=True)} + elif fmt == "nvfp4": + values = { + "scale": tensor("weight_scale_2", required=True), + "block_scale": tensor("weight_scale", dtype=torch.float8_e4m3fn, required=True), + } + elif fmt == "int8_tensorwise": + values = {"scale": tensor("weight_scale", required=True)} + if config.get("convrot", params_conf.get("convrot", False)): + values.update({ + "convrot": True, + "convrot_groupsize": int(config.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256))), + }) + elif fmt == "convrot_w4a4": + values = { + "scale": tensor("weight_scale", required=True), + "convrot_groupsize": int(config.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256))), + "quant_group_size": 64, + "linear_dtype": config.get("linear_dtype", params_conf.get("linear_dtype", "int4")), + } + elif fmt == "asym_w4a8_int8": + if f"{stem}.weight_correction" in tensors: + raise ValueError( + f"[{self.label}] Core's W4A8 loader does not apply weight_correction for {key}" + ) + scale_value = tensors.get(f"{stem}.weight_s_rel") + values = { + "scale": tensor("weight_s_rel", dtype=torch.float8_e4m3fn if scale_value is not None and scale_value.dtype == torch.uint8 else None, required=True), + "s_channel": tensor("weight_s_channel", required=True), + "codebook": tensor("weight_codebook"), + "group_size": int(config.get("group_size", params_conf.get("group_size", 16))), + "convrot_groupsize": int(config.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256))), + } + else: + raise ValueError(f"[{self.label}] Unsupported ComfyUI quantization format {fmt!r} for {key}") + + layout_type = QUANT_ALGOS[fmt]["comfy_tensor_layout"] + layout = get_layout_class(layout_type) + params = layout.Params(**values, orig_dtype=compute_dtype, orig_shape=shape) + qdata = tensors[key].to(device=device, dtype=QUANT_ALGOS[fmt]["storage_t"]) + return QuantizedTensor(qdata, layout_type, params).dequantize() + + def output_metadata(self) -> dict: + metadata = dict(self._metadata) + metadata.pop("_quantization_metadata", None) + return metadata + + def _header(handler): if getattr(handler, "_header", None) is not None: return handler._header diff --git a/tests/test_adaptive_extraction_pipeline.py b/tests/test_adaptive_extraction_pipeline.py index 076c0af..f4f4eba 100644 --- a/tests/test_adaptive_extraction_pipeline.py +++ b/tests/test_adaptive_extraction_pipeline.py @@ -122,6 +122,48 @@ def test_paired_stream_is_bounded_ordered_and_releases_once(): assert all(stream.closed for stream in handler_a.streams + handler_b.streams) +def test_paired_stream_decodes_quantized_logical_keys_per_work_unit(): + streaming = _load_module("nodes.extraction_stream") + handler_a = _FakeHandler({ + "layer1.code": torch.tensor(2), + "layer1.scale": torch.tensor(3), + "layer2.code": torch.tensor(4), + "layer2.scale": torch.tensor(5), + }) + handler_b = _FakeHandler({"layer1": torch.tensor(1), "layer2": torch.tensor(2)}) + + class _Quantization: + quantized_keys = {"layer1", "layer2"} + + def required_keys(self, key): + return (f"{key}.code", f"{key}.scale") + + def decode(self, key, tensors, compute_dtype, device="cpu"): + assert set(tensors) == {f"{key}.code", f"{key}.scale"} + return tensors[f"{key}.code"] * tensors[f"{key}.scale"] + + quantization = _Quantization() + observed = list(streaming.paired_async_tensors( + handler_a, + handler_b, + [("layer1", "layer1", "layer1"), ("layer2", "layer2", "layer2")], + pin_memory=False, + quantization_a=quantization, + )) + + assert [(key, tensor_a.item(), tensor_b.item()) for key, tensor_a, tensor_b in observed] == [ + ("layer1", 6, 1), + ("layer2", 20, 2), + ] + assert handler_a.calls[0][0] == ( + "layer1.code", "layer1.scale", "layer2.code", "layer2.scale" + ) + assert Counter(handler_a.marked) == Counter({ + "layer1.code": 1, "layer1.scale": 1, "layer2.code": 1, "layer2.scale": 1, + }) + assert all(stream.closed for stream in handler_a.streams + handler_b.streams) + + def test_cuda_oom_compute_retries_on_cpu(monkeypatch): streaming = _load_module("nodes.extraction_stream") monkeypatch.setattr(streaming, "release_cuda_oom", lambda: None) diff --git a/tests/test_lora_model_analysis.py b/tests/test_lora_model_analysis.py index d9fd754..72d5f3d 100644 --- a/tests/test_lora_model_analysis.py +++ b/tests/test_lora_model_analysis.py @@ -1,6 +1,7 @@ import csv import importlib from io import StringIO +import json import logging from pathlib import Path import sys @@ -67,6 +68,21 @@ def test_reports_actual_applied_delta_without_saving(monkeypatch, tmp_path, anal assert set(tmp_path.iterdir()) == set(paths.values()) +def test_quantized_base_is_dequantized_and_sidecars_are_not_analyzed(monkeypatch, tmp_path, analysis): + prefix = "model.diffusion_model.layer" + quant = torch.tensor(list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8) + setup(monkeypatch, tmp_path, analysis, pair(), { + f"{prefix}.weight": torch.ones((2, 2), dtype=torch.int8), + f"{prefix}.weight_scale": torch.tensor(2.0), + f"{prefix}.comfy_quant": quant, + }) + comparison, _, metrics_csv, _ = analysis.analyze_lora_on_model("adapter", "base", params()) + rows = list(csv.DictReader(StringIO(metrics_csv))) + assert [row["key"] for row in rows] == [f"{prefix}.weight"] + assert float(rows[0]["mae"]) == pytest.approx(2.25) + assert "weight_scale" not in comparison + + @pytest.mark.parametrize("suffixes", [(".lora_down.weight", ".lora_up.weight"), (".lora_A.default.weight", ".lora_B.default.weight"), (".lora_A", ".lora_B")]) def test_pair_formats_and_flattened_reference_mapping(monkeypatch, tmp_path, analysis, suffixes): setup(monkeypatch, tmp_path, analysis, { diff --git a/tests/test_lora_quantization_guard.py b/tests/test_lora_quantization_guard.py index fdbb15b..424e927 100644 --- a/tests/test_lora_quantization_guard.py +++ b/tests/test_lora_quantization_guard.py @@ -436,3 +436,152 @@ def test_merge_to_model_preserves_low_bit_base_and_skips_low_bit_adapter_layer( assert "diffusion_model.guarded.lora_down.weight" in report assert "torch.int8" in report assert "model.diffusion_model.base_quant.weight" in report + + +def test_merge_to_model_dequantizes_base_and_omits_sidecars(monkeypatch, tmp_path, modules): + resize = modules["resize"] + _patch_output(monkeypatch, resize, tmp_path) + base = tmp_path / "quantized_base.safetensors" + adapter = tmp_path / "quantized_adapter.safetensors" + quant = torch.tensor(list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8) + prefix = "model.diffusion_model.layer" + save_file({ + f"{prefix}.weight": torch.ones((2, 2), dtype=torch.int8), + f"{prefix}.weight_scale": torch.tensor(2.0), + f"{prefix}.comfy_quant": quant, + }, str(base)) + save_file({"diffusion_model.layer.diff": torch.ones((2, 2))}, str(adapter)) + + output = resize.merge_loras_to_model( + [str(adapter)], [1.0], str(base), "cpu", torch.float16, + "dequantized_merge", verbose=False, + ) + result = load_file(output) + assert list(result) == [f"{prefix}.weight"] + torch.testing.assert_close(result[f"{prefix}.weight"].float(), torch.full((2, 2), 3.0)) + + +def test_diffusion_decoder_matches_core_quantized_tensor(tmp_path, modules): + from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor, get_layout_class + + path = tmp_path / "core_int8.safetensors" + key = "layer.weight" + qdata = torch.tensor([[1, -2], [3, 4]], dtype=torch.int8) + scale = torch.tensor(0.5) + save_file({ + key: qdata, + "layer.weight_scale": scale, + "layer.comfy_quant": torch.tensor( + list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8 + ), + }, str(path)) + handler = modules["resize"].MemoryEfficientSafeOpen(str(path), low_memory=True) + try: + decoder = modules["guard"].DiffusionQuantization(handler, str(path), "Core comparison") + keys = decoder.required_keys(key) + assert set(keys) == {key, "layer.weight_scale", "layer.comfy_quant"} + assert decoder.data_keys == (key,) + stream = handler.async_stream(keys, batch_size=1, prefetch_batches=1, pin_memory=False) + tensors = {} + try: + for expected in keys: + batch = next(stream) + assert len(batch) == 1 and batch[0][0] == expected + tensors[expected] = batch[0][1] + actual = decoder.decode(key, tensors, torch.float32) + finally: + for source_key in tensors: + handler.mark_processed(source_key) + stream.close() + + layout_name = QUANT_ALGOS["int8_tensorwise"]["comfy_tensor_layout"] + layout = get_layout_class(layout_name) + expected = QuantizedTensor( + qdata, layout_name, + layout.Params(scale=scale, orig_dtype=torch.float32, orig_shape=(2, 2)), + ).dequantize() + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(actual, qdata.float() * scale) + finally: + handler.__exit__(None, None, None) + + +def test_padded_shape_comes_from_core_meta_architecture(monkeypatch, modules): + import comfy.model_detection + + class HeaderOnly: + def keys(self): + return ["model.layer.weight", "model.layer.weight_scale", "model.layer.weight_scale_2"] + + def metadata(self): + return {"_quantization_metadata": json.dumps({"layers": {"model.layer": {"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] + + def get_dtype(self, key): + return torch.uint8 if key == "model.layer.weight" else torch.float32 + + class Config: + def get_model(self, state, prefix, device): + assert device == torch.device("meta") + assert all(tensor.device.type == "meta" for tensor in state.values()) + model = torch.nn.Module() + model.diffusion_model = torch.nn.Module() + model.diffusion_model.layer = torch.nn.Module() + model.diffusion_model.layer._orig_shape = (3, 5) + return model + + 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) + + +def test_nvfp4_stream_decode_matches_core_full_dequantization(monkeypatch, tmp_path, modules): + import comfy.model_detection + from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor, get_layout_class + + layout_name = QUANT_ALGOS["nvfp4"]["comfy_tensor_layout"] + layout = get_layout_class(layout_name) + original = torch.arange(15, dtype=torch.float16).reshape(3, 5) / 10 + qdata, params = layout.quantize(original) + path = tmp_path / "nvfp4.safetensors" + metadata = {"_quantization_metadata": json.dumps({"layers": {"model.layer": {"format": "nvfp4"}}})} + save_file({ + "model.layer.weight": qdata, + "model.layer.weight_scale": params.block_scale, + "model.layer.weight_scale_2": params.scale, + }, str(path), metadata=metadata) + + class Config: + def get_model(self, state, prefix, device): + model = torch.nn.Module() + model.diffusion_model = torch.nn.Module() + model.diffusion_model.layer = torch.nn.Module() + model.diffusion_model.layer._orig_shape = tuple(original.shape) + return model + + monkeypatch.setattr(comfy.model_detection, "model_config_from_unet", lambda *args, **kwargs: Config()) + handler = modules["resize"].MemoryEfficientSafeOpen(str(path), low_memory=True) + try: + decoder = modules["guard"].DiffusionQuantization(handler, str(path), "NVFP4 test") + key = "model.layer.weight" + assert decoder.logical_shape(key) == tuple(original.shape) + keys = decoder.required_keys(key) + stream = handler.async_stream(keys, batch_size=1, prefetch_batches=1, pin_memory=False) + tensors = {} + try: + for source_key in keys: + batch = next(stream) + assert len(batch) == 1 and batch[0][0] == source_key + tensors[source_key] = batch[0][1] + actual = decoder.decode(key, tensors, torch.float16) + finally: + for source_key in tensors: + handler.mark_processed(source_key) + stream.close() + expected = QuantizedTensor(qdata, layout_name, params).dequantize() + torch.testing.assert_close(actual, expected) + assert actual.shape == original.shape + finally: + handler.__exit__(None, None, None) diff --git a/tests/test_merger_filters.py b/tests/test_merger_filters.py index 7b5aa08..7af35da 100644 --- a/tests/test_merger_filters.py +++ b/tests/test_merger_filters.py @@ -1,10 +1,12 @@ import importlib +import json import sys import types from pathlib import Path import pytest import torch +from safetensors import safe_open from safetensors.torch import load_file, save_file @@ -130,3 +132,47 @@ def test_lora_factor_group_uses_one_normalized_layer_name(monkeypatch, tmp_path, assert "diffusion_model.layer.lora_A.weight" in result assert "diffusion_model.layer.lora_B.weight" in result assert not any(key.endswith(".alpha") for key in result) + + +def test_diffusion_merge_decodes_tensorwise_quantization_and_drops_sidecars( + monkeypatch, tmp_path, modules, +): + merger, operations = modules + paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")} + quant = torch.tensor(list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8) + save_file({ + "model.diffusion_model.quant.weight": torch.tensor([1, 2], dtype=torch.int8), + "model.diffusion_model.quant.weight_scale": torch.tensor(2.0), + "model.diffusion_model.quant.comfy_quant": quant, + "dense.weight": torch.tensor([2.0, 4.0]), + }, paths["a"]) + save_file({ + "model.diffusion_model.quant.weight": torch.tensor([3, 6], dtype=torch.int8), + "model.diffusion_model.quant.weight_scale": torch.tensor(2.0), + "model.diffusion_model.quant.comfy_quant": quant, + "dense.weight": torch.tensor([6.0, 8.0]), + }, paths["b"]) + + 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("quantized_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_path = tmp_path / "diffusion_models" / "quantized_merge.safetensors" + result = load_file(str(output_path)) + key = "model.diffusion_model.quant.weight" + assert set(result) == {key, "dense.weight"} + assert result[key].tolist() == pytest.approx([4.0, 8.0]) + assert result[key].dtype == torch.float32 + assert result["dense.weight"].tolist() == pytest.approx([4.0, 6.0]) + with safe_open(str(output_path), framework="pt", device="cpu") as output: + metadata = output.metadata() + assert metadata in (None, {}) diff --git a/tests/test_model_analysis.py b/tests/test_model_analysis.py index 6856844..fd4c88b 100644 --- a/tests/test_model_analysis.py +++ b/tests/test_model_analysis.py @@ -1,6 +1,7 @@ import importlib import csv from io import StringIO +import json import math import sys import types @@ -331,6 +332,51 @@ def test_include_filter_limits_analysis_inventory( assert sorted(streamed) == sorted(expected_keys * 2) +def test_diffusion_analysis_decodes_int8_tensorwise_and_skips_sidecars( + monkeypatch, tmp_path, analysis, +): + prefix = "model.diffusion_model.layer" + quant = torch.tensor(list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8) + paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")} + save_file({ + f"{prefix}.weight": torch.tensor([1, 2], dtype=torch.int8), + f"{prefix}.weight_scale": torch.tensor(2.0), + f"{prefix}.comfy_quant": quant, + "dense.weight": torch.tensor([2.0, 4.0]), + }, paths["a"]) + save_file({ + f"{prefix}.weight": torch.tensor([3, 6], dtype=torch.int8), + f"{prefix}.weight_scale": torch.tensor(2.0), + f"{prefix}.comfy_quant": quant, + "dense.weight": torch.tensor([4.0, 8.0]), + }, 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) + streamed = [] + original_stream = analysis.MemoryEfficientSafeOpen.async_stream + + def track_stream(handler, keys, *args, **kwargs): + streamed.append((handler.filename, tuple(keys))) + return original_stream(handler, keys, *args, **kwargs) + + monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", track_stream) + _, _, metrics_csv, _ = analysis.ModelAnalysisLogic.execute( + "a", "b", "diffusion_models", _params(), + ) + + 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 float(rows["dense.weight"]["MAE"]) == pytest.approx(3.0) + for path, keys in streamed: + if path.endswith("a.safetensors") or path.endswith("b.safetensors"): + if f"{prefix}.weight" in keys: + assert f"{prefix}.weight_scale" in keys + assert f"{prefix}.comfy_quant" in keys + assert sum(f"{prefix}.weight" in keys or "dense.weight" in keys for _, keys in streamed) == 4 + + 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]])