From e7e49a5f7e356181ba87f49a9fbdcf2d5cb35bc0 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Mon, 24 Aug 2026 16:24:09 +0200 Subject: [PATCH] feat(dtype): map dtypes from reference --- nodes/dtype_conversion.py | 41 ++++++++++++++++++++++++++++++++- tests/test_dtype_conversion.py | 42 ++++++++++++++++++++++++++++++++++ 2 files changed, 82 insertions(+), 1 deletion(-) diff --git a/nodes/dtype_conversion.py b/nodes/dtype_conversion.py index 7206517..c20a92d 100644 --- a/nodes/dtype_conversion.py +++ b/nodes/dtype_conversion.py @@ -5,6 +5,7 @@ from __future__ import annotations import os import re from collections import Counter +from contextlib import nullcontext import comfy.utils import folder_paths @@ -53,6 +54,7 @@ def convert_diffusion_model_dtype( target_dtype: str, exclude_patterns: str, output_filename: str, + reference_model: str | None = None, ) -> tuple[str, str]: try: destination_dtype = TARGET_DTYPES[target_dtype] @@ -71,10 +73,19 @@ def convert_diffusion_model_dtype( exclusions = _compile_exclusions(exclude_patterns) counts = Counter() original_dtypes = Counter() + reference_path = None + if reference_model and reference_model != "None": + reference_path = folder_paths.get_full_path_or_raise( + "diffusion_models", reference_model + ) try: - with MemoryEfficientSafeOpen(source_path, low_memory=True) as loader: + with MemoryEfficientSafeOpen(source_path, low_memory=True) as loader, ( + MemoryEfficientSafeOpen(reference_path, low_memory=True) + if reference_path else nullcontext() + ) as reference_loader: metadata = (loader.metadata() or {}).copy() + reference_keys = set(reference_loader.keys()) if reference_loader else set() with atomic_uel_writer(output_path, metadata) as writer: keys = list(loader.keys()) progress = comfy.utils.ProgressBar(len(keys)) @@ -97,6 +108,22 @@ def convert_diffusion_model_dtype( counts["excluded"] += 1 elif not source.is_floating_point(): counts["non_floating"] += 1 + elif reference_loader is not None and key in reference_keys: + reference_dtype = reference_loader.get_dtype(key) + if not reference_dtype.is_floating_point: + raise ValueError( + f"Reference tensor '{key}' is {reference_dtype}, but the input tensor is floating." + ) + if source.dtype == reference_dtype: + counts["reference_already_matched"] += 1 + else: + output = source.to(dtype=reference_dtype).contiguous() + counts["reference_mapped"] += 1 + elif reference_loader is not None: + output = source if source.dtype == destination_dtype else source.to( + dtype=destination_dtype + ).contiguous() + counts["reference_missing_fallback"] += 1 elif source.dtype == destination_dtype: counts["already_target"] += 1 else: @@ -121,12 +148,16 @@ def convert_diffusion_model_dtype( ( "DIFFUSION MODEL DTYPE CONVERSION", f"Input: {model_name}", + f"Reference model: {reference_model or 'None'}", f"Output: {output_name}", f"Target dtype: {target_dtype}", f"Total tensors: {total}", f"Converted floating tensors: {counts['converted']}", f"Excluded tensors preserved: {counts['excluded']}", f"Already target dtype: {counts['already_target']}", + f"Reference-mapped floating tensors: {counts['reference_mapped']}", + f"Reference dtype already matched: {counts['reference_already_matched']}", + f"Reference-missing fallback tensors: {counts['reference_missing_fallback']}", f"Non-floating tensors preserved: {counts['non_floating']}", f"Original dtype counts: {dtype_summary or 'None'}", ) @@ -156,6 +187,12 @@ class DiffusionModelDtypeConversion(io.ComfyNode): "target_dtype", options=["fp32", "fp16", "bf16"], tooltip="Target dtype for floating tensors not matched by an exclusion pattern.", ), + io.Combo.Input( + "reference_model", + options=["None", *folder_paths.get_filename_list("diffusion_models")], + default="None", + tooltip="Optional authoritative dtype map. Matching floating tensor keys use this model's header dtype without loading its tensors; missing keys use Target Dtype.", + ), io.String.Input( "exclude_patterns", default="", @@ -182,6 +219,7 @@ class DiffusionModelDtypeConversion(io.ComfyNode): cls, model_name: str, target_dtype: str, + reference_model: str, exclude_patterns: str, output_filename: str, ) -> io.NodeOutput: @@ -190,5 +228,6 @@ class DiffusionModelDtypeConversion(io.ComfyNode): target_dtype, exclude_patterns, output_filename, + reference_model, ) return io.NodeOutput(output_path, report) diff --git a/tests/test_dtype_conversion.py b/tests/test_dtype_conversion.py index b913e12..deb5325 100644 --- a/tests/test_dtype_conversion.py +++ b/tests/test_dtype_conversion.py @@ -65,6 +65,7 @@ def test_schema_is_diffusion_model_scoped(conversion): assert [item.id for item in schema.inputs] == [ "model_name", "target_dtype", + "reference_model", "exclude_patterns", "output_filename", ] @@ -75,6 +76,47 @@ def test_schema_is_diffusion_model_scoped(conversion): ] +def test_reference_model_dtypes_override_target_without_tensor_loading( + monkeypatch, tmp_path, conversion +): + source = tmp_path / "source.safetensors" + reference = tmp_path / "reference.safetensors" + _write_uel(source, { + "mapped_bf16": torch.ones(2, dtype=torch.float32), + "mapped_fp32": torch.ones(2, dtype=torch.float16), + "fallback": torch.ones(2, dtype=torch.float32), + "nonfloating": torch.tensor([1], dtype=torch.int64), + }) + _write_uel(reference, { + "mapped_bf16": torch.ones(2, dtype=torch.bfloat16), + "mapped_fp32": torch.ones(2, dtype=torch.float32), + "nonfloating": torch.tensor([1], dtype=torch.int64), + }) + monkeypatch.setattr(conversion.folder_paths, "models_dir", str(tmp_path)) + monkeypatch.setattr( + conversion.folder_paths, + "get_full_path_or_raise", + lambda _, name: {"source": str(source), "reference": str(reference)}[name], + ) + monkeypatch.setattr(conversion, "cleanup_after_operation", lambda: None) + monkeypatch.setattr( + conversion.MemoryEfficientSafeOpen, + "get_tensor", + lambda *args, **kwargs: pytest.fail("synchronous reference loading was used"), + ) + + output_path, report = conversion.convert_diffusion_model_dtype( + "source", "fp16", "", "reference_mapped", "reference" + ) + result, _ = _read_uel(tmp_path / "diffusion_models" / output_path) + assert result["mapped_bf16"].dtype == torch.bfloat16 + assert result["mapped_fp32"].dtype == torch.float32 + assert result["fallback"].dtype == torch.float16 + assert result["nonfloating"].dtype == torch.int64 + assert "Reference-mapped floating tensors: 2" in report + assert "Reference-missing fallback tensors: 1" in report + + @pytest.mark.parametrize( ("target_name", "expected_dtype"), [("fp32", torch.float32), ("fp16", torch.float16), ("bf16", torch.bfloat16)],