




import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyInpaintPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()





class     SCREAMING_SNAKE_CASE__		(UpperCAmelCase_    ,					unittest.TestCase			):
     lowercase_ :      Optional[Any]				= KandinskyInpaintPipeline
     lowercase_ :      Dict				= ["prompt", "image_embeds", "negative_image_embeds", "image", "mask_image"]
     lowercase_ :      Any				= [
         "prompt",
         "negative_prompt",
         "image_embeds",
         "negative_image_embeds",
         "image",
         "mask_image",
     ]
     lowercase_ :      Tuple				= [
         "generator",
         "height",
         "width",
         "latents",
         "guidance_scale",
         "negative_prompt",
         "num_inference_steps",
         "return_dict",
         "guidance_scale",
         "num_images_per_prompt",
         "output_type",
         "return_dict",
     ]
     lowercase_ :      Dict				= False
     @property
     def        A__      (						self	:		Optional[Any]						):



       """simple docstring"""
       return 32
     @property
     def        A__      (						self	:		int						):



       """simple docstring"""
       return 32
     @property
     def        A__      (						self	:		List[Any]						):



       """simple docstring"""
       return self.time_input_dim
     @property
     def        A__      (						self	:		str						):



       """simple docstring"""
       return self.time_input_dim * 4
     @property
     def        A__      (						self	:		List[str]						):



       """simple docstring"""
       return 1_00
     @property
     def        A__      (						self	:		Any						):



       """simple docstring"""
       lowerCAmelCase__										=      XLMRobertaTokenizerFast.from_pretrained('''YiYiXu/tiny-random-mclip-base'''						)
       return tokenizer
     @property
     def        A__      (						self	:		str						):



       """simple docstring"""
       torch.manual_seed(0						)
       lowerCAmelCase__										=      MCLIPConfig(
           numDims=self.cross_attention_dim       ,						transformerDimensions=self.text_embedder_hidden_size       ,						hidden_size=self.text_embedder_hidden_size       ,						intermediate_size=37       ,						num_attention_heads=4       ,						num_hidden_layers=5       ,						vocab_size=10_05       ,						)

       lowerCAmelCase__										=      MultilingualCLIP(_snake_case						)
       lowerCAmelCase__										=      text_encoder.eval()

       return text_encoder
     @property
     def        A__      (						self	:		Dict						):



       """simple docstring"""
       torch.manual_seed(0						)

       lowerCAmelCase__										=      {
           '''in_channels''': 9,
           # Out channels is double in channels because predicts mean and variance
           '''out_channels''': 8,
           '''addition_embed_type''': '''text_image''',
           '''down_block_types''': ('''ResnetDownsampleBlock2D''', '''SimpleCrossAttnDownBlock2D'''),
           '''up_block_types''': ('''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D'''),
           '''mid_block_type''': '''UNetMidBlock2DSimpleCrossAttn''',
           '''block_out_channels''': (self.block_out_channels_a, self.block_out_channels_a * 2),
           '''layers_per_block''': 1,
           '''encoder_hid_dim''': self.text_embedder_hidden_size,
           '''encoder_hid_dim_type''': '''text_image_proj''',
           '''cross_attention_dim''': self.cross_attention_dim,
           '''attention_head_dim''': 4,
           '''resnet_time_scale_shift''': '''scale_shift''',
           '''class_embed_type''': None,
       }

       lowerCAmelCase__										=      UNetaDConditionModel(**_snake_case						)
       return model
     @property
     def        A__      (						self	:		int						):



       """simple docstring"""
       return {
           "block_out_channels": [32, 64],
           "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
           "in_channels": 3,
           "latent_channels": 4,
           "layers_per_block": 1,
           "norm_num_groups": 8,
           "norm_type": "spatial",
           "num_vq_embeddings": 12,
           "out_channels": 3,
           "up_block_types": [
               "AttnUpDecoderBlock2D",
               "UpDecoderBlock2D",
           ],
           "vq_embed_dim": 4,
       }
     @property
     def        A__      (						self	:		Optional[Any]						):



       """simple docstring"""
       torch.manual_seed(0						)
       lowerCAmelCase__										=      VQModel(**self.dummy_movq_kwargs						)
       return model
     def        A__      (						self	:		List[Any]						):



       """simple docstring"""
       lowerCAmelCase__										=      self.dummy_text_encoder
       lowerCAmelCase__										=      self.dummy_tokenizer
       lowerCAmelCase__										=      self.dummy_unet
       lowerCAmelCase__										=      self.dummy_movq

       lowerCAmelCase__										=      DDIMScheduler(
           num_train_timesteps=10_00       ,						beta_schedule='''linear'''       ,						beta_start=0.0_0085       ,						beta_end=0.012       ,						clip_sample=_snake_case       ,						set_alpha_to_one=_snake_case       ,						steps_offset=1       ,						prediction_type='''epsilon'''       ,						thresholding=_snake_case       ,						)

       lowerCAmelCase__										=      {
           '''text_encoder''': text_encoder,
           '''tokenizer''': tokenizer,
           '''unet''': unet,
           '''scheduler''': scheduler,
           '''movq''': movq,
       }

       return components
     def        A__      (						self	:		str       ,						__lowerCamelCase	:		int       ,						__lowerCamelCase	:		List[str]=0						):



       """simple docstring"""
       lowerCAmelCase__										=      floats_tensor((1, self.cross_attention_dim)       ,						rng=random.Random(_snake_case						)						).to(_snake_case						)
       lowerCAmelCase__										=      floats_tensor((1, self.cross_attention_dim)       ,						rng=random.Random(seed + 1						)						).to(_snake_case						)
       # create init_image
       lowerCAmelCase__										=      floats_tensor((1, 3, 64, 64)       ,						rng=random.Random(_snake_case						)						).to(_snake_case						)
       lowerCAmelCase__										=      image.cpu().permute(0       ,						2       ,						3       ,						1						)[0]
       lowerCAmelCase__										=      Image.fromarray(np.uinta(_snake_case						)						).convert('''RGB'''						).resize((2_56, 2_56)						)
       # create mask
       lowerCAmelCase__										=      np.ones((64, 64)       ,						dtype=np.floataa						)
       lowerCAmelCase__										=      0

       if str(_snake_case						).startswith('''mps'''						):
         lowerCAmelCase__										=      torch.manual_seed(_snake_case						)
       else:
         lowerCAmelCase__										=      torch.Generator(device=_snake_case						).manual_seed(_snake_case						)
       lowerCAmelCase__										=      {
           '''prompt''': '''horse''',
           '''image''': init_image,
           '''mask_image''': mask,
           '''image_embeds''': image_embeds,
           '''negative_image_embeds''': negative_image_embeds,
           '''generator''': generator,
           '''height''': 64,
           '''width''': 64,
           '''num_inference_steps''': 2,
           '''guidance_scale''': 4.0,
           '''output_type''': '''np''',
       }
       return inputs
     def        A__      (						self	:		List[str]						):



       """simple docstring"""
       lowerCAmelCase__										=      '''cpu'''

       lowerCAmelCase__										=      self.get_dummy_components()

       lowerCAmelCase__										=      self.pipeline_class(**_snake_case						)
       lowerCAmelCase__										=      pipe.to(_snake_case						)

       pipe.set_progress_bar_config(disable=_snake_case						)

       lowerCAmelCase__										=      pipe(**self.get_dummy_inputs(_snake_case						)						)
       lowerCAmelCase__										=      output.images

       lowerCAmelCase__										=      pipe(
           **self.get_dummy_inputs(_snake_case						)       ,						return_dict=_snake_case       ,						)[0]

       lowerCAmelCase__										=      image[0, -3:, -3:, -1]
       lowerCAmelCase__										=      image_from_tuple[0, -3:, -3:, -1]

       print(F"""image.shape {image.shape}"""						)

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

       lowerCAmelCase__										=      np.array(
           [0.832_6919, 0.7379_0467, 0.2091_8581, 0.930_9612, 0.551_1791, 0.4371_3328, 0.551_3321, 0.4992_2934, 0.5949_7786]						)

       assert (
           np.abs(image_slice.flatten() - expected_slice						).max() < 1e-2
       ), F""" expected_slice {expected_slice}, but got {image_slice.flatten()}"""
       assert (
           np.abs(image_from_tuple_slice.flatten() - expected_slice						).max() < 1e-2
       ), F""" expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}"""



     def        A__      (						self	:		Dict						):



       """simple docstring"""
       super().test_inference_batch_single_identical(expected_max_diff=3e-3						)





