






"""simple docstring"""






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 _SCREAMING_SNAKE_CASE				(		self  :						Dict):




          '''simple docstring'''


          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()
      @property
      def _SCREAMING_SNAKE_CASE				(		self  :						List[str]):




          '''simple docstring'''


          torch.manual_seed(0)
          SCREAMING_SNAKE_CASE_   :						List[str]							=	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 _SCREAMING_SNAKE_CASE				(		self  :						Tuple):




          '''simple docstring'''


          torch.manual_seed(0)
          SCREAMING_SNAKE_CASE_   :						Optional[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 _SCREAMING_SNAKE_CASE				(		self  :						str):




          '''simple docstring'''


          torch.manual_seed(0)
          SCREAMING_SNAKE_CASE_   :						Dict							=	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''')		,			)
          SCREAMING_SNAKE_CASE_   :						int							=	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 _SCREAMING_SNAKE_CASE				(		self  :						List[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Tuple							=	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
          SCREAMING_SNAKE_CASE_   :						List[str]							=	Mel(
              x_res=self.dummy_unet.config.sample_size[1]		,			y_res=self.dummy_unet.config.sample_size[0]		,			)

          SCREAMING_SNAKE_CASE_   :						List[str]							=	DDPMScheduler()
          SCREAMING_SNAKE_CASE_   :						Dict							=	AudioDiffusionPipeline(vqvae=_a		,			unet=self.dummy_unet		,			mel=_a		,			scheduler=_a)
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	pipe.to(_a)
          pipe.set_progress_bar_config(disable=_a)

          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	torch.Generator(device=_a).manual_seed(42)
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	pipe(generator=_a		,			steps=4)
          SCREAMING_SNAKE_CASE_   :						int							=	output.audios[0]
          SCREAMING_SNAKE_CASE_   :						Dict							=	output.images[0]

          SCREAMING_SNAKE_CASE_   :						Any							=	torch.Generator(device=_a).manual_seed(42)
          SCREAMING_SNAKE_CASE_   :						Tuple							=	pipe(generator=_a		,			steps=4		,			return_dict=_a)
          SCREAMING_SNAKE_CASE_   :						List[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]
          )
          SCREAMING_SNAKE_CASE_   :						List[str]							=	np.frombuffer(image.tobytes()		,			dtype='''uint8''')[:10]
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	np.frombuffer(image_from_tuple.tobytes()		,			dtype='''uint8''')[:10]
          SCREAMING_SNAKE_CASE_   :						List[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

          SCREAMING_SNAKE_CASE_   :						str							=	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]		,			)

          SCREAMING_SNAKE_CASE_   :						Any							=	DDIMScheduler()
          SCREAMING_SNAKE_CASE_   :						Tuple							=	self.dummy_vqvae_and_unet
          SCREAMING_SNAKE_CASE_   :						List[str]							=	AudioDiffusionPipeline(
              vqvae=self.dummy_vqvae_and_unet[0]		,			unet=dummy_vqvae_and_unet[1]		,			mel=_a		,			scheduler=_a)
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	pipe.to(_a)
          pipe.set_progress_bar_config(disable=_a)

          np.random.seed(0)
          SCREAMING_SNAKE_CASE_   :						Dict							=	np.random.uniform(-1		,			1		,			((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,))
          SCREAMING_SNAKE_CASE_   :						int							=	torch.Generator(device=_a).manual_seed(42)
          SCREAMING_SNAKE_CASE_   :						int							=	pipe(raw_audio=_a		,			generator=_a		,			start_step=5		,			steps=10)
          SCREAMING_SNAKE_CASE_   :						Dict							=	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]
          )
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	np.frombuffer(image.tobytes()		,			dtype='''uint8''')[:10]
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	np.array([120, 117, 110, 109, 138, 167, 138, 148, 132, 121])

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

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

          np.random.seed(0)
          SCREAMING_SNAKE_CASE_   :						List[str]							=	torch.rand((1, 1, 10))
          SCREAMING_SNAKE_CASE_   :						Tuple							=	pipe(generator=_a		,			encoding=_a)
          SCREAMING_SNAKE_CASE_   :						Tuple							=	output.images[0]
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	np.frombuffer(image.tobytes()		,			dtype='''uint8''')[:10]
          SCREAMING_SNAKE_CASE_   :						int							=	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 _SCREAMING_SNAKE_CASE				(		self  :						List[str]):




          '''simple docstring'''


          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()



      def _SCREAMING_SNAKE_CASE				(		self  :						Any):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Tuple							=	torch_device

          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	DiffusionPipeline.from_pretrained('''teticio/audio-diffusion-ddim-256''')
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	pipe.to(_a)
          pipe.set_progress_bar_config(disable=_a)

          SCREAMING_SNAKE_CASE_   :						List[str]							=	torch.Generator(device=_a).manual_seed(42)
          SCREAMING_SNAKE_CASE_   :						Dict							=	pipe(generator=_a)
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	output.audios[0]
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	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]
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	np.frombuffer(image.tobytes()		,			dtype='''uint8''')[:10]
          SCREAMING_SNAKE_CASE_   :						Any							=	np.array([151, 167, 154, 144, 122, 134, 121, 105, 70, 26])

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










"""simple docstring"""






from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


UpperCAmelCase_     :				List[str]         =				logging.get_logger(__name__)

UpperCAmelCase_     :				Dict         =				{
    """facebook/data2vec-text-base""": """https://huggingface.co/data2vec/resolve/main/config.json""",
}




class 							lowerCAmelCase__       (    UpperCAmelCase__				):






      '''simple docstring'''

      __UpperCamelCase							=						"data2vec-text"
      def __init__(		self  :						Any		,			lowercase_  :						Any=30522		,			lowercase_  :						Any=768		,			lowercase_  :						Union[str, Any]=12		,			lowercase_  :						Dict=12		,			lowercase_  :						List[Any]=3072		,			lowercase_  :						str="gelu"		,			lowercase_  :						int=0.1		,			lowercase_  :						Dict=0.1		,			lowercase_  :						str=512		,			lowercase_  :						Optional[int]=2		,			lowercase_  :						int=0.02		,			lowercase_  :						int=1e-12		,			lowercase_  :						Any=1		,			lowercase_  :						Any=0		,			lowercase_  :						List[Any]=2		,			lowercase_  :						Tuple="absolute"		,			lowercase_  :						Optional[int]=True		,			lowercase_  :						int=None		,			**lowercase_  :						Tuple		,			):




          '''simple docstring'''


          super().__init__(pad_token_id=lowercase_		,			bos_token_id=lowercase_		,			eos_token_id=lowercase_		,			**lowercase_)

          SCREAMING_SNAKE_CASE_   :						List[Any]							=	vocab_size
          SCREAMING_SNAKE_CASE_   :						Tuple							=	hidden_size
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	num_hidden_layers
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	num_attention_heads
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	hidden_act
          SCREAMING_SNAKE_CASE_   :						Any							=	intermediate_size
          SCREAMING_SNAKE_CASE_   :						Tuple							=	hidden_dropout_prob
          SCREAMING_SNAKE_CASE_   :						Tuple							=	attention_probs_dropout_prob
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	max_position_embeddings
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	type_vocab_size
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	initializer_range
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	layer_norm_eps
          SCREAMING_SNAKE_CASE_   :						str							=	position_embedding_type
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	use_cache
          SCREAMING_SNAKE_CASE_   :						str							=	classifier_dropout




class 							lowerCAmelCase__       (    UpperCAmelCase__				):






      '''simple docstring'''

      @property
      def _SCREAMING_SNAKE_CASE				(		self  :						int):




          '''simple docstring'''


          if self.task == "multiple-choice":
              SCREAMING_SNAKE_CASE_   :						List[Any]							=	{0: '''batch''', 1: '''choice''', 2: '''sequence'''}
          else:
              SCREAMING_SNAKE_CASE_   :						Any							=	{0: '''batch''', 1: '''sequence'''}
          return OrderedDict(
              [
                  ('''input_ids''', dynamic_axis),
                  ('''attention_mask''', dynamic_axis),
              ])


