






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			:    Dict					  =logging.get_logger(__name__)

enable_full_determinism()

class _lowercase			(       _lowercase       ,   _lowercase       ,   unittest.TestCase	):
							a          = UNetaDModel
							a          = """sample"""




							@property
							def 	lowerCamelCase_       (							self:       str  ):
														lowerCamelCase__      :			Union[str, Any] 						= 4
														lowerCamelCase__      :			Union[str, Any] 						= 3
														lowerCamelCase__      :			int 						= (32, 32)

														lowerCamelCase__      :			Union[str, Any] 						= floats_tensor((batch_size, num_channels) + sizes  ).to(UpperCamelCase__  )
														lowerCamelCase__      :			Union[str, Any] 						= torch.tensor([10]  ).to(UpperCamelCase__  )

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




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




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




							def 	lowerCamelCase_       (							self:       Optional[int]  ):
														lowerCamelCase__      :			Dict 						= {
														    """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,
														}
														lowerCamelCase__      :			List[Any] 						= self.dummy_input
														return init_dict, inputs_dict



class _lowercase			(       _lowercase       ,   _lowercase       ,   unittest.TestCase	):
							a          = UNetaDModel
							a          = """sample"""




							@property
							def 	lowerCamelCase_       (							self:       Optional[int]  ):
														lowerCamelCase__      :			List[str] 						= 4
														lowerCamelCase__      :			str 						= 4
														lowerCamelCase__      :			Any 						= (32, 32)

														lowerCamelCase__      :			Tuple 						= floats_tensor((batch_size, num_channels) + sizes  ).to(UpperCamelCase__  )
														lowerCamelCase__      :			List[Any] 						= torch.tensor([10]  ).to(UpperCamelCase__  )

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




							@property
							def 	lowerCamelCase_       (							self:       List[str]  ):
														return (4, 32, 32)




							@property
							def 	lowerCamelCase_       (							self:       Dict  ):
														return (4, 32, 32)




							def 	lowerCamelCase_       (							self:       List[Any]  ):
														lowerCamelCase__      :			Tuple 						= {
														    """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"""),
														}
														lowerCamelCase__      :			str 						= self.dummy_input
														return init_dict, inputs_dict




							def 	lowerCamelCase_       (							self:       List[str]  ):
														lowerCamelCase__      :			int 						= UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""    ,							output_loading_info=UpperCamelCase__  )

														self.assertIsNotNone(UpperCamelCase__  )
														self.assertEqual(len(loading_info["""missing_keys"""]  )    ,							0  )

														model.to(UpperCamelCase__  )
														lowerCamelCase__      :			List[Any] 						= 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 	lowerCamelCase_       (							self:       List[str]  ):
														lowerCamelCase__      :			int 						= UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""    ,							output_loading_info=UpperCamelCase__  )
														model.to(UpperCamelCase__  )
														lowerCamelCase__      :			Dict 						= 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 	lowerCamelCase_       (							self:       List[str]  ):
														# by defautl model loading will use accelerate as `low_cpu_mem_usage=True`
														lowerCamelCase__      :			List[Any] 						= UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""    ,							output_loading_info=UpperCamelCase__  )
														model_accelerate.to(UpperCamelCase__  )
														model_accelerate.eval()

														lowerCamelCase__      :			List[str] 						= torch.randn(
														    1    ,							model_accelerate.config.in_channels    ,							model_accelerate.config.sample_size    ,							model_accelerate.config.sample_size    ,							generator=torch.manual_seed(0  )    ,							)
														lowerCamelCase__      :			Tuple 						= noise.to(UpperCamelCase__  )
														lowerCamelCase__      :			List[Any] 						= torch.tensor([10] * noise.shape[0]  ).to(UpperCamelCase__  )

														lowerCamelCase__      :			Optional[Any] 						= model_accelerate(UpperCamelCase__    ,							UpperCamelCase__  )["""sample"""]

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

														lowerCamelCase__      :			List[str] 						= UNetaDModel.from_pretrained(
														    """fusing/unet-ldm-dummy-update"""    ,							output_loading_info=UpperCamelCase__    ,							low_cpu_mem_usage=UpperCamelCase__  )
														model_normal_load.to(UpperCamelCase__  )
														model_normal_load.eval()
														lowerCamelCase__      :			Union[str, Any] 						= model_normal_load(UpperCamelCase__    ,							UpperCamelCase__  )["""sample"""]

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




							def 	lowerCamelCase_       (							self:       Tuple  ):
														lowerCamelCase__      :			Any 						= UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""  )
														model.eval()
														model.to(UpperCamelCase__  )

														lowerCamelCase__      :			List[Any] 						= torch.randn(
														    1    ,							model.config.in_channels    ,							model.config.sample_size    ,							model.config.sample_size    ,							generator=torch.manual_seed(0  )    ,							)
														lowerCamelCase__      :			Dict 						= noise.to(UpperCamelCase__  )
														lowerCamelCase__      :			List[Any] 						= torch.tensor([10] * noise.shape[0]  ).to(UpperCamelCase__  )

														with torch.no_grad():
																					lowerCamelCase__      :			List[Any] 						= model(UpperCamelCase__    ,							UpperCamelCase__  ).sample

														lowerCamelCase__      :			int 						= output[0, -1, -3:, -3:].flatten().cpu()
														# fmt: off
														lowerCamelCase__      :			int 						= torch.tensor([-13.3_258, -20.1_100, -15.9_873, -17.6_617, -23.0_596, -17.9_419, -13.3_675, -16.1_889, -12.3_800]  )
														# fmt: on

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



class _lowercase			(       _lowercase       ,   _lowercase       ,   unittest.TestCase	):
							a          = UNetaDModel
							a          = """sample"""




							@property
							def 	lowerCamelCase_       (							self:       Optional[Any]    ,							UpperCamelCase__:       Dict=(32, 32)  ):
														lowerCamelCase__      :			Dict 						= 4
														lowerCamelCase__      :			Tuple 						= 3

														lowerCamelCase__      :			Dict 						= floats_tensor((batch_size, num_channels) + sizes  ).to(UpperCamelCase__  )
														lowerCamelCase__      :			Any 						= torch.tensor(batch_size * [10]  ).to(dtype=torch.intaa    ,							device=UpperCamelCase__  )

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




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




							@property
							def 	lowerCamelCase_       (							self:       Dict  ):
														return (3, 32, 32)




							def 	lowerCamelCase_       (							self:       str  ):
														lowerCamelCase__      :			List[Any] 						= {
														    """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""",
														    ],
														}
														lowerCamelCase__      :			int 						= self.dummy_input
														return init_dict, inputs_dict




							@slow
							def 	lowerCamelCase_       (							self:       List[Any]  ):
														lowerCamelCase__      :			str 						= UNetaDModel.from_pretrained("""google/ncsnpp-celebahq-256"""    ,							output_loading_info=UpperCamelCase__  )
														self.assertIsNotNone(UpperCamelCase__  )
														self.assertEqual(len(loading_info["""missing_keys"""]  )    ,							0  )

														model.to(UpperCamelCase__  )
														lowerCamelCase__      :			Dict 						= self.dummy_input
														lowerCamelCase__      :			List[Any] 						= floats_tensor((4, 3) + (256, 256)  ).to(UpperCamelCase__  )
														lowerCamelCase__      :			List[Any] 						= noise
														lowerCamelCase__      :			Union[str, Any] 						= model(**UpperCamelCase__  )

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




							@slow
							def 	lowerCamelCase_       (							self:       Optional[int]  ):
														lowerCamelCase__      :			Union[str, Any] 						= UNetaDModel.from_pretrained("""google/ncsnpp-celebahq-256"""  )
														model.to(UpperCamelCase__  )

														lowerCamelCase__      :			Union[str, Any] 						= 4
														lowerCamelCase__      :			str 						= 3
														lowerCamelCase__      :			Optional[int] 						= (256, 256)

														lowerCamelCase__      :			Union[str, Any] 						= torch.ones((batch_size, num_channels) + sizes  ).to(UpperCamelCase__  )
														lowerCamelCase__      :			List[Any] 						= torch.tensor(batch_size * [1e-4]  ).to(UpperCamelCase__  )

														with torch.no_grad():
																					lowerCamelCase__      :			Tuple 						= model(UpperCamelCase__    ,							UpperCamelCase__  ).sample

														lowerCamelCase__      :			List[str] 						= output[0, -3:, -3:, -1].flatten().cpu()
														# fmt: off
														lowerCamelCase__      :			int 						= torch.tensor([-4_842.8_691, -6_499.6_631, -3_800.1_953, -7_978.2_686, -10_980.7_129, -20_028.8_535, 8_148.2_822, 2_342.2_905, 567.7_608]  )
														# fmt: on

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




							def 	lowerCamelCase_       (							self:       Tuple  ):
														lowerCamelCase__      :			Dict 						= UNetaDModel.from_pretrained("""fusing/ncsnpp-ffhq-ve-dummy-update"""  )
														model.to(UpperCamelCase__  )

														lowerCamelCase__      :			int 						= 4
														lowerCamelCase__      :			int 						= 3
														lowerCamelCase__      :			Any 						= (32, 32)

														lowerCamelCase__      :			List[Any] 						= torch.ones((batch_size, num_channels) + sizes  ).to(UpperCamelCase__  )
														lowerCamelCase__      :			List[str] 						= torch.tensor(batch_size * [1e-4]  ).to(UpperCamelCase__  )

														with torch.no_grad():
																					lowerCamelCase__      :			List[str] 						= model(UpperCamelCase__    ,							UpperCamelCase__  ).sample

														lowerCamelCase__      :			Optional[int] 						= output[0, -3:, -3:, -1].flatten().cpu()
														# fmt: off
														lowerCamelCase__      :			Any 						= torch.tensor([-0.0_325, -0.0_900, -0.0_869, -0.0_332, -0.0_725, -0.0_270, -0.0_101, 0.0_227, 0.0_256]  )
														# fmt: on

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




							def 	lowerCamelCase_       (							self:       Union[str, Any]  ):
														# not required for this model
														pass









'''simple docstring'''







from ...configuration_utils import PretrainedConfig
from ...utils import logging


_A			:    Union[str, Any]					  =logging.get_logger(__name__)

_A			:    List[str]					  ={
    '''MIT/ast-finetuned-audioset-10-10-0.4593''': (
        '''https://huggingface.co/MIT/ast-finetuned-audioset-10-10-0.4593/resolve/main/config.json'''
    ),
}

class _lowercase			(       _lowercase	):
							a          = """audio-spectrogram-transformer"""




							def __init__(							self:       str    ,							UpperCamelCase__:       Any=768    ,							UpperCamelCase__:       Union[str, Any]=12    ,							UpperCamelCase__:       List[Any]=12    ,							UpperCamelCase__:       int=3_072    ,							UpperCamelCase__:       Optional[Any]="gelu"    ,							UpperCamelCase__:       Optional[int]=0.0    ,							UpperCamelCase__:       Tuple=0.0    ,							UpperCamelCase__:       Union[str, Any]=0.02    ,							UpperCamelCase__:       Dict=1e-12    ,							UpperCamelCase__:       List[str]=16    ,							UpperCamelCase__:       List[str]=True    ,							UpperCamelCase__:       Any=10    ,							UpperCamelCase__:       List[str]=10    ,							UpperCamelCase__:       Any=1_024    ,							UpperCamelCase__:       Optional[Any]=128    ,							**UpperCamelCase__:       Union[str, Any]    ,							):
														super().__init__(**UpperCamelCase__  )

														lowerCamelCase__      :			List[Any] 						= hidden_size
														lowerCamelCase__      :			int 						= num_hidden_layers
														lowerCamelCase__      :			List[str] 						= num_attention_heads
														lowerCamelCase__      :			Optional[int] 						= intermediate_size
														lowerCamelCase__      :			List[Any] 						= hidden_act
														lowerCamelCase__      :			List[Any] 						= hidden_dropout_prob
														lowerCamelCase__      :			str 						= attention_probs_dropout_prob
														lowerCamelCase__      :			Dict 						= initializer_range
														lowerCamelCase__      :			List[str] 						= layer_norm_eps
														lowerCamelCase__      :			List[Any] 						= patch_size
														lowerCamelCase__      :			List[str] 						= qkv_bias
														lowerCamelCase__      :			Dict 						= frequency_stride
														lowerCamelCase__      :			List[Any] 						= time_stride
														lowerCamelCase__      :			str 						= max_length
														lowerCamelCase__      :			Dict 						= num_mel_bins