@slow
@require_torch_gpu
class     SCREAMING_SNAKE_CASE__		(unittest.TestCase			):
     def        A__      (						self	:		Optional[int]						):



       """simple docstring"""
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()



     def        A__      (						self	:		Union[str, Any]						):



       """simple docstring"""
       lowerCAmelCase__										=      load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/kandinsky/kandinsky_inpaint_cat_with_hat_fp16.npy'''						)

       lowerCAmelCase__										=      load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/kandinsky/cat.png'''						)
       lowerCAmelCase__										=      np.ones((7_68, 7_68)       ,						dtype=np.floataa						)
       lowerCAmelCase__										=      0

       lowerCAmelCase__										=      '''a hat'''

       lowerCAmelCase__										=      KandinskyPriorPipeline.from_pretrained(
           '''kandinsky-community/kandinsky-2-1-prior'''       ,						torch_dtype=torch.floataa						)
       pipe_prior.to(_snake_case						)

       lowerCAmelCase__										=      KandinskyInpaintPipeline.from_pretrained(
           '''kandinsky-community/kandinsky-2-1-inpaint'''       ,						torch_dtype=torch.floataa						)
       lowerCAmelCase__										=      pipeline.to(_snake_case						)
       pipeline.set_progress_bar_config(disable=_snake_case						)

       lowerCAmelCase__										=      torch.Generator(device='''cpu'''						).manual_seed(0						)
       lowerCAmelCase__				,      lowerCAmelCase__										=      pipe_prior(
           _snake_case       ,						generator=_snake_case       ,						num_inference_steps=5       ,						negative_prompt=''''''       ,						).to_tuple()

       lowerCAmelCase__										=      pipeline(
           _snake_case       ,						image=_snake_case       ,						mask_image=_snake_case       ,						image_embeds=_snake_case       ,						negative_image_embeds=_snake_case       ,						generator=_snake_case       ,						num_inference_steps=1_00       ,						height=7_68       ,						width=7_68       ,						output_type='''np'''       ,						)

       lowerCAmelCase__										=      output.images[0]

       assert image.shape == (7_68, 7_68, 3)

       assert_mean_pixel_difference(_snake_case       ,						_snake_case						)






