274 lines
11 KiB
Python
274 lines
11 KiB
Python
import importlib.util
|
|
from pathlib import Path
|
|
import sys
|
|
from types import SimpleNamespace
|
|
import unittest
|
|
|
|
import torch
|
|
|
|
|
|
PACKAGE_ROOT = Path(__file__).resolve().parents[1]
|
|
COMFY_ROOT = PACKAGE_ROOT.parents[1]
|
|
if str(COMFY_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(COMFY_ROOT))
|
|
|
|
spec = importlib.util.spec_from_file_location("comfyui_precision_detector", PACKAGE_ROOT / "__init__.py")
|
|
precision_detector = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(precision_detector)
|
|
|
|
|
|
class FakeRuntimeModel(torch.nn.Module):
|
|
def __init__(self, weight_dtype=torch.float32, manual_cast_dtype=None):
|
|
super().__init__()
|
|
self.weight = torch.nn.Parameter(torch.empty(2, dtype=weight_dtype), requires_grad=False)
|
|
self.manual_cast_dtype = manual_cast_dtype
|
|
|
|
def get_dtype(self):
|
|
return self.weight.dtype
|
|
|
|
def get_dtype_inference(self):
|
|
return self.manual_cast_dtype or self.get_dtype()
|
|
|
|
|
|
class FakePatcher:
|
|
def __init__(self, model, object_patches=None):
|
|
self.model = model
|
|
self.object_patches = object_patches or {}
|
|
self.load_device = torch.device("cuda:0")
|
|
self.offload_device = torch.device("cpu")
|
|
|
|
def model_dtype(self):
|
|
return self.model.get_dtype()
|
|
|
|
|
|
class FakeCLIP:
|
|
def __init__(self, dtype=torch.float16, compute_dtype=torch.float32):
|
|
self.cond_stage_model = FakeRuntimeModel(dtype)
|
|
self.patcher = FakePatcher(
|
|
self.cond_stage_model,
|
|
{"manual_cast_dtype": compute_dtype},
|
|
)
|
|
|
|
|
|
class FakeVAE:
|
|
def __init__(self, dtype):
|
|
self.vae_dtype = dtype
|
|
self.first_stage_model = FakeRuntimeModel(dtype)
|
|
self.device = torch.device("cuda:0")
|
|
self.patcher = FakePatcher(self.first_stage_model)
|
|
|
|
def vae_output_dtype(self):
|
|
return torch.float32
|
|
|
|
|
|
class MixedRuntimeModel(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.fp16_weight = torch.nn.Parameter(torch.empty(1, dtype=torch.float16), requires_grad=False)
|
|
self.fp32_weight = torch.nn.Parameter(torch.empty(1, dtype=torch.float32), requires_grad=False)
|
|
self.manual_cast_dtype = torch.bfloat16
|
|
|
|
def get_dtype(self):
|
|
return self.fp16_weight.dtype
|
|
|
|
def get_dtype_inference(self):
|
|
return self.manual_cast_dtype
|
|
|
|
|
|
class QuantizedRuntimeModel(FakeRuntimeModel):
|
|
def __init__(self):
|
|
super().__init__(torch.bfloat16, torch.bfloat16)
|
|
self.quant_format = "int8_tensorwise"
|
|
self.layout_type = "TensorWiseINT8Layout"
|
|
|
|
|
|
class LayoutTaggedRuntimeModel(FakeRuntimeModel):
|
|
def __init__(self, quant_format, layout_type, storage_dtype, linear_dtype=None):
|
|
super().__init__(torch.bfloat16, torch.bfloat16)
|
|
self.quant_format = quant_format
|
|
self.layout_type = layout_type
|
|
self.weight._layout_cls = layout_type
|
|
self.weight.storage_dtype = storage_dtype
|
|
if linear_dtype is not None:
|
|
self.weight.params = SimpleNamespace(linear_dtype=linear_dtype)
|
|
|
|
|
|
class PrecisionReportTests(unittest.TestCase):
|
|
def test_model_reports_fp32_without_manual_cast(self):
|
|
report = precision_detector.model_precision_report(FakePatcher(FakeRuntimeModel(torch.float32)))
|
|
self.assertTrue(report.startswith("Active compute: FP32\n"))
|
|
self.assertIn("Runtime weights: FP32", report)
|
|
self.assertIn("Manual cast: None", report)
|
|
|
|
def test_model_reports_fp16(self):
|
|
report = precision_detector.model_precision_report(FakePatcher(FakeRuntimeModel(torch.float16)))
|
|
self.assertTrue(report.startswith("Active compute: FP16\n"))
|
|
|
|
def test_model_reports_bf16_manual_cast(self):
|
|
report = precision_detector.model_precision_report(
|
|
FakePatcher(FakeRuntimeModel(torch.float32, torch.bfloat16))
|
|
)
|
|
self.assertTrue(report.startswith("Active compute: BF16\n"))
|
|
self.assertIn("Runtime weights: FP32", report)
|
|
self.assertIn("Manual cast: BF16", report)
|
|
|
|
@unittest.skipUnless(hasattr(torch, "float8_e4m3fn"), "PyTorch build has no FP8 dtype")
|
|
def test_model_reports_fp8(self):
|
|
report = precision_detector.model_precision_report(
|
|
FakePatcher(FakeRuntimeModel(torch.float8_e4m3fn))
|
|
)
|
|
self.assertTrue(report.startswith("Active compute: FP8 E4M3FN\n"))
|
|
self.assertIn("Runtime weights: FP8 E4M3FN", report)
|
|
|
|
def test_model_reports_int8_quantization_separately(self):
|
|
report = precision_detector.model_precision_report(FakePatcher(QuantizedRuntimeModel()))
|
|
self.assertTrue(report.startswith("Active compute: BF16\n"))
|
|
self.assertIn("INT8", report)
|
|
self.assertIn("TensorWiseINT8Layout", report)
|
|
|
|
def test_all_native_comfyui_quantization_layouts(self):
|
|
from comfy.quant_ops import QUANT_ALGOS
|
|
|
|
expected = {
|
|
"float8_e4m3fn": ("FP8 E4M3FN", None),
|
|
"float8_e5m2": ("FP8 E5M2", None),
|
|
"nvfp4": ("NVFP4", None),
|
|
"mxfp8": ("MXFP8", None),
|
|
"int8_tensorwise": ("INT8", None),
|
|
"convrot_w4a4": ("INT4 / ConvRot W4A4", "int4"),
|
|
}
|
|
self.assertEqual(set(QUANT_ALGOS), set(expected))
|
|
|
|
for quant_format, config in QUANT_ALGOS.items():
|
|
expected_label, linear_dtype = expected[quant_format]
|
|
model = LayoutTaggedRuntimeModel(
|
|
quant_format,
|
|
config["comfy_tensor_layout"],
|
|
config["storage_t"],
|
|
linear_dtype,
|
|
)
|
|
report = precision_detector.model_precision_report(FakePatcher(model))
|
|
self.assertIn(f"Runtime weights: {expected_label}", report)
|
|
self.assertIn(f"Quantization: {expected_label}", report)
|
|
self.assertIn(config["comfy_tensor_layout"], report)
|
|
self.assertNotIn("Quantization: None", report)
|
|
|
|
def test_quantized_layout_detection_is_shared_by_model_clip_and_vae(self):
|
|
runtime_model = LayoutTaggedRuntimeModel(
|
|
"nvfp4",
|
|
"TensorCoreNVFP4Layout",
|
|
torch.uint8,
|
|
)
|
|
patcher = FakePatcher(runtime_model, {"manual_cast_dtype": torch.bfloat16})
|
|
clip = SimpleNamespace(cond_stage_model=runtime_model, patcher=patcher)
|
|
vae = SimpleNamespace(
|
|
first_stage_model=runtime_model,
|
|
patcher=patcher,
|
|
vae_dtype=torch.bfloat16,
|
|
device=torch.device("cuda:0"),
|
|
vae_output_dtype=lambda: torch.float32,
|
|
)
|
|
|
|
reports = (
|
|
precision_detector.model_precision_report(patcher),
|
|
precision_detector.clip_precision_report(clip),
|
|
precision_detector.vae_precision_report(vae),
|
|
)
|
|
for report in reports:
|
|
self.assertIn("Runtime weights: NVFP4", report)
|
|
self.assertIn("layout TensorCoreNVFP4Layout", report)
|
|
self.assertIn("storage UINT8", report)
|
|
|
|
def test_all_installed_pytorch_dtypes_have_readable_labels(self):
|
|
dtypes = {
|
|
value
|
|
for name in dir(torch)
|
|
if isinstance((value := getattr(torch, name)), torch.dtype)
|
|
}
|
|
for dtype in dtypes:
|
|
label = precision_detector._dtype_label(dtype)
|
|
self.assertIsInstance(label, str)
|
|
self.assertTrue(label)
|
|
self.assertNotIn("TORCH.", label)
|
|
|
|
if hasattr(torch, "float4_e2m1fn_x2"):
|
|
self.assertEqual(
|
|
precision_detector._dtype_label(torch.float4_e2m1fn_x2),
|
|
"FP4 E2M1FN (packed x2)",
|
|
)
|
|
for name in ("int1", "int4", "int7", "uint1", "uint4", "uint7"):
|
|
if hasattr(torch, name):
|
|
self.assertEqual(
|
|
precision_detector._dtype_label(getattr(torch, name)),
|
|
name.upper(),
|
|
)
|
|
|
|
def test_model_reports_mixed_runtime_weights(self):
|
|
report = precision_detector.model_precision_report(FakePatcher(MixedRuntimeModel()))
|
|
self.assertIn("Runtime weights: Mixed (FP16, FP32)", report)
|
|
|
|
def test_patcher_compute_override_wins(self):
|
|
model = FakeRuntimeModel(torch.float32, torch.bfloat16)
|
|
report = precision_detector.model_precision_report(
|
|
FakePatcher(model, {"manual_cast_dtype": torch.float16})
|
|
)
|
|
self.assertTrue(report.startswith("Active compute: FP16\n"))
|
|
|
|
def test_clip_reports_comfyui_fp32_compute_override(self):
|
|
report = precision_detector.clip_precision_report(FakeCLIP())
|
|
self.assertTrue(report.startswith("Active compute: FP32\n"))
|
|
self.assertIn("Runtime weights: FP16", report)
|
|
|
|
def test_vae_reports_selected_compute_and_output_dtype(self):
|
|
report = precision_detector.vae_precision_report(FakeVAE(torch.bfloat16))
|
|
self.assertTrue(report.startswith("Active compute: BF16\n"))
|
|
self.assertIn("Output tensor: FP32", report)
|
|
|
|
def test_unknown_objects_are_reported_without_guessing(self):
|
|
model_report = precision_detector.model_precision_report(object())
|
|
clip_report = precision_detector.clip_precision_report(object())
|
|
vae_report = precision_detector.vae_precision_report(object())
|
|
self.assertTrue(model_report.startswith("Active compute: UNKNOWN"))
|
|
self.assertTrue(clip_report.startswith("Active compute: UNKNOWN"))
|
|
self.assertTrue(vae_report.startswith("Active compute: UNKNOWN"))
|
|
|
|
|
|
class NodeContractTests(unittest.TestCase):
|
|
def test_exactly_three_node_schemas(self):
|
|
classes = (
|
|
precision_detector.RuntimePrecisionModel,
|
|
precision_detector.RuntimePrecisionCLIP,
|
|
precision_detector.RuntimePrecisionVAE,
|
|
)
|
|
schemas = [node.define_schema() for node in classes]
|
|
self.assertEqual(
|
|
[schema.node_id for schema in schemas],
|
|
["RuntimePrecisionModel", "RuntimePrecisionCLIP", "RuntimePrecisionVAE"],
|
|
)
|
|
self.assertTrue(all(schema.is_output_node for schema in schemas))
|
|
self.assertTrue(all(len(schema.inputs) == 1 for schema in schemas))
|
|
self.assertTrue(all(len(schema.outputs) == 1 for schema in schemas))
|
|
|
|
def test_string_result_matches_visible_text(self):
|
|
output = precision_detector.RuntimePrecisionVAE.execute(FakeVAE(torch.float16))
|
|
report = output.result[0]
|
|
self.assertEqual(output.ui.as_dict()["text"], (report,))
|
|
self.assertTrue(report.startswith("Active compute: FP16\n"))
|
|
|
|
def test_frontend_extension_displays_reports_in_all_detector_nodes(self):
|
|
self.assertEqual(precision_detector.WEB_DIRECTORY, "./web")
|
|
script = (PACKAGE_ROOT / "web" / "precision_display.js").read_text(encoding="utf-8")
|
|
for node_id in (
|
|
"RuntimePrecisionModel",
|
|
"RuntimePrecisionCLIP",
|
|
"RuntimePrecisionVAE",
|
|
):
|
|
self.assertIn(node_id, script)
|
|
self.assertIn("ComfyWidgets.STRING", script)
|
|
self.assertIn("message?.text", script)
|
|
self.assertIn("readOnly = true", script)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|