import unittest

import numpy as np
import torch

from diffusers import DDIMPipeline, DDIMScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, slow, torch_device

from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


class      lowerCamelCase_					(       UpperCAmelCase_						,     unittest.TestCase			):






    '''simple docstring'''






    a__     :						Optional[Any]      =				DDIMPipeline
    a__     :						str      =				UNCONDITIONAL_IMAGE_GENERATION_PARAMS
    a__     :						Any      =				PipelineTesterMixin.required_optional_params - {
        """num_images_per_prompt""",
        """latents""",
        """callback""",
        """callback_steps""",
    }
    a__     :						str      =				UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS
    a__     :						Optional[Any]      =				False







    def       UpperCamelCase__       (     self)	->     Tuple:
        torch.manual_seed(0)
        __UpperCamelCase					:Optional[int]        =		UNetaDModel(
            block_out_channels=(32, 64)			,			layers_per_block=2			,			sample_size=32			,			in_channels=3			,			out_channels=3			,			down_block_types=('''DownBlock2D''', '''AttnDownBlock2D''')			,			up_block_types=('''AttnUpBlock2D''', '''UpBlock2D''')			,			)
        __UpperCamelCase					:int        =		DDIMScheduler()
        __UpperCamelCase					:Any        =		{'''unet''': unet, '''scheduler''': scheduler}
        return components







    def       UpperCamelCase__       (     self			,			__lowercase			,			__lowercase=0)	->     Optional[Any]:
        if str(__lowercase).startswith('''mps'''):
            __UpperCamelCase					:List[Any]        =		torch.manual_seed(__lowercase)
        else:
            __UpperCamelCase					:List[str]        =		torch.Generator(device=__lowercase).manual_seed(__lowercase)
        __UpperCamelCase					:List[Any]        =		{
            '''batch_size''': 1,
            '''generator''': generator,
            '''num_inference_steps''': 2,
            '''output_type''': '''numpy''',
        }
        return inputs







    def       UpperCamelCase__       (     self)	->     Optional[Any]:
        __UpperCamelCase					:List[Any]        =		'''cpu'''

        __UpperCamelCase					:int        =		self.get_dummy_components()
        __UpperCamelCase					:List[Any]        =		self.pipeline_class(**__lowercase)
        pipe.to(__lowercase)
        pipe.set_progress_bar_config(disable=__lowercase)

        __UpperCamelCase					:Any        =		self.get_dummy_inputs(__lowercase)
        __UpperCamelCase					:Union[str, Any]        =		pipe(**__lowercase).images
        __UpperCamelCase					:List[Any]        =		image[0, -3:, -3:, -1]

        self.assertEqual(image.shape			,			(1, 32, 32, 3))
        __UpperCamelCase					:int        =		np.array(
            [1.0_0_0E0_0, 5.7_1_7E-0_1, 4.7_1_7E-0_1, 1.0_0_0E0_0, 0.0_0_0E0_0, 1.0_0_0E0_0, 3.0_0_0E-0_4, 0.0_0_0E0_0, 9.0_0_0E-0_4])
        __UpperCamelCase					:str        =		np.abs(image_slice.flatten() - expected_slice).max()
        self.assertLessEqual(__lowercase			,			1E-3)







    def       UpperCamelCase__       (     self)	->     List[Any]:
        super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3)







    def       UpperCamelCase__       (     self)	->     Optional[int]:
        super().test_save_load_local(expected_max_difference=3E-3)







    def       UpperCamelCase__       (     self)	->     Tuple:
        super().test_save_load_optional_components(expected_max_difference=3E-3)







    def       UpperCamelCase__       (     self)	->     int:
        super().test_inference_batch_single_identical(expected_max_diff=3E-3)


