feat(dtype): map dtypes from reference

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)

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