diffusers/tests/pipelines/testing_utils/quantization.py

# coding=utf-8
# Copyright 2026 The HuggingFace Team Inc.
#
# 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.
"""Pipeline-level quantization tests.

Model-level quantization tests live in `tests/models/testing_utils/quantization.py` and are wired
into the individual model test files. Backend-level tests (config validation, loading error paths,
kernels) live in `tests/quantization/`. This module only covers behavior that needs a pipeline.

The module name intentionally does not match pytest's `test_*.py` discovery pattern: these tests
only run when this file is passed to pytest explicitly, as the nightly quantization CI jobs do.
Every class is marked with the `quantization` marker plus its backend marker (`bitsandbytes`,
`torchao`, `gguf`, `modelopt`) so CI can select per-backend subsets with `pytest -m`.
"""

import gc
import json
import tempfile

import numpy as np
import pytest
import torch
from huggingface_hub import hf_hub_download
from parameterized import parameterized
from PIL import Image

from diffusers import (
    AuraFlowPipeline,
    AuraFlowTransformer2DModel,
    AutoencoderKL,
    BitsAndBytesConfig,
    DiffusionPipeline,
    FlowMatchEulerDiscreteScheduler,
    FluxControlPipeline,
    FluxPipeline,
    FluxTransformer2DModel,
    GGUFQuantizationConfig,
    NVIDIAModelOptConfig,
    SD3Transformer2DModel,
    StableDiffusion3Pipeline,
    TorchAoConfig,
)
from diffusers.quantizers import PipelineQuantizationConfig
from diffusers.utils import is_accelerate_version, load_image, logging

from ...testing_utils import (
    CaptureLogger,
    Expectations,
    backend_empty_cache,
    backend_reset_peak_memory_stats,
    backend_synchronize,
    enable_full_determinism,
    is_bitsandbytes,
    is_bitsandbytes_available,
    is_gguf,
    is_gguf_available,
    is_modelopt,
    is_quantization,
    is_torchao,
    is_torchao_available,
    is_transformers_available,
    nightly,
    numpy_cosine_similarity_distance,
    require_accelerate,
    require_big_accelerator,
    require_bitsandbytes_version_greater,
    require_gguf_version_greater_or_equal,
    require_modelopt_version_greater_or_equal,
    require_peft_backend,
    require_peft_version_greater,
    require_torch,
    require_torch_accelerator,
    require_torch_version_greater,
    require_torch_version_greater_equal,
    require_torchao_version_greater_or_equal,
    require_transformers_version_greater,
    slow,
    torch_device,
)


if is_transformers_available():
    from transformers import AutoTokenizer, CLIPTextModel, CLIPTokenizer, T5EncoderModel
    from transformers import BitsAndBytesConfig as TranBitsAndBytesConfig
else:
    TranBitsAndBytesConfig = None

if is_bitsandbytes_available():
    pass

if is_torchao_available():
    from torchao.quantization import (
        Float8WeightOnlyConfig,
        Int4WeightOnlyConfig,
        Int8DynamicActivationInt8WeightConfig,
        Int8DynamicActivationIntxWeightConfig,
        Int8Tensor,
        Int8WeightOnlyConfig,
        IntxWeightOnlyConfig,
    )
    from torchao.utils import TorchAOBaseTensor

if is_gguf_available():
    pass


enable_full_determinism()


# ======================== Shared compile base ========================


@is_quantization
@require_torch_accelerator
@slow
class QuantCompileTests:
    @property
    def quantization_config(self):
        raise NotImplementedError(
            "This property should be implemented in the subclass to return the appropriate quantization config."
        )

    @pytest.fixture(autouse=True)
    def _cleanup(self):
        gc.collect()
        backend_empty_cache(torch_device)
        torch.compiler.reset()
        yield
        gc.collect()
        backend_empty_cache(torch_device)
        torch.compiler.reset()

    def _init_pipeline(self, quantization_config, torch_dtype):
        pipe = DiffusionPipeline.from_pretrained(
            "stabilityai/stable-diffusion-3-medium-diffusers",
            quantization_config=quantization_config,
            torch_dtype=torch_dtype,
        )
        return pipe

    def _test_torch_compile_with_cpu_offload(self, torch_dtype=torch.bfloat16):
        pipe = self._init_pipeline(self.quantization_config, torch_dtype)
        pipe.enable_model_cpu_offload()
        # regional compilation is better for offloading.
        # see: https://pytorch.org/blog/torch-compile-and-diffusers-a-hands-on-guide-to-peak-performance/
        if getattr(pipe.transformer, "_repeated_blocks"):
            pipe.transformer.compile_repeated_blocks(fullgraph=True)
        else:
            pipe.transformer.compile()

        # small resolutions to ensure speedy execution.
        pipe("a dog", num_inference_steps=2, max_sequence_length=16, height=256, width=256)

    def test_torch_compile_with_cpu_offload(self):
        self._test_torch_compile_with_cpu_offload()


# ======================== PipelineQuantizationConfig ========================


