"""simple docstring"""







from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Mapping, Optional, Union

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


if TYPE_CHECKING:
	from ...feature_extraction_utils import FeatureExtractionMixin
	from ...tokenization_utils_base import PreTrainedTokenizerBase
	from ...utils import TensorType

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

_a					:					Optional[int]=			{
    "openai/whisper-base": "https://huggingface.co/openai/whisper-base/resolve/main/config.json",
}

# fmt: off
_a					:					str=			[
    1, 2, 7, 8, 9, 10, 14, 25,
    26, 27, 28, 29, 31, 58, 59, 60, 61, 62,
    63, 90, 91, 92, 93, 357, 366, 438, 532, 685,
    705, 796, 930, 1_058, 1_220, 1_267, 1_279, 1_303, 1_343, 1_377,
    1_391, 1_635, 1_782, 1_875, 2_162, 2_361, 2_488, 3_467, 4_008, 4_211,
    4_600, 4_808, 5_299, 5_855, 6_329, 7_203, 9_609, 9_959, 10_563, 10_786,
    11_420, 11_709, 11_907, 13_163, 13_697, 13_700, 14_808, 15_306, 16_410, 16_791,
    17_992, 19_203, 19_510, 20_724, 22_305, 22_935, 27_007, 30_109, 30_420, 33_409,
    34_949, 40_283, 40_493, 40_549, 47_282, 49_146, 50_257, 50_359, 50_360, 50_361
]
_a					:					List[str]=			[
    1, 2, 7, 8, 9, 10, 14, 25,
    26, 27, 28, 29, 31, 58, 59, 60, 61, 62,
    63, 90, 91, 92, 93, 359, 503, 522, 542, 873,
    893, 902, 918, 922, 931, 1_350, 1_853, 1_982, 2_460, 2_627,
    3_246, 3_253, 3_268, 3_536, 3_846, 3_961, 4_183, 4_667, 6_585, 6_647,
    7_273, 9_061, 9_383, 10_428, 10_929, 11_938, 12_033, 12_331, 12_562, 13_793,
    14_157, 14_635, 15_265, 15_618, 16_553, 16_604, 18_362, 18_956, 20_075, 21_675,
    22_520, 26_130, 26_161, 26_435, 28_279, 29_464, 31_650, 32_302, 32_470, 36_865,
    42_863, 47_425, 49_870, 50_254, 50_258, 50_360, 50_361, 50_362
]