import warnings
from typing import Any, Dict, List, Optional, Union

import numpy as np

from ...audio_utils import mel_filter_bank, optimal_fft_length, spectrogram, window_function
from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
from ...feature_extraction_utils import BatchFeature
from ...utils import PaddingStrategy, TensorType, logging


SCREAMING_SNAKE_CASE__   			=   logging.get_logger(__name__)


class   __lowerCAmelCase  (		UpperCAmelCase_   ):
					"""simple docstring"""






					A__					:		Optional[Any]      							=     ["input_values", "attention_mask"]



					def __init__(  self       :							str					,       _snake_case       :							int = 1					,       _snake_case       :							int = 1_60_00					,       _snake_case       :							float = 0.0					,       _snake_case       :							bool = False					,       _snake_case       :							int = 80					,       _snake_case       :							int = 16					,       _snake_case       :							int = 64					,       _snake_case       :							str = "hann_window"					,       _snake_case       :							float = 1.0					,       _snake_case       :							float = 80					,       _snake_case       :							float = 76_00					,       _snake_case       :							float = 1E-10					,       _snake_case       :							int = 2					,       _snake_case       :							bool = True					,       **_snake_case       :							Union[str, Any]					,       ):






								"""simple docstring"""


								super().__init__(feature_size=_snake_case					,       sampling_rate=_snake_case					,       padding_value=_snake_case					,       **_snake_case       )
								A__	  =     do_normalize
								A__	  =     return_attention_mask

								A__	  =     num_mel_bins
								A__	  =     hop_length
								A__	  =     win_length
								A__	  =     win_function
								A__	  =     frame_signal_scale
								A__	  =     fmin
								A__	  =     fmax
								A__	  =     mel_floor
								A__	  =     reduction_factor

								A__	  =     win_length * sampling_rate // 10_00
								A__	  =     hop_length * sampling_rate // 10_00
								A__	  =     optimal_fft_length(self.sample_size       )
								A__	  =     (self.n_fft // 2) + 1

								A__	  =     window_function(window_length=self.sample_size					,       name=self.win_function					,       periodic=_snake_case       )

								A__	  =     mel_filter_bank(
								    num_frequency_bins=self.n_freqs					,       num_mel_filters=self.num_mel_bins					,       min_frequency=self.fmin					,       max_frequency=self.fmax					,       sampling_rate=self.sampling_rate					,       norm='slaney'					,       mel_scale='slaney'					,       )

								if frame_signal_scale != 1.0:
											warnings.warn(
											    'The argument `frame_signal_scale` is deprecated and will be removed in version 4.30.0 of Transformers'					,       _snake_case					,       )
								if reduction_factor != 2.0:
											warnings.warn(
											    'The argument `reduction_factor` is deprecated and will be removed in version 4.30.0 of Transformers'					,       _snake_case					,       )



					@staticmethod
					# Copied from transformers.models.wav2vec2.feature_extraction_wav2vec2.Wav2Vec2FeatureExtractor.zero_mean_unit_var_norm
					def      _a				(  _snake_case       :							List[np.ndarray]					,       _snake_case       :							List[np.ndarray]					,       _snake_case       :							float = 0.0       ):






								"""simple docstring"""


								if attention_mask is not None:
											A__	  =     np.array(_snake_case					,       np.intaa       )
											A__	  =     []

											for vector, length in zip(_snake_case					,       attention_mask.sum(-1       )       ):
														A__	  =     (vector - vector[:length].mean()) / np.sqrt(vector[:length].var() + 1E-7       )
														if length < normed_slice.shape[0]:
																	A__	  =     padding_value

														normed_input_values.append(_snake_case       )
								else:
											A__	  =     [(x - x.mean()) / np.sqrt(x.var() + 1E-7       ) for x in input_values]

								return normed_input_values



					def      _a				(  self       :							Tuple					,       _snake_case       :							np.ndarray					,       ):






								"""simple docstring"""


								A__	  =     spectrogram(
								    _snake_case					,       window=self.window					,       frame_length=self.sample_size					,       hop_length=self.sample_stride					,       fft_length=self.n_fft					,       mel_filters=self.mel_filters					,       mel_floor=self.mel_floor					,       log_mel='log10'					,       )
								return log_mel_spec.T



					def __call__(  self       :							List[str]					,       _snake_case       :							Optional[Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]] = None					,       _snake_case       :							Optional[Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]] = None					,       _snake_case       :							Union[bool, str, PaddingStrategy] = False					,       _snake_case       :							Optional[int] = None					,       _snake_case       :							bool = False					,       _snake_case       :							Optional[int] = None					,       _snake_case       :							Optional[bool] = None					,       _snake_case       :							Optional[Union[str, TensorType]] = None					,       _snake_case       :							Optional[int] = None					,       **_snake_case       :							Tuple					,       ):






								"""simple docstring"""


								if audio is None and audio_target is None:
											raise ValueError('You must provide either `audio` or `audio_target` values.'       )

								if sampling_rate is not None:
											if sampling_rate != self.sampling_rate:
														raise ValueError(
														    F'''The model corresponding to this feature extractor: {self} was trained using a sampling rate of'''
														    F''' {self.sampling_rate}. Please make sure that the provided audio input was sampled with'''
														    F''' {self.sampling_rate} and not {sampling_rate}.'''       )
								else:
											logger.warning(
											    'It is strongly recommended to pass the ``sampling_rate`` argument to this function. '
											    'Failing to do so can result in silent errors that might be hard to debug.'       )

								if audio is not None:
											A__	  =     self._process_audio(
											    _snake_case					,       _snake_case					,       _snake_case					,       _snake_case					,       _snake_case					,       _snake_case					,       _snake_case					,       _snake_case					,       **_snake_case					,       )
								else:
											A__	  =     None

								if audio_target is not None:
											A__	  =     self._process_audio(
											    _snake_case					,       _snake_case					,       _snake_case					,       _snake_case					,       _snake_case					,       _snake_case					,       _snake_case					,       _snake_case					,       **_snake_case					,       )

											if inputs is None:
														return inputs_target
											else:
														A__	  =     inputs_target['input_values']
														A__	  =     inputs_target.get('attention_mask'       )
														if decoder_attention_mask is not None:
																	A__	  =     decoder_attention_mask

								return inputs



					def      _a				(  self       :							Tuple					,       _snake_case       :							Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]					,       _snake_case       :							bool = False					,       _snake_case       :							Union[bool, str, PaddingStrategy] = False					,       _snake_case       :							Optional[int] = None					,       _snake_case       :							bool = False					,       _snake_case       :							Optional[int] = None					,       _snake_case       :							Optional[bool] = None					,       _snake_case       :							Optional[Union[str, TensorType]] = None					,       **_snake_case       :							Tuple					,       ):






								"""simple docstring"""


								A__	  =     isinstance(_snake_case					,       np.ndarray       ) and len(speech.shape       ) > 1
								if is_batched_numpy and len(speech.shape       ) > 2:
											raise ValueError(F'''Only mono-channel audio is supported for input to {self}'''       )
								A__	  =     is_batched_numpy or (
								    isinstance(_snake_case					,       (list, tuple)       ) and (isinstance(speech[0]					,       (np.ndarray, tuple, list)       ))
								)

								if is_batched:
											A__	  =     [np.asarray(_snake_case					,       dtype=np.floataa       ) for speech in speech]
								elif not is_batched and not isinstance(_snake_case					,       np.ndarray       ):
											A__	  =     np.asarray(_snake_case					,       dtype=np.floataa       )
								elif isinstance(_snake_case					,       np.ndarray       ) and speech.dtype is np.dtype(np.floataa       ):
											A__	  =     speech.astype(np.floataa       )

								# always return batch
								if not is_batched:
											A__	  =     [speech]

								# needed to make pad() work on spectrogram inputs
								A__	  =     self.feature_size

								# convert into correct format for padding
								if is_target:
											A__	  =     [self._extract_mel_features(_snake_case       ) for waveform in speech]
											A__	  =     BatchFeature({'input_values': features}       )
											A__	  =     self.num_mel_bins
								else:
											A__	  =     BatchFeature({'input_values': speech}       )

								A__	  =     self.pad(
								    _snake_case					,       padding=_snake_case					,       max_length=_snake_case					,       truncation=_snake_case					,       pad_to_multiple_of=_snake_case					,       return_attention_mask=_snake_case					,       **_snake_case					,       )

								A__	  =     feature_size_hack

								# convert input values to correct format
								A__	  =     padded_inputs['input_values']
								if not isinstance(input_values[0]					,       np.ndarray       ):
											A__	  =     [np.asarray(_snake_case					,       dtype=np.floataa       ) for array in input_values]
								elif (
								    not isinstance(_snake_case					,       np.ndarray       )
								    and isinstance(input_values[0]					,       np.ndarray       )
								    and input_values[0].dtype is np.dtype(np.floataa       )
								):
											A__	  =     [array.astype(np.floataa       ) for array in input_values]
								elif isinstance(_snake_case					,       np.ndarray       ) and input_values.dtype is np.dtype(np.floataa       ):
											A__	  =     input_values.astype(np.floataa       )

								# convert attention_mask to correct format
								A__	  =     padded_inputs.get('attention_mask'       )
								if attention_mask is not None:
											A__	  =     [np.asarray(_snake_case					,       dtype=np.intaa       ) for array in attention_mask]

								# zero-mean and unit-variance normalization
								if not is_target and self.do_normalize:
											A__	  =     (
											    attention_mask
											    if self._get_padding_strategies(_snake_case					,       max_length=_snake_case       ) is not PaddingStrategy.DO_NOT_PAD
											    else None
											)
											A__	  =     self.zero_mean_unit_var_norm(
											    padded_inputs['input_values']					,       attention_mask=_snake_case					,       padding_value=self.padding_value       )

								if return_tensors is not None:
											A__	  =     padded_inputs.convert_to_tensors(_snake_case       )

								return padded_inputs



					def      _a				(  self       :							Optional[Any]       ):






								"""simple docstring"""


								A__	  =     super().to_dict()

								# Don't serialize these as they are derived from the other properties.
								A__	  =     ['window', 'mel_filters', 'sample_size', 'sample_stride', 'n_fft', 'n_freqs']
								for name in names:
											if name in output:
														del output[name]

								return output




