432 lines
15 KiB
Python
432 lines
15 KiB
Python
# SimpleSyrup - workflow-focused ComfyUI extensions for image generation
|
|
# Copyright (C) 2026 Artificial Sweetener and contributors
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
|
|
"""Tests for streaming ComfyUI checkpoint quantization."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import importlib
|
|
import sys
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from types import ModuleType
|
|
from typing import Any
|
|
|
|
import pytest
|
|
import torch
|
|
from safetensors import safe_open
|
|
from safetensors.torch import load_file, save_file
|
|
|
|
from simple_syrup.domain.anima_quantization import (
|
|
MXFP8_PROFILE,
|
|
NVFP4_MIXED_PROFILE,
|
|
AnimaQuantizationRecipe,
|
|
)
|
|
from simple_syrup.domain.model_quantization import QuantizationFormat
|
|
from simple_syrup.domain.quant_cache import SourceCheckpointIdentity
|
|
from simple_syrup.runtime.checkpoint_quantizer import SafetensorsCheckpointQuantizer
|
|
from simple_syrup.runtime.comfy_safetensors_dtypes import (
|
|
ComfySafetensorsDtypeRegistry,
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class RecordingProgress:
|
|
"""Record absolute quantization progress updates."""
|
|
|
|
updates: list[tuple[int, int]] = field(default_factory=list)
|
|
|
|
def start(self, label: str, total: int) -> None:
|
|
"""Record no start because the resolver owns operation setup."""
|
|
|
|
del label, total
|
|
|
|
def advance(self, current: int, total: int) -> None:
|
|
"""Record one absolute update."""
|
|
|
|
self.updates.append((current, total))
|
|
|
|
def finish(self) -> None:
|
|
"""Record no finish because the resolver owns completion."""
|
|
|
|
|
|
class FakeQuantizedTensor:
|
|
"""Serialize deterministic fake quant tensors through Comfy's boundary."""
|
|
|
|
def __init__(self, source: torch.Tensor) -> None:
|
|
"""Retain only the source shape needed by the fake serializer."""
|
|
|
|
self._shape = source.shape
|
|
|
|
@classmethod
|
|
def from_float(
|
|
cls,
|
|
tensor: torch.Tensor,
|
|
layout_name: str,
|
|
**kwargs: object,
|
|
) -> FakeQuantizedTensor:
|
|
"""Accept the same call shape as ComfyUI's QuantizedTensor."""
|
|
|
|
assert layout_name == "FakeLayout"
|
|
assert kwargs == {"scale": "recalculate"}
|
|
return cls(tensor)
|
|
|
|
def state_dict(self, prefix: str) -> dict[str, torch.Tensor]:
|
|
"""Return representative quantized weight and scale tensors."""
|
|
|
|
return {
|
|
prefix: torch.zeros(self._shape, dtype=torch.uint8),
|
|
f"{prefix}_scale": torch.ones((), dtype=torch.float32),
|
|
}
|
|
|
|
|
|
class RecordingDtypeRegistry(ComfySafetensorsDtypeRegistry):
|
|
"""Record header and marker validation performed before publication."""
|
|
|
|
def __init__(self, *, supports_markers: bool = True) -> None:
|
|
"""Create an empty validation record."""
|
|
|
|
self._supports_markers = supports_markers
|
|
self.validated_paths: list[Path] = []
|
|
self.validated_formats: list[QuantizationFormat] = []
|
|
|
|
def validate_checkpoint_header(self, checkpoint_path: Path) -> None:
|
|
"""Record the generated checkpoint whose header was validated."""
|
|
|
|
self.validated_paths.append(checkpoint_path)
|
|
|
|
def supports_format(self, quantization_format: QuantizationFormat) -> bool:
|
|
"""Record and accept one generated Comfy marker format."""
|
|
|
|
self.validated_formats.append(quantization_format)
|
|
return self._supports_markers
|
|
|
|
|
|
def test_quantizer_streams_eligible_tensors_and_preserves_anima_policy(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Only eligible matrices receive Comfy layer markers and scale tensors."""
|
|
|
|
source_path = tmp_path / "anima.safetensors"
|
|
destination = tmp_path / "quantized.safetensors"
|
|
save_file(
|
|
{
|
|
"net.blocks.0.attn.q_proj.weight": torch.randn(4, 4, dtype=torch.bfloat16),
|
|
"net.blocks.2.attn.q_proj.weight": torch.randn(4, 4, dtype=torch.bfloat16),
|
|
"net.blocks.2.attn.q_proj.bias": torch.randn(4, dtype=torch.bfloat16),
|
|
},
|
|
str(source_path),
|
|
metadata={"source": "test"},
|
|
)
|
|
quant_ops = ModuleType("comfy.quant_ops")
|
|
quant_ops.QUANT_ALGOS = { # type: ignore[attr-defined]
|
|
"float8_e4m3fn": {"comfy_tensor_layout": "FakeLayout"},
|
|
"nvfp4": {"comfy_tensor_layout": "FakeLayout"},
|
|
}
|
|
quant_ops.QuantizedTensor = FakeQuantizedTensor # type: ignore[attr-defined]
|
|
model_management = ModuleType("comfy.model_management")
|
|
model_management.get_torch_device = lambda: torch.device("cpu") # type: ignore[attr-defined]
|
|
monkeypatch.setitem(sys.modules, "comfy.quant_ops", quant_ops)
|
|
monkeypatch.setitem(sys.modules, "comfy.model_management", model_management)
|
|
progress = RecordingProgress()
|
|
source = _identity(source_path)
|
|
dtype_registry = RecordingDtypeRegistry()
|
|
|
|
result = SafetensorsCheckpointQuantizer(dtype_registry).quantize(
|
|
source=source,
|
|
destination_path=destination,
|
|
profile=NVFP4_MIXED_PROFILE,
|
|
recipe=AnimaQuantizationRecipe(),
|
|
progress=progress,
|
|
progress_base=source.size_bytes,
|
|
progress_total=source.size_bytes * 2,
|
|
)
|
|
|
|
assert result.quantized_tensor_count == 1
|
|
assert result.preserved_tensor_count == 2
|
|
assert progress.updates
|
|
assert dtype_registry.validated_paths == [destination]
|
|
assert dtype_registry.validated_formats == [QuantizationFormat.NVFP4]
|
|
with safe_open(str(destination), framework="pt", device="cpu") as checkpoint:
|
|
assert set(checkpoint.keys()) == {
|
|
"net.blocks.0.attn.q_proj.weight",
|
|
"net.blocks.2.attn.q_proj.bias",
|
|
"net.blocks.2.attn.q_proj.comfy_quant",
|
|
"net.blocks.2.attn.q_proj.weight",
|
|
"net.blocks.2.attn.q_proj.weight_scale",
|
|
}
|
|
assert (
|
|
checkpoint.get_tensor("net.blocks.0.attn.q_proj.weight").dtype
|
|
is torch.bfloat16
|
|
)
|
|
assert checkpoint.metadata() == {
|
|
"source": "test",
|
|
"simple_syrup.derived_model": "true",
|
|
"simple_syrup.model_family": "Anima",
|
|
"simple_syrup.profile_label": "NVFP4 (Mixed)",
|
|
"simple_syrup.profile_version": "3",
|
|
"simple_syrup.quantization_formats": "float8_e4m3fn,nvfp4",
|
|
"simple_syrup.quantization_profile": "nvfp4-mixed",
|
|
"simple_syrup.recipe_version": "2",
|
|
"simple_syrup.source_model": "Anima/anima.safetensors",
|
|
"simple_syrup.source_sha256": source.sha256,
|
|
}
|
|
|
|
|
|
def test_quantizer_rejects_a_quantized_source(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Derived checkpoints cannot silently become lossy quantization sources."""
|
|
|
|
source_path = tmp_path / "already-quantized.safetensors"
|
|
save_file(
|
|
{
|
|
"net.blocks.2.linear.weight": torch.ones(4, 4),
|
|
"net.blocks.2.linear.comfy_quant": torch.tensor(
|
|
list(b'{"format":"nvfp4"}'), dtype=torch.uint8
|
|
),
|
|
},
|
|
str(source_path),
|
|
)
|
|
quant_ops = ModuleType("comfy.quant_ops")
|
|
quant_ops.QUANT_ALGOS = { # type: ignore[attr-defined]
|
|
"float8_e4m3fn": {"comfy_tensor_layout": "FakeLayout"},
|
|
"nvfp4": {"comfy_tensor_layout": "FakeLayout"},
|
|
}
|
|
quant_ops.QuantizedTensor = FakeQuantizedTensor # type: ignore[attr-defined]
|
|
model_management = ModuleType("comfy.model_management")
|
|
model_management.get_torch_device = lambda: torch.device("cpu") # type: ignore[attr-defined]
|
|
monkeypatch.setitem(sys.modules, "comfy.quant_ops", quant_ops)
|
|
monkeypatch.setitem(sys.modules, "comfy.model_management", model_management)
|
|
|
|
with pytest.raises(ValueError, match="already contains ComfyUI quantization"):
|
|
SafetensorsCheckpointQuantizer().quantize(
|
|
source=_identity(source_path),
|
|
destination_path=tmp_path / "output.safetensors",
|
|
profile=NVFP4_MIXED_PROFILE,
|
|
recipe=AnimaQuantizationRecipe(),
|
|
progress=RecordingProgress(),
|
|
progress_base=0,
|
|
progress_total=1,
|
|
)
|
|
|
|
|
|
def test_validation_rejects_inconsistent_profile_metadata(tmp_path: Path) -> None:
|
|
"""A generated artifact cannot publish under mismatched profile metadata."""
|
|
|
|
destination = tmp_path / "mismatched-profile.safetensors"
|
|
save_file(
|
|
{
|
|
"net.blocks.2.attn.q_proj.weight": torch.zeros(4, 4, dtype=torch.uint8),
|
|
"net.blocks.2.attn.q_proj.comfy_quant": torch.tensor(
|
|
list(b'{"format":"nvfp4"}'), dtype=torch.uint8
|
|
),
|
|
},
|
|
str(destination),
|
|
metadata={
|
|
"simple_syrup.quantization_profile": "wrong-profile",
|
|
"simple_syrup.profile_version": "3",
|
|
"simple_syrup.quantization_formats": "float8_e4m3fn,nvfp4",
|
|
},
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="inconsistent profile metadata"):
|
|
SafetensorsCheckpointQuantizer(RecordingDtypeRegistry())._validate_output(
|
|
destination,
|
|
NVFP4_MIXED_PROFILE,
|
|
1,
|
|
)
|
|
|
|
|
|
def test_validation_rejects_marker_unsupported_by_active_loader(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""Publication fails when the active loader cannot read a marker format."""
|
|
|
|
destination = tmp_path / "unsupported-marker.safetensors"
|
|
save_file(
|
|
{
|
|
"net.blocks.2.attn.q_proj.weight": torch.zeros(4, 4, dtype=torch.uint8),
|
|
"net.blocks.2.attn.q_proj.comfy_quant": torch.tensor(
|
|
list(b'{"format":"nvfp4"}'), dtype=torch.uint8
|
|
),
|
|
},
|
|
str(destination),
|
|
metadata={
|
|
"simple_syrup.quantization_profile": "nvfp4-mixed",
|
|
"simple_syrup.profile_version": "3",
|
|
"simple_syrup.quantization_formats": "float8_e4m3fn,nvfp4",
|
|
},
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="unsupported by the active Comfy loader"):
|
|
SafetensorsCheckpointQuantizer(
|
|
RecordingDtypeRegistry(supports_markers=False)
|
|
)._validate_output(
|
|
destination,
|
|
NVFP4_MIXED_PROFILE,
|
|
1,
|
|
)
|
|
|
|
|
|
def test_native_nvfp4_generation_uses_comfy_scale_layout(tmp_path: Path) -> None:
|
|
"""The installed ComfyUI runtime produces its native NVFP4 checkpoint shape."""
|
|
|
|
try:
|
|
import comfy.model_management as model_management
|
|
except ImportError:
|
|
pytest.skip("ComfyUI runtime is unavailable.")
|
|
if not model_management.supports_nvfp4_compute(model_management.get_torch_device()):
|
|
pytest.skip("NVFP4 compute is unavailable on this GPU.")
|
|
|
|
source_path = tmp_path / "native-anima.safetensors"
|
|
destination = tmp_path / "native-nvfp4.safetensors"
|
|
save_file(
|
|
{"net.blocks.2.attn.q_proj.weight": torch.randn(16, 16, dtype=torch.bfloat16)},
|
|
str(source_path),
|
|
)
|
|
source = _identity(source_path)
|
|
|
|
SafetensorsCheckpointQuantizer().quantize(
|
|
source=source,
|
|
destination_path=destination,
|
|
profile=NVFP4_MIXED_PROFILE,
|
|
recipe=AnimaQuantizationRecipe(),
|
|
progress=RecordingProgress(),
|
|
progress_base=source.size_bytes,
|
|
progress_total=source.size_bytes * 2,
|
|
)
|
|
|
|
with safe_open(str(destination), framework="pt", device="cpu") as checkpoint:
|
|
keys = set(checkpoint.keys())
|
|
assert "net.blocks.2.attn.q_proj.comfy_quant" in keys
|
|
assert "net.blocks.2.attn.q_proj.weight_scale" in keys
|
|
assert "net.blocks.2.attn.q_proj.weight_scale_2" in keys
|
|
assert (
|
|
checkpoint.get_tensor("net.blocks.2.attn.q_proj.weight").dtype
|
|
is torch.uint8
|
|
)
|
|
|
|
import comfy.ops
|
|
import comfy.quant_ops
|
|
|
|
state_dict = {
|
|
key.removeprefix("net.blocks.2.attn.q_proj."): value
|
|
for key, value in load_file(str(destination)).items()
|
|
}
|
|
operations = comfy.ops.mixed_precision_ops(
|
|
{"mixed_ops": True}, compute_dtype=torch.bfloat16
|
|
)
|
|
layer = operations.Linear(
|
|
16,
|
|
16,
|
|
bias=False,
|
|
device=torch.device("cpu"),
|
|
dtype=torch.bfloat16,
|
|
)
|
|
|
|
layer.load_state_dict(state_dict, strict=True)
|
|
|
|
assert isinstance(layer.weight, comfy.quant_ops.QuantizedTensor)
|
|
assert layer.quant_format == "nvfp4"
|
|
|
|
|
|
def test_native_mxfp8_generation_reloads_through_aimdo(tmp_path: Path) -> None:
|
|
"""E8M0 scale tensors reload through Comfy's active AIMDO mmap path."""
|
|
|
|
import comfy.model_management as model_management
|
|
import comfy.utils as comfy_utils
|
|
|
|
aimdo_control: Any = importlib.import_module("comfy_aimdo.control")
|
|
aimdo_model_mmap: Any = importlib.import_module("comfy_aimdo.model_mmap")
|
|
|
|
device = model_management.get_torch_device()
|
|
if not model_management.supports_mxfp8_compute(device):
|
|
pytest.skip("MXFP8 compute is unavailable on this GPU.")
|
|
source_path = tmp_path / "native-mxfp8-source.safetensors"
|
|
destination = tmp_path / "native-mxfp8.safetensors"
|
|
save_file(
|
|
{"net.blocks.14.attn.q_proj.weight": torch.randn(32, 32, dtype=torch.bfloat16)},
|
|
str(source_path),
|
|
)
|
|
source = _identity(source_path)
|
|
|
|
SafetensorsCheckpointQuantizer().quantize(
|
|
source=source,
|
|
destination_path=destination,
|
|
profile=MXFP8_PROFILE,
|
|
recipe=AnimaQuantizationRecipe(),
|
|
progress=RecordingProgress(),
|
|
progress_base=source.size_bytes,
|
|
progress_total=source.size_bytes * 2,
|
|
)
|
|
if not aimdo_control.init():
|
|
pytest.skip("comfy-aimdo is unavailable.")
|
|
importlib.reload(aimdo_model_mmap)
|
|
state_dict, metadata = comfy_utils.load_safetensors(str(destination))
|
|
|
|
assert metadata["simple_syrup.quantization_profile"] == "mxfp8"
|
|
assert any(tensor.dtype is torch.float8_e8m0fnu for tensor in state_dict.values())
|
|
|
|
|
|
def test_native_recommended_profile_serializes_both_formats(tmp_path: Path) -> None:
|
|
"""Recommended profile emits native NVFP4 and FP8 markers in one checkpoint."""
|
|
|
|
import comfy.model_management as model_management
|
|
|
|
device = model_management.get_torch_device()
|
|
if not model_management.supports_nvfp4_compute(device):
|
|
pytest.skip("NVFP4 compute is unavailable on this GPU.")
|
|
source_path = tmp_path / "native-mixed-source.safetensors"
|
|
destination = tmp_path / "native-mixed.safetensors"
|
|
save_file(
|
|
{
|
|
"net.blocks.14.attn.q_proj.weight": torch.randn(
|
|
32, 32, dtype=torch.bfloat16
|
|
),
|
|
"net.blocks.14.attn.v_proj.weight": torch.randn(
|
|
32, 32, dtype=torch.bfloat16
|
|
),
|
|
},
|
|
str(source_path),
|
|
)
|
|
source = _identity(source_path)
|
|
|
|
SafetensorsCheckpointQuantizer().quantize(
|
|
source=source,
|
|
destination_path=destination,
|
|
profile=NVFP4_MIXED_PROFILE,
|
|
recipe=AnimaQuantizationRecipe(),
|
|
progress=RecordingProgress(),
|
|
progress_base=source.size_bytes,
|
|
progress_total=source.size_bytes * 2,
|
|
)
|
|
|
|
with safe_open(str(destination), framework="pt", device="cpu") as checkpoint:
|
|
q_marker = bytes(
|
|
checkpoint.get_tensor("net.blocks.14.attn.q_proj.comfy_quant").tolist()
|
|
)
|
|
v_marker = bytes(
|
|
checkpoint.get_tensor("net.blocks.14.attn.v_proj.comfy_quant").tolist()
|
|
)
|
|
assert b'"format":"nvfp4"' in q_marker
|
|
assert b'"format":"float8_e4m3fn"' in v_marker
|
|
|
|
|
|
def _identity(path: Path) -> SourceCheckpointIdentity:
|
|
"""Create a complete source identity for a tiny test checkpoint."""
|
|
|
|
content = path.read_bytes()
|
|
stat = path.stat()
|
|
return SourceCheckpointIdentity(
|
|
display_name="Anima/anima.safetensors",
|
|
path=path.resolve(),
|
|
size_bytes=stat.st_size,
|
|
modified_ns=stat.st_mtime_ns,
|
|
sha256=hashlib.sha256(content).hexdigest(),
|
|
)
|