@is_quantization
@require_bitsandbytes_version_greater("0.43.2")
@require_torchao_version_greater_or_equal("0.16.0")
@require_accelerate
@require_torch
@require_torch_accelerator
@slow
class TestPipelineQuantization:
    model_name = "hf-internal-testing/tiny-flux-pipe"
    prompt = "a beautiful sunset amidst the mountains."
    num_inference_steps = 10
    seed = 0

    def test_quant_config_set_correctly_through_kwargs(self):
        components_to_quantize = ["transformer", "text_encoder_2"]
        quant_config = PipelineQuantizationConfig(
            quant_backend="bitsandbytes_4bit",
            quant_kwargs={
                "load_in_4bit": True,
                "bnb_4bit_quant_type": "nf4",
                "bnb_4bit_compute_dtype": torch.bfloat16,
            },
            components_to_quantize=components_to_quantize,
        )
        pipe = DiffusionPipeline.from_pretrained(
            self.model_name,
            quantization_config=quant_config,
            torch_dtype=torch.bfloat16,
        ).to(torch_device)
        for name, component in pipe.components.items():
            if name in components_to_quantize:
                assert getattr(component.config, "quantization_config", None) is not None
                quantization_config = component.config.quantization_config
                assert quantization_config.load_in_4bit
                assert quantization_config.quant_method == "bitsandbytes"

        _ = pipe(self.prompt, num_inference_steps=self.num_inference_steps)

    def test_quant_config_set_correctly_through_granular(self):
        quant_config = PipelineQuantizationConfig(
            quant_mapping={
                "transformer": TorchAoConfig(Int8WeightOnlyConfig(version=2)),
                "text_encoder_2": TranBitsAndBytesConfig(load_in_4bit=True, compute_dtype=torch.bfloat16),
            }
        )
        components_to_quantize = list(quant_config.quant_mapping.keys())
        pipe = DiffusionPipeline.from_pretrained(
            self.model_name,
            quantization_config=quant_config,
            torch_dtype=torch.bfloat16,
        ).to(torch_device)
        for name, component in pipe.components.items():
            if name in components_to_quantize:
                assert getattr(component.config, "quantization_config", None) is not None
                quantization_config = component.config.quantization_config

                if name == "text_encoder_2":
                    assert quantization_config.load_in_4bit
                    assert quantization_config.quant_method == "bitsandbytes"
                else:
                    assert quantization_config.quant_method == "torchao"

        _ = pipe(self.prompt, num_inference_steps=self.num_inference_steps)

    def test_raises_error_for_invalid_config(self):
        with pytest.raises(ValueError) as err_context:
            _ = PipelineQuantizationConfig(
                quant_mapping={
                    "transformer": TorchAoConfig(Int8WeightOnlyConfig(version=2)),
                    "text_encoder_2": TranBitsAndBytesConfig(load_in_4bit=True, compute_dtype=torch.bfloat16),
                },
                quant_backend="bitsandbytes_4bit",
            )

        assert (
            str(err_context.value) == "Both `quant_backend` and `quant_mapping` cannot be specified at the same time."
        )

    def test_validation_for_kwargs(self):
        components_to_quantize = ["transformer", "text_encoder_2"]
        with pytest.raises(ValueError) as err_context:
            _ = PipelineQuantizationConfig(
                quant_backend="torchao",
                quant_kwargs={"quant_type": Int8WeightOnlyConfig(version=2)},
                components_to_quantize=components_to_quantize,
            )

        assert "The signatures of the __init__ methods of the quantization config classes" in str(err_context.value)

    def test_raises_error_for_wrong_config_class(self):
        quant_config = {
            "transformer": TorchAoConfig(Int8WeightOnlyConfig(version=2)),
            "text_encoder_2": TranBitsAndBytesConfig(load_in_4bit=True, compute_dtype=torch.bfloat16),
        }
        with pytest.raises(ValueError) as err_context:
            _ = DiffusionPipeline.from_pretrained(
                self.model_name,
                quantization_config=quant_config,
                torch_dtype=torch.bfloat16,
            )
        assert str(err_context.value) == "`quantization_config` must be an instance of `PipelineQuantizationConfig`."

    def test_validation_for_mapping(self):
        with pytest.raises(ValueError) as err_context:
            _ = PipelineQuantizationConfig(
                quant_mapping={
                    "transformer": DiffusionPipeline(),
                    "text_encoder_2": TranBitsAndBytesConfig(load_in_4bit=True, compute_dtype=torch.bfloat16),
                }
            )

        assert "Provided config for module_name=transformer could not be found" in str(err_context.value)

    def test_saving_loading(self):
        quant_config = PipelineQuantizationConfig(
            quant_mapping={
                "transformer": TorchAoConfig(Int8WeightOnlyConfig(version=2)),
                "text_encoder_2": TranBitsAndBytesConfig(load_in_4bit=True, compute_dtype=torch.bfloat16),
            }
        )
        components_to_quantize = list(quant_config.quant_mapping.keys())
        pipe = DiffusionPipeline.from_pretrained(
            self.model_name,
            quantization_config=quant_config,
            torch_dtype=torch.bfloat16,
        ).to(torch_device)

        pipe_inputs = {"prompt": self.prompt, "num_inference_steps": self.num_inference_steps, "output_type": "latent"}
        output_1 = pipe(**pipe_inputs, generator=torch.manual_seed(self.seed)).images

        with tempfile.TemporaryDirectory() as tmpdir:
            pipe.save_pretrained(tmpdir)
            loaded_pipe = DiffusionPipeline.from_pretrained(tmpdir, torch_dtype=torch.bfloat16).to(torch_device)
        for name, component in loaded_pipe.components.items():
            if name in components_to_quantize:
                assert getattr(component.config, "quantization_config", None) is not None
                quantization_config = component.config.quantization_config

                if name == "text_encoder_2":
                    assert quantization_config.load_in_4bit
                    assert quantization_config.quant_method == "bitsandbytes"
                else:
                    assert quantization_config.quant_method == "torchao"

        output_2 = loaded_pipe(**pipe_inputs, generator=torch.manual_seed(self.seed)).images

        assert torch.allclose(output_1, output_2)

    @parameterized.expand(["quant_kwargs", "quant_mapping"])
    def test_warn_invalid_component(self, method):
        invalid_component = "foo"
        if method == "quant_kwargs":
            components_to_quantize = ["transformer", invalid_component]
            quant_config = PipelineQuantizationConfig(
                quant_backend="bitsandbytes_8bit",
                quant_kwargs={"load_in_8bit": True},
                components_to_quantize=components_to_quantize,
            )
        else:
            quant_config = PipelineQuantizationConfig(
                quant_mapping={
                    "transformer": TorchAoConfig(Int8WeightOnlyConfig(version=2)),
                    invalid_component: TranBitsAndBytesConfig(load_in_8bit=True),
                }
            )

        logger = logging.get_logger("diffusers.pipelines.pipeline_loading_utils")
        logger.setLevel(logging.WARNING)
        with CaptureLogger(logger) as cap_logger:
            _ = DiffusionPipeline.from_pretrained(
                self.model_name,
                quantization_config=quant_config,
                torch_dtype=torch.bfloat16,
            )
        assert invalid_component in cap_logger.out

    @parameterized.expand(["quant_kwargs", "quant_mapping"])
    def test_no_quantization_for_all_invalid_components(self, method):
        invalid_component = "foo"
        if method == "quant_kwargs":
            components_to_quantize = [invalid_component]
            quant_config = PipelineQuantizationConfig(
                quant_backend="bitsandbytes_8bit",
                quant_kwargs={"load_in_8bit": True},
                components_to_quantize=components_to_quantize,
            )
        else:
            quant_config = PipelineQuantizationConfig(
                quant_mapping={invalid_component: TranBitsAndBytesConfig(load_in_8bit=True)}
            )

        pipe = DiffusionPipeline.from_pretrained(
            self.model_name,
            quantization_config=quant_config,
            torch_dtype=torch.bfloat16,
        )
        for name, component in pipe.components.items():
            if isinstance(component, torch.nn.Module):
                assert not hasattr(component.config, "quantization_config")

    @parameterized.expand(["quant_kwargs", "quant_mapping"])
    def test_quant_config_repr(self, method):
        component_name = "transformer"
        if method == "quant_kwargs":
            components_to_quantize = [component_name]
            quant_config = PipelineQuantizationConfig(
                quant_backend="bitsandbytes_8bit",
                quant_kwargs={"load_in_8bit": True},
                components_to_quantize=components_to_quantize,
            )
        else:
            quant_config = PipelineQuantizationConfig(
                quant_mapping={component_name: BitsAndBytesConfig(load_in_8bit=True)}
            )

        pipe = DiffusionPipeline.from_pretrained(
            self.model_name,
            quantization_config=quant_config,
            torch_dtype=torch.bfloat16,
        )
        assert getattr(pipe, "quantization_config", None) is not None
        retrieved_config = pipe.quantization_config
        expected_config = """
transformer BitsAndBytesConfig {
  "_load_in_4bit": false,
  "_load_in_8bit": true,
  "bnb_4bit_compute_dtype": "float32",
  "bnb_4bit_quant_storage": "uint8",
  "bnb_4bit_quant_type": "fp4",
  "bnb_4bit_use_double_quant": false,
  "llm_int8_enable_fp32_cpu_offload": false,
  "llm_int8_has_fp16_weight": false,
  "llm_int8_skip_modules": null,
  "llm_int8_threshold": 6.0,
  "load_in_4bit": false,
  "load_in_8bit": true,
  "quant_method": "bitsandbytes"
}

"""
        expected_data = self._parse_config_string(expected_config)
        actual_data = self._parse_config_string(str(retrieved_config))
        assert actual_data == expected_data

    def _parse_config_string(self, config_string: str) -> tuple[str, dict]:
        first_brace = config_string.find("{")
        if first_brace == -1:
            raise ValueError("Could not find opening brace '{' in the string.")

        json_part = config_string[first_brace:]
        data = json.loads(json_part)

        return data

    def test_single_component_to_quantize(self):
        component_to_quantize = "transformer"
        quant_config = PipelineQuantizationConfig(
            quant_backend="bitsandbytes_8bit",
            quant_kwargs={"load_in_8bit": True},
            components_to_quantize=component_to_quantize,
        )
        pipe = DiffusionPipeline.from_pretrained(
            self.model_name,
            quantization_config=quant_config,
            torch_dtype=torch.bfloat16,
        )
        for name, component in pipe.components.items():
            if name == component_to_quantize:
                assert hasattr(component.config, "quantization_config")


# ======================== BitsAndBytes ========================