class 		UpperCamelCase      (   lowercase						):
				UpperCAmelCase    :				int   =		"""whisper"""
				UpperCAmelCase    :				Any   =		["""past_key_values"""]
				UpperCAmelCase    :				Optional[int]   =		{"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}



				def __init__(self    :				List[str]     ,  _A    :				str=5_18_65     ,  _A    :				Optional[int]=80     ,  _A    :				Optional[Any]=6     ,  _A    :				Any=4     ,  _A    :				Any=6     ,  _A    :				Optional[Any]=4     ,  _A    :				Optional[Any]=15_36     ,  _A    :				Any=15_36     ,  _A    :				List[str]=0.0     ,  _A    :				int=0.0     ,  _A    :				Tuple=5_02_57     ,  _A    :				int=True     ,  _A    :				int=True     ,  _A    :				str="gelu"     ,  _A    :				Union[str, Any]=2_56     ,  _A    :				Tuple=0.0     ,  _A    :				Tuple=0.0     ,  _A    :				Union[str, Any]=0.0     ,  _A    :				Optional[Any]=0.02     ,  _A    :				Any=False     ,  _A    :				Optional[Any]=15_00     ,  _A    :				int=4_48     ,  _A    :				Dict=5_02_56     ,  _A    :				Any=5_02_56     ,  _A    :				int=5_02_56     ,  _A    :				Tuple=None     ,  _A    :				Dict=[2_20, 5_02_56]     ,  _A    :				Optional[int]=False     ,  _A    :				List[str]=2_56     ,  _A    :				Any=False     ,  _A    :				Union[str, Any]=0.05     ,  _A    :				Optional[Any]=10     ,  _A    :				Optional[Any]=2     ,  _A    :				List[Any]=0.0     ,  _A    :				int=10     ,  _A    :				str=0     ,  _A    :				Tuple=7     ,  **_A    :				Any     ,  )   ->     int:
					__snake_case      :							List[str]						=			vocab_size
					__snake_case      :							Tuple						=			num_mel_bins
					__snake_case      :							Optional[Any]						=			d_model
					__snake_case      :							Dict						=			encoder_layers
					__snake_case      :							List[str]						=			encoder_attention_heads
					__snake_case      :							Dict						=			decoder_layers
					__snake_case      :							int						=			decoder_attention_heads
					__snake_case      :							int						=			decoder_ffn_dim
					__snake_case      :							Any						=			encoder_ffn_dim
					__snake_case      :							Any						=			dropout
					__snake_case      :							Any						=			attention_dropout
					__snake_case      :							Tuple						=			activation_dropout
					__snake_case      :							Any						=			activation_function
					__snake_case      :							Union[str, Any]						=			init_std
					__snake_case      :							Any						=			encoder_layerdrop
					__snake_case      :							int						=			decoder_layerdrop
					__snake_case      :							List[Any]						=			use_cache
					__snake_case      :							int						=			encoder_layers
					__snake_case      :							Tuple						=			scale_embedding  # scale factor will be sqrt(d_model) if True
					__snake_case      :							Optional[int]						=			max_source_positions
					__snake_case      :							List[str]						=			max_target_positions

					# Audio Classification-specific parameters. Feel free to ignore for other classes.
					__snake_case      :							Union[str, Any]						=			classifier_proj_size
					__snake_case      :							List[str]						=			use_weighted_layer_sum

					# fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
					__snake_case      :							Optional[Any]						=			apply_spec_augment
					__snake_case      :							Dict						=			mask_time_prob
					__snake_case      :							List[Any]						=			mask_time_length
					__snake_case      :							Union[str, Any]						=			mask_time_min_masks
					__snake_case      :							str						=			mask_feature_prob
					__snake_case      :							List[str]						=			mask_feature_length
					__snake_case      :							Optional[int]						=			mask_feature_min_masks

					__snake_case      :							Dict						=			median_filter_width

					super().__init__(
					    pad_token_id=_A     ,  bos_token_id=_A     ,  eos_token_id=_A     ,  is_encoder_decoder=_A     ,  decoder_start_token_id=_A     ,  suppress_tokens=_A     ,  begin_suppress_tokens=_A     ,  **_A     ,  )







class 		UpperCamelCase      (   lowercase						):



				@property
				def  _lowercase					(self    :				Optional[int])   ->     Mapping[str, Mapping[int, str]]:
					__snake_case      :							Any						=			OrderedDict(
					    [
					        ('input_features', {0: 'batch', 1: 'feature_size', 2: 'encoder_sequence'}),
					    ])
					if self.use_past:
						__snake_case      :							List[Any]						=			{0: 'batch'}
					else:
						__snake_case      :							Tuple						=			{0: 'batch', 1: 'decoder_sequence'}

					if self.use_past:
						self.fill_with_past_key_values_(_A     ,  direction='inputs')

					return common_inputs



				def  _lowercase					(self    :				Dict     ,  _A    :				Union["PreTrainedTokenizerBase", "FeatureExtractionMixin"]     ,  _A    :				int = -1     ,  _A    :				int = -1     ,  _A    :				bool = False     ,  _A    :				Optional["TensorType"] = None     ,  _A    :				int = 2_20_50     ,  _A    :				float = 5.0     ,  _A    :				int = 2_20     ,  )   ->     Mapping[str, Any]:
					__snake_case      :							Any						=			OrderedDict()
					__snake_case      :							Optional[int]						=			OnnxConfig.generate_dummy_inputs(
					    self     ,  preprocessor=preprocessor.feature_extractor     ,  batch_size=_A     ,  framework=_A     ,  sampling_rate=_A     ,  time_duration=_A     ,  frequency=_A     ,  )
					__snake_case      :							Optional[Any]						=			encoder_inputs['input_features'].shape[2]
					__snake_case      :							List[str]						=			encoder_sequence_length // 2 if self.use_past else seq_length

					__snake_case      :							str						=			super().generate_dummy_inputs(
					    preprocessor.tokenizer     ,  _A     ,  _A     ,  _A     ,  _A)

					__snake_case      :							List[Any]						=			encoder_inputs.pop('input_features')
					__snake_case      :							Dict						=			decoder_inputs.pop('decoder_input_ids')

					if "past_key_values" in decoder_inputs:
						__snake_case      :							Union[str, Any]						=			decoder_inputs.pop('past_key_values')

					return dummy_inputs






				@property
				def  _lowercase					(self    :				Optional[int])   ->     float:
					return 1E-3






"""simple docstring"""







import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


_a					:					List[str]=			False

class 		UpperCamelCase      (   unittest.TestCase						):



				def  _lowercase					(self    :				List[str])   ->     List[Any]:
					# clean up the VRAM after each test
					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()



				@property
				def  _lowercase					(self    :				Optional[Any])   ->     Union[str, Any]:
					return 12



				@property
				def  _lowercase					(self    :				Dict)   ->     Union[str, Any]:
					return 12



				@property
				def  _lowercase					(self    :				int)   ->     Tuple:
					return 32



				@property
				def  _lowercase					(self    :				Optional[int])   ->     Dict:
					torch.manual_seed(0)
					__snake_case      :							Any						=			VQModel(
					    block_out_channels=[32, 64]     ,  in_channels=3     ,  out_channels=3     ,  down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']     ,  up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']     ,  latent_channels=3     ,  num_vq_embeddings=self.num_embed     ,  vq_embed_dim=3     ,  )
					return model



				@property
				def  _lowercase					(self    :				List[Any])   ->     Optional[int]:
					__snake_case      :							Dict						=			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')
					return tokenizer



				@property
				def  _lowercase					(self    :				Union[str, Any])   ->     Optional[int]:
					torch.manual_seed(0)
					__snake_case      :							Any						=			CLIPTextConfig(
					    bos_token_id=0     ,  eos_token_id=2     ,  hidden_size=self.text_embedder_hidden_size     ,  intermediate_size=37     ,  layer_norm_eps=1E-05     ,  num_attention_heads=4     ,  num_hidden_layers=5     ,  pad_token_id=1     ,  vocab_size=10_00     ,  )
					return CLIPTextModel(_A)



				@property
				def  _lowercase					(self    :				Union[str, Any])   ->     Dict:
					torch.manual_seed(0)

					__snake_case      :							Any						=			12
					__snake_case      :							int						=			12

					__snake_case      :							List[Any]						=			{
					    'attention_bias': True,
					    'cross_attention_dim': 32,
					    'attention_head_dim': height * width,
					    'num_attention_heads': 1,
					    'num_vector_embeds': self.num_embed,
					    'num_embeds_ada_norm': self.num_embeds_ada_norm,
					    'norm_num_groups': 32,
					    'sample_size': width,
					    'activation_fn': 'geglu-approximate',
					}

					__snake_case      :							Union[str, Any]						=			TransformeraDModel(**_A)
					return model



				def  _lowercase					(self    :				Union[str, Any])   ->     Dict:
					__snake_case      :							Tuple						=			'cpu'

					__snake_case      :							List[str]						=			self.dummy_vqvae
					__snake_case      :							str						=			self.dummy_text_encoder
					__snake_case      :							Optional[Any]						=			self.dummy_tokenizer
					__snake_case      :							Dict						=			self.dummy_transformer
					__snake_case      :							Optional[int]						=			VQDiffusionScheduler(self.num_embed)
					__snake_case      :							List[Any]						=			LearnedClassifierFreeSamplingEmbeddings(learnable=_A)

					__snake_case      :							List[Any]						=			VQDiffusionPipeline(
					    vqvae=_A     ,  text_encoder=_A     ,  tokenizer=_A     ,  transformer=_A     ,  scheduler=_A     ,  learned_classifier_free_sampling_embeddings=_A     ,  )
					__snake_case      :							List[Any]						=			pipe.to(_A)
					pipe.set_progress_bar_config(disable=_A)

					__snake_case      :							Optional[Any]						=			'teddy bear playing in the pool'

					__snake_case      :							str						=			torch.Generator(device=_A).manual_seed(0)
					__snake_case      :							Union[str, Any]						=			pipe([prompt]     ,  generator=_A     ,  num_inference_steps=2     ,  output_type='np')
					__snake_case      :							Optional[int]						=			output.images

					__snake_case      :							int						=			torch.Generator(device=_A).manual_seed(0)
					__snake_case      :							Tuple						=			pipe(
					    [prompt]     ,  generator=_A     ,  output_type='np'     ,  return_dict=_A     ,  num_inference_steps=2)[0]

					__snake_case      :							str						=			image[0, -3:, -3:, -1]
					__snake_case      :							Union[str, Any]						=			image_from_tuple[0, -3:, -3:, -1]

					assert image.shape == (1, 24, 24, 3)

					__snake_case      :							str						=			np.array([0.6_551, 0.6_168, 0.5_008, 0.5_676, 0.5_659, 0.4_295, 0.6_073, 0.5_599, 0.4_992])

					assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2
					assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1E-2






				def  _lowercase					(self    :				Tuple)   ->     Optional[int]:
					__snake_case      :							Optional[Any]						=			'cpu'

					__snake_case      :							Optional[int]						=			self.dummy_vqvae
					__snake_case      :							List[str]						=			self.dummy_text_encoder
					__snake_case      :							Optional[int]						=			self.dummy_tokenizer
					__snake_case      :							Optional[Any]						=			self.dummy_transformer
					__snake_case      :							Union[str, Any]						=			VQDiffusionScheduler(self.num_embed)
					__snake_case      :							Optional[int]						=			LearnedClassifierFreeSamplingEmbeddings(
					    learnable=_A     ,  hidden_size=self.text_embedder_hidden_size     ,  length=tokenizer.model_max_length)

					__snake_case      :							Union[str, Any]						=			VQDiffusionPipeline(
					    vqvae=_A     ,  text_encoder=_A     ,  tokenizer=_A     ,  transformer=_A     ,  scheduler=_A     ,  learned_classifier_free_sampling_embeddings=_A     ,  )
					__snake_case      :							Union[str, Any]						=			pipe.to(_A)
					pipe.set_progress_bar_config(disable=_A)

					__snake_case      :							Union[str, Any]						=			'teddy bear playing in the pool'

					__snake_case      :							Optional[int]						=			torch.Generator(device=_A).manual_seed(0)
					__snake_case      :							Tuple						=			pipe([prompt]     ,  generator=_A     ,  num_inference_steps=2     ,  output_type='np')
					__snake_case      :							Optional[Any]						=			output.images

					__snake_case      :							str						=			torch.Generator(device=_A).manual_seed(0)
					__snake_case      :							Dict						=			pipe(
					    [prompt]     ,  generator=_A     ,  output_type='np'     ,  return_dict=_A     ,  num_inference_steps=2)[0]

					__snake_case      :							int						=			image[0, -3:, -3:, -1]
					__snake_case      :							Tuple						=			image_from_tuple[0, -3:, -3:, -1]

					assert image.shape == (1, 24, 24, 3)

					__snake_case      :							Optional[Any]						=			np.array([0.6_693, 0.6_075, 0.4_959, 0.5_701, 0.5_583, 0.4_333, 0.6_171, 0.5_684, 0.4_988])

					assert np.abs(image_slice.flatten() - expected_slice).max() < 2.0
					assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1E-2







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



				def  _lowercase					(self    :				Any)   ->     Union[str, Any]:
					# clean up the VRAM after each test
					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()






				def  _lowercase					(self    :				Tuple)   ->     Optional[int]:
					__snake_case      :							List[Any]						=			load_numpy(
					    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
					    '/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy')

					__snake_case      :							Union[str, Any]						=			VQDiffusionPipeline.from_pretrained('microsoft/vq-diffusion-ithq')
					__snake_case      :							Tuple						=			pipeline.to(_A)
					pipeline.set_progress_bar_config(disable=_A)

					# requires GPU generator for gumbel softmax
					# don't use GPU generator in tests though
					__snake_case      :							Optional[int]						=			torch.Generator(device=_A).manual_seed(0)
					__snake_case      :							int						=			pipeline(
					    'teddy bear playing in the pool'     ,  num_images_per_prompt=1     ,  generator=_A     ,  output_type='np'     ,  )

					__snake_case      :							int						=			output.images[0]

					assert image.shape == (2_56, 2_56, 3)
					assert np.abs(expected_image - image).max() < 2.0





