update

Author: DN6

docs/source/en/api/quantization.md

@@ -31,6 +31,11 @@ Learn how to quantize models in the [Quantization](../quantization/overview) gui
 ## GGUFQuantizationConfig
 
 [[autodoc]] GGUFQuantizationConfig
+
+## QuantoConfig
+
+[[autodoc]] QuantoConfig
+
 ## TorchAoConfig
 
 [[autodoc]] TorchAoConfig

docs/source/en/quantization/quanto.md

@@ -33,7 +33,7 @@ import torch
 from diffusers import FluxTransformer2DModel, QuantoConfig
 
 model_id = "black-forest-labs/FLUX.1-dev"
-quantization_config = QuantoConfig(weights="float8")
+quantization_config = QuantoConfig(weights_dtype="float8")
 transformer = FluxTransformer2DModel.from_pretrained(model_id, quantization_config=quantization_config, torch_dtype=torch.bfloat16)
 
 pipe = FluxPipeline.from_pretrained(model_id, transformer=transformer, torch_dtype=torch_dtype)
@@ -55,7 +55,7 @@ import torch
 from diffusers import FluxTransformer2DModel, QuantoConfig
 
 model_id = "black-forest-labs/FLUX.1-dev"
-quantization_config = QuantoConfig(weights="float8", modules_to_not_convert=["proj_out"])
+quantization_config = QuantoConfig(weights_dtype="float8", modules_to_not_convert=["proj_out"])
 transformer = FluxTransformer2DModel.from_pretrained(model_id, quantization_config=quantization_config, torch_dtype=torch.bfloat16)
 ```
 
@@ -66,7 +66,7 @@ import torch
 from diffusers import FluxTransformer2DModel, QuantoConfig
 
 ckpt_path = "https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/flux1-dev.safetensors"
-quantization_config = QuantoConfig(weights="float8")
+quantization_config = QuantoConfig(weights_dtype="float8")
 transformer = FluxTransformer2DModel.from_single_file(ckpt_path, quantization_config=quantization_config, torch_dtype=torch.bfloat16)
 ```
 
@@ -79,7 +79,7 @@ import torch
 from diffusers import FluxTransformer2DModel, QuantoConfig
 
 model_id = "black-forest-labs/FLUX.1-dev"
-quantization_config = QuantoConfig(weights="float8")
+quantization_config = QuantoConfig(weights_dtype="float8")
 transformer = FluxTransformer2DModel.from_pretrained(model_id, quantization_config=quantization_config, torch_dtype=torch.bfloat16)
 
 # save quantized model to reuse
@@ -100,7 +100,7 @@ import torch
 from diffusers import FluxPipeline, FluxTransformer2DModel, QuantoConfig
 
 model_id = "black-forest-labs/FLUX.1-dev"
