[Quantization] Add support for Comfy-Kitchen / Comfy-Quants

docs/source/en/_toctree.yml

@@ -97,6 +97,8 @@
       title: bitsandbytes
     - local: quantization/gguf
       title: gguf
+    - local: quantization/comfy_quant
+      title: Comfy Quant
     - local: quantization/nunchaku
       title: Nunchaku Lite
     - local: quantization/torchao

docs/source/en/quantization/comfy_quant.md

@@ -0,0 +1,77 @@
+<!--Copyright 2026 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+-->
+
+# Comfy Quant
+
+The [Comfy Quant](https://github.com/Comfy-Org/comfy-quants) toolkit provides state-of-the-art quantization techniques. While `comfy-quants` is used for exporting and quantizing models, Diffusers natively supports running inference on these models using the [comfy-kitchen](https://github.com/Comfy-Org/comfy-kitchen) library. 
+
+`comfy-kitchen` provides highly optimized GPU kernels that allow you to seamlessly run quantized layers. By passing a `ComfyQuantConfig` to Diffusers, the library will dynamically intercept parameters and wrap them in a `QuantizedTensor` that maps directly to the optimized `comfy-kitchen` layouts.
+
+Before starting, please install `comfy-kitchen` in your environment:
+
+```shell
+pip install comfy-kitchen
+```
+
+## Loading a Comfy Quant Model
+
+To load a model prequantized with Comfy Quant, use the [`~FromSingleFileMixin.from_single_file`] method and pass in the [`ComfyQuantConfig`]. 
+
+The configuration requires you to specify the `quant_format` that the model was quantized in, and the `compute_dtype` for active inference calculations.
+
+The following example demonstrates how to load a quantized FLUX transformer:
+
+```python
+import torch
+from diffusers import FluxPipeline, FluxTransformer2DModel, ComfyQuantConfig
+
+ckpt_path = "path/to/comfy_quant_checkpoint.safetensors"
+
+# Initialize the config with your desired format and compute dtype
+quantization_config = ComfyQuantConfig(
+    quant_format="fp8",
+    compute_dtype=torch.bfloat16
+)
+
+# Load the transformer directly from the safetensors file
+transformer = FluxTransformer2DModel.from_single_file(
+    ckpt_path,
+    quantization_config=quantization_config,
+    dtype=torch.bfloat16,
+)
+
+# Pass the quantized transformer into the pipeline
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev",
+    transformer=transformer,
+    dtype=torch.bfloat16,
+)
+pipe.enable_model_cpu_offload()
+
+prompt = "A cat holding a sign that says hello world"
+image = pipe(prompt, generator=torch.manual_seed(0)).images[0]
+image.save("flux-comfy-quant.png")
+```
+
+## Supported Quantization Formats
+
+Diffusers currently maps the following Comfy Quant formats to `comfy-kitchen` layouts:
+
+- **FP8** (`fp8`): Maps to `TensorCoreFP8Layout` (E4M3/E5M2)
+- **INT8** (`int8`): Maps to `TensorWiseINT8Layout` (W8A8, tensorwise)
+- **MXFP8** (`mxfp8`): Maps to `TensorCoreMXFP8Layout`
+- **NVFP4** (`nvfp4`): Maps to `TensorCoreNVFP4Layout`
+- **INT4 SVD** (`int4_svd`): Maps to `TensorCoreSVDQuantW4A4Layout` (SVDQuant W4A4)
+- **INT4 AWQ** (`int4_awq`): Maps to `TensorCoreAWQW4A16Layout` (AWQ W4A16)
+
+When using optimized layouts, `comfy-kitchen` automatically dispatches the operations to the best available backend (HIP, CUDA, Triton, or Eager).

src/diffusers/__init__.py

@@ -9,6 +9,7 @@
     is_accelerate_available,
     is_auto_round_available,
     is_bitsandbytes_available,
+    is_comfy_kitchen_available,
     is_gguf_available,
     is_librosa_available,
     is_note_seq_available,
@@ -47,6 +48,7 @@
     "schedulers": [],
     "utils": [
         "OptionalDependencyNotAvailable",
+        "is_comfy_kitchen_available",
         "is_inflect_available",
         "is_invisible_watermark_available",
         "is_librosa_available",
@@ -158,6 +160,18 @@
 else:
     _import_structure["quantizers.quantization_config"].append("SDNQConfig")
 
+try:
+    if not is_comfy_kitchen_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    from .utils import dummy_comfy_kitchen_objects
+
+    _import_structure["utils.dummy_comfy_kitchen_objects"] = [
+        name for name in dir(dummy_comfy_kitchen_objects) if not name.startswith("_")
+    ]
+else:
+    _import_structure["quantizers.quantization_config"].append("ComfyQuantConfig")
+
 try:
     if not is_onnx_available():
         raise OptionalDependencyNotAvailable()
@@ -1053,6 +1067,14 @@
     else:
         from .quantizers.quantization_config import SDNQConfig
 
+    try:
+        if not is_comfy_kitchen_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from .utils.dummy_comfy_kitchen_objects import *
+    else:
+        from .quantizers.quantization_config import ComfyQuantConfig
+
     try:
         if not is_onnx_available():
             raise OptionalDependencyNotAvailable()

src/diffusers/quantizers/auto.py

@@ -20,12 +20,14 @@
 
 from .autoround import AutoRoundQuantizer
 from .bitsandbytes import BnB4BitDiffusersQuantizer, BnB8BitDiffusersQuantizer
+from .comfy_quant import ComfyQuantizer
 from .gguf import GGUFQuantizer
 from .modelopt import NVIDIAModelOptQuantizer
 from .nunchaku import NunchakuLiteQuantizer
 from .quantization_config import (
     AutoRoundConfig,
     BitsAndBytesConfig,
+    ComfyQuantConfig,
     GGUFQuantizationConfig,
     NunchakuLiteQuantizationConfig,
     NVIDIAModelOptConfig,
@@ -43,6 +45,7 @@
 AUTO_QUANTIZER_MAPPING = {
     "bitsandbytes_4bit": BnB4BitDiffusersQuantizer,
     "bitsandbytes_8bit": BnB8BitDiffusersQuantizer,
+    "comfy_quant": ComfyQuantizer,
     "gguf": GGUFQuantizer,
     "quanto": QuantoQuantizer,
     "torchao": TorchAoHfQuantizer,
@@ -55,6 +58,7 @@
 AUTO_QUANTIZATION_CONFIG_MAPPING = {
     "bitsandbytes_4bit": BitsAndBytesConfig,
     "bitsandbytes_8bit": BitsAndBytesConfig,
+    "comfy_quant": ComfyQuantConfig,
     "gguf": GGUFQuantizationConfig,
     "quanto": QuantoConfig,
     "torchao": TorchAoConfig,

src/diffusers/quantizers/comfy_quant/__init__.py

@@ -0,0 +1 @@
+from .comfy_quantizer import ComfyQuantizer

src/diffusers/quantizers/comfy_quant/comfy_quantizer.py

@@ -0,0 +1,121 @@
+from typing import TYPE_CHECKING, Any
+
+from ...utils import get_module_from_name, is_comfy_kitchen_available, is_torch_available, logging
+from ..base import DiffusersQuantizer
+
+
+if is_comfy_kitchen_available():
+    import comfy_kitchen.tensor as ck_tensor
+
+if TYPE_CHECKING:
+    from ...models.modeling_utils import ModelMixin
+
+if is_torch_available():
+    import torch
+
+
+logger = logging.get_logger(__name__)
+
+
+class ComfyQuantizer(DiffusersQuantizer):
+    """
+    Quantizer for comfy-kitchen formats (FP8, INT8, etc.).
+    """
+
+    use_keep_in_fp32_modules = True
+    requires_calibration = False
+    required_packages = ["comfy_kitchen"]
+
+    def __init__(self, quantization_config, **kwargs):
+        super().__init__(quantization_config, **kwargs)
+        self.quant_format = getattr(quantization_config, "quant_format", "fp8")
+        self.compute_dtype = quantization_config.compute_dtype
+        self.modules_to_not_convert = quantization_config.modules_to_not_convert or []
+        if not isinstance(self.modules_to_not_convert, list):
+            self.modules_to_not_convert = [self.modules_to_not_convert]
+
+        if is_comfy_kitchen_available():
+            # Resolve the layout class once since quant_format is constant.
+            layout_map = {
+                "fp8": getattr(ck_tensor, "TensorCoreFP8Layout", None),
+                "nvfp4": getattr(ck_tensor, "TensorCoreNVFP4Layout", None),
+                "mxfp8": getattr(ck_tensor, "TensorCoreMXFP8Layout", None),
+                "int8": getattr(ck_tensor, "TensorWiseINT8Layout", None),
+                "int4_svd": getattr(ck_tensor, "TensorCoreSVDQuantW4A4Layout", None),
+                "int4_awq": getattr(ck_tensor, "TensorCoreAWQW4A16Layout", None),
+            }
+            self.layout = layout_map.get(self.quant_format.lower())
+            if self.layout is None:
+                supported = list(layout_map.keys())
+                raise ValueError(
+                    f"The layout for '{self.quant_format}' was not found in `comfy_kitchen`. "
+                    f"Supported formats are: {supported}."
+                )
+        else:
+            self.layout = None
+
+    def validate_environment(self, *args, **kwargs):
+        if not is_comfy_kitchen_available():
+            raise ImportError(
+                "Loading Comfy Quant weights requires `comfy-kitchen`. "
+                "Please install it with: `pip install comfy-kitchen`."
+            )
+
+    def check_if_quantized_param(
+        self,
+        model: "ModelMixin",
+        param_value: "torch.Tensor",
+        param_name: str,
+        state_dict: dict[str, Any],
+        **kwargs,
+    ) -> bool:
+        # Based on comfy_kitchen, we will likely wrap tensors based on some config layout.
+        # For now, we assume all linear weights that aren't excluded are quantized.
+        # This will be refined based on comfy-kitchen's actual detection logic.
+        if any(m in param_name.split(".") for m in self.modules_to_not_convert):
+            return False
+        return True
+
+    def create_quantized_param(
+        self,
+        model: "ModelMixin",
+        param_value: "torch.Tensor",
+        param_name: str,
+        target_device: "torch.device",
+        state_dict: dict[str, Any] | None = None,
+        unexpected_keys: list[str] | None = None,
+        **kwargs,
+    ):
+        module, tensor_name = get_module_from_name(model, param_name)
+
+        quantized_weight = ck_tensor.QuantizedTensor.from_float(param_value.to(target_device), self.layout.__name__)
+
+        if tensor_name in module._parameters:
+            module._parameters[tensor_name] = quantized_weight.to(target_device)
+        if tensor_name in module._buffers:
+            module._buffers[tensor_name] = quantized_weight.to(target_device)
+
+    def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype":
+        if torch_dtype is None:
+            torch_dtype = self.compute_dtype
+        return torch_dtype
+
+    def _process_model_before_weight_loading(
+        self,
+        model: "ModelMixin",
+        device_map,
+        keep_in_fp32_modules: list[str] = [],
+        **kwargs,
+    ):
+        pass
+
+    def _process_model_after_weight_loading(self, model, **kwargs):
+        pass
+
+    @property
+    def is_serializable(self):
+        return False
+
+    @property
+    def is_trainable(self):
+        return False

src/diffusers/quantizers/quantization_config.py

@@ -51,6 +51,7 @@ class QuantizationMethod(str, Enum):
     MODELOPT = "modelopt"
     AUTOROUND = "auto-round"
     SDNQ = "sdnq"
+    COMFY_QUANT = "comfy_quant"
 
 
 @dataclass
@@ -994,3 +995,37 @@ def __new__(cls, *args, **kwargs):
         from sdnq import SDNQConfig as SDNQLibConfig
 
         return SDNQLibConfig(*args, **kwargs)
+
+
+@dataclass
+class ComfyQuantConfig(QuantizationConfigMixin):
+    """
+    This is a wrapper class about all possible attributes and features that you can play with for a model that has been
+    quantized using comfy-kitchen.
+
+    Args:
+        quant_format (`str`, *optional*, defaults to `"fp8"`):
+            The quantization format. Supported values include `"fp8"`, `"int8"`, `"mxfp8"`, `"nvfp4"`, `"int4_svd"`,
+            and `"int4_awq"`.
+        compute_dtype (`torch.dtype`, *optional*):
+            The target dtype for the compute operations.
+        modules_to_not_convert (`list[str]`, *optional*, defaults to `None`):
+            The list of modules to skip during quantization.
+    """
+
+    def __init__(
+        self,
+        quant_format: str = "fp8",
+        compute_dtype: Any = None,
+        modules_to_not_convert: list[str] | None = None,
+        **kwargs,
+    ):
+        supported_formats = {"fp8", "nvfp4", "mxfp8", "int8", "int4_svd", "int4_awq"}
+        if quant_format not in supported_formats:
+            raise ValueError(
+                f"Unsupported quant_format: {quant_format}. Supported formats are: {sorted(supported_formats)}"
+            )
+        self.quant_method = QuantizationMethod.COMFY_QUANT
+        self.quant_format = quant_format
+        self.compute_dtype = compute_dtype
+        self.modules_to_not_convert = modules_to_not_convert

src/diffusers/utils/__init__.py

@@ -72,6 +72,7 @@
     is_bitsandbytes_available,
     is_bitsandbytes_version,
     is_bs4_available,
+    is_comfy_kitchen_available,
     is_cosmos_guardrail_available,
     is_flash_attn_3_available,
     is_flash_attn_available,

src/diffusers/utils/dummy_comfy_kitchen_objects.py

@@ -0,0 +1,17 @@
+# This file is autogenerated by the command `make fix-copies`, do not edit.
+from ..utils import DummyObject, requires_backends
+
+
+class ComfyQuantConfig(metaclass=DummyObject):
+    _backends = ["comfy_kitchen"]
+
+    def __init__(self, *args, **kwargs):
+        requires_backends(self, ["comfy_kitchen"])
+
+    @classmethod
+    def from_config(cls, *args, **kwargs):
+        requires_backends(cls, ["comfy_kitchen"])
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        requires_backends(cls, ["comfy_kitchen"])