# Model-level BitsAndBytes tests live in `tests/models/testing_utils/quantization.py`
# (`BitsAndBytesTesterMixin` / `BitsAndBytesCompileTesterMixin`), wired into model test files via
# concrete classes (e.g. `TestFluxTransformerBitsAndBytes`). Only pipeline-level coverage remains here.
@is_quantization
@is_bitsandbytes
@require_bitsandbytes_version_greater("0.43.2")
@require_accelerate
@require_torch
@require_torch_accelerator
@slow
class Base4bitTests:
    # We need to test on relatively large models (aka >1b parameters otherwise the quantiztion may not work as expected)
    # Therefore here we use only SD3 to test our module
    model_name = "stabilityai/stable-diffusion-3-medium-diffusers"

    prompt = "a beautiful sunset amidst the mountains."
    num_inference_steps = 10
    seed = 0

    @pytest.fixture(autouse=True, scope="class")
    def _toggle_determinism(self):
        was_enabled = torch.are_deterministic_algorithms_enabled()
        if not was_enabled:
            torch.use_deterministic_algorithms(True)
        yield
        if not was_enabled:
            torch.use_deterministic_algorithms(False)


@require_transformers_version_greater("4.44.0")
class TestSlowBnb4Bit(Base4bitTests):
    @pytest.fixture(autouse=True)
    def _setup_slow(self):
        gc.collect()
        backend_empty_cache(torch_device)

        nf4_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.float16,
        )
        model_4bit = SD3Transformer2DModel.from_pretrained(
            self.model_name, subfolder="transformer", quantization_config=nf4_config, device_map=torch_device
        )
        self.pipeline_4bit = DiffusionPipeline.from_pretrained(
            self.model_name, transformer=model_4bit, torch_dtype=torch.float16
        )
        self.pipeline_4bit.enable_model_cpu_offload()
        yield
        del self.pipeline_4bit

        gc.collect()
        backend_empty_cache(torch_device)

    def test_quality(self):
        output = self.pipeline_4bit(
            prompt=self.prompt,
            num_inference_steps=self.num_inference_steps,
            generator=torch.manual_seed(self.seed),
            output_type="np",
        ).images

        out_slice = output[0, -3:, -3:, -1].flatten()
        expected_slice = np.array([0.1123, 0.1296, 0.1609, 0.1042, 0.1230, 0.1274, 0.0928, 0.1165, 0.1216])

        max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice)
        assert max_diff < 1e-2

    def test_generate_quality_dequantize(self):
        r"""
        Test that loading the model and unquantize it produce correct results.
        """
        self.pipeline_4bit.transformer.dequantize()
        output = self.pipeline_4bit(
            prompt=self.prompt,
            num_inference_steps=self.num_inference_steps,
            generator=torch.manual_seed(self.seed),
            output_type="np",
        ).images

        out_slice = output[0, -3:, -3:, -1].flatten()
        expected_slice = np.array([0.1216, 0.1387, 0.1584, 0.1152, 0.1318, 0.1282, 0.1062, 0.1226, 0.1228])
        max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice)
        assert max_diff < 1e-3

        # Since we offloaded the `pipeline_4bit.transformer` to CPU (result of `enable_model_cpu_offload()), check
        # the following.
        assert self.pipeline_4bit.transformer.device.type == "cpu"
        # calling it again shouldn't be a problem
        _ = self.pipeline_4bit(
            prompt=self.prompt,
            num_inference_steps=2,
            generator=torch.manual_seed(self.seed),
            output_type="np",
        ).images

    def test_moving_to_cpu_throws_warning(self):
        nf4_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.float16,
        )
        model_4bit = SD3Transformer2DModel.from_pretrained(
            self.model_name, subfolder="transformer", quantization_config=nf4_config, device_map=torch_device
        )

        logger = logging.get_logger("diffusers.pipelines.pipeline_utils")
        logger.setLevel(30)
        with CaptureLogger(logger) as cap_logger:
            # Because `model.dtype` will return torch.float16 as SD3 transformer has
            # a conv layer as the first layer.
            _ = DiffusionPipeline.from_pretrained(
                self.model_name, transformer=model_4bit, torch_dtype=torch.float16
            ).to("cpu")

        assert "Pipelines loaded with `dtype=torch.float16`" in cap_logger.out

    @pytest.mark.xfail(
        condition=is_accelerate_version("<=", "1.1.1"),
        reason="Test will pass after https://github.com/huggingface/accelerate/pull/3223 is in a release.",
        strict=True,
    )
    def test_pipeline_cuda_placement_works_with_nf4(self):
        transformer_nf4_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.float16,
        )
        transformer_4bit = SD3Transformer2DModel.from_pretrained(
            self.model_name,
            subfolder="transformer",
            quantization_config=transformer_nf4_config,
            torch_dtype=torch.float16,
            device_map=torch_device,
        )
        text_encoder_3_nf4_config = TranBitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.float16,
        )
        text_encoder_3_4bit = T5EncoderModel.from_pretrained(
            self.model_name,
            subfolder="text_encoder_3",
            quantization_config=text_encoder_3_nf4_config,
            torch_dtype=torch.float16,
            device_map=torch_device,
        )
        # CUDA device placement works.
        pipeline_4bit = DiffusionPipeline.from_pretrained(
            self.model_name,
            transformer=transformer_4bit,
            text_encoder_3=text_encoder_3_4bit,
            torch_dtype=torch.float16,
        ).to(torch_device)

        # Check if inference works.
        _ = pipeline_4bit(self.prompt, max_sequence_length=20, num_inference_steps=2)

        del pipeline_4bit


@require_transformers_version_greater("4.44.0")
class TestSlowBnb4BitFlux(Base4bitTests):
    @pytest.fixture(autouse=True)
    def _setup_flux(self):
        gc.collect()
        backend_empty_cache(torch_device)

        model_id = "hf-internal-testing/flux.1-dev-nf4-pkg"
        t5_4bit = T5EncoderModel.from_pretrained(model_id, subfolder="text_encoder_2")
        transformer_4bit = FluxTransformer2DModel.from_pretrained(model_id, subfolder="transformer")
        self.pipeline_4bit = DiffusionPipeline.from_pretrained(
            "black-forest-labs/FLUX.1-dev",
            text_encoder_2=t5_4bit,
            transformer=transformer_4bit,
            torch_dtype=torch.float16,
        )
        self.pipeline_4bit.enable_model_cpu_offload()
        yield
        del self.pipeline_4bit

        gc.collect()
        backend_empty_cache(torch_device)

    def test_quality(self):
        # keep the resolution and max tokens to a lower number for faster execution.
        output = self.pipeline_4bit(
            prompt=self.prompt,
            num_inference_steps=self.num_inference_steps,
            generator=torch.manual_seed(self.seed),
            height=256,
            width=256,
            max_sequence_length=64,
            output_type="np",
        ).images

        out_slice = output[0, -3:, -3:, -1].flatten()
        expected_slice = np.array([0.0583, 0.0586, 0.0632, 0.0815, 0.0813, 0.0947, 0.1040, 0.1145, 0.1265])

        max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice)
        assert max_diff < 1e-3

    @require_peft_backend
    def test_lora_loading(self):
        self.pipeline_4bit.load_lora_weights(
            hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"), adapter_name="hyper-sd"
        )
        self.pipeline_4bit.set_adapters("hyper-sd", adapter_weights=0.125)

        output = self.pipeline_4bit(
            prompt=self.prompt,
            height=256,
            width=256,
            max_sequence_length=64,
            output_type="np",
            num_inference_steps=8,
            generator=torch.Generator().manual_seed(42),
        ).images
        out_slice = output[0, -3:, -3:, -1].flatten()
        expected_slice = np.array([0.5347, 0.5342, 0.5283, 0.5093, 0.4988, 0.5093, 0.5044, 0.5015, 0.4946])

        max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice)
        assert max_diff < 1e-3


