




import gc
import unittest

import numpy as np
import torch

from diffusers import (
    AudioDiffusionPipeline,
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    DiffusionPipeline,
    Mel,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()


class    lowerCamelCase    (		unittest.TestCase						):






   '''simple docstring'''





   def 							UpperCAmelCase__       ( self			:   Dict  )						->   Optional[int]:
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()

   @property
   def 							UpperCAmelCase__       ( self			:   str  )						->   List[Any]:
        torch.manual_seed(0  )
        __magic_name__							:	Optional[int]    							=							UNetaDModel(
            sample_size=(32, 64)       ,						in_channels=1       ,						out_channels=1       ,						layers_per_block=2       ,						block_out_channels=(128, 128)       ,						down_block_types=('''AttnDownBlock2D''', '''DownBlock2D''')       ,						up_block_types=('''UpBlock2D''', '''AttnUpBlock2D''')       ,						)
        return model

   @property
   def 							UpperCAmelCase__       ( self			:   Optional[int]  )						->   Dict:
        torch.manual_seed(0  )
        __magic_name__							:	Any    							=							UNetaDConditionModel(
            sample_size=(64, 32)       ,						in_channels=1       ,						out_channels=1       ,						layers_per_block=2       ,						block_out_channels=(128, 128)       ,						down_block_types=('''CrossAttnDownBlock2D''', '''DownBlock2D''')       ,						up_block_types=('''UpBlock2D''', '''CrossAttnUpBlock2D''')       ,						cross_attention_dim=10       ,						)
        return model

   @property
   def 							UpperCAmelCase__       ( self			:   int  )						->   Union[str, Any]:
        torch.manual_seed(0  )
        __magic_name__							:	Tuple    							=							AutoencoderKL(
            sample_size=(128, 64)       ,						in_channels=1       ,						out_channels=1       ,						latent_channels=1       ,						layers_per_block=2       ,						block_out_channels=(128, 128)       ,						down_block_types=('''DownEncoderBlock2D''', '''DownEncoderBlock2D''')       ,						up_block_types=('''UpDecoderBlock2D''', '''UpDecoderBlock2D''')       ,						)
        __magic_name__							:	Dict    							=							UNetaDModel(
            sample_size=(64, 32)       ,						in_channels=1       ,						out_channels=1       ,						layers_per_block=2       ,						block_out_channels=(128, 128)       ,						down_block_types=('''AttnDownBlock2D''', '''DownBlock2D''')       ,						up_block_types=('''UpBlock2D''', '''AttnUpBlock2D''')       ,						)
        return vqvae, unet


   @slow
   def 							UpperCAmelCase__       ( self			:   Optional[Any]  )						->   Tuple:
        __magic_name__							:	Optional[Any]    							=							"""cpu"""  # ensure determinism for the device-dependent torch.Generator
        __magic_name__							:	List[str]    							=							Mel(
            x_res=self.dummy_unet.config.sample_size[1]       ,						y_res=self.dummy_unet.config.sample_size[0]       ,						)

        __magic_name__							:	Any    							=							DDPMScheduler()
        __magic_name__							:	Optional[Any]    							=							AudioDiffusionPipeline(vqvae=__a       ,						unet=self.dummy_unet       ,						mel=__a       ,						scheduler=__a  )
        __magic_name__							:	Any    							=							pipe.to(__a  )
        pipe.set_progress_bar_config(disable=__a  )

        __magic_name__							:	Optional[Any]    							=							torch.Generator(device=__a  ).manual_seed(42  )
        __magic_name__							:	Tuple    							=							pipe(generator=__a       ,						steps=4  )
        __magic_name__							:	str    							=							output.audios[0]
        __magic_name__							:	Tuple    							=							output.images[0]

        __magic_name__							:	Optional[int]    							=							torch.Generator(device=__a  ).manual_seed(42  )
        __magic_name__							:	Any    							=							pipe(generator=__a       ,						steps=4       ,						return_dict=__a  )
        __magic_name__							:	Any    							=							output[0][0]

        assert audio.shape == (1, (self.dummy_unet.config.sample_size[1] - 1) * mel.hop_length)
        assert (
            image.height == self.dummy_unet.config.sample_size[0]
            and image.width == self.dummy_unet.config.sample_size[1]
        )
        __magic_name__							:	List[Any]    							=							np.frombuffer(image.tobytes()       ,						dtype='''uint8'''  )[:10]
        __magic_name__							:	Optional[int]    							=							np.frombuffer(image_from_tuple.tobytes()       ,						dtype='''uint8'''  )[:10]
        __magic_name__							:	Any    							=							np.array([69, 255, 255, 255, 0, 0, 77, 181, 12, 127]  )

        assert np.abs(image_slice.flatten() - expected_slice  ).max() == 0
        assert np.abs(image_from_tuple_slice.flatten() - expected_slice  ).max() == 0

        __magic_name__							:	Optional[Any]    							=							Mel(
            x_res=self.dummy_vqvae_and_unet[0].config.sample_size[1]       ,						y_res=self.dummy_vqvae_and_unet[0].config.sample_size[0]       ,						)

        __magic_name__							:	int    							=							DDIMScheduler()
        __magic_name__							:	int    							=							self.dummy_vqvae_and_unet
        __magic_name__							:	Dict    							=							AudioDiffusionPipeline(
            vqvae=self.dummy_vqvae_and_unet[0]       ,						unet=dummy_vqvae_and_unet[1]       ,						mel=__a       ,						scheduler=__a  )
        __magic_name__							:	List[str]    							=							pipe.to(__a  )
        pipe.set_progress_bar_config(disable=__a  )

        np.random.seed(0  )
        __magic_name__							:	List[Any]    							=							np.random.uniform(-1       ,						1       ,						((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)  )
        __magic_name__							:	List[str]    							=							torch.Generator(device=__a  ).manual_seed(42  )
        __magic_name__							:	Optional[int]    							=							pipe(raw_audio=__a       ,						generator=__a       ,						start_step=5       ,						steps=10  )
        __magic_name__							:	int    							=							output.images[0]

        assert (
            image.height == self.dummy_vqvae_and_unet[0].config.sample_size[0]
            and image.width == self.dummy_vqvae_and_unet[0].config.sample_size[1]
        )
        __magic_name__							:	Optional[Any]    							=							np.frombuffer(image.tobytes()       ,						dtype='''uint8'''  )[:10]
        __magic_name__							:	str    							=							np.array([120, 117, 110, 109, 138, 167, 138, 148, 132, 121]  )

        assert np.abs(image_slice.flatten() - expected_slice  ).max() == 0

        __magic_name__							:	List[str]    							=							self.dummy_unet_condition
        __magic_name__							:	Union[str, Any]    							=							AudioDiffusionPipeline(
            vqvae=self.dummy_vqvae_and_unet[0]       ,						unet=__a       ,						mel=__a       ,						scheduler=__a  )
        __magic_name__							:	Union[str, Any]    							=							pipe.to(__a  )
        pipe.set_progress_bar_config(disable=__a  )

        np.random.seed(0  )
        __magic_name__							:	Tuple    							=							torch.rand((1, 1, 10)  )
        __magic_name__							:	Any    							=							pipe(generator=__a       ,						encoding=__a  )
        __magic_name__							:	int    							=							output.images[0]
        __magic_name__							:	str    							=							np.frombuffer(image.tobytes()       ,						dtype='''uint8'''  )[:10]
        __magic_name__							:	List[str]    							=							np.array([107, 103, 120, 127, 142, 122, 113, 122, 97, 111]  )

        assert np.abs(image_slice.flatten() - expected_slice  ).max() == 0






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






   '''simple docstring'''





   def 							UpperCAmelCase__       ( self			:   Dict  )						->   Dict:
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()


   def 							UpperCAmelCase__       ( self			:   Optional[int]  )						->   Optional[Any]:
        __magic_name__							:	Tuple    							=							torch_device

        __magic_name__							:	List[str]    							=							DiffusionPipeline.from_pretrained('''teticio/audio-diffusion-ddim-256'''  )
        __magic_name__							:	List[Any]    							=							pipe.to(__a  )
        pipe.set_progress_bar_config(disable=__a  )

        __magic_name__							:	Tuple    							=							torch.Generator(device=__a  ).manual_seed(42  )
        __magic_name__							:	int    							=							pipe(generator=__a  )
        __magic_name__							:	int    							=							output.audios[0]
        __magic_name__							:	int    							=							output.images[0]

        assert audio.shape == (1, (pipe.unet.config.sample_size[1] - 1) * pipe.mel.hop_length)
        assert image.height == pipe.unet.config.sample_size[0] and image.width == pipe.unet.config.sample_size[1]
        __magic_name__							:	Optional[Any]    							=							np.frombuffer(image.tobytes()       ,						dtype='''uint8'''  )[:10]
        __magic_name__							:	List[str]    							=							np.array([151, 167, 154, 144, 122, 134, 121, 105, 70, 26]  )

        assert np.abs(image_slice.flatten() - expected_slice  ).max() == 0









from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


__lowerCamelCase   :						Optional[int]				=			logging.get_logger(__name__)

__lowerCamelCase   :						Tuple				=			{
    '''shi-labs/nat-mini-in1k-224''': '''https://huggingface.co/shi-labs/nat-mini-in1k-224/resolve/main/config.json''',
    # See all Nat models at https://huggingface.co/models?filter=nat
}


class    lowerCamelCase    (		_lowerCamelCase				,_lowerCamelCase						):






			'''simple docstring'''





			UpperCamelCase__								='''nat'''

			UpperCamelCase__								={
			    '''num_attention_heads''': '''num_heads''',
			    '''num_hidden_layers''': '''num_layers''',
			}

			def __init__( self			:   Union[str, Any]       ,						lowerCamelCase_			:   Dict=4       ,						lowerCamelCase_			:   Union[str, Any]=3       ,						lowerCamelCase_			:   List[str]=64       ,						lowerCamelCase_			:   Union[str, Any]=[3, 4, 6, 5]       ,						lowerCamelCase_			:   List[Any]=[2, 4, 8, 16]       ,						lowerCamelCase_			:   List[Any]=7       ,						lowerCamelCase_			:   Union[str, Any]=3.0       ,						lowerCamelCase_			:   int=True       ,						lowerCamelCase_			:   str=0.0       ,						lowerCamelCase_			:   int=0.0       ,						lowerCamelCase_			:   Dict=0.1       ,						lowerCamelCase_			:   Optional[int]="gelu"       ,						lowerCamelCase_			:   Tuple=0.0_2       ,						lowerCamelCase_			:   Tuple=1E-5       ,						lowerCamelCase_			:   List[Any]=0.0       ,						lowerCamelCase_			:   str=None       ,						lowerCamelCase_			:   Tuple=None       ,						**lowerCamelCase_			:   Any       ,						)						->   Any:
								super().__init__(**lowerCamelCase_  )

								__magic_name__							:	List[Any]    							=							patch_size
								__magic_name__							:	str    							=							num_channels
								__magic_name__							:	Union[str, Any]    							=							embed_dim
								__magic_name__							:	Dict    							=							depths
								__magic_name__							:	str    							=							len(lowerCamelCase_  )
								__magic_name__							:	List[str]    							=							num_heads
								__magic_name__							:	Optional[int]    							=							kernel_size
								__magic_name__							:	Any    							=							mlp_ratio
								__magic_name__							:	Any    							=							qkv_bias
								__magic_name__							:	Union[str, Any]    							=							hidden_dropout_prob
								__magic_name__							:	Any    							=							attention_probs_dropout_prob
								__magic_name__							:	List[str]    							=							drop_path_rate
								__magic_name__							:	str    							=							hidden_act
								__magic_name__							:	Tuple    							=							layer_norm_eps
								__magic_name__							:	Any    							=							initializer_range
								# we set the hidden_size attribute in order to make Nat work with VisionEncoderDecoderModel
								# this indicates the channel dimension after the last stage of the model
								__magic_name__							:	str    							=							int(embed_dim * 2 ** (len(lowerCamelCase_  ) - 1)  )
								__magic_name__							:	str    							=							layer_scale_init_value
								__magic_name__							:	Dict    							=							['''stem'''] + [F'''stage{idx}''' for idx in range(1       ,						len(lowerCamelCase_  ) + 1  )]
								__magic_name__							,    __magic_name__							:	int    							=							get_aligned_output_features_output_indices(
								    out_features=lowerCamelCase_       ,						out_indices=lowerCamelCase_       ,						stage_names=self.stage_names  )