@slow
@require_torch_gpu
class      lowerCamelCase_					(       unittest.TestCase			):






    '''simple docstring'''






    def       UpperCamelCase__       (     self)	->     Any:
        __UpperCamelCase					:List[str]        =		'''google/ddpm-cifar10-32'''

        __UpperCamelCase					:Optional[Any]        =		UNetaDModel.from_pretrained(__lowercase)
        __UpperCamelCase					:Dict        =		DDIMScheduler()

        __UpperCamelCase					:Optional[Any]        =		DDIMPipeline(unet=__lowercase			,			scheduler=__lowercase)
        ddim.to(__lowercase)
        ddim.set_progress_bar_config(disable=__lowercase)

        __UpperCamelCase					:List[str]        =		torch.manual_seed(0)
        __UpperCamelCase					:Optional[int]        =		ddim(generator=__lowercase			,			eta=0.0			,			output_type='''numpy''').images

        __UpperCamelCase					:Optional[Any]        =		image[0, -3:, -3:, -1]

        assert image.shape == (1, 32, 32, 3)
        __UpperCamelCase					:Union[str, Any]        =		np.array([0.17_23, 0.16_17, 0.16_00, 0.16_26, 0.14_97, 0.15_13, 0.15_05, 0.14_42, 0.14_53])

        assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2







    def       UpperCamelCase__       (     self)	->     List[str]:
        __UpperCamelCase					:Optional[int]        =		'''google/ddpm-ema-bedroom-256'''

        __UpperCamelCase					:str        =		UNetaDModel.from_pretrained(__lowercase)
        __UpperCamelCase					:int        =		DDIMScheduler.from_pretrained(__lowercase)

        __UpperCamelCase					:Any        =		DDIMPipeline(unet=__lowercase			,			scheduler=__lowercase)
        ddpm.to(__lowercase)
        ddpm.set_progress_bar_config(disable=__lowercase)

        __UpperCamelCase					:Dict        =		torch.manual_seed(0)
        __UpperCamelCase					:Any        =		ddpm(generator=__lowercase			,			output_type='''numpy''').images

        __UpperCamelCase					:List[str]        =		image[0, -3:, -3:, -1]

        assert image.shape == (1, 256, 256, 3)
        __UpperCamelCase					:Dict        =		np.array([0.00_60, 0.02_01, 0.03_44, 0.00_24, 0.00_18, 0.00_02, 0.00_22, 0.00_00, 0.00_69])

        assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2






# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import copy
import importlib.metadata
import json
import os
from dataclasses import dataclass
from typing import Any, Dict, Union

from packaging import version

from ..utils import is_torch_available, logging


if is_torch_available():
    import torch


__lowercase										= logging.get_logger(__name__)


