Support quantized diffusion model inputs

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:

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, "

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)

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())

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

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)
 

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

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)

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, {

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)

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, {})

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]])