-quantization_config = QuantoConfig(weights="int8")
+quantization_config = QuantoConfig(weights_dtype="int8")
 transformer = FluxTransformer2DModel.from_pretrained(
     model_id,
     subfolder="transformer",

setup.py

@@ -241,7 +241,7 @@ def run(self):
 
 extras["bitsandbytes"] = deps_list("bitsandbytes", "accelerate")
 extras["gguf"] = deps_list("gguf", "accelerate")
-extras["quanto"] = deps_list("quanto", "accelerate")
+extras["quanto"] = deps_list("optimum_quanto", "accelerate")
 extras["torchao"] = deps_list("torchao", "accelerate")
 
 if os.name == "nt":  # windows

src/diffusers/models/model_loading_utils.py

@@ -259,6 +259,9 @@ def load_model_dict_into_meta(
             ):
                 param = param.to(torch.float32)
                 set_module_kwargs["dtype"] = torch.float32
+            # For quantizers have save weights using torch.float8_e4m3fn
+            elif hf_quantizer is not None and param.dtype == getattr(torch, "float8_e4m3fn", None):
+                pass
             else:
                 param = param.to(dtype)
                 set_module_kwargs["dtype"] = dtype
@@ -306,7 +309,9 @@ def load_model_dict_into_meta(
         elif is_quantized and (
             hf_quantizer.check_if_quantized_param(model, param, param_name, state_dict, param_device=param_device)
         ):
-            hf_quantizer.create_quantized_param(model, param, param_name, param_device, state_dict, unexpected_keys)
+            hf_quantizer.create_quantized_param(
+                model, param, param_name, param_device, state_dict, unexpected_keys, dtype=dtype
+            )
         else:
             set_module_tensor_to_device(model, param_name, param_device, value=param, **set_module_kwargs)
 

src/diffusers/quantizers/quanto/quanto_quantizer.py

@@ -1,5 +1,7 @@
 from typing import TYPE_CHECKING, Any, Dict, List, Union
 
+from diffusers.utils.import_utils import is_optimum_quanto_version
+
 from ...utils import (
     get_module_from_name,
     is_accelerate_available,
@@ -44,6 +46,12 @@ def validate_environment(self, *args, **kwargs):
             raise ImportError(
                 "Loading an optimum-quanto quantized model requires optimum-quanto library (`pip install optimum-quanto`)"
             )
+        if not is_optimum_quanto_version(">=", "0.2.6"):
+            raise ImportError(
+                "Loading an optimum-quanto quantized model requires `optimum-quanto>=0.2.6`. "
+                "Please upgrade your installation with `pip install --upgrade optimum-quanto"
+            )
+
         if not is_accelerate_available():
             raise ImportError(
                 "Loading an optimum-quanto quantized model requires accelerate library (`pip install accelerate`)"

tests/quantization/quanto/test_quanto.py

@@ -1,19 +1,18 @@
-import tempfile
 import gc
+import tempfile
 import unittest
 
-import torch
-
-from diffusers import QuantoConfig, FluxTransformer2DModel, FluxPipeline
-from diffusers.utils import is_torch_available, is_optimum_quanto_available
+from diffusers import FluxPipeline, FluxTransformer2DModel, QuantoConfig
+from diffusers.models.attention_processor import Attention
+from diffusers.utils import is_optimum_quanto_available, is_torch_available
 from diffusers.utils.testing_utils import (
     nightly,
     numpy_cosine_similarity_distance,
     require_accelerate,
     require_big_gpu_with_torch_cuda,
     torch_device,
 )
-from diffusers.models.attention_processor import Attention
+
 
 if is_optimum_quanto_available():
     from optimum.quanto import QLinear
@@ -192,7 +191,11 @@ def test_torch_compile(self):
         with torch.no_grad():
             compiled_model_output = compiled_model(**self.get_dummy_inputs()).sample
 
-        assert torch.allclose(model_output, compiled_model_output, rtol=1e-2, atol=1e-3)
+        model_output = model_output.detach().float().cpu().numpy()
+        compiled_model_output = compiled_model_output.detach().float().cpu().numpy()
+
+        max_diff = numpy_cosine_similarity_distance(model_output.flatten(), compiled_model_output.flatten())
+        assert max_diff < 1e-3
 
 
 class FluxTransformerQuantoMixin(QuantoBaseTesterMixin):
@@ -275,7 +278,7 @@ def test_model_cpu_offload(self):
             "hf-internal-testing/tiny-flux-pipe", transformer=transformer, torch_dtype=torch.bfloat16
         )
         pipe.enable_model_cpu_offload(device=torch_device)
-        images = pipe("a cat holding a sign that says hello", num_inference_steps=2)
+        _ = pipe("a cat holding a sign that says hello", num_inference_steps=2)
 
     def test_training(self):
         quantization_config = QuantoConfig(**self.get_dummy_init_kwargs())
@@ -311,7 +314,6 @@ def test_training(self):
 
 class FluxTransformerFloat8WeightsTest(FluxTransformerQuantoMixin, unittest.TestCase):
     expected_memory_reduction = 0.3
-    _test_torch_compile = True
 
     def get_dummy_init_kwargs(self):
         return {"weights_dtype": "float8"}
@@ -341,7 +343,6 @@ def get_dummy_init_kwargs(self):
 
 class FluxTransformerInt4WeightsTest(FluxTransformerQuantoMixin, unittest.TestCase):
     expected_memory_reduction = 0.55
-    _test_torch_compile = True
 
     def get_dummy_init_kwargs(self):
         return {"weights_dtype": "int4"}