@dataclass
class      lowerCamelCase_					:






    '''simple docstring'''






    def __init__(     self			,			__lowercase=False			,			__lowercase=False			,			__lowercase=6.0			,			__lowercase=None			,			__lowercase=False			,			__lowercase=False			,			__lowercase=None			,			__lowercase="fp4"			,			__lowercase=False			,			**__lowercase			,			)	->     Tuple:
        __UpperCamelCase					:List[str]        =		load_in_abit
        __UpperCamelCase					:Union[str, Any]        =		load_in_abit
        __UpperCamelCase					:str        =		llm_inta_threshold
        __UpperCamelCase					:List[str]        =		llm_inta_skip_modules
        __UpperCamelCase					:Any        =		llm_inta_enable_fpaa_cpu_offload
        __UpperCamelCase					:List[Any]        =		llm_inta_has_fpaa_weight
        __UpperCamelCase					:str        =		bnb_abit_quant_type
        __UpperCamelCase					:Optional[int]        =		bnb_abit_use_double_quant

        if bnb_abit_compute_dtype is None:
            __UpperCamelCase					:Tuple        =		torch.floataa
        elif isinstance(__lowercase			,			__lowercase):
            __UpperCamelCase					:Union[str, Any]        =		getattr(__lowercase			,			__lowercase)
        elif isinstance(__lowercase			,			torch.dtype):
            __UpperCamelCase					:int        =		bnb_abit_compute_dtype
        else:
            raise ValueError('''bnb_4bit_compute_dtype must be a string or a torch.dtype''')

        self.post_init()







    def       UpperCamelCase__       (     self)	->     Union[str, Any]:
        if not isinstance(self.llm_inta_threshold			,			__lowercase):
            raise ValueError('''llm_int8_threshold must be a float''')

        if self.llm_inta_skip_modules is not None and not isinstance(self.llm_inta_skip_modules			,			__lowercase):
            raise ValueError('''llm_int8_skip_modules must be a list of strings''')
        if not isinstance(self.llm_inta_enable_fpaa_cpu_offload			,			__lowercase):
            raise ValueError('''llm_int8_enable_fp32_cpu_offload must be a boolean''')

        if not isinstance(self.llm_inta_has_fpaa_weight			,			__lowercase):
            raise ValueError('''llm_int8_has_fp16_weight must be a boolean''')

        if self.bnb_abit_compute_dtype is not None and not isinstance(self.bnb_abit_compute_dtype			,			torch.dtype):
            raise ValueError('''bnb_4bit_compute_dtype must be torch.dtype''')

        if not isinstance(self.bnb_abit_quant_type			,			__lowercase):
            raise ValueError('''bnb_4bit_quant_type must be a string''')

        if not isinstance(self.bnb_abit_use_double_quant			,			__lowercase):
            raise ValueError('''bnb_4bit_use_double_quant must be a boolean''')

        if self.load_in_abit and not version.parse(importlib.metadata.version('''bitsandbytes''')) >= version.parse(
            '''0.39.0'''):
            raise ValueError(
                '''4 bit quantization requires bitsandbytes>=0.39.0 - please upgrade your bitsandbytes version''')







    def       UpperCamelCase__       (     self)	->     Any:
        return self.load_in_abit or self.load_in_abit







    def       UpperCamelCase__       (     self)	->     List[Any]:
        if self.load_in_abit:
            return "llm_int8"
        elif self.load_in_abit and self.bnb_abit_quant_type == "fp4":
            return "fp4"
        elif self.load_in_abit and self.bnb_abit_quant_type == "nf4":
            return "nf4"
        else:
            return None







    @classmethod
    def       UpperCamelCase__       (     cls			,			__lowercase			,			__lowercase			,			**__lowercase)	->     List[str]:
        __UpperCamelCase					:Optional[int]        =		cls(**__lowercase)

        __UpperCamelCase					:Optional[Any]        =		[]
        for key, value in kwargs.items():
            if hasattr(__lowercase			,			__lowercase):
                setattr(__lowercase			,			__lowercase			,			__lowercase)
                to_remove.append(__lowercase)
        for key in to_remove:
            kwargs.pop(__lowercase			,			__lowercase)

        if return_unused_kwargs:
            return config, kwargs
        else:
            return config







    def       UpperCamelCase__       (     self			,			__lowercase)	->     Union[str, Any]:
        with open(__lowercase			,			'''w'''			,			encoding='''utf-8''') as writer:
            __UpperCamelCase					:Optional[int]        =		self.to_dict()
            __UpperCamelCase					:Optional[int]        =		json.dumps(__lowercase			,			indent=2			,			sort_keys=__lowercase) + '''\n'''

            writer.write(__lowercase)







    def       UpperCamelCase__       (     self)	->     Dict[str, Any]:
        __UpperCamelCase					:Optional[Any]        =		copy.deepcopy(self.__dict__)
        __UpperCamelCase					:Optional[int]        =		str(output['''bnb_4bit_compute_dtype''']).split('''.''')[1]

        return output







    def __repr__(     self)	->     Dict:
        return f"""{self.__class__.__name__} {self.to_json_string()}"""







    def       UpperCamelCase__       (     self			,			__lowercase = True)	->     str:
        if use_diff is True:
            __UpperCamelCase					:Union[str, Any]        =		self.to_diff_dict()
        else:
            __UpperCamelCase					:Dict        =		self.to_dict()
        return json.dumps(__lowercase			,			indent=2			,			sort_keys=__lowercase) + "\n"







    def       UpperCamelCase__       (     self)	->     Dict[str, Any]:
        __UpperCamelCase					:Union[str, Any]        =		self.to_dict()

        # get the default config dict
        __UpperCamelCase					:Optional[Any]        =		BitsAndBytesConfig().to_dict()

        __UpperCamelCase					:str        =		{}

        # only serialize values that differ from the default config
        for key, value in config_dict.items():
            if value != default_config_dict[key]:
                __UpperCamelCase					:str        =		value

        return serializable_config_dict