@require_transformers_version_greater("4.44.0")
@require_peft_backend
class TestSlowBnb4BitFluxControlWithLora(Base4bitTests):
    @pytest.fixture(autouse=True)
    def _setup_flux_control(self):
        gc.collect()
        backend_empty_cache(torch_device)

        self.pipeline_4bit = FluxControlPipeline.from_pretrained("eramth/flux-4bit", torch_dtype=torch.float16)
        self.pipeline_4bit.enable_model_cpu_offload()
        yield
        del self.pipeline_4bit

        gc.collect()
        backend_empty_cache(torch_device)

    def test_lora_loading(self):
        self.pipeline_4bit.load_lora_weights("black-forest-labs/FLUX.1-Canny-dev-lora")

        output = self.pipeline_4bit(
            prompt=self.prompt,
            control_image=Image.new(mode="RGB", size=(256, 256)),
            height=256,
            width=256,
            max_sequence_length=64,
            output_type="np",
            num_inference_steps=8,
            generator=torch.Generator().manual_seed(42),
        ).images
        out_slice = output[0, -3:, -3:, -1].flatten()
        expected_slice = np.array([0.1636, 0.1675, 0.1982, 0.1743, 0.1809, 0.1936, 0.1743, 0.2095, 0.2139])

        max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice)
        assert max_diff < 1e-3, f"{out_slice=} != {expected_slice=}"


@is_quantization
@is_bitsandbytes
@require_torch_version_greater("2.7.1")
@require_bitsandbytes_version_greater("0.45.5")
class TestBnb4BitCompile(QuantCompileTests):
    @property
    def quantization_config(self):
        return PipelineQuantizationConfig(
            quant_backend="bitsandbytes_4bit",
            quant_kwargs={
                "load_in_4bit": True,
                "bnb_4bit_quant_type": "nf4",
                "bnb_4bit_compute_dtype": torch.bfloat16,
            },
            components_to_quantize=["transformer", "text_encoder_2"],
        )


@is_quantization
@is_bitsandbytes
@require_bitsandbytes_version_greater("0.43.2")
@require_accelerate
@require_torch
@require_torch_accelerator
@slow
class Base8bitTests:
    # We need to test on relatively large models (aka >1b parameters otherwise the quantiztion may not work as expected)
    # Therefore here we use only SD3 to test our module
    model_name = "stabilityai/stable-diffusion-3-medium-diffusers"

    prompt = "a beautiful sunset amidst the mountains."
    num_inference_steps = 10
    seed = 0

    @pytest.fixture(autouse=True, scope="class")
    def _toggle_determinism(self):
        was_enabled = torch.are_deterministic_algorithms_enabled()
        if not was_enabled:
            torch.use_deterministic_algorithms(True)
        yield
        if not was_enabled:
            torch.use_deterministic_algorithms(False)


@require_transformers_version_greater("4.44.0")
class TestSlowBnb8bit(Base8bitTests):
    @pytest.fixture(autouse=True)
    def _setup_slow(self):
        gc.collect()
        backend_empty_cache(torch_device)

        mixed_int8_config = BitsAndBytesConfig(load_in_8bit=True)
        model_8bit = SD3Transformer2DModel.from_pretrained(
            self.model_name, subfolder="transformer", quantization_config=mixed_int8_config, device_map=torch_device
        )
        self.pipeline_8bit = DiffusionPipeline.from_pretrained(
            self.model_name, transformer=model_8bit, torch_dtype=torch.float16
        )
        self.pipeline_8bit.enable_model_cpu_offload()
        yield
        del self.pipeline_8bit

        gc.collect()
        backend_empty_cache(torch_device)

    def test_quality(self):
        output = self.pipeline_8bit(
            prompt=self.prompt,
            num_inference_steps=self.num_inference_steps,
            generator=torch.manual_seed(self.seed),
            output_type="np",
        ).images
        out_slice = output[0, -3:, -3:, -1].flatten()
        expected_slice = np.array([0.0674, 0.0623, 0.0364, 0.0632, 0.0671, 0.0430, 0.0317, 0.0493, 0.0583])

        max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice)
        assert max_diff < 1e-2

    def test_model_cpu_offload_raises_warning(self):
        model_8bit = SD3Transformer2DModel.from_pretrained(
            self.model_name,
            subfolder="transformer",
            quantization_config=BitsAndBytesConfig(load_in_8bit=True),
            device_map=torch_device,
        )
        pipeline_8bit = DiffusionPipeline.from_pretrained(
            self.model_name, transformer=model_8bit, torch_dtype=torch.float16
        )
        logger = logging.get_logger("diffusers.pipelines.pipeline_utils")
        logger.setLevel(30)

        with CaptureLogger(logger) as cap_logger:
            pipeline_8bit.enable_model_cpu_offload()

        assert "has been loaded in `bitsandbytes` 8bit" in cap_logger.out

    def test_moving_to_cpu_throws_warning(self):
        model_8bit = SD3Transformer2DModel.from_pretrained(
            self.model_name,
            subfolder="transformer",
            quantization_config=BitsAndBytesConfig(load_in_8bit=True),
            device_map=torch_device,
        )
        logger = logging.get_logger("diffusers.pipelines.pipeline_utils")
        logger.setLevel(30)

        with CaptureLogger(logger) as cap_logger:
            # Because `model.dtype` will return torch.float16 as SD3 transformer has
            # a conv layer as the first layer.
            _ = DiffusionPipeline.from_pretrained(
                self.model_name, transformer=model_8bit, torch_dtype=torch.float16
            ).to("cpu")

        assert "Pipelines loaded with `dtype=torch.float16`" in cap_logger.out

    def test_generate_quality_dequantize(self):
        r"""
        Test that loading the model and unquantize it produce correct results.
        """
        self.pipeline_8bit.transformer.dequantize()
        output = self.pipeline_8bit(
            prompt=self.prompt,
            num_inference_steps=self.num_inference_steps,
            generator=torch.manual_seed(self.seed),
            output_type="np",
        ).images

        out_slice = output[0, -3:, -3:, -1].flatten()
        expected_slice = np.array([0.0266, 0.0264, 0.0271, 0.0110, 0.0310, 0.0098, 0.0078, 0.0256, 0.0208])
        max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice)
        assert max_diff < 1e-2

        # 8bit models cannot be offloaded to CPU.
        assert self.pipeline_8bit.transformer.device.type == torch_device
        # calling it again shouldn't be a problem
        _ = self.pipeline_8bit(
            prompt=self.prompt,
            num_inference_steps=2,
            generator=torch.manual_seed(self.seed),
            output_type="np",
        ).images

    @pytest.mark.xfail(
        condition=is_accelerate_version("<=", "1.1.1"),
        reason="Test will pass after https://github.com/huggingface/accelerate/pull/3223 is in a release.",
        strict=True,
    )
    def test_pipeline_cuda_placement_works_with_mixed_int8(self):
        transformer_8bit_config = BitsAndBytesConfig(load_in_8bit=True)
        transformer_8bit = SD3Transformer2DModel.from_pretrained(
            self.model_name,
            subfolder="transformer",
            quantization_config=transformer_8bit_config,
            torch_dtype=torch.float16,
            device_map=torch_device,
        )
        text_encoder_3_8bit_config = TranBitsAndBytesConfig(load_in_8bit=True)
        text_encoder_3_8bit = T5EncoderModel.from_pretrained(
            self.model_name,
            subfolder="text_encoder_3",
            quantization_config=text_encoder_3_8bit_config,
            torch_dtype=torch.float16,
            device_map=torch_device,
        )

        # CUDA device placement works.
        device = torch_device if torch_device != "rocm" else "cuda"
        pipeline_8bit = DiffusionPipeline.from_pretrained(
            self.model_name,
            transformer=transformer_8bit,
            text_encoder_3=text_encoder_3_8bit,
            torch_dtype=torch.float16,
        ).to(device)

        # Check if inference works.
        _ = pipeline_8bit(self.prompt, max_sequence_length=20, num_inference_steps=2)

        del pipeline_8bit


