

import gc
import math
import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, logging, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


__A      		=     logging.get_logger(__name__)

enable_full_determinism()



class   lowercase_   (  __lowercase       ,				__lowercase       ,				unittest.TestCase					):
      UpperCamelCase_      :		Optional[Any]										=    UNetaDModel
      UpperCamelCase_      :		List[Any]										=    "sample"



      @property
      def 						UpperCamelCase_		(      self				:    Dict		)   ->					Dict:
        _snake_case						=							4
        _snake_case						=							3
        _snake_case						=							(32, 32)

        _snake_case						=							floats_tensor((batch_size, num_channels) + sizes		).to(A__		)
        _snake_case						=							torch.tensor([10]		).to(A__		)

        return {"sample": noise, "timestep": time_step}



      @property
      def 						UpperCamelCase_		(      self				:    int		)   ->					List[Any]:
        return (3, 32, 32)



      @property
      def 						UpperCamelCase_		(      self				:    Optional[Any]		)   ->					str:
        return (3, 32, 32)







      def 						UpperCamelCase_		(      self				:    Union[str, Any]		)   ->					Tuple:
        _snake_case						=							{
            '''block_out_channels''': (32, 64),
            '''down_block_types''': ('''DownBlock2D''', '''AttnDownBlock2D'''),
            '''up_block_types''': ('''AttnUpBlock2D''', '''UpBlock2D'''),
            '''attention_head_dim''': 3,
            '''out_channels''': 3,
            '''in_channels''': 3,
            '''layers_per_block''': 2,
            '''sample_size''': 32,
        }
        _snake_case						=							self.dummy_input
        return init_dict, inputs_dict



class   lowercase_   (  __lowercase       ,				__lowercase       ,				unittest.TestCase					):
      UpperCamelCase_      :		Optional[int]										=    UNetaDModel
      UpperCamelCase_      :		List[str]										=    "sample"



      @property
      def 						UpperCamelCase_		(      self				:    List[Any]		)   ->					int:
        _snake_case						=							4
        _snake_case						=							4
        _snake_case						=							(32, 32)

        _snake_case						=							floats_tensor((batch_size, num_channels) + sizes		).to(A__		)
        _snake_case						=							torch.tensor([10]		).to(A__		)

        return {"sample": noise, "timestep": time_step}



      @property
      def 						UpperCamelCase_		(      self				:    str		)   ->					Optional[Any]:
        return (4, 32, 32)



      @property
      def 						UpperCamelCase_		(      self				:    Any		)   ->					Union[str, Any]:
        return (4, 32, 32)



      def 						UpperCamelCase_		(      self				:    Any		)   ->					List[Any]:
        _snake_case						=							{
            '''sample_size''': 32,
            '''in_channels''': 4,
            '''out_channels''': 4,
            '''layers_per_block''': 2,
            '''block_out_channels''': (32, 64),
            '''attention_head_dim''': 32,
            '''down_block_types''': ('''DownBlock2D''', '''DownBlock2D'''),
            '''up_block_types''': ('''UpBlock2D''', '''UpBlock2D'''),
        }
        _snake_case						=							self.dummy_input
        return init_dict, inputs_dict



      def 						UpperCamelCase_		(      self				:    List[str]		)   ->					int:
        _snake_case,					_snake_case						=							UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''      , output_loading_info=A__		)

        self.assertIsNotNone(A__		)
        self.assertEqual(len(loading_info['''missing_keys''']		)      , 0		)

        model.to(A__		)
        _snake_case						=							model(**self.dummy_input		).sample

        assert image is not None, "Make sure output is not None"



      @unittest.skipIf(torch_device != '''cuda'''      , '''This test is supposed to run on GPU'''		)
      def 						UpperCamelCase_		(      self				:    str		)   ->					int:
        _snake_case,					_snake_case						=							UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''      , output_loading_info=A__		)
        model.to(A__		)
        _snake_case						=							model(**self.dummy_input		).sample

        assert image is not None, "Make sure output is not None"



      @unittest.skipIf(torch_device != '''cuda'''      , '''This test is supposed to run on GPU'''		)
      def 						UpperCamelCase_		(      self				:    List[Any]		)   ->					Optional[int]:
        # by defautl model loading will use accelerate as `low_cpu_mem_usage=True`
        _snake_case,					_snake_case						=							UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''      , output_loading_info=A__		)
        model_accelerate.to(A__		)
        model_accelerate.eval()

        _snake_case						=							torch.randn(
            1      , model_accelerate.config.in_channels      , model_accelerate.config.sample_size      , model_accelerate.config.sample_size      , generator=torch.manual_seed(0		)      , )
        _snake_case						=							noise.to(A__		)
        _snake_case						=							torch.tensor([10] * noise.shape[0]		).to(A__		)

        _snake_case						=							model_accelerate(A__      , A__		)['''sample''']

        # two models don't need to stay in the device at the same time
        del model_accelerate
        torch.cuda.empty_cache()
        gc.collect()

        _snake_case,					_snake_case						=							UNetaDModel.from_pretrained(
            '''fusing/unet-ldm-dummy-update'''      , output_loading_info=A__      , low_cpu_mem_usage=A__		)
        model_normal_load.to(A__		)
        model_normal_load.eval()
        _snake_case						=							model_normal_load(A__      , A__		)['''sample''']

        assert torch_all_close(A__      , A__      , rtol=1e-3		)







      def 						UpperCamelCase_		(      self				:    Any		)   ->					Union[str, Any]:
        _snake_case						=							UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''		)
        model.eval()
        model.to(A__		)

        _snake_case						=							torch.randn(
            1      , model.config.in_channels      , model.config.sample_size      , model.config.sample_size      , generator=torch.manual_seed(0		)      , )
        _snake_case						=							noise.to(A__		)
        _snake_case						=							torch.tensor([10] * noise.shape[0]		).to(A__		)

        with torch.no_grad():
          _snake_case						=							model(A__      , A__		).sample

        _snake_case						=							output[0, -1, -3:, -3:].flatten().cpu()
        # fmt: off
        _snake_case						=							torch.tensor([-13.3258, -20.1100, -15.9873, -17.6617, -23.0596, -17.9419, -13.3675, -16.1889, -12.3800]		)
        # fmt: on

        self.assertTrue(torch_all_close(A__      , A__      , rtol=1e-3		)		)



