




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


__SCREAMING_SNAKE_CASE 						=logging.get_logger(__name__)

enable_full_determinism()






class  __magic_name__						(  __a    ,  __a    ,  unittest.TestCase):
    '''simple docstring'''







    SCREAMING_SNAKE_CASE__					:							Any	      =					UNetaDModel
    SCREAMING_SNAKE_CASE__					:							List[Any]	      =					'''sample'''


    @property
    def     _A							(							self:   str		):
           SCREAMING_SNAKE_CASE_   	=    4
           SCREAMING_SNAKE_CASE_   	=    3
           SCREAMING_SNAKE_CASE_   	=    (32, 32)

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

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


    @property
    def     _A							(							self:   List[str]		):
           return (3, 32, 32)


    @property
    def     _A							(							self:   Optional[Any]		):
           return (3, 32, 32)





    def     _A							(							self:   int		):
           SCREAMING_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,
           }
           SCREAMING_SNAKE_CASE_   	=    self.dummy_input
           return init_dict, inputs_dict






class  __magic_name__						(  __a    ,  __a    ,  unittest.TestCase):
    '''simple docstring'''







    SCREAMING_SNAKE_CASE__					:							List[Any]	      =					UNetaDModel
    SCREAMING_SNAKE_CASE__					:							Optional[int]	      =					'''sample'''


    @property
    def     _A							(							self:   Optional[int]		):
           SCREAMING_SNAKE_CASE_   	=    4
           SCREAMING_SNAKE_CASE_   	=    4
           SCREAMING_SNAKE_CASE_   	=    (32, 32)

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

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


    @property
    def     _A							(							self:   Optional[Any]		):
           return (4, 32, 32)


    @property
    def     _A							(							self:   List[Any]		):
           return (4, 32, 32)


    def     _A							(							self:   Union[str, Any]		):
           SCREAMING_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'''),
           }
           SCREAMING_SNAKE_CASE_   	=    self.dummy_input
           return init_dict, inputs_dict


    def     _A							(							self:   Optional[Any]		):
           SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_   	=    UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''		,					output_loading_info=snake_case__		)

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

           model.to(snake_case__		)
           SCREAMING_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     _A							(							self:   str		):
           SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_   	=    UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''		,					output_loading_info=snake_case__		)
           model.to(snake_case__		)
           SCREAMING_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     _A							(							self:   int		):
           # by defautl model loading will use accelerate as `low_cpu_mem_usage=True`
           SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_   	=    UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''		,					output_loading_info=snake_case__		)
           model_accelerate.to(snake_case__		)
           model_accelerate.eval()

           SCREAMING_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		)		,					)
           SCREAMING_SNAKE_CASE_   	=    noise.to(snake_case__		)
           SCREAMING_SNAKE_CASE_   	=    torch.tensor([10] * noise.shape[0]		).to(snake_case__		)

           SCREAMING_SNAKE_CASE_   	=    model_accelerate(snake_case__		,					snake_case__		)['''sample''']

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

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

           assert torch_all_close(snake_case__		,					snake_case__		,					rtol=1E-3		)





    def     _A							(							self:   List[str]		):
           SCREAMING_SNAKE_CASE_   	=    UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''		)
           model.eval()
           model.to(snake_case__		)

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

           with torch.no_grad():
                  SCREAMING_SNAKE_CASE_   	=    model(snake_case__		,					snake_case__		).sample

           SCREAMING_SNAKE_CASE_   	=    output[0, -1, -3:, -3:].flatten().cpu()
           # fmt: off
           SCREAMING_SNAKE_CASE_   	=    torch.tensor([-13.32_58, -20.11_00, -15.98_73, -17.66_17, -23.05_96, -17.94_19, -13.36_75, -16.18_89, -12.38_00]		)
           # fmt: on

           self.assertTrue(torch_all_close(snake_case__		,					snake_case__		,					rtol=1E-3		)		)






class  __magic_name__						(  __a    ,  __a    ,  unittest.TestCase):
    '''simple docstring'''







    SCREAMING_SNAKE_CASE__					:							int	      =					UNetaDModel
    SCREAMING_SNAKE_CASE__					:							str	      =					'''sample'''


    @property
    def     _A							(							self:   List[str]		,					_lowerCamelCase:   Optional[Any]=(32, 32)		):
           SCREAMING_SNAKE_CASE_   	=    4
           SCREAMING_SNAKE_CASE_   	=    3

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

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


    @property
    def     _A							(							self:   Optional[Any]		):
           return (3, 32, 32)


    @property
    def     _A							(							self:   int		):
           return (3, 32, 32)


    def     _A							(							self:   List[str]		):
           SCREAMING_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''',
               ],
           }
           SCREAMING_SNAKE_CASE_   	=    self.dummy_input
           return init_dict, inputs_dict


    @slow
    def     _A							(							self:   Optional[Any]		):
           SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_   	=    UNetaDModel.from_pretrained('''google/ncsnpp-celebahq-256'''		,					output_loading_info=snake_case__		)
           self.assertIsNotNone(snake_case__		)
           self.assertEqual(len(loading_info['''missing_keys''']		)		,					0		)

           model.to(snake_case__		)
           SCREAMING_SNAKE_CASE_   	=    self.dummy_input
           SCREAMING_SNAKE_CASE_   	=    floats_tensor((4, 3) + (2_56, 2_56)		).to(snake_case__		)
           SCREAMING_SNAKE_CASE_   	=    noise
           SCREAMING_SNAKE_CASE_   	=    model(**snake_case__		)

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


    @slow
    def     _A							(							self:   Dict		):
           SCREAMING_SNAKE_CASE_   	=    UNetaDModel.from_pretrained('''google/ncsnpp-celebahq-256'''		)
           model.to(snake_case__		)

           SCREAMING_SNAKE_CASE_   	=    4
           SCREAMING_SNAKE_CASE_   	=    3
           SCREAMING_SNAKE_CASE_   	=    (2_56, 2_56)

           SCREAMING_SNAKE_CASE_   	=    torch.ones((batch_size, num_channels) + sizes		).to(snake_case__		)
           SCREAMING_SNAKE_CASE_   	=    torch.tensor(batch_size * [1E-4]		).to(snake_case__		)

           with torch.no_grad():
                  SCREAMING_SNAKE_CASE_   	=    model(snake_case__		,					snake_case__		).sample

           SCREAMING_SNAKE_CASE_   	=    output[0, -3:, -3:, -1].flatten().cpu()
           # fmt: off
           SCREAMING_SNAKE_CASE_   	=    torch.tensor([-48_42.86_91, -64_99.66_31, -38_00.19_53, -79_78.26_86, -1_09_80.71_29, -2_00_28.85_35, 81_48.28_22, 23_42.29_05, 5_67.76_08]		)
           # fmt: on

           self.assertTrue(torch_all_close(snake_case__		,					snake_case__		,					rtol=1E-2		)		)


    def     _A							(							self:   Optional[int]		):
           SCREAMING_SNAKE_CASE_   	=    UNetaDModel.from_pretrained('''fusing/ncsnpp-ffhq-ve-dummy-update'''		)
           model.to(snake_case__		)

           SCREAMING_SNAKE_CASE_   	=    4
           SCREAMING_SNAKE_CASE_   	=    3
           SCREAMING_SNAKE_CASE_   	=    (32, 32)

           SCREAMING_SNAKE_CASE_   	=    torch.ones((batch_size, num_channels) + sizes		).to(snake_case__		)
           SCREAMING_SNAKE_CASE_   	=    torch.tensor(batch_size * [1E-4]		).to(snake_case__		)

           with torch.no_grad():
                  SCREAMING_SNAKE_CASE_   	=    model(snake_case__		,					snake_case__		).sample

           SCREAMING_SNAKE_CASE_   	=    output[0, -3:, -3:, -1].flatten().cpu()
           # fmt: off
           SCREAMING_SNAKE_CASE_   	=    torch.tensor([-0.03_25, -0.09_00, -0.08_69, -0.03_32, -0.07_25, -0.02_70, -0.01_01, 0.02_27, 0.02_56]		)
           # fmt: on

           self.assertTrue(torch_all_close(snake_case__		,					snake_case__		,					rtol=1E-2		)		)





    def     _A							(							self:   Optional[int]		):
           # not required for this model
           pass