@require_transformers_version_greater("4.44.0")
@require_big_accelerator
class TestSlowBnb8bitFlux(Base8bitTests):
    @pytest.fixture(autouse=True)
    def _setup_slow_flux(self):
        gc.collect()
        backend_empty_cache(torch_device)

        model_id = "hf-internal-testing/flux.1-dev-int8-pkg"
        t5_8bit = T5EncoderModel.from_pretrained(model_id, subfolder="text_encoder_2")
        transformer_8bit = FluxTransformer2DModel.from_pretrained(model_id, subfolder="transformer")
        self.pipeline_8bit = DiffusionPipeline.from_pretrained(
            "black-forest-labs/FLUX.1-dev",
            text_encoder_2=t5_8bit,
            transformer=transformer_8bit,
            torch_dtype=torch.float16,
        )
        self.pipeline_8bit.enable_model_cpu_offload()
        yield
        del self.pipeline_8bit

        gc.collect()
        backend_empty_cache(torch_device)

    def test_quality(self):
        # keep the resolution and max tokens to a lower number for faster execution.
        output = self.pipeline_8bit(
            prompt=self.prompt,
            num_inference_steps=self.num_inference_steps,
            generator=torch.manual_seed(self.seed),
            height=256,
            width=256,
            max_sequence_length=64,
            output_type="np",
        ).images
        out_slice = output[0, -3:, -3:, -1].flatten()
        expected_slice = np.array([0.0574, 0.0554, 0.0581, 0.0686, 0.0676, 0.0759, 0.0757, 0.0803, 0.0930])

        max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice)
        assert max_diff < 1e-3

    @require_peft_version_greater("0.14.0")
    def test_lora_loading(self):
        self.pipeline_8bit.load_lora_weights(
            hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"), adapter_name="hyper-sd"
        )
        self.pipeline_8bit.set_adapters("hyper-sd", adapter_weights=0.125)

        output = self.pipeline_8bit(
            prompt=self.prompt,
            height=256,
            width=256,
            max_sequence_length=64,
            output_type="np",
            num_inference_steps=8,
            generator=torch.manual_seed(42),
        ).images
        out_slice = output[0, -3:, -3:, -1].flatten()

        expected_slice = np.array([0.3916, 0.3916, 0.3887, 0.4243, 0.4155, 0.4233, 0.4570, 0.4531, 0.4248])

        max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice)
        assert max_diff < 1e-3


@require_transformers_version_greater("4.44.0")
@require_peft_backend
class TestSlowBnb8bitFluxControlWithLora(Base8bitTests):
    @pytest.fixture(autouse=True)
    def _setup_flux_control_lora(self):
        gc.collect()
        backend_empty_cache(torch_device)

        self.pipeline_8bit = FluxControlPipeline.from_pretrained(
            "black-forest-labs/FLUX.1-dev",
            quantization_config=PipelineQuantizationConfig(
                quant_backend="bitsandbytes_8bit",
                quant_kwargs={"load_in_8bit": True},
                components_to_quantize=["transformer", "text_encoder_2"],
            ),
            torch_dtype=torch.float16,
        )
        self.pipeline_8bit.enable_model_cpu_offload()
        yield
        del self.pipeline_8bit

        gc.collect()
        backend_empty_cache(torch_device)

    def test_lora_loading(self):
        self.pipeline_8bit.load_lora_weights("black-forest-labs/FLUX.1-Canny-dev-lora")

        output = self.pipeline_8bit(
            prompt=self.prompt,
            control_image=Image.new(mode="RGB", size=(256, 256)),
            height=256,
            width=256,
            max_sequence_length=64,
            output_type="np",
            num_inference_steps=8,
            generator=torch.Generator().manual_seed(42),
        ).images
        out_slice = output[0, -3:, -3:, -1].flatten()
        # Hardware-dependent: the Control LoRA dequantizes and expands `x_embedder`, and the error
        # accumulates over the 8 denoising steps enough that even different CUDA GPUs disagree, so
        # reference slices are stored per accelerator backend.
        expected_slices = Expectations(
            {
                (None, None): np.array([0.2029, 0.2136, 0.2268, 0.1921, 0.1997, 0.2185, 0.2021, 0.2183, 0.2292]),
                ("xpu", 5): np.array([0.0955, 0.1223, 0.1509, 0.0872, 0.1155, 0.1890, 0.0754, 0.1028, 0.2178]),
            }
        )
        expected_slice = expected_slices.get_expectation()

        max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice)
        assert max_diff < 1e-3, f"{out_slice=} != {expected_slice=}"


@is_quantization
@is_bitsandbytes
@require_torch_version_greater_equal("2.6.0")
@require_bitsandbytes_version_greater("0.48.0")
class TestBnb8BitCompile(QuantCompileTests):
    @property
    def quantization_config(self):
        return PipelineQuantizationConfig(
            quant_backend="bitsandbytes_8bit",
            quant_kwargs={"load_in_8bit": True},
            components_to_quantize=["transformer", "text_encoder_2"],
        )

    def test_torch_compile_with_cpu_offload(self):
        super()._test_torch_compile_with_cpu_offload(torch_dtype=torch.float16)


# ======================== TorchAO ========================


def _is_xpu_or_cuda_capability_atleast_8_9() -> bool:
    if torch.cuda.is_available():
        major, minor = torch.cuda.get_device_capability()
        if major == 8:
            return minor >= 9
        return major >= 9
    elif torch.xpu.is_available():
        return True
    return False