class   lowercase_   (  __lowercase       ,				__lowercase       ,				unittest.TestCase					):
      UpperCamelCase_      :		Union[str, Any]										=    UNetaDModel
      UpperCamelCase_      :		int										=    "sample"



      @property
      def 						UpperCamelCase_		(      self				:    str      , A__				:    Optional[Any]=(32, 32)		)   ->					Tuple:
        _snake_case						=							4
        _snake_case						=							3

        _snake_case						=							floats_tensor((batch_size, num_channels) + sizes		).to(A__		)
        _snake_case						=							torch.tensor(batch_size * [10]		).to(dtype=torch.intaa      , device=A__		)

        return {"sample": noise, "timestep": time_step}



      @property
      def 						UpperCamelCase_		(      self				:    Any		)   ->					List[str]:
        return (3, 32, 32)



      @property
      def 						UpperCamelCase_		(      self				:    str		)   ->					List[str]:
        return (3, 32, 32)



      def 						UpperCamelCase_		(      self				:    List[str]		)   ->					int:
        _snake_case						=							{
            '''block_out_channels''': [32, 64, 64, 64],
            '''in_channels''': 3,
            '''layers_per_block''': 1,
            '''out_channels''': 3,
            '''time_embedding_type''': '''fourier''',
            '''norm_eps''': 1e-6,
            '''mid_block_scale_factor''': math.sqrt(2.0		),
            '''norm_num_groups''': None,
            '''down_block_types''': [
                '''SkipDownBlock2D''',
                '''AttnSkipDownBlock2D''',
                '''SkipDownBlock2D''',
                '''SkipDownBlock2D''',
            ],
            '''up_block_types''': [
                '''SkipUpBlock2D''',
                '''SkipUpBlock2D''',
                '''AttnSkipUpBlock2D''',
                '''SkipUpBlock2D''',
            ],
        }
        _snake_case						=							self.dummy_input
        return init_dict, inputs_dict



      @slow
      def 						UpperCamelCase_		(      self				:    Optional[Any]		)   ->					int:
        _snake_case,					_snake_case						=							UNetaDModel.from_pretrained('''google/ncsnpp-celebahq-256'''      , output_loading_info=A__		)
        self.assertIsNotNone(A__		)
        self.assertEqual(len(loading_info['''missing_keys''']		)      , 0		)

        model.to(A__		)
        _snake_case						=							self.dummy_input
        _snake_case						=							floats_tensor((4, 3) + (256, 256)		).to(A__		)
        _snake_case						=							noise
        _snake_case						=							model(**A__		)

        assert image is not None, "Make sure output is not None"



      @slow
      def 						UpperCamelCase_		(      self				:    List[str]		)   ->					int:
        _snake_case						=							UNetaDModel.from_pretrained('''google/ncsnpp-celebahq-256'''		)
        model.to(A__		)

        _snake_case						=							4
        _snake_case						=							3
        _snake_case						=							(256, 256)

        _snake_case						=							torch.ones((batch_size, num_channels) + sizes		).to(A__		)
        _snake_case						=							torch.tensor(batch_size * [1e-4]		).to(A__		)

        with torch.no_grad():
          _snake_case						=							model(A__      , A__		).sample

        _snake_case						=							output[0, -3:, -3:, -1].flatten().cpu()
        # fmt: off
        _snake_case						=							torch.tensor([-4842.8691, -6499.6631, -3800.1953, -7978.2686, -10980.7129, -20028.8535, 8148.2822, 2342.2905, 567.7608]		)
        # fmt: on

        self.assertTrue(torch_all_close(A__      , A__      , rtol=1e-2		)		)



      def 						UpperCamelCase_		(      self				:    Optional[int]		)   ->					List[Any]:
        _snake_case						=							UNetaDModel.from_pretrained('''fusing/ncsnpp-ffhq-ve-dummy-update'''		)
        model.to(A__		)

        _snake_case						=							4
        _snake_case						=							3
        _snake_case						=							(32, 32)

        _snake_case						=							torch.ones((batch_size, num_channels) + sizes		).to(A__		)
        _snake_case						=							torch.tensor(batch_size * [1e-4]		).to(A__		)

        with torch.no_grad():
          _snake_case						=							model(A__      , A__		).sample

        _snake_case						=							output[0, -3:, -3:, -1].flatten().cpu()
        # fmt: off
        _snake_case						=							torch.tensor([-0.0325, -0.0900, -0.0869, -0.0332, -0.0725, -0.0270, -0.0101, 0.0227, 0.0256]		)
        # fmt: on

        self.assertTrue(torch_all_close(A__      , A__      , rtol=1e-2		)		)







      def 						UpperCamelCase_		(      self				:    Optional[int]		)   ->					Union[str, Any]:
        # not required for this model
        pass




__A      		=     {
    '''Pillow''': '''Pillow<10.0.0''',
    '''accelerate''': '''accelerate>=0.20.3''',
    '''av''': '''av==9.2.0''',
    '''beautifulsoup4''': '''beautifulsoup4''',
    '''black''': '''black~=23.1''',
    '''codecarbon''': '''codecarbon==1.2.0''',
    '''cookiecutter''': '''cookiecutter==1.7.3''',
    '''dataclasses''': '''dataclasses''',
    '''datasets''': '''datasets!=2.5.0''',
    '''decord''': '''decord==0.6.0''',
    '''deepspeed''': '''deepspeed>=0.9.3''',
    '''diffusers''': '''diffusers''',
    '''dill''': '''dill<0.3.5''',
    '''evaluate''': '''evaluate>=0.2.0''',
    '''fairscale''': '''fairscale>0.3''',
    '''faiss-cpu''': '''faiss-cpu''',
    '''fastapi''': '''fastapi''',
    '''filelock''': '''filelock''',
    '''flax''': '''flax>=0.4.1,<=0.7.0''',
    '''ftfy''': '''ftfy''',
    '''fugashi''': '''fugashi>=1.0''',
    '''GitPython''': '''GitPython<3.1.19''',
    '''hf-doc-builder''': '''hf-doc-builder>=0.3.0''',
    '''huggingface-hub''': '''huggingface-hub>=0.14.1,<1.0''',
    '''importlib_metadata''': '''importlib_metadata''',
    '''ipadic''': '''ipadic>=1.0.0,<2.0''',
    '''isort''': '''isort>=5.5.4''',
    '''jax''': '''jax>=0.2.8,!=0.3.2,<=0.4.13''',
    '''jaxlib''': '''jaxlib>=0.1.65,<=0.4.13''',
    '''jieba''': '''jieba''',
    '''kenlm''': '''kenlm''',
    '''keras-nlp''': '''keras-nlp>=0.3.1''',
    '''librosa''': '''librosa''',
    '''nltk''': '''nltk''',
    '''natten''': '''natten>=0.14.6''',
    '''numpy''': '''numpy>=1.17''',
    '''onnxconverter-common''': '''onnxconverter-common''',
    '''onnxruntime-tools''': '''onnxruntime-tools>=1.4.2''',
    '''onnxruntime''': '''onnxruntime>=1.4.0''',
    '''opencv-python''': '''opencv-python''',
    '''optuna''': '''optuna''',
    '''optax''': '''optax>=0.0.8,<=0.1.4''',
    '''packaging''': '''packaging>=20.0''',
    '''parameterized''': '''parameterized''',
    '''phonemizer''': '''phonemizer''',
    '''protobuf''': '''protobuf''',
    '''psutil''': '''psutil''',
    '''pyyaml''': '''pyyaml>=5.1''',
    '''pydantic''': '''pydantic<2''',
    '''pytest''': '''pytest>=7.2.0''',
    '''pytest-timeout''': '''pytest-timeout''',
    '''pytest-xdist''': '''pytest-xdist''',
    '''python''': '''python>=3.8.0''',
    '''ray[tune]''': '''ray[tune]''',
    '''regex''': '''regex!=2019.12.17''',
    '''requests''': '''requests''',
    '''rhoknp''': '''rhoknp>=1.1.0,<1.3.1''',
    '''rjieba''': '''rjieba''',
    '''rouge-score''': '''rouge-score!=0.0.7,!=0.0.8,!=0.1,!=0.1.1''',
    '''ruff''': '''ruff>=0.0.241,<=0.0.259''',
    '''sacrebleu''': '''sacrebleu>=1.4.12,<2.0.0''',
    '''sacremoses''': '''sacremoses''',
    '''safetensors''': '''safetensors>=0.3.1''',
    '''sagemaker''': '''sagemaker>=2.31.0''',
    '''scikit-learn''': '''scikit-learn''',
    '''sentencepiece''': '''sentencepiece>=0.1.91,!=0.1.92''',
    '''sigopt''': '''sigopt''',
    '''starlette''': '''starlette''',
    '''sudachipy''': '''sudachipy>=0.6.6''',
    '''sudachidict_core''': '''sudachidict_core>=20220729''',
    '''tensorflow-cpu''': '''tensorflow-cpu>=2.6,<2.14''',
    '''tensorflow''': '''tensorflow>=2.6,<2.14''',
    '''tensorflow-text''': '''tensorflow-text<2.14''',
    '''tf2onnx''': '''tf2onnx''',
    '''timeout-decorator''': '''timeout-decorator''',
    '''timm''': '''timm''',
    '''tokenizers''': '''tokenizers>=0.11.1,!=0.11.3,<0.14''',
    '''torch''': '''torch>=1.9,!=1.12.0''',
    '''torchaudio''': '''torchaudio''',
    '''torchvision''': '''torchvision''',
    '''pyctcdecode''': '''pyctcdecode>=0.4.0''',
    '''tqdm''': '''tqdm>=4.27''',
    '''unidic''': '''unidic>=1.0.2''',
    '''unidic_lite''': '''unidic_lite>=1.0.7''',
    '''urllib3''': '''urllib3<2.0.0''',
    '''uvicorn''': '''uvicorn''',
}

