From 2dd589f06a08a9c98904d1592a829ee5ff1d2299 Mon Sep 17 00:00:00 2001 From: PrakshaaleJain Date: Wed, 9 Sep 2026 10:31:53 +0530 Subject: [PATCH 1/9] adding comfy_quant methods --- .../quantizers/comfy_quant/__init__.py | 1 + .../quantizers/comfy_quant/comfy_quantizer.py | 109 ++++++++++++++++++ 2 files changed, 110 insertions(+) create mode 100644 src/diffusers/quantizers/comfy_quant/__init__.py create mode 100644 src/diffusers/quantizers/comfy_quant/comfy_quantizer.py diff --git a/src/diffusers/quantizers/comfy_quant/__init__.py b/src/diffusers/quantizers/comfy_quant/__init__.py new file mode 100644 index 000000000000..bbfe42225cde --- /dev/null +++ b/src/diffusers/quantizers/comfy_quant/__init__.py @@ -0,0 +1 @@ +from .comfy_quantizer import ComfyQuantizer diff --git a/src/diffusers/quantizers/comfy_quant/comfy_quantizer.py b/src/diffusers/quantizers/comfy_quant/comfy_quantizer.py new file mode 100644 index 000000000000..2d94420ab003 --- /dev/null +++ b/src/diffusers/quantizers/comfy_quant/comfy_quantizer.py @@ -0,0 +1,109 @@ +from typing import TYPE_CHECKING, Any + +from ...utils import ( + get_module_from_name, + is_torch_available, + logging, +) +from ..base import DiffusersQuantizer + + +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.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] + + def validate_environment(self, *args, **kwargs): + from ...utils.import_utils import is_comfy_kitchen_available + + 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) + + # Defaulting to an example layout for now. Ideally this is pulled from config or metadata. + # Since ComfyQuantConfig can store the exact layout/format, we'd use it here. + # quantized_weight = ck_tensor.QuantizedTensor.from_float(param_value, ck_tensor.TensorCoreFP8Layout) + + # Since we don't have the exact layout detection in this PR snippet, we do a basic wrap. + # This is a placeholder for the actual comfy-kitchen wrapping logic. + quantized_weight = param_value + + 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 From 78054ff022504644905e0d7bfa7e88b1e7271d44 Mon Sep 17 00:00:00 2001 From: PrakshaaleJain Date: Wed, 9 Sep 2026 11:38:23 +0530 Subject: [PATCH 2/9] feature: add aupport for comfy-kitchen quantization --- src/diffusers/__init__.py | 14 +++++++++ src/diffusers/pipelines/ltx2/__init__.py | 2 +- src/diffusers/quantizers/auto.py | 4 +++ .../quantizers/quantization_config.py | 25 ++++++++++++++++ src/diffusers/utils/__init__.py | 1 + .../utils/dummy_comfy_kitchen_objects.py | 29 +++++++++++++++++++ src/diffusers/utils/import_utils.py | 11 +++++++ 7 files changed, 85 insertions(+), 1 deletion(-) create mode 100644 src/diffusers/utils/dummy_comfy_kitchen_objects.py diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 1fc34e6bdbf6..a46ae675d5c0 100644 --- a/src/diffusers/__init__.py +++ b/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_torch_available() and not is_accelerate_available() and 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() diff --git a/src/diffusers/pipelines/ltx2/__init__.py b/src/diffusers/pipelines/ltx2/__init__.py index d4aa35127403..d48d890f4cb6 100644 --- a/src/diffusers/pipelines/ltx2/__init__.py +++ b/src/diffusers/pipelines/ltx2/__init__.py @@ -30,12 +30,12 @@ _import_structure["pipeline_ltx2_condition"] = ["LTX2ConditionPipeline", "LTX2VideoCondition"] _import_structure["pipeline_ltx2_dfr"] = ["LTX2DFRPipeline"] _import_structure["pipeline_ltx2_dfr_temporal_refine"] = ["LTX2DFRTemporalRefinePipeline"] - _import_structure["pipeline_output"] = ["LTX2DFRPipelineOutput", "LTX2PipelineOutput", "LTX2VideoDecodeOutput"] _import_structure["pipeline_ltx2_diffusion_decode"] = ["LTX2VideoDiffusionDecodePipeline"] _import_structure["pipeline_ltx2_hdr_lora"] = ["LTX2HDRPipeline", "LTX2HDRReferenceCondition"] _import_structure["pipeline_ltx2_ic_lora"] = ["LTX2InContextPipeline", "LTX2ReferenceCondition"] _import_structure["pipeline_ltx2_image2video"] = ["LTX2ImageToVideoPipeline"] _import_structure["pipeline_ltx2_latent_upsample"] = ["LTX2LatentUpsamplePipeline"] + _import_structure["pipeline_output"] = ["LTX2DFRPipelineOutput", "LTX2PipelineOutput", "LTX2VideoDecodeOutput"] _import_structure["vocoder"] = ["LTX2Vocoder", "LTX2VocoderWithBWE"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: diff --git a/src/diffusers/quantizers/auto.py b/src/diffusers/quantizers/auto.py index ea6caf91ab80..4b4358fecf30 100644 --- a/src/diffusers/quantizers/auto.py +++ b/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, diff --git a/src/diffusers/quantizers/quantization_config.py b/src/diffusers/quantizers/quantization_config.py index 30e89f53f906..d22e63819bc1 100644 --- a/src/diffusers/quantizers/quantization_config.py +++ b/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,27 @@ 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: + 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, + compute_dtype: Any = None, + modules_to_not_convert: list[str] | None = None, + **kwargs, + ): + self.quant_method = QuantizationMethod.COMFY_QUANT + self.compute_dtype = compute_dtype + self.modules_to_not_convert = modules_to_not_convert diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py index 5c63a4bc7661..a1ec33bcf3fb 100644 --- a/src/diffusers/utils/__init__.py +++ b/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, diff --git a/src/diffusers/utils/dummy_comfy_kitchen_objects.py b/src/diffusers/utils/dummy_comfy_kitchen_objects.py new file mode 100644 index 000000000000..f6bc79109b83 --- /dev/null +++ b/src/diffusers/utils/dummy_comfy_kitchen_objects.py @@ -0,0 +1,29 @@ +# 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. + +from ..utils import DummyObject, requires_backends + + +class ComfyQuantConfig(metaclass=DummyObject): + _backends = ["comfy_kitchen"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["comfy_kitchen"]) + + +class ComfyQuantizer(metaclass=DummyObject): + _backends = ["comfy_kitchen"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["comfy_kitchen"]) diff --git a/src/diffusers/utils/import_utils.py b/src/diffusers/utils/import_utils.py index d2cf394cd9a7..d29d137251bc 100644 --- a/src/diffusers/utils/import_utils.py +++ b/src/diffusers/utils/import_utils.py @@ -220,6 +220,7 @@ def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> tuple[b _sdnq_available, _sdnq_version = _is_package_available("sdnq") _flashpack_available, _flashpack_version = _is_package_available("flashpack") _av_available, _av_version = _is_package_available("av") +_comfy_kitchen_available, _comfy_kitchen_version = _is_package_available("comfy_kitchen") def is_torch_available(): @@ -422,6 +423,10 @@ def is_av_available(): return _av_available +def is_comfy_kitchen_available(): + return _comfy_kitchen_available + + # docstyle-ignore INFLECT_IMPORT_ERROR = """ {0} requires the inflect library but it was not found in your environment. You can install it with pip: `pip install @@ -562,6 +567,11 @@ def is_av_available(): {0} requires the sdnq library but it was not found in your environment. You can install it with pip: `pip install sdnq` """ +COMFY_KITCHEN_IMPORT_ERROR = """ +{0} requires the comfy-kitchen library but it was not found in your environment. You can install it with pip: `pip +install comfy-kitchen` +""" + # docstyle-ignore PYTORCH_RETINAFACE_IMPORT_ERROR = """ {0} requires the pytorch_retinaface library but it was not found in your environment. You can install it with pip: `pip install pytorch_retinaface` @@ -613,6 +623,7 @@ def is_av_available(): ("pytorch_retinaface", (is_pytorch_retinaface_available, PYTORCH_RETINAFACE_IMPORT_ERROR)), ("better_profanity", (is_better_profanity_available, BETTER_PROFANITY_IMPORT_ERROR)), ("nltk", (is_nltk_available, NLTK_IMPORT_ERROR)), + ("comfy_kitchen", (is_comfy_kitchen_available, COMFY_KITCHEN_IMPORT_ERROR)), ("torch_neuronx", (is_torch_neuronx_available, TORCH_NEURONX_IMPORT_ERROR)), ] ) From eb4811cc1fd8ecac7e1f97b31422b5fcf0848b2a Mon Sep 17 00:00:00 2001 From: PrakshaaleJain Date: Wed, 9 Sep 2026 18:12:25 +0530 Subject: [PATCH 3/9] feat: add support for INT8 and INT4 formats to comfy-kitchen quantizer --- .../quantizers/comfy_quant/comfy_quantizer.py | 34 +++++++++++++++---- .../quantizers/quantization_config.py | 5 +++ 2 files changed, 32 insertions(+), 7 deletions(-) diff --git a/src/diffusers/quantizers/comfy_quant/comfy_quantizer.py b/src/diffusers/quantizers/comfy_quant/comfy_quantizer.py index 2d94420ab003..c76ec13ca79e 100644 --- a/src/diffusers/quantizers/comfy_quant/comfy_quantizer.py +++ b/src/diffusers/quantizers/comfy_quant/comfy_quantizer.py @@ -29,6 +29,7 @@ class ComfyQuantizer(DiffusersQuantizer): 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): @@ -70,13 +71,32 @@ def create_quantized_param( ): module, tensor_name = get_module_from_name(model, param_name) - # Defaulting to an example layout for now. Ideally this is pulled from config or metadata. - # Since ComfyQuantConfig can store the exact layout/format, we'd use it here. - # quantized_weight = ck_tensor.QuantizedTensor.from_float(param_value, ck_tensor.TensorCoreFP8Layout) - - # Since we don't have the exact layout detection in this PR snippet, we do a basic wrap. - # This is a placeholder for the actual comfy-kitchen wrapping logic. - quantized_weight = param_value + import comfy_kitchen.tensor as ck_tensor + + layout_map = { + "fp8": getattr(ck_tensor, "TensorCoreFP8Layout", None), + "nvfp4": getattr(ck_tensor, "TensorCoreNVFP4Layout", None), + "mxfp8": getattr(ck_tensor, "TensorCoreMXFP8Layout", None), + "int8": getattr(ck_tensor, "Int8Layout", None), + "int4_svd": getattr(ck_tensor, "SVDQuantW4A4Layout", None), + "int4_awq": getattr(ck_tensor, "AWQW4A16Layout", None), + } + + # Check if it's already a QuantizedTensor (e.g., if loaded directly from a custom loader) + if hasattr(param_value, "layout") and isinstance(param_value.layout, getattr(ck_tensor, "BaseLayout", type)): + quantized_weight = param_value + else: + layout = layout_map.get(self.quant_format.lower()) + if layout is None: + raise ValueError( + f"The layout for '{self.quant_format}' was not found in `comfy_kitchen`. " + f"Make sure you have the latest version installed that supports this format." + ) + + # comfy-kitchen natively handles wrapping standard float tensors via from_float + # If the tensor is pre-quantized raw bytes, comfy-kitchen exposes `.from_quantized(...)` or similar internally, + # but `.from_float` guarantees we intercept float weights (e.g. standard safetensors float weights). + quantized_weight = ck_tensor.QuantizedTensor.from_float(param_value.to(target_device), layout) if tensor_name in module._parameters: module._parameters[tensor_name] = quantized_weight.to(target_device) diff --git a/src/diffusers/quantizers/quantization_config.py b/src/diffusers/quantizers/quantization_config.py index d22e63819bc1..6cc1c92b45f5 100644 --- a/src/diffusers/quantizers/quantization_config.py +++ b/src/diffusers/quantizers/quantization_config.py @@ -1004,6 +1004,9 @@ class ComfyQuantConfig(QuantizationConfigMixin): 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`): @@ -1012,10 +1015,12 @@ class ComfyQuantConfig(QuantizationConfigMixin): def __init__( self, + quant_format: str = "fp8", compute_dtype: Any = None, modules_to_not_convert: list[str] | None = None, **kwargs, ): 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 From b185a13a172d8788134c865c74162464de5912aa Mon Sep 17 00:00:00 2001 From: PrakshaaleJain Date: Wed, 9 Sep 2026 18:40:38 +0530 Subject: [PATCH 4/9] fix: register comfy-kitchen dummy objects correctly for check_dummies.py --- src/diffusers/__init__.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index a46ae675d5c0..7581614e6c7e 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -1060,6 +1060,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() From a8b2e12096c6276b50adf7a0df15cbcb60c9d9dd Mon Sep 17 00:00:00 2001 From: PrakshaaleJain Date: Wed, 9 Sep 2026 18:55:48 +0530 Subject: [PATCH 5/9] docfix: deleted suggestion for installing latest version of comfy-kitchen --- src/diffusers/__init__.py | 4 ++-- src/diffusers/quantizers/comfy_quant/comfy_quantizer.py | 3 +-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 7581614e6c7e..82d801d09eda 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -1066,8 +1066,8 @@ except OptionalDependencyNotAvailable: from .utils.dummy_comfy_kitchen_objects import * else: - from .quantizers.quantization_config import ComfyQuantConfig - + from .quantizers.quantization_config import Com + try: if not is_onnx_available(): raise OptionalDependencyNotAvailable() diff --git a/src/diffusers/quantizers/comfy_quant/comfy_quantizer.py b/src/diffusers/quantizers/comfy_quant/comfy_quantizer.py index c76ec13ca79e..d7659d52f79b 100644 --- a/src/diffusers/quantizers/comfy_quant/comfy_quantizer.py +++ b/src/diffusers/quantizers/comfy_quant/comfy_quantizer.py @@ -89,8 +89,7 @@ def create_quantized_param( layout = layout_map.get(self.quant_format.lower()) if layout is None: raise ValueError( - f"The layout for '{self.quant_format}' was not found in `comfy_kitchen`. " - f"Make sure you have the latest version installed that supports this format." + f"The layout for '{self.quant_format}' was not found in `comfy_kitchen`." ) # comfy-kitchen natively handles wrapping standard float tensors via from_float From 03178749acd1e57bb12b205e6ffce998de8587c1 Mon Sep 17 00:00:00 2001 From: PrakshaaleJain Date: Wed, 9 Sep 2026 19:13:07 +0530 Subject: [PATCH 6/9] test: add comfy-quant unit tests and fix layout passing --- src/diffusers/__init__.py | 2 +- .../quantizers/comfy_quant/comfy_quantizer.py | 8 +-- .../utils/dummy_comfy_kitchen_objects.py | 28 +++------ tests/quantization/comfy_quant/__init__.py | 0 .../comfy_quant/test_comfy_quant.py | 62 +++++++++++++++++++ 5 files changed, 74 insertions(+), 26 deletions(-) create mode 100644 tests/quantization/comfy_quant/__init__.py create mode 100644 tests/quantization/comfy_quant/test_comfy_quant.py diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 82d801d09eda..460edbeb325f 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -1067,7 +1067,7 @@ from .utils.dummy_comfy_kitchen_objects import * else: from .quantizers.quantization_config import Com - + try: if not is_onnx_available(): raise OptionalDependencyNotAvailable() diff --git a/src/diffusers/quantizers/comfy_quant/comfy_quantizer.py b/src/diffusers/quantizers/comfy_quant/comfy_quantizer.py index d7659d52f79b..622eb82550fe 100644 --- a/src/diffusers/quantizers/comfy_quant/comfy_quantizer.py +++ b/src/diffusers/quantizers/comfy_quant/comfy_quantizer.py @@ -83,19 +83,17 @@ def create_quantized_param( } # Check if it's already a QuantizedTensor (e.g., if loaded directly from a custom loader) - if hasattr(param_value, "layout") and isinstance(param_value.layout, getattr(ck_tensor, "BaseLayout", type)): + if isinstance(param_value, ck_tensor.QuantizedTensor): quantized_weight = param_value else: layout = layout_map.get(self.quant_format.lower()) if layout is None: - raise ValueError( - f"The layout for '{self.quant_format}' was not found in `comfy_kitchen`." - ) + raise ValueError(f"The layout for '{self.quant_format}' was not found in `comfy_kitchen`.") # comfy-kitchen natively handles wrapping standard float tensors via from_float # If the tensor is pre-quantized raw bytes, comfy-kitchen exposes `.from_quantized(...)` or similar internally, # but `.from_float` guarantees we intercept float weights (e.g. standard safetensors float weights). - quantized_weight = ck_tensor.QuantizedTensor.from_float(param_value.to(target_device), layout) + quantized_weight = ck_tensor.QuantizedTensor.from_float(param_value.to(target_device), layout.__name__) if tensor_name in module._parameters: module._parameters[tensor_name] = quantized_weight.to(target_device) diff --git a/src/diffusers/utils/dummy_comfy_kitchen_objects.py b/src/diffusers/utils/dummy_comfy_kitchen_objects.py index f6bc79109b83..cda3dedbb573 100644 --- a/src/diffusers/utils/dummy_comfy_kitchen_objects.py +++ b/src/diffusers/utils/dummy_comfy_kitchen_objects.py @@ -1,29 +1,17 @@ -# 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. - +# This file is autogenerated by the command `make fix-copies`, do not edit. from ..utils import DummyObject, requires_backends -class ComfyQuantConfig(metaclass=DummyObject): +class Com(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"]) -class ComfyQuantizer(metaclass=DummyObject): - _backends = ["comfy_kitchen"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["comfy_kitchen"]) + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["comfy_kitchen"]) diff --git a/tests/quantization/comfy_quant/__init__.py b/tests/quantization/comfy_quant/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/quantization/comfy_quant/test_comfy_quant.py b/tests/quantization/comfy_quant/test_comfy_quant.py new file mode 100644 index 000000000000..36f58f738ee2 --- /dev/null +++ b/tests/quantization/comfy_quant/test_comfy_quant.py @@ -0,0 +1,62 @@ +import pytest +import torch +import torch.nn as nn + +from diffusers import ComfyQuantConfig +from diffusers.utils import is_comfy_kitchen_available + +from ...testing_utils import require_torch + + +if is_comfy_kitchen_available(): + import comfy_kitchen.tensor as ck_tensor + + from diffusers.quantizers.comfy_quant.comfy_quantizer import ComfyQuantizer + +device = "cuda" if torch.cuda.is_available() else "cpu" + + +@require_torch +@pytest.mark.skipif(not is_comfy_kitchen_available(), reason="comfy-kitchen is not available") +class TestComfyQuantizer: + def test_create_quantized_param_fp8(self): + config = ComfyQuantConfig(quant_format="fp8") + quantizer = ComfyQuantizer(config) + + class DummyModel(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(16, 16) + + model = DummyModel() + param_value = torch.randn(16, 16, dtype=torch.float32) + + quantizer.create_quantized_param( + model=model, + param_value=param_value, + param_name="linear.weight", + target_device=torch.device(device), + ) + + assert isinstance(model.linear.weight, ck_tensor.QuantizedTensor) + assert model.linear.weight._layout_cls == "TensorCoreFP8Layout" + + def test_create_quantized_param_invalid_format(self): + config = ComfyQuantConfig(quant_format="non_existent_format") + quantizer = ComfyQuantizer(config) + + class DummyModel(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(16, 16) + + model = DummyModel() + param_value = torch.randn(16, 16, dtype=torch.float32) + + with pytest.raises(ValueError, match="not found in `comfy_kitchen`"): + quantizer.create_quantized_param( + model=model, + param_value=param_value, + param_name="linear.weight", + target_device=torch.device(device), + ) From 0c97d8dd090f369207a87fc89b6668df04fc5672 Mon Sep 17 00:00:00 2001 From: PrakshaaleJain Date: Wed, 9 Sep 2026 21:11:02 +0530 Subject: [PATCH 7/9] test: refine invalid format test --- tests/quantization/comfy_quant/test_comfy_quant.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/quantization/comfy_quant/test_comfy_quant.py b/tests/quantization/comfy_quant/test_comfy_quant.py index 36f58f738ee2..a40835e90673 100644 --- a/tests/quantization/comfy_quant/test_comfy_quant.py +++ b/tests/quantization/comfy_quant/test_comfy_quant.py @@ -42,7 +42,7 @@ def __init__(self): assert model.linear.weight._layout_cls == "TensorCoreFP8Layout" def test_create_quantized_param_invalid_format(self): - config = ComfyQuantConfig(quant_format="non_existent_format") + config = ComfyQuantConfig(quant_format="INT_42") #INT_42 is an non-existent format quantizer = ComfyQuantizer(config) class DummyModel(nn.Module): From e38bf7faf85b7dfd4de23f1c115d38be446d6418 Mon Sep 17 00:00:00 2001 From: PrakshaaleJain Date: Wed, 9 Sep 2026 21:13:51 +0530 Subject: [PATCH 8/9] style fixes --- tests/quantization/comfy_quant/test_comfy_quant.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/quantization/comfy_quant/test_comfy_quant.py b/tests/quantization/comfy_quant/test_comfy_quant.py index a40835e90673..af94ab99018a 100644 --- a/tests/quantization/comfy_quant/test_comfy_quant.py +++ b/tests/quantization/comfy_quant/test_comfy_quant.py @@ -42,7 +42,7 @@ def __init__(self): assert model.linear.weight._layout_cls == "TensorCoreFP8Layout" def test_create_quantized_param_invalid_format(self): - config = ComfyQuantConfig(quant_format="INT_42") #INT_42 is an non-existent format + config = ComfyQuantConfig(quant_format="INT_42") # INT_42 is an non-existent format quantizer = ComfyQuantizer(config) class DummyModel(nn.Module): From 2452aa8628b516a44925282c61c9ee37ebdf97c7 Mon Sep 17 00:00:00 2001 From: PrakshaaleJain Date: Wed, 9 Sep 2026 21:21:18 +0530 Subject: [PATCH 9/9] docs: add comfy_quant usage tutorial --- docs/source/en/_toctree.yml | 2 + docs/source/en/quantization/comfy_quant.md | 77 ++++++++++++++++++++++ 2 files changed, 79 insertions(+) create mode 100644 docs/source/en/quantization/comfy_quant.md diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index f05667986f11..04aa4b355a19 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -172,6 +172,8 @@ title: bitsandbytes - local: quantization/gguf title: gguf + - local: quantization/comfy_quant + title: Comfy Quant - local: quantization/nunchaku title: Nunchaku Lite - local: quantization/torchao diff --git a/docs/source/en/quantization/comfy_quant.md b/docs/source/en/quantization/comfy_quant.md new file mode 100644 index 000000000000..4e96bf955396 --- /dev/null +++ b/docs/source/en/quantization/comfy_quant.md @@ -0,0 +1,77 @@ + + +# 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 `TensorCoreInt8Layout` (W8A8, tensorwise) +- **MXFP8** (`mxfp8`): Maps to `TensorCoreMXFP8Layout` +- **NVFP4** (`nvfp4`): Maps to `TensorCoreNVFP4Layout` +- **INT4 SVD** (`int4_svd`): Maps to `SVDQuantW4A4Layout` (SVDQuant W4A4) +- **INT4 AWQ** (`int4_awq`): Maps to `AWQW4A16Layout` (AWQ W4A16) + +When using optimized layouts, `comfy-kitchen` automatically dispatches the operations to the best available backend (HIP, CUDA, Triton, or Eager).