# Model-level TorchAO tests live in `tests/models/testing_utils/quantization.py`
# (`TorchAoTesterMixin` / `TorchAoCompileTesterMixin`), wired into model test files via concrete
# classes (e.g. `TestFluxTransformerTorchAo`). Only pipeline-level coverage remains here.
# Slices for these tests have been obtained on our aws-g6e-xlarge-plus runners
@is_quantization
@is_torchao
@require_torch
@require_torch_accelerator
@require_torchao_version_greater_or_equal("0.15.0")
class TestTorchAo:
    @pytest.fixture(autouse=True)
    def _setup_torchao(self):
        yield
        gc.collect()
        backend_empty_cache(torch_device)

    def get_dummy_components(
        self, quantization_config: TorchAoConfig, model_id: str = "hf-internal-testing/tiny-flux-pipe"
    ):
        transformer = FluxTransformer2DModel.from_pretrained(
            model_id,
            subfolder="transformer",
            quantization_config=quantization_config,
            torch_dtype=torch.bfloat16,
        )
        text_encoder = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder", torch_dtype=torch.bfloat16)
        text_encoder_2 = T5EncoderModel.from_pretrained(
            model_id, subfolder="text_encoder_2", torch_dtype=torch.bfloat16
        )
        tokenizer = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer")
        tokenizer_2 = AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer_2")
        vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.bfloat16)
        scheduler = FlowMatchEulerDiscreteScheduler()

        return {
            "scheduler": scheduler,
            "text_encoder": text_encoder,
            "text_encoder_2": text_encoder_2,
            "tokenizer": tokenizer,
            "tokenizer_2": tokenizer_2,
            "transformer": transformer,
            "vae": vae,
        }

    def get_dummy_inputs(self, device: torch.device, seed: int = 0):
        if str(device).startswith("mps"):
            generator = torch.manual_seed(seed)
        else:
            generator = torch.Generator().manual_seed(seed)

        inputs = {
            "prompt": "an astronaut riding a horse in space",
            "height": 32,
            "width": 32,
            "num_inference_steps": 2,
            "output_type": "np",
            "generator": generator,
        }

        return inputs

    def _test_quant_type(self, quantization_config: TorchAoConfig, expected_slice: list[float], model_id: str):
        components = self.get_dummy_components(quantization_config, model_id)
        pipe = FluxPipeline(**components)
        pipe.to(device=torch_device)

        inputs = self.get_dummy_inputs(torch_device)
        output = pipe(**inputs)[0]
        output_slice = output[-1, -1, -3:, -3:].flatten()

        assert np.allclose(output_slice, expected_slice, atol=1e-3, rtol=1e-3)

    def test_quantization(self):
        for model_id in ["hf-internal-testing/tiny-flux-pipe", "hf-internal-testing/tiny-flux-sharded"]:
            # fmt: off
            QUANTIZATION_TYPES_TO_TEST = [
                (Int4WeightOnlyConfig(version=2), np.array([0.4648, 0.5234, 0.5547, 0.4219, 0.4414, 0.6445, 0.4336, 0.4531, 0.5625])),
                (Int8DynamicActivationIntxWeightConfig(version=2), np.array([0.4688, 0.5195, 0.5547, 0.418, 0.4414, 0.6406, 0.4336, 0.4531, 0.5625])),
                (Int8WeightOnlyConfig(version=2), np.array([0.4648, 0.5195, 0.5547, 0.4199, 0.4414, 0.6445, 0.4316, 0.4531, 0.5625])),
                (Int8DynamicActivationInt8WeightConfig(version=2), np.array([0.4648, 0.5195, 0.5547, 0.4199, 0.4414, 0.6445, 0.4316, 0.4531, 0.5625])),
                (IntxWeightOnlyConfig(dtype=torch.uint4, group_size=16, version=2), np.array([0.4609, 0.5234, 0.5508, 0.4199, 0.4336, 0.6406, 0.4316, 0.4531, 0.5625])),
                (IntxWeightOnlyConfig(dtype=torch.uint7, group_size=16, version=2), np.array([0.4648, 0.5195, 0.5547, 0.4219, 0.4414, 0.6445, 0.4316, 0.4531, 0.5625])),
            ]

            if _is_xpu_or_cuda_capability_atleast_8_9():
                QUANTIZATION_TYPES_TO_TEST.extend([
                    (Float8WeightOnlyConfig(weight_dtype=torch.float8_e5m2), np.array([0.4590, 0.5273, 0.5547, 0.4219, 0.4375, 0.6406, 0.4316, 0.4512, 0.5625])),
                    (Float8WeightOnlyConfig(weight_dtype=torch.float8_e4m3fn), np.array([0.4648, 0.5234, 0.5547, 0.4219, 0.4414, 0.6406, 0.4316, 0.4531, 0.5625])),
                ])
            # fmt: on

            for quant_config, expected_slice in QUANTIZATION_TYPES_TO_TEST:
                quantization_config = TorchAoConfig(quant_type=quant_config, modules_to_not_convert=["x_embedder"])
                self._test_quant_type(quantization_config, expected_slice, model_id)

    def test_sequential_cpu_offload(self):
        r"""
        A test that checks if inference runs as expected when sequential cpu offloading is enabled.
        """
        quantization_config = TorchAoConfig(Int8WeightOnlyConfig())
        components = self.get_dummy_components(quantization_config)
        pipe = FluxPipeline(**components)
        pipe.enable_sequential_cpu_offload()

        inputs = self.get_dummy_inputs(torch_device)
        _ = pipe(**inputs)

    @require_torchao_version_greater_or_equal("0.15.0")
    def test_aobase_config(self):
        quantization_config = TorchAoConfig(Int8WeightOnlyConfig())
        components = self.get_dummy_components(quantization_config)
        pipe = FluxPipeline(**components).to(torch_device)

        inputs = self.get_dummy_inputs(torch_device)
        _ = pipe(**inputs)


@is_quantization
@is_torchao
@require_torchao_version_greater_or_equal("0.15.0")
class TestTorchAoCompile(QuantCompileTests):
    @property
    def quantization_config(self):
        return PipelineQuantizationConfig(
            quant_mapping={"transformer": TorchAoConfig(Int8WeightOnlyConfig())},
        )

    def test_torch_compile_with_cpu_offload(self):
        pipe = self._init_pipeline(self.quantization_config, torch.bfloat16)
        pipe.enable_model_cpu_offload()
        # No compilation because it fails with:
        # RuntimeError: _apply(): Couldn't swap Linear.weight

        # small resolutions to ensure speedy execution.
        pipe("a dog", num_inference_steps=2, max_sequence_length=16, height=256, width=256)


# Slices for these tests have been obtained on our aws-g6e-xlarge-plus runners
@is_quantization
@is_torchao
@require_torch
@require_torch_accelerator
@require_torchao_version_greater_or_equal("0.15.0")
@slow
@nightly
class TestSlowTorchAo:
    @pytest.fixture(autouse=True)
    def _setup_slow_torchao(self):
        yield
        gc.collect()
        backend_empty_cache(torch_device)

    def get_dummy_components(self, quantization_config: TorchAoConfig):
        # This is just for convenience, so that we can modify it at one place for custom environments and locally testing
        cache_dir = None
        model_id = "black-forest-labs/FLUX.1-dev"
        transformer = FluxTransformer2DModel.from_pretrained(
            model_id,
            subfolder="transformer",
            quantization_config=quantization_config,
            torch_dtype=torch.bfloat16,
            cache_dir=cache_dir,
        )
        text_encoder = CLIPTextModel.from_pretrained(
            model_id, subfolder="text_encoder", torch_dtype=torch.bfloat16, cache_dir=cache_dir
        )
        text_encoder_2 = T5EncoderModel.from_pretrained(
            model_id, subfolder="text_encoder_2", torch_dtype=torch.bfloat16, cache_dir=cache_dir
        )
        tokenizer = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer", cache_dir=cache_dir)
        tokenizer_2 = AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer_2", cache_dir=cache_dir)
        vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.bfloat16, cache_dir=cache_dir)
        scheduler = FlowMatchEulerDiscreteScheduler()

        return {
            "scheduler": scheduler,
            "text_encoder": text_encoder,
            "text_encoder_2": text_encoder_2,
            "tokenizer": tokenizer,
            "tokenizer_2": tokenizer_2,
            "transformer": transformer,
            "vae": vae,
        }

    def get_dummy_inputs(self, device: torch.device, seed: int = 0):
        if str(device).startswith("mps"):
            generator = torch.manual_seed(seed)
        else:
            generator = torch.Generator().manual_seed(seed)

        inputs = {
            "prompt": "an astronaut riding a horse in space",
            "height": 512,
            "width": 512,
            "num_inference_steps": 20,
            "output_type": "np",
            "generator": generator,
        }

        return inputs

    def _test_quant_type(self, quantization_config, expected_slice):
        components = self.get_dummy_components(quantization_config)
        pipe = FluxPipeline(**components)
        pipe.enable_model_cpu_offload()

        weight = pipe.transformer.transformer_blocks[0].ff.net[2].weight
        assert isinstance(weight, TorchAOBaseTensor)

        inputs = self.get_dummy_inputs(torch_device)
        output = pipe(**inputs)[0].flatten()
        output_slice = np.concatenate((output[:16], output[-16:]))
        assert np.allclose(output_slice, expected_slice, atol=1e-3, rtol=1e-3)

    def test_quantization(self):
        # fmt: off
        QUANTIZATION_TYPES_TO_TEST = [
            (Int8WeightOnlyConfig(), np.array([0.0505, 0.0742, 0.1367, 0.0429, 0.0585, 0.1386, 0.0585, 0.0703, 0.1367, 0.0566, 0.0703, 0.1464, 0.0546, 0.0703, 0.1425, 0.0546, 0.3535, 0.7578, 0.5000, 0.4062, 0.7656, 0.5117, 0.4121, 0.7656, 0.5117, 0.3984, 0.7578, 0.5234, 0.4023, 0.7382, 0.5390, 0.4570])),
            (Int8DynamicActivationInt8WeightConfig(), np.array([0.0546, 0.0761, 0.1386, 0.0488, 0.0644, 0.1425, 0.0605, 0.0742, 0.1406, 0.0625, 0.0722, 0.1523, 0.0625, 0.0742, 0.1503, 0.0605, 0.3886, 0.7968, 0.5507, 0.4492, 0.7890, 0.5351, 0.4316, 0.8007, 0.5390, 0.4179, 0.8281, 0.5820, 0.4531, 0.7812, 0.5703, 0.4921])),
        ]

        if _is_xpu_or_cuda_capability_atleast_8_9():
            QUANTIZATION_TYPES_TO_TEST.extend([
                (Float8WeightOnlyConfig(weight_dtype=torch.float8_e4m3fn), np.array([0.0546, 0.0722, 0.1328, 0.0468, 0.0585, 0.1367, 0.0605, 0.0703, 0.1328, 0.0625, 0.0703, 0.1445, 0.0585, 0.0703, 0.1406, 0.0605, 0.3496, 0.7109, 0.4843, 0.4042, 0.7226, 0.5000, 0.4160, 0.7031, 0.4824, 0.3886, 0.6757, 0.4667, 0.3710, 0.6679, 0.4902, 0.4238])),
            ])
        # fmt: on

        for quant_config, expected_slice in QUANTIZATION_TYPES_TO_TEST:
            quantization_config = TorchAoConfig(quant_type=quant_config, modules_to_not_convert=["x_embedder"])
            self._test_quant_type(quantization_config, expected_slice)
            gc.collect()
            backend_empty_cache(torch_device)
            backend_synchronize(torch_device)

    def test_serialization_int8wo(self):
        quantization_config = TorchAoConfig(Int8WeightOnlyConfig())
        components = self.get_dummy_components(quantization_config)
        pipe = FluxPipeline(**components)
        pipe.enable_model_cpu_offload()

        weight = pipe.transformer.x_embedder.weight
        assert isinstance(weight, Int8Tensor)

        inputs = self.get_dummy_inputs(torch_device)
        output = pipe(**inputs)[0].flatten()[:128]

        with tempfile.TemporaryDirectory() as tmp_dir:
            pipe.transformer.save_pretrained(tmp_dir, safe_serialization=False)
            pipe.remove_all_hooks()
            del pipe.transformer
            gc.collect()
            backend_empty_cache(torch_device)
            backend_synchronize(torch_device)
            transformer = FluxTransformer2DModel.from_pretrained(
                tmp_dir, torch_dtype=torch.bfloat16, use_safetensors=False
            )
            pipe.transformer = transformer
            pipe.enable_model_cpu_offload()

        weight = transformer.x_embedder.weight
        assert isinstance(weight, Int8Tensor)

        loaded_output = pipe(**inputs)[0].flatten()[:128]
        # Seems to require higher tolerance depending on which machine it is being run.
        # A difference of 0.06 in normalized pixel space (-1 to 1), corresponds to a difference of
        # 0.06 / 2 * 255 = 7.65 in pixel space (0 to 255). On our CI runners, the difference is about 0.04,
        # on DGX it is 0.06, and on audace it is 0.037. So, we are using a tolerance of 0.06 here.
        assert np.allclose(output, loaded_output, atol=0.06)


