83 lines
2.3 KiB
Python
83 lines
2.3 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
|
|
|
|
"""Define model-independent checkpoint quantization profile contracts."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from enum import StrEnum
|
|
from typing import Protocol
|
|
|
|
|
|
class QuantizationFormat(StrEnum):
|
|
"""Identify one reusable ComfyUI tensor quantization format."""
|
|
|
|
FP8_E4M3 = "float8_e4m3fn"
|
|
FP8_E5M2 = "float8_e5m2"
|
|
NVFP4 = "nvfp4"
|
|
MXFP8 = "mxfp8"
|
|
|
|
@property
|
|
def label(self) -> str:
|
|
"""Return the concise format label used in diagnostics."""
|
|
|
|
return {
|
|
QuantizationFormat.FP8_E4M3: "FP8 E4M3",
|
|
QuantizationFormat.FP8_E5M2: "FP8 E5M2",
|
|
QuantizationFormat.NVFP4: "NVFP4",
|
|
QuantizationFormat.MXFP8: "MXFP8",
|
|
}[self]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class QuantizationProfile:
|
|
"""Describe one workflow-facing, versioned per-tensor policy profile."""
|
|
|
|
profile_id: str
|
|
label: str
|
|
version: int
|
|
required_formats: frozenset[QuantizationFormat]
|
|
|
|
@property
|
|
def is_original(self) -> bool:
|
|
"""Return whether this profile loads the source checkpoint unchanged."""
|
|
|
|
return not self.required_formats
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TensorDescriptor:
|
|
"""Describe a checkpoint tensor without coupling policy to PyTorch."""
|
|
|
|
name: str
|
|
shape: tuple[int, ...]
|
|
dtype_name: str
|
|
|
|
|
|
class ModelQuantizationRecipe(Protocol):
|
|
"""Assign model-specific per-tensor formats for named profiles."""
|
|
|
|
@property
|
|
def model_family(self) -> str:
|
|
"""Return the stable family identifier used in cache identity."""
|
|
|
|
@property
|
|
def version(self) -> int:
|
|
"""Return the recipe version used in cache invalidation."""
|
|
|
|
@property
|
|
def profiles(self) -> tuple[QuantizationProfile, ...]:
|
|
"""Return deterministic workflow profiles owned by this recipe."""
|
|
|
|
def profile_from_selection(self, selection: str) -> QuantizationProfile:
|
|
"""Parse a workflow selection into a recipe-owned profile."""
|
|
|
|
def policy_for(
|
|
self,
|
|
tensor: TensorDescriptor,
|
|
profile: QuantizationProfile,
|
|
) -> QuantizationFormat | None:
|
|
"""Return the tensor format or ``None`` to preserve source precision."""
|