# pylint: disable=redefined-builtin,no-member

from typing import Any, Dict, List, Tuple, Optional, Union
from dataclasses import dataclass
from enum import Enum
import os
import torch
from diffusers.quantizers.base import DiffusersQuantizer
from diffusers.quantizers.quantization_config import QuantizationConfigMixin
from diffusers.utils import get_module_from_name
from accelerate.utils import CustomDtype
from modules import devices, shared


debug = os.environ.get('SD_QUANT_DEBUG', None) is not None

dtype_dict = {
    "int8": {"min": -128, "max": 127, "num_bits": 8, "target_dtype": torch.int8, "torch_dtype": torch.int8, "storage_dtype": torch.int8, "is_unsigned": False, "is_integer": True},
    "uint8": {"min": 0, "max": 255, "num_bits": 8, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
    "int4": {"min": -8, "max": 7, "num_bits": 4, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
    "uint4": {"min": 0, "max": 15, "num_bits": 4, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
    "float8_e4m3fn": {"min": -448, "max": 448, "num_bits": 8, "target_dtype": torch.float8_e4m3fn, "torch_dtype": torch.float8_e4m3fn, "storage_dtype": torch.float8_e4m3fn, "is_unsigned": False, "is_integer": False},
    "float8_e5m2": {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": torch.float8_e5m2, "torch_dtype": torch.float8_e5m2, "storage_dtype": torch.float8_e5m2, "is_unsigned": False, "is_integer": False},
    "float8_e4m3fnuz": {"min": -240, "max": 240, "num_bits": 8, "target_dtype": CustomDtype.FP8, "torch_dtype": torch.float8_e4m3fnuz, "storage_dtype": torch.float8_e4m3fnuz, "is_unsigned": False, "is_integer": False},
    "float8_e5m2fnuz": {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": CustomDtype.FP8, "torch_dtype": torch.float8_e5m2fnuz, "storage_dtype": torch.float8_e5m2fnuz, "is_unsigned": False, "is_integer": False},
}
if hasattr(torch, "float8_e8m0fnu"):
    dtype_dict["float8_e8m0fnu"] = {"min": 5.87747e-39, "max": 1.70141e+38, "num_bits": 8, "target_dtype": CustomDtype.FP8, "torch_dtype": torch.float8_e8m0fnu, "storage_dtype": torch.float8_e8m0fnu, "is_unsigned": True, "is_integer": False}

quantized_matmul_dtypes = ("int8", "int4") # todo: float8_e4m3fn

linear_types = ("Linear",)
conv_types = ("Conv1d", "Conv2d", "Conv3d")
conv_transpose_types = ("ConvTranspose1d", "ConvTranspose2d", "ConvTranspose3d")
allowed_types = linear_types + conv_types + conv_transpose_types


class QuantizationMethod(str, Enum):
    SDNQ = "sdnq"


def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, param_name=None): # pylint: disable=unused-argument
    layer_class_name = layer.__class__.__name__
    if layer_class_name in allowed_types:
        is_conv_type = False
        is_conv_transpose_type = False
        is_linear_type = False
        result_shape = None
        if torch_dtype is None:
            torch_dtype = devices.dtype

        if layer_class_name in conv_types:
            if not quant_conv:
                return layer
            reduction_axes = [i for i in range(layer.weight.ndim) if i != 0]
            use_quantized_matmul = False
            is_conv_type = True
        elif layer_class_name in conv_transpose_types:
            if not quant_conv:
                return layer
            reduction_axes = [i for i in range(layer.weight.ndim) if i != 1]
            use_quantized_matmul = False
            is_conv_transpose_type = True
        else:
            is_linear_type = True
            reduction_axes = -1
            channel_size = layer.weight.shape[-1]
            use_quantized_matmul = use_quantized_matmul and weights_dtype in quantized_matmul_dtypes and channel_size >= 32 and layer.weight.shape[0] >= 32

            if not use_quantized_matmul and (group_size > 0 or (dtype_dict[weights_dtype]["num_bits"] == 4 and group_size != -1)):
                if group_size == 0:
                    group_size = 64
                    num_of_groups = channel_size // group_size

                if group_size >= channel_size:
                    group_size = channel_size
                    num_of_groups = 1
                else:
                    num_of_groups = channel_size // group_size
                    while channel_size % group_size != 0: # find something divisible
                        num_of_groups -= 1
                        if num_of_groups <= 1:
                            group_size = channel_size
                            num_of_groups = 1
                            break
                        group_size = channel_size / num_of_groups

                if num_of_groups > 1:
                    result_shape = layer.weight.shape
                    new_shape = list(result_shape)
                    last_dim_index = layer.weight.ndim
                    new_shape[last_dim_index - 1 : last_dim_index] = (int(num_of_groups), int(group_size))
                    layer.weight.data = layer.weight.reshape(new_shape)

        layer.weight.requires_grad = False
        if shared.opts.diffusers_offload_mode in {"none", "model"}:
            return_device = devices.device
        else:
            return_device = devices.cpu
        layer.weight.data = layer.weight.data.to(devices.device, dtype=torch.float32)

        if dtype_dict[weights_dtype]["is_unsigned"]:
            scale, zero_point = get_scale_asymmetric(layer.weight, reduction_axes, weights_dtype)
        else:
            scale = get_scale_symmetric(layer.weight, reduction_axes, weights_dtype)
            zero_point = None
        layer.weight.data = quantize_weight(layer.weight, scale, zero_point, weights_dtype)

        if not shared.opts.sdnq_decompress_fp32:
            scale = scale.to(torch_dtype)
            if zero_point is not None:
                zero_point = zero_point.to(torch_dtype)

        if use_quantized_matmul:
            scale = scale.squeeze(-1)
            if dtype_dict[weights_dtype]["num_bits"] == 8:
                layer.weight.data = layer.weight.data.transpose(0,1)

        layer.sdnq_decompressor = decompressor_dict[weights_dtype](
            scale=scale,
            zero_point=zero_point,
            compressed_weight_shape=layer.weight.data.shape,
            result_dtype=torch_dtype,
            result_shape=result_shape,
            use_quantized_matmul=use_quantized_matmul,
            weights_dtype=weights_dtype,
        ).to(return_device)
        layer.weight.data = layer.sdnq_decompressor.pack_weight(layer.weight.data).to(return_device)

        if is_linear_type:
            if use_quantized_matmul:
                layer.forward = quantized_linear_forward_int8_matmul
            else:
                layer.forward = quantized_linear_forward
        elif is_conv_type:
            layer.forward = quantized_conv_forward
        elif is_conv_transpose_type:
            if layer_class_name.endswith("1d"):
                layer.forward = quantized_conv_transpose_1d_forward
            elif layer_class_name.endswith("2d"):
                layer.forward = quantized_conv_transpose_2d_forward
            elif layer_class_name.endswith("3d"):
                layer.forward = quantized_conv_transpose_3d_forward
        layer.forward = layer.forward.__get__(layer, layer.__class__)
        devices.torch_gc(force=False, reason=f"SDNQ param_name: {param_name}")
    return layer


def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, param_name=None):
    has_children = list(model.children())
    if not has_children:
        return model
    for module_param_name, module in model.named_children():
        if hasattr(module, "weight") and module.weight is not None:
            module = sdnq_quantize_layer(
                module,
                weights_dtype=weights_dtype,
                torch_dtype=torch_dtype,
                group_size=group_size,
                quant_conv=quant_conv,
                use_quantized_matmul=use_quantized_matmul,
                param_name=module_param_name,
            )
        module = apply_sdnq_to_module(
                module,
                weights_dtype=weights_dtype,
                torch_dtype=torch_dtype,
                group_size=group_size,
                quant_conv=quant_conv,
                use_quantized_matmul=use_quantized_matmul,
                param_name=module_param_name,
            )
    return model


def get_scale_asymmetric(weight: torch.FloatTensor, reduction_axes: List[int], weights_dtype: str) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
    zero_point = torch.amin(weight, dim=reduction_axes, keepdims=True)
    scale = torch.amax(weight, dim=reduction_axes, keepdims=True).sub_(zero_point).div_(dtype_dict[weights_dtype]["max"])
    eps = torch.finfo(scale.dtype).eps # prevent divison by 0
    scale = torch.where(torch.abs(scale) < eps, eps, scale)
    if dtype_dict[weights_dtype]["min"] != 0:
        zero_point.add_(dtype_dict[weights_dtype]["min"])
    return scale, zero_point


def get_scale_symmetric(weight: torch.FloatTensor, reduction_axes: List[int], weights_dtype: str) -> torch.FloatTensor:
    abs_min_values = torch.amin(weight, dim=reduction_axes, keepdims=True).abs_()
    max_values = torch.amax(weight, dim=reduction_axes, keepdims=True)
    scale = torch.where(abs_min_values >= max_values, abs_min_values, -max_values).div_(dtype_dict[weights_dtype]["max"])
    eps = torch.finfo(scale.dtype).eps # prevent divison by 0
    scale = torch.where(torch.abs(scale) < eps, eps, scale)
    return scale


def quantize_weight(weight: torch.FloatTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, weights_dtype: str) -> torch.ByteTensor:
    if zero_point is not None:
        compressed_weight = torch.sub(weight, zero_point).div_(scale)
    else:
        compressed_weight = torch.div(weight, scale)
    if dtype_dict[weights_dtype]["is_integer"]:
        compressed_weight.round_()
    compressed_weight = compressed_weight.clamp_(dtype_dict[weights_dtype]["min"], dtype_dict[weights_dtype]["max"]).to(dtype_dict[weights_dtype]["torch_dtype"])
    return compressed_weight


def decompress_asymmetric(input: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.Tensor:
    result = torch.addcmul(zero_point, input.to(dtype=scale.dtype), scale).to(dtype=dtype)
    if result_shape is not None:
        result = result.reshape(result_shape)
    return result


def decompress_symmetric(input: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype, result_shape: torch.Size, skip_int8_matmul: bool = False) -> torch.Tensor:
    if skip_int8_matmul:
        result = input.transpose(0,1).to(dtype=scale.dtype).mul_(scale.unsqueeze(-1)).to(dtype=dtype)
    else:
        result = input.to(dtype=scale.dtype).mul_(scale).to(dtype=dtype)
    if result_shape is not None:
        result = result.reshape(result_shape)
    return result


def decompress_int4_asymmetric(input: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size) -> torch.Tensor:
    return decompress_asymmetric(unpack_uint4(input, shape), scale, zero_point, dtype, result_shape)


def decompress_int4_symmetric(input: torch.Tensor, scale: torch.Tensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, skip_int8_matmul: bool = False) -> torch.Tensor:
    if skip_int8_matmul:
        return decompress_symmetric(unpack_int4(input, shape, dtype=scale.dtype), scale.unsqueeze(-1), dtype, result_shape)
    else:
        return decompress_symmetric(unpack_int4(input, shape, dtype=scale.dtype), scale, dtype, result_shape)


def pack_uint4(tensor: torch.Tensor) -> torch.Tensor:
    if tensor.dtype != torch.uint8:
        raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.")
    packed_tensor = tensor.contiguous().reshape(-1, 2)
    packed_tensor = torch.bitwise_or(torch.bitwise_and(packed_tensor[:, 0], 15), torch.bitwise_left_shift(packed_tensor[:, 1], 4))
    return packed_tensor


def pack_int4(tensor: torch.Tensor) -> torch.Tensor:
    if tensor.dtype != torch.int8:
        raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.int8 type is supported.")
    return pack_uint4((tensor + 8).to(dtype=torch.uint8))


def unpack_uint4(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
    result = torch.stack((torch.bitwise_and(packed_tensor, 15), torch.bitwise_right_shift(packed_tensor, 4)), dim=-1).reshape(shape)
    return result


def unpack_int4(packed_tensor: torch.Tensor, shape: torch.Size, dtype: Optional[torch.dtype] = torch.int8, transpose: Optional[bool] = False) -> torch.Tensor:
    result = unpack_uint4(packed_tensor, shape).to(dtype=dtype).sub_(8)
    if transpose:
        result = result.transpose(0,1)
    return result


def quantize_int8_matmul_input(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.ByteTensor, torch.FloatTensor]:
    input_scale = torch.div(input.abs().max(dim=-1).values, 127).unsqueeze(-1)
    input = torch.div(input, input_scale).round_().clamp_(-128, 127).to(torch.int8).flatten(0,-2).contiguous()
    scale = torch.mul(input_scale, scale).flatten(0,-2).contiguous()
    if scale.dtype == torch.float16: # fp16 will overflow
        scale = scale.to(dtype=torch.float32)
    return input, scale


def int8_matmul(
    input: torch.FloatTensor,
    weight: torch.Tensor,
    bias: torch.FloatTensor,
    scale: torch.FloatTensor,
    compressed_weight_shape: torch.Size,
) -> torch.FloatTensor:
    if compressed_weight_shape is not None:
        weight = unpack_int4_compiled(weight, compressed_weight_shape, transpose=True)
    return_dtype = input.dtype
    output_shape = list(input.shape)
    output_shape[-1] = weight.shape[-1]
    input, scale = quantize_int8_matmul_input_compiled(input, scale)
    result = decompress_symmetric_compiled(torch._int_mm(input, weight), scale, return_dtype, output_shape) # pylint: disable=protected-access
    if bias is not None:
        result.add_(bias)
    return result


def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
    if torch.numel(input) / input.shape[-1] < 32:
        return torch.nn.functional.linear(input, self.sdnq_decompressor(self.weight, skip_int8_matmul=True), self.bias)
    return int8_matmul(input, self.weight, self.bias, self.sdnq_decompressor.scale, getattr(self.sdnq_decompressor, "compressed_weight_shape", None))


def quantized_linear_forward(self, input: torch.FloatTensor) -> torch.FloatTensor:
    return torch.nn.functional.linear(input, self.sdnq_decompressor(self.weight), self.bias)


def quantized_conv_forward(self, input) -> torch.FloatTensor:
    return self._conv_forward(input, self.sdnq_decompressor(self.weight), self.bias)


def quantized_conv_transpose_1d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor:
    output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 1, self.dilation)
    return torch.nn.functional.conv_transpose1d(input, self.sdnq_decompressor(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)


def quantized_conv_transpose_2d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor:
    output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 2, self.dilation)
    return torch.nn.functional.conv_transpose2d(input, self.sdnq_decompressor(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)


def quantized_conv_transpose_3d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor:
    output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 3, self.dilation)
    return torch.nn.functional.conv_transpose3d(input, self.sdnq_decompressor(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)


class AsymmetricWeightsDecompressor(torch.nn.Module):
    def __init__(
        self,
        scale: torch.Tensor,
        zero_point: torch.Tensor,
        result_dtype: torch.dtype,
        result_shape: torch.Size,
        weights_dtype: torch.dtype,
        **kwargs,
    ):
        super().__init__()
        self.weights_dtype = weights_dtype
        self.use_quantized_matmul = False
        self.scale = scale
        self.zero_point = zero_point
        self.result_dtype = result_dtype
        self.result_shape = result_shape

    def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
        return weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"])

    def forward(self, weight, **kwargs):
        return decompress_asymmetric_compiled(weight, self.scale, self.zero_point, self.result_dtype, self.result_shape)


class SymmetricWeightsDecompressor(torch.nn.Module):
    def __init__(
        self,
        scale: torch.Tensor,
        result_dtype: torch.dtype,
        result_shape: torch.Size,
        weights_dtype: torch.dtype,
        use_quantized_matmul: bool = False,
        **kwargs,
    ):
        super().__init__()
        self.weights_dtype = weights_dtype
        self.use_quantized_matmul = use_quantized_matmul
        self.scale = scale
        self.result_dtype = result_dtype
        self.result_shape = result_shape

    def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
        return weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"])

    def forward(self, weight, skip_int8_matmul=False, **kwargs):
        return decompress_symmetric_compiled(weight, self.scale, self.result_dtype, self.result_shape, skip_int8_matmul=skip_int8_matmul)


class INT4AsymmetricWeightsDecompressor(torch.nn.Module):
    def __init__(
        self,
        scale: torch.Tensor,
        zero_point: torch.Tensor,
        compressed_weight_shape: torch.Size,
        result_dtype: torch.dtype,
        result_shape: torch.Size,
        **kwargs,
    ):
        super().__init__()
        self.weights_dtype = "uint4"
        self.use_quantized_matmul = False
        self.scale = scale
        self.zero_point = zero_point
        self.compressed_weight_shape = compressed_weight_shape
        self.result_dtype = result_dtype
        self.result_shape = result_shape

    def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
        return pack_uint4(weight.to(dtype=torch.uint8))

    def forward(self, weight, **kwargs):
        return decompress_int4_asymmetric_compiled(weight, self.scale, self.zero_point, self.compressed_weight_shape, self.result_dtype, self.result_shape)


class INT4SymmetricWeightsDecompressor(torch.nn.Module):
    def __init__(
        self,
        scale: torch.Tensor,
        compressed_weight_shape: torch.Size,
        result_dtype: torch.dtype,
        result_shape: torch.Size,
        use_quantized_matmul: bool = False,
        **kwargs,
    ):
        super().__init__()
        self.weights_dtype = "int4"
        self.use_quantized_matmul = use_quantized_matmul
        self.scale = scale
        self.compressed_weight_shape = compressed_weight_shape
        self.result_dtype = result_dtype
        self.result_shape = result_shape

    def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
        return pack_int4(weight.to(dtype=torch.int8))

    def forward(self, weight, skip_int8_matmul=False, **kwargs):
        return decompress_int4_symmetric_compiled(weight, self.scale, self.compressed_weight_shape, self.result_dtype, self.result_shape, skip_int8_matmul=skip_int8_matmul)


decompressor_dict = {
    "int8": SymmetricWeightsDecompressor,
    "uint8": AsymmetricWeightsDecompressor,
    "int4": INT4SymmetricWeightsDecompressor,
    "uint4": INT4AsymmetricWeightsDecompressor,
    "float8_e4m3fn": SymmetricWeightsDecompressor,
    "float8_e4m3fnuz": SymmetricWeightsDecompressor,
    "float8_e5m2": SymmetricWeightsDecompressor,
    "float8_e5m2fnuz": SymmetricWeightsDecompressor,
    "float8_e8m0fnu": AsymmetricWeightsDecompressor,
}


class SDNQQuantizer(DiffusersQuantizer):
    r"""
    Diffusers Quantizer for SDNQ
    """

    requires_parameters_quantization = True
    use_keep_in_fp32_modules = True
    requires_calibration = False
    required_packages = None
    torch_dtype = None

    def __init__(self, quantization_config, **kwargs): # pylint: disable=useless-parent-delegation
        super().__init__(quantization_config, **kwargs)

    def check_if_quantized_param(
        self,
        model,
        param_value: "torch.Tensor",
        param_name: str,
        state_dict: Dict[str, Any],
        **kwargs,
    ):
        if param_name.endswith(".weight"):
            split_param_name = param_name.split(".")
            if param_name not in self.modules_to_not_convert and not any(param in split_param_name for param in self.modules_to_not_convert):
                layer, _ = get_module_from_name(model, param_name)
                return layer.__class__.__name__ in allowed_types
        return False

    def check_quantized_param(self, *args, **kwargs) -> bool:
        """
        needed for transformers compatibilty, returns self.check_if_quantized_param
        """
        return self.check_if_quantized_param(*args, **kwargs)

    def create_quantized_param( # pylint: disable=arguments-differ
        self,
        model,
        param_value: "torch.Tensor",
        param_name: str,
        target_device: "torch.device",
        state_dict: Dict[str, Any], # pylint: disable=unused-argument
        unexpected_keys: List[str], # pylint: disable=unused-argument
        **kwargs,
    ):
        # load the model params to target_device first
        layer, tensor_name = get_module_from_name(model, param_name)
        layer.weight = torch.nn.Parameter(param_value.to(device=target_device), requires_grad=False)
        layer = sdnq_quantize_layer(
            layer,
            weights_dtype=self.quantization_config.weights_dtype,
            torch_dtype=self.torch_dtype,
            group_size=self.quantization_config.group_size,
            quant_conv=self.quantization_config.quant_conv,
            use_quantized_matmul=self.quantization_config.use_quantized_matmul,
            param_name=param_name,
        )

    def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]:
        max_memory = {key: val * 0.70 for key, val in max_memory.items()}
        return max_memory

    def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype": # pylint: disable=unused-argument,arguments-renamed
        return dtype_dict[self.quantization_config.weights_dtype]["target_dtype"]

    def update_torch_dtype(self, torch_dtype: "torch.dtype" = None) -> "torch.dtype":
        if torch_dtype is None:
            torch_dtype = devices.dtype
        self.torch_dtype = torch_dtype
        return torch_dtype

    def _process_model_before_weight_loading( # pylint: disable=arguments-differ
        self,
        model,
        device_map, # pylint: disable=unused-argument
        keep_in_fp32_modules: List[str] = [],
        **kwargs,
    ):
        model.config.quantization_config = self.quantization_config
        self.modules_to_not_convert = self.quantization_config.modules_to_not_convert
        if not isinstance(self.modules_to_not_convert, list):
            self.modules_to_not_convert = [self.modules_to_not_convert]
        if keep_in_fp32_modules is not None:
            self.modules_to_not_convert.extend(keep_in_fp32_modules)

    def _process_model_after_weight_loading(self, model, **kwargs):
        if shared.opts.diffusers_offload_mode == "model":
            model = model.to(devices.cpu)
        devices.torch_gc(force=True)
        return model

    def update_tp_plan(self, config):
        """
        needed for transformers compatibilty, no-op function
        """
        return config

    def update_unexpected_keys(self, model, unexpected_keys: List[str], prefix: str) -> List[str]: # pylint: disable=unused-argument
        """
        needed for transformers compatibilty, no-op function
        """
        return unexpected_keys

    def update_missing_keys_after_loading(self, model, missing_keys: List[str], prefix: str) -> List[str]: # pylint: disable=unused-argument
        """
        needed for transformers compatibilty, no-op function
        """
        return missing_keys

    def update_expected_keys(self, model, expected_keys: List[str], loaded_keys: List[str]) -> List[str]: # pylint: disable=unused-argument
        """
        needed for transformers compatibilty, no-op function
        """
        return expected_keys

    @property
    def is_trainable(self):
        return False

    @property
    def is_serializable(self):
        return False


@dataclass
class SDNQConfig(QuantizationConfigMixin):
    """
    This is a wrapper class about all possible attributes and features that you can play with a model that has been
    loaded using `sdnq`.

    Args:
        weights_dtype (`str`, *optional*, defaults to `"int8"`):
            The target dtype for the weights after quantization. Supported values are:
            ("int8", "uint8", "int4", "uint4", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz", "float8_e8m0fnu")
       modules_to_not_convert (`list`, *optional*, default to `None`):
            The list of modules to not quantize, useful for quantizing models that explicitly require to have some
            modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers).
    """

    def __init__( # pylint: disable=super-init-not-called
        self,
        weights_dtype: str = "int8",
        group_size: int = 0,
        quant_conv: bool = False,
        use_quantized_matmul: bool = False,
        modules_to_not_convert: Optional[List[str]] = None,
        **kwargs, # pylint: disable=unused-argument
    ):
        self.weights_dtype = weights_dtype
        self.quant_method = QuantizationMethod.SDNQ
        self.group_size = group_size
        self.quant_conv = quant_conv
        self.use_quantized_matmul = use_quantized_matmul
        self.modules_to_not_convert = modules_to_not_convert
        self.post_init()
        self.is_integer = dtype_dict[self.weights_dtype]["is_integer"]

    def post_init(self):
        r"""
        Safety checker that arguments are correct
        """
        accepted_weights = ["int8", "uint8", "int4", "uint4", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz", "float8_e8m0fnu"]
        if self.weights_dtype not in accepted_weights:
            raise ValueError(f"Only support weights in {accepted_weights} but found {self.weights_dtype}")


class SDNQ_T5DenseGatedActDense(torch.nn.Module): # forward can't find what self is without creating a class
    def __init__(self, T5DenseGatedActDense, dtype):
        super().__init__()
        self.wi_0 = T5DenseGatedActDense.wi_0
        self.wi_1 = T5DenseGatedActDense.wi_1
        self.wo = T5DenseGatedActDense.wo
        self.dropout = T5DenseGatedActDense.dropout
        self.act = T5DenseGatedActDense.act
        self.torch_dtype = dtype

    def forward(self, hidden_states):
        hidden_gelu = self.act(self.wi_0(hidden_states))
        hidden_linear = self.wi_1(hidden_states)
        hidden_states = hidden_gelu * hidden_linear
        hidden_states = self.dropout(hidden_states)
        hidden_states = hidden_states.to(self.torch_dtype) # this line needs to be forced
        hidden_states = self.wo(hidden_states)
        return hidden_states


if shared.opts.sdnq_decompress_compile:
    try:
        torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit) # pylint: disable=protected-access
        decompress_asymmetric_compiled = torch.compile(decompress_asymmetric, fullgraph=True)
        decompress_symmetric_compiled = torch.compile(decompress_symmetric, fullgraph=True)
        decompress_int4_asymmetric_compiled = torch.compile(decompress_int4_asymmetric, fullgraph=True)
        decompress_int4_symmetric_compiled = torch.compile(decompress_int4_symmetric, fullgraph=True)
        if devices.backend != "ipex": # pytorch uses the cpu device in torch._int_mm op with ipex + torch.compile
            quantize_int8_matmul_input_compiled = quantize_int8_matmul_input
            unpack_int4_compiled = unpack_int4
            int8_matmul = torch.compile(int8_matmul, fullgraph=True)
        else:
            quantize_int8_matmul_input_compiled = torch.compile(quantize_int8_matmul_input, fullgraph=True)
            unpack_int4_compiled = torch.compile(unpack_int4, fullgraph=True)
    except Exception as e:
        shared.log.warning(f"Quantization: type=sdnq Decompress using torch.compile is not available: {e}")
        decompress_asymmetric_compiled = decompress_asymmetric
        decompress_symmetric_compiled = decompress_symmetric
        decompress_int4_asymmetric_compiled = decompress_int4_asymmetric
        decompress_int4_symmetric_compiled = decompress_int4_symmetric
        quantize_int8_matmul_input_compiled = quantize_int8_matmul_input
        unpack_int4_compiled = unpack_int4
else:
    decompress_asymmetric_compiled = decompress_asymmetric
    decompress_symmetric_compiled = decompress_symmetric
    decompress_int4_asymmetric_compiled = decompress_int4_asymmetric
    decompress_int4_symmetric_compiled = decompress_int4_symmetric
    quantize_int8_matmul_input_compiled = quantize_int8_matmul_input
    unpack_int4_compiled = unpack_int4