# ======================== GGUF ========================


# Model-level GGUF tests live in `tests/models/testing_utils/quantization.py`
# (`GGUFTesterMixin` / `GGUFCompileTesterMixin`) and backend-level ones (quantized parameter/layer
# inspection, memory use, CUDA kernels) in `tests/quantization/gguf/`. Only pipeline-level coverage
# remains here.
@is_quantization
@is_gguf
@nightly
@require_big_accelerator
@require_accelerate
@require_gguf_version_greater_or_equal("0.10.0")
class GGUFPipelineTests:
    @pytest.fixture(autouse=True)
    def _cleanup(self):
        gc.collect()
        backend_empty_cache(torch_device)
        yield
        gc.collect()
        backend_empty_cache(torch_device)


class TestFluxGGUFPipeline(GGUFPipelineTests):
    ckpt_path = "https://huggingface.co/city96/FLUX.1-dev-gguf/blob/main/flux1-dev-Q2_K.gguf"
    model_cls = FluxTransformer2DModel
    torch_dtype = torch.bfloat16

    def test_pipeline_inference(self):
        quantization_config = GGUFQuantizationConfig(compute_dtype=self.torch_dtype)
        transformer = self.model_cls.from_single_file(
            self.ckpt_path, quantization_config=quantization_config, torch_dtype=self.torch_dtype
        )
        pipe = FluxPipeline.from_pretrained(
            "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=self.torch_dtype
        )
        pipe.enable_model_cpu_offload()

        prompt = "a cat holding a sign that says hello"
        output = pipe(
            prompt=prompt, num_inference_steps=2, generator=torch.Generator("cpu").manual_seed(0), output_type="np"
        ).images[0]
        output_slice = output[:3, :3, :].flatten()
        expected_slice = np.array(
            [
                0.47265625,
                0.43359375,
                0.359375,
                0.47070312,
                0.421875,
                0.34375,
                0.46875,
                0.421875,
                0.34765625,
                0.46484375,
                0.421875,
                0.34179688,
                0.47070312,
                0.42578125,
                0.34570312,
                0.46875,
                0.42578125,
                0.3515625,
                0.45507812,
                0.4140625,
                0.33984375,
                0.4609375,
                0.41796875,
                0.34375,
                0.45898438,
                0.41796875,
                0.34375,
            ]
        )
        max_diff = numpy_cosine_similarity_distance(expected_slice, output_slice)
        assert max_diff < 1e-4


class TestSD35LargeGGUFPipeline(GGUFPipelineTests):
    ckpt_path = "https://huggingface.co/city96/stable-diffusion-3.5-large-gguf/blob/main/sd3.5_large-Q4_0.gguf"
    model_cls = SD3Transformer2DModel
    torch_dtype = torch.bfloat16

    def test_pipeline_inference(self):
        quantization_config = GGUFQuantizationConfig(compute_dtype=self.torch_dtype)
        transformer = self.model_cls.from_single_file(
            self.ckpt_path, quantization_config=quantization_config, torch_dtype=self.torch_dtype
        )
        pipe = StableDiffusion3Pipeline.from_pretrained(
            "stabilityai/stable-diffusion-3.5-large", transformer=transformer, torch_dtype=self.torch_dtype
        )
        pipe.enable_model_cpu_offload()

        prompt = "a cat holding a sign that says hello"
        output = pipe(
            prompt=prompt,
            num_inference_steps=2,
            generator=torch.Generator("cpu").manual_seed(0),
            output_type="np",
        ).images[0]
        output_slice = output[:3, :3, :].flatten()
        expected_slices = Expectations(
            {
                ("xpu", 3): np.array(
                    [
                        0.16796875,
                        0.27929688,
                        0.28320312,
                        0.11328125,
                        0.27539062,
                        0.26171875,
                        0.10742188,
                        0.26367188,
                        0.26171875,
                        0.1484375,
                        0.2734375,
                        0.296875,
                        0.13476562,
                        0.2890625,
                        0.30078125,
                        0.1171875,
                        0.28125,
                        0.28125,
                        0.16015625,
                        0.31445312,
                        0.30078125,
                        0.15625,
                        0.32421875,
                        0.296875,
                        0.14453125,
                        0.30859375,
                        0.2890625,
                    ]
                ),
                ("cuda", 7): np.array(
                    [
                        0.17578125,
                        0.27539062,
                        0.27734375,
                        0.11914062,
                        0.26953125,
                        0.25390625,
                        0.109375,
                        0.25390625,
                        0.25,
                        0.15039062,
                        0.26171875,
                        0.28515625,
                        0.13671875,
                        0.27734375,
                        0.28515625,
                        0.12109375,
                        0.26757812,
                        0.265625,
                        0.16210938,
                        0.29882812,
                        0.28515625,
                        0.15625,
                        0.30664062,
                        0.27734375,
                        0.14648438,
                        0.29296875,
                        0.26953125,
                    ]
                ),
            }
        )
        expected_slice = expected_slices.get_expectation()
        max_diff = numpy_cosine_similarity_distance(expected_slice, output_slice)
        assert max_diff < 1e-4


