ThereforeGames
722dadcbc5
Rebrand the project as "ComfyUI Quantization Toolkit" and introduce native INT4 (ConvRot W4A4) support, mixed INT4/INT8 modes, and expanded quantization policies. Add an INT4 torch.compile compatibility shim, INT4-aware lazy-compile changes, richer LoRA stack nodes with autogrow, and a Quantized LoRA entry/apply pair. Replace ad-hoc prints with structured logging, bump package metadata to 2.0.0, and add documentation (advanced usage, checkpoints, policies) plus many unit tests for compile, lazy-compile, LoRA, adapter caching, saving, and modes. Misc: kernel config, model save/load, and numerous runtime/adapter refactors and robustness fixes.
127 lines
4.8 KiB
Python
127 lines
4.8 KiB
Python
"""
|
|
ComfyUI Quantization Toolkit - native INT4 and Toolkit INT8 support
|
|
|
|
Provides:
|
|
- Int8TensorwiseOps: Custom operations for direct int8 weight loading
|
|
- OTUNetLoaderW8A8: Backward-compatible loader ID for mixed INT4/INT8 diffusion models
|
|
- INT8LoraLoader / INT8LoraLoaderStack: Unified LoRA loaders with mode switching
|
|
- INT8ModelAdapter: Backward-compatible adapter ID for quantizing an existing MODEL
|
|
|
|
Uses torch._int_mm for fast inference.
|
|
"""
|
|
|
|
import logging
|
|
import torch
|
|
|
|
# =============================================================================
|
|
# Layout Registration
|
|
# =============================================================================
|
|
|
|
def _register_layouts():
|
|
"""
|
|
Register the Int8Tensorwise layout with ComfyUI's model management.
|
|
"""
|
|
try:
|
|
from comfy.quant_ops import QUANT_ALGOS, register_layout_class, QuantizedLayout
|
|
|
|
class Int8TensorwiseLayout(QuantizedLayout):
|
|
"""Minimal layout class to satisfy ComfyUI's registry requirements."""
|
|
class Params:
|
|
def __init__(self, scale=None, orig_dtype=None, orig_shape=None, **kwargs):
|
|
self.scale = scale
|
|
self.orig_dtype = orig_dtype
|
|
self.orig_shape = orig_shape
|
|
|
|
def clone(self):
|
|
return Int8TensorwiseLayout.Params(
|
|
scale=self.scale.clone() if isinstance(self.scale, torch.Tensor) else self.scale,
|
|
orig_dtype=self.orig_dtype,
|
|
orig_shape=self.orig_shape
|
|
)
|
|
|
|
@classmethod
|
|
def state_dict_tensors(cls, qdata, params):
|
|
return {"": qdata, "weight_scale": params.scale}
|
|
|
|
@classmethod
|
|
def dequantize(cls, qdata, params):
|
|
return qdata.float() * params.scale
|
|
|
|
# Register the class
|
|
register_layout_class("Int8TensorwiseLayout", Int8TensorwiseLayout)
|
|
|
|
# Register the Algo Config
|
|
QUANT_ALGOS.setdefault(
|
|
"int8_tensorwise",
|
|
{
|
|
"storage_t": torch.int8,
|
|
# We include input_scale here so ComfyUI extracts it from checkpoints if present,
|
|
# even though our LinearW8A8 implementation explicitly ignores it.
|
|
"parameters": {"weight_scale", "input_scale"},
|
|
"comfy_tensor_layout": "Int8TensorwiseLayout",
|
|
}
|
|
)
|
|
|
|
except ImportError:
|
|
logging.warning("Quantization Toolkit: ComfyUI quantization support was not found; update ComfyUI.")
|
|
except Exception as e:
|
|
logging.error(f"Quantization Toolkit: failed to register quantization layouts ({e}).")
|
|
|
|
|
|
# =============================================================================
|
|
# Module Initialization
|
|
# =============================================================================
|
|
|
|
# 1. Register Layouts
|
|
_register_layouts()
|
|
|
|
# 2. Export Custom Ops (for external use)
|
|
try:
|
|
from .int8_quant import Int8TensorwiseOps
|
|
except ImportError:
|
|
Int8TensorwiseOps = None
|
|
|
|
# 3. Node Mappings
|
|
# Wrap imports in try/except to prevent total failure if dependencies are missing
|
|
try:
|
|
from .int8_unet_loader import UNetLoaderINTW8A8
|
|
from .int8_lora import INT8LoraLoader, INT8LoraLoaderStack, QuantizedLoraConfig, QuantizedLoraPatcher
|
|
from .int8_kernel_config_node import INT8KernelConfigTuner
|
|
from .int8_model_adapter import INT8ModelAdapter
|
|
from .int8_model_save import INT8ModelSave
|
|
from .int8_lazy_compile import INT8LazyTorchCompile
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"OTUNetLoaderW8A8": UNetLoaderINTW8A8,
|
|
"INT8ModelAdapter": INT8ModelAdapter,
|
|
"INT8ModelSave": INT8ModelSave,
|
|
"INT8LazyTorchCompile": INT8LazyTorchCompile,
|
|
"INT8LoraLoader": INT8LoraLoader,
|
|
"INT8LoraLoaderStack": INT8LoraLoaderStack,
|
|
"QuantizedLoraConfig": QuantizedLoraConfig,
|
|
"QuantizedLoraPatcher": QuantizedLoraPatcher,
|
|
"INT8KernelConfigTuner": INT8KernelConfigTuner,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"OTUNetLoaderW8A8": "Load Diffusion Model Quantized",
|
|
"INT8ModelAdapter": "Enable Quantization on MODEL",
|
|
"INT8ModelSave": "Save Quantized Model (DynamicVRAM Safe)",
|
|
"INT8LazyTorchCompile": "Quantized Lazy Torch Compile",
|
|
"INT8LoraLoader": "Load LoRA (Quantized)",
|
|
"INT8LoraLoaderStack": "Load LoRA Stack (Quantized)",
|
|
"QuantizedLoraConfig": "LoRA Stack Entry (Quantized)",
|
|
"QuantizedLoraPatcher": "Apply LoRA Stack (Quantized)",
|
|
"INT8KernelConfigTuner": "INT8 Kernel Config",
|
|
}
|
|
except ImportError as e:
|
|
logging.error(f"Quantization Toolkit: failed to import nodes ({e}).")
|
|
NODE_CLASS_MAPPINGS = {}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
|
|
__all__ = [
|
|
"NODE_CLASS_MAPPINGS",
|
|
"NODE_DISPLAY_NAME_MAPPINGS",
|
|
"Int8TensorwiseOps",
|
|
]
|