diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 525734a2e4cb..aa4e7521540c 100644 --- a/docs/source/en/_toctree.yml +++ b/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 diff --git a/docs/source/en/quantization/comfy_quant.md b/docs/source/en/quantization/comfy_quant.md new file mode 100644 index 000000000000..19c6da128437 --- /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 `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). diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 2825e9888c98..7424d4d571d1 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_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() 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/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..3420e65cfa01 --- /dev/null +++ b/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 diff --git a/src/diffusers/quantizers/quantization_config.py b/src/diffusers/quantizers/quantization_config.py index 30e89f53f906..d6bd11ad98d7 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,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 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..e00c405c2bb2 --- /dev/null +++ b/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"]) 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)), ] ) diff --git a/tests/models/testing_utils/__init__.py b/tests/models/testing_utils/__init__.py index 2d7d5ae23257..8673533508bb 100644 --- a/tests/models/testing_utils/__init__.py +++ b/tests/models/testing_utils/__init__.py @@ -31,6 +31,8 @@ BitsAndBytesCompileTesterMixin, BitsAndBytesConfigMixin, BitsAndBytesTesterMixin, + ComfyQuantConfigMixin, + ComfyQuantTesterMixin, GGUFCompileTesterMixin, GGUFConfigMixin, GGUFTesterMixin, @@ -86,6 +88,8 @@ "ModelOptCompileTesterMixin", "ModelOptConfigMixin", "ModelOptTesterMixin", + "ComfyQuantConfigMixin", + "ComfyQuantTesterMixin", "ModelTesterMixin", "NunchakuLiteCompileTesterMixin", "NunchakuLiteConfigMixin", diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py index 65f904bd4948..5b99b10d8262 100644 --- a/tests/models/testing_utils/quantization.py +++ b/tests/models/testing_utils/quantization.py @@ -43,6 +43,7 @@ backend_empty_cache, is_autoround, is_bitsandbytes, + is_comfy_kitchen, is_gguf, is_modelopt, is_quantization, @@ -53,6 +54,7 @@ require_accelerator, require_auto_round_version_greater_or_equal, require_bitsandbytes_version_greater, + require_comfy_kitchen, require_gguf_version_greater_or_equal, require_modelopt_version_greater_or_equal, require_sdnq, @@ -1823,3 +1825,133 @@ def test_autoround_torch_compile(self): def test_autoround_torch_compile_with_group_offload(self): self._test_torch_compile_with_group_offload(self.config_dict) + + +@is_quantization +@is_comfy_kitchen +@require_comfy_kitchen +@require_accelerator +@require_accelerate +class ComfyQuantConfigMixin: + """ + Base mixin providing ComfyQuant quantization config and model creation. + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained + """ + + COMFY_QUANT_CONFIGS = { + "fp8": {"quant_format": "fp8"}, + "nvfp4": {"quant_format": "nvfp4"}, + "mxfp8": {"quant_format": "mxfp8"}, + "int8": {"quant_format": "int8"}, + "int4_svd": {"quant_format": "int4_svd"}, + "int4_awq": {"quant_format": "int4_awq"}, + } + + COMFY_QUANT_EXPECTED_MEMORY_REDUCTIONS = { + "fp8": 1.2, + "nvfp4": 1.2, + "mxfp8": 1.2, + "int8": 1.2, + "int4_svd": 1.2, + "int4_awq": 1.2, + } + + def _create_quantized_model(self, config_kwargs, **extra_kwargs): + from diffusers.quantizers.quantization_config import ComfyQuantConfig + + config = ComfyQuantConfig(**config_kwargs) + kwargs = getattr(self, "pretrained_model_kwargs", {}).copy() + kwargs["quantization_config"] = config + kwargs["device_map"] = str(torch_device) + kwargs.update(extra_kwargs) + return self.model_class.from_pretrained(self.pretrained_model_name_or_path, **kwargs) + + def _verify_if_layer_quantized(self, name, module, config_kwargs): + import comfy_kitchen.tensor as ck_tensor + + assert isinstance(module.weight, ck_tensor.QuantizedTensor), f"Layer {name} is not a ck_tensor.QuantizedTensor" + + +@is_comfy_kitchen +@require_comfy_kitchen +@require_accelerate +@require_accelerator +class ComfyQuantTesterMixin(ComfyQuantConfigMixin, QuantizationTesterMixin): + """ + Mixin class for testing ComfyQuant quantization on models. + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained (e.g., {"subfolder": "transformer"}) + + Expected methods to be implemented by subclasses: + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Optional class attributes: + - COMFY_QUANT_CONFIGS: Dict of config name -> ComfyQuantConfig kwargs to test + + Pytest mark: comfy_kitchen + Use `pytest -m "not comfy_kitchen"` to skip these tests + """ + + @pytest.mark.parametrize("config_name", ["fp8"], ids=["fp8"]) + def test_comfy_quant_quantization_num_parameters(self, config_name): + self._test_quantization_num_parameters(ComfyQuantConfigMixin.COMFY_QUANT_CONFIGS[config_name]) + + @pytest.mark.parametrize( + "config_name", + ["fp8", "int8", "int4_svd"], + ids=["fp8", "int8", "int4_svd"], + ) + def test_comfy_quant_quantization_memory_footprint(self, config_name): + expected = ComfyQuantConfigMixin.COMFY_QUANT_EXPECTED_MEMORY_REDUCTIONS.get(config_name, 1.2) + self._test_quantization_memory_footprint( + ComfyQuantConfigMixin.COMFY_QUANT_CONFIGS[config_name], expected_memory_reduction=expected + ) + + @pytest.mark.parametrize( + "config_name", + ["fp8", "int8", "int4_svd"], + ids=["fp8", "int8", "int4_svd"], + ) + def test_comfy_quant_quantization_inference(self, config_name): + self._test_quantization_inference(ComfyQuantConfigMixin.COMFY_QUANT_CONFIGS[config_name]) + + @pytest.mark.parametrize("config_name", ["fp8"], ids=["fp8"]) + def test_comfy_quant_quantization_dtype_assignment(self, config_name): + self._test_quantization_dtype_assignment(ComfyQuantConfigMixin.COMFY_QUANT_CONFIGS[config_name]) + + @pytest.mark.parametrize("config_name", ["fp8"], ids=["fp8"]) + def test_comfy_quant_quantization_lora_inference(self, config_name): + self._test_quantization_lora_inference(ComfyQuantConfigMixin.COMFY_QUANT_CONFIGS[config_name]) + + @pytest.mark.parametrize("config_name", ["fp8"], ids=["fp8"]) + def test_comfy_quant_quantization_serialization(self, config_name, tmp_path): + self._test_quantization_serialization(ComfyQuantConfigMixin.COMFY_QUANT_CONFIGS[config_name], tmp_path) + + @pytest.mark.parametrize("config_name", ["fp8"], ids=["fp8"]) + def test_comfy_quant_quantized_layers(self, config_name): + self._test_quantized_layers(ComfyQuantConfigMixin.COMFY_QUANT_CONFIGS[config_name]) + + def test_comfy_quant_modules_to_not_convert(self): + """Test that modules_to_not_convert parameter works correctly.""" + modules_to_exclude = getattr(self, "modules_to_not_convert_for_test", None) + if modules_to_exclude is None: + pytest.skip("modules_to_not_convert_for_test not defined for this model") + + self._test_quantization_modules_to_not_convert( + ComfyQuantConfigMixin.COMFY_QUANT_CONFIGS["fp8"], modules_to_exclude + ) + + def test_comfy_quant_device_map(self): + """Test that device_map='auto' works correctly with quantization.""" + self._test_quantization_device_map(ComfyQuantConfigMixin.COMFY_QUANT_CONFIGS["fp8"]) + + def test_comfy_quant_dequantize(self): + """Test that dequantize() works correctly.""" + self._test_dequantize(ComfyQuantConfigMixin.COMFY_QUANT_CONFIGS["fp8"]) 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..c74a8090c137 --- /dev/null +++ b/tests/quantization/comfy_quant/test_comfy_quant.py @@ -0,0 +1,46 @@ +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): + with pytest.raises(ValueError, match="Unsupported quant_format"): + ComfyQuantConfig(quant_format="INT_42") # INT_42 is an non-existent format diff --git a/tests/testing_utils.py b/tests/testing_utils.py index 9da89e626198..20cd37f1273a 100644 --- a/tests/testing_utils.py +++ b/tests/testing_utils.py @@ -34,6 +34,7 @@ is_accelerate_available, is_auto_round_available, is_bitsandbytes_available, + is_comfy_kitchen_available, is_compel_available, is_flashpack_available, is_gguf_available, @@ -492,6 +493,14 @@ def is_cache(test_case): return pytest.mark.cache(test_case) +def is_comfy_kitchen(test_case): + """ + Decorator marking a test as a Comfy Kitchen test. These tests can be filtered using: + pytest -m "not comfy_kitchen" to skip pytest -m comfy_kitchen to run only these tests + """ + return pytest.mark.comfy_kitchen(test_case) + + def require_torch(test_case): """ Decorator marking a test that requires PyTorch. These tests are skipped when PyTorch isn't installed. @@ -508,6 +517,18 @@ def require_torch_2(test_case): )(test_case) +def require_comfy_kitchen(test_case): + """ + Decorator marking a test that requires comfy_kitchen. + """ + import pytest + + if not is_comfy_kitchen_available(): + return pytest.mark.skipif(not is_comfy_kitchen_available(), reason="test requires comfy-kitchen")(test_case) + + return test_case + + def require_torch_version_greater_equal(torch_version): """Decorator marking a test that requires torch with a specific version or greater."""