class TestSD35MediumGGUFPipeline(GGUFPipelineTests):
    ckpt_path = "https://huggingface.co/city96/stable-diffusion-3.5-medium-gguf/blob/main/sd3.5_medium-Q3_K_M.gguf"
    model_cls = SD3Transformer2DModel
    torch_dtype = torch.bfloat16

    def test_pipeline_inference(self):
        quantization_config = GGUFQuantizationConfig(compute_dtype=self.torch_dtype)
        transformer = self.model_cls.from_single_file(
            self.ckpt_path, quantization_config=quantization_config, torch_dtype=self.torch_dtype
        )
        pipe = StableDiffusion3Pipeline.from_pretrained(
            "stabilityai/stable-diffusion-3.5-medium", transformer=transformer, torch_dtype=self.torch_dtype
        )
        pipe.enable_model_cpu_offload()

        prompt = "a cat holding a sign that says hello"
        output = pipe(
            prompt=prompt, num_inference_steps=2, generator=torch.Generator("cpu").manual_seed(0), output_type="np"
        ).images[0]
        output_slice = output[:3, :3, :].flatten()
        expected_slice = np.array(
            [
                0.625,
                0.6171875,
                0.609375,
                0.65625,
                0.65234375,
                0.640625,
                0.6484375,
                0.640625,
                0.625,
                0.6484375,
                0.63671875,
                0.6484375,
                0.66796875,
                0.65625,
                0.65234375,
                0.6640625,
                0.6484375,
                0.6328125,
                0.6640625,
                0.6484375,
                0.640625,
                0.67578125,
                0.66015625,
                0.62109375,
                0.671875,
                0.65625,
                0.62109375,
            ]
        )
        max_diff = numpy_cosine_similarity_distance(expected_slice, output_slice)
        assert max_diff < 1e-4


class TestAuraFlowGGUFPipeline(GGUFPipelineTests):
    ckpt_path = "https://huggingface.co/city96/AuraFlow-v0.3-gguf/blob/main/aura_flow_0.3-Q2_K.gguf"
    model_cls = AuraFlowTransformer2DModel
    torch_dtype = torch.bfloat16

    def test_pipeline_inference(self):
        quantization_config = GGUFQuantizationConfig(compute_dtype=self.torch_dtype)
        transformer = self.model_cls.from_single_file(
            self.ckpt_path, quantization_config=quantization_config, torch_dtype=self.torch_dtype
        )
        pipe = AuraFlowPipeline.from_pretrained(
            "fal/AuraFlow-v0.3", transformer=transformer, torch_dtype=self.torch_dtype
        )
        pipe.enable_model_cpu_offload()

        prompt = "a pony holding a sign that says hello"
        output = pipe(
            prompt=prompt, num_inference_steps=2, generator=torch.Generator("cpu").manual_seed(0), output_type="np"
        ).images[0]
        output_slice = output[:3, :3, :].flatten()
        expected_slice = np.array(
            [
                0.46484375,
                0.546875,
                0.64453125,
                0.48242188,
                0.53515625,
                0.59765625,
                0.47070312,
                0.5078125,
                0.5703125,
                0.42773438,
                0.50390625,
                0.5703125,
                0.47070312,
                0.515625,
                0.57421875,
                0.45898438,
                0.48632812,
                0.53515625,
                0.4453125,
                0.5078125,
                0.56640625,
                0.47851562,
                0.5234375,
                0.57421875,
                0.48632812,
                0.5234375,
                0.56640625,
            ]
        )
        max_diff = numpy_cosine_similarity_distance(expected_slice, output_slice)
        assert max_diff < 1e-4


@is_quantization
@is_gguf
@require_peft_backend
@nightly
@require_big_accelerator
@require_accelerate
@require_gguf_version_greater_or_equal("0.10.0")
class TestFluxControlLoRAGGUF:
    def test_lora_loading(self):
        ckpt_path = "https://huggingface.co/city96/FLUX.1-dev-gguf/blob/main/flux1-dev-Q2_K.gguf"
        transformer = FluxTransformer2DModel.from_single_file(
            ckpt_path,
            quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
            torch_dtype=torch.bfloat16,
        )
        pipe = FluxControlPipeline.from_pretrained(
            "black-forest-labs/FLUX.1-dev",
            transformer=transformer,
            torch_dtype=torch.bfloat16,
        ).to(torch_device)
        pipe.load_lora_weights("black-forest-labs/FLUX.1-Canny-dev-lora")

        prompt = "A robot made of exotic candies and chocolates of different kinds. The background is filled with confetti and celebratory gifts."
        control_image = load_image(
            "https://huggingface.co/datasets/sayakpaul/sample-datasets/resolve/main/control_image_robot_canny.png"
        )

        output = pipe(
            prompt=prompt,
            control_image=control_image,
            height=256,
            width=256,
            num_inference_steps=10,
            guidance_scale=30.0,
            output_type="np",
            generator=torch.manual_seed(0),
        ).images

        out_slice = output[0, -3:, -3:, -1].flatten()
        expected_slice = np.array([0.8047, 0.8359, 0.8711, 0.6875, 0.7070, 0.7383, 0.5469, 0.5820, 0.6641])

        max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice)
        assert max_diff < 1e-3


@is_quantization
@is_gguf
@require_torch_version_greater("2.7.1")
class TestGGUFCompile(QuantCompileTests):
    torch_dtype = torch.bfloat16
    gguf_ckpt = "https://huggingface.co/city96/FLUX.1-dev-gguf/blob/main/flux1-dev-Q2_K.gguf"

    @property
    def quantization_config(self):
        return GGUFQuantizationConfig(compute_dtype=self.torch_dtype)

    def _init_pipeline(self, *args, **kwargs):
        transformer = FluxTransformer2DModel.from_single_file(
            self.gguf_ckpt, quantization_config=self.quantization_config, torch_dtype=self.torch_dtype
        )
        pipe = DiffusionPipeline.from_pretrained(
            "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=self.torch_dtype
        )
        return pipe


# ======================== NVIDIA ModelOpt ========================


# Model-level ModelOpt tests live in `tests/models/testing_utils/quantization.py`
# (`ModelOptTesterMixin` / `ModelOptCompileTesterMixin`), wired into model test files via concrete
# classes (e.g. `TestSD3TransformerModelOpt`). Only pipeline-level coverage remains here.
@is_quantization
@is_modelopt
@nightly
@require_big_accelerator
@require_accelerate
@require_modelopt_version_greater_or_equal("0.33.1")
class TestModelOptFP8:
    model_id = "hf-internal-testing/tiny-sd3-pipe"

    @pytest.fixture(autouse=True)
    def _setup(self):
        backend_reset_peak_memory_stats(torch_device)
        backend_empty_cache(torch_device)
        gc.collect()
        yield
        backend_reset_peak_memory_stats(torch_device)
        backend_empty_cache(torch_device)
        gc.collect()

    def test_model_cpu_offload(self):
        transformer = SD3Transformer2DModel.from_pretrained(
            self.model_id,
            quantization_config=NVIDIAModelOptConfig(quant_type="FP8"),
            subfolder="transformer",
            torch_dtype=torch.bfloat16,
        )
        pipe = StableDiffusion3Pipeline.from_pretrained(
            self.model_id, transformer=transformer, torch_dtype=torch.bfloat16
        )
        pipe.enable_model_cpu_offload(device=torch_device)
        _ = pipe("a cat holding a sign that says hello", num_inference_steps=2)