import unittest

from diffusers import FlaxAutoencoderKL
from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import require_flax

from .test_modeling_common_flax import FlaxModelTesterMixin


if is_flax_available():
			import jax






@require_flax
class  __magic_name__						(  __UpperCAmelCase    ,  unittest.TestCase):
				'''simple docstring'''







				SCREAMING_SNAKE_CASE__					:							str	      =					FlaxAutoencoderKL


				@property
				def     _A							(							self:   Any		):
											SCREAMING_SNAKE_CASE_   	=    4
											SCREAMING_SNAKE_CASE_   	=    3
											SCREAMING_SNAKE_CASE_   	=    (32, 32)

											SCREAMING_SNAKE_CASE_   	=    jax.random.PRNGKey(0		)
											SCREAMING_SNAKE_CASE_   	=    jax.random.uniform(_lowerCamelCase		,					((batch_size, num_channels) + sizes)		)

											return {"sample": image, "prng_key": prng_key}





				def     _A							(							self:   List[str]		):
											SCREAMING_SNAKE_CASE_   	=    {
											    '''block_out_channels''': [32, 64],
											    '''in_channels''': 3,
											    '''out_channels''': 3,
											    '''down_block_types''': ['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],
											    '''up_block_types''': ['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],
											    '''latent_channels''': 4,
											}
											SCREAMING_SNAKE_CASE_   	=    self.dummy_input
											return init_dict, inputs_dict


