
"""simple docstring"""




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       __magic_name__							(							SCREAMING_SNAKE_CASE__ ,       unittest.TestCase			):
       UpperCamelCase_ =		KandinskyInpaintPipeline
       UpperCamelCase_ =		['''prompt''', '''image_embeds''', '''negative_image_embeds''', '''image''', '''mask_image''']
       UpperCamelCase_ =		[
           '''prompt''',
           '''negative_prompt''',
           '''image_embeds''',
           '''negative_image_embeds''',
           '''image''',
           '''mask_image''',
       ]
       UpperCamelCase_ =		[
           '''generator''',
           '''height''',
           '''width''',
           '''latents''',
           '''guidance_scale''',
           '''negative_prompt''',
           '''num_inference_steps''',
           '''return_dict''',
           '''guidance_scale''',
           '''num_images_per_prompt''',
           '''output_type''',
           '''return_dict''',
       ]
       UpperCamelCase_ =		False
       @property
       def 				lowercase_				(  self							)  ->						Tuple:





         """simple docstring"""




         return 32
       @property
       def 				lowercase_				(  self							)  ->						Union[str, Any]:





         """simple docstring"""




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





         """simple docstring"""




         return self.time_input_dim
       @property
       def 				lowercase_				(  self							)  ->						Optional[Any]:





         """simple docstring"""




         return self.time_input_dim * 4
       @property
       def 				lowercase_				(  self							)  ->						List[str]:





         """simple docstring"""




         return 100
       @property
       def 				lowercase_				(  self							)  ->						Dict:





         """simple docstring"""




         _lowercase:      Optional[Any]   			=  XLMRobertaTokenizerFast.from_pretrained('''YiYiXu/tiny-random-mclip-base'''							)
         return tokenizer
       @property
       def 				lowercase_				(  self							)  ->						List[Any]:





         """simple docstring"""




         torch.manual_seed(0							)
         _lowercase:      Dict   			=  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=1005			,					)

         _lowercase:      Optional[Any]   			=  MultilingualCLIP(A_							)
         _lowercase:      str   			=  text_encoder.eval()

         return text_encoder
       @property
       def 				lowercase_				(  self							)  ->						Dict:





         """simple docstring"""




         torch.manual_seed(0							)

         _lowercase:      Union[str, Any]   			=  {
             '''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,
         }

         _lowercase:      Dict   			=  UNetaDConditionModel(**A_							)
         return model
       @property
       def 				lowercase_				(  self							)  ->						List[Any]:





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





         """simple docstring"""




         torch.manual_seed(0							)
         _lowercase:      Any   			=  VQModel(**self.dummy_movq_kwargs							)
         return model
       def 				lowercase_				(  self							)  ->						Union[str, Any]:





         """simple docstring"""




         _lowercase:      List[Any]   			=  self.dummy_text_encoder
         _lowercase:      Tuple   			=  self.dummy_tokenizer
         _lowercase:      List[str]   			=  self.dummy_unet
         _lowercase:      Any   			=  self.dummy_movq

         _lowercase:      Optional[Any]   			=  DDIMScheduler(
             num_train_timesteps=1000			,					beta_schedule='''linear'''			,					beta_start=0.0_00_85			,					beta_end=0.0_12			,					clip_sample=A_			,					set_alpha_to_one=A_			,					steps_offset=1			,					prediction_type='''epsilon'''			,					thresholding=A_			,					)

         _lowercase:      Optional[Any]   			=  {
             '''text_encoder''': text_encoder,
             '''tokenizer''': tokenizer,
             '''unet''': unet,
             '''scheduler''': scheduler,
             '''movq''': movq,
         }

         return components
       def 				lowercase_				(  self			,					A_			,					A_=0							)  ->						Dict:





         """simple docstring"""




         _lowercase:      int   			=  floats_tensor((1, self.cross_attention_dim)			,					rng=random.Random(A_							)							).to(A_							)
         _lowercase:      List[Any]   			=  floats_tensor((1, self.cross_attention_dim)			,					rng=random.Random(seed + 1							)							).to(A_							)
         # create init_image
         _lowercase:      Union[str, Any]   			=  floats_tensor((1, 3, 64, 64)			,					rng=random.Random(A_							)							).to(A_							)
         _lowercase:      Tuple   			=  image.cpu().permute(0			,					2			,					3			,					1							)[0]
         _lowercase:      List[Any]   			=  Image.fromarray(np.uinta(A_							)							).convert('''RGB'''							).resize((256, 256)							)
         # create mask
         _lowercase:      Union[str, Any]   			=  np.ones((64, 64)			,					dtype=np.floataa							)
         _lowercase:      Any   			=  0

         if str(A_							).startswith('''mps'''							):
           _lowercase:      int   			=  torch.manual_seed(A_							)
         else:
           _lowercase:      List[Any]   			=  torch.Generator(device=A_							).manual_seed(A_							)
         _lowercase:      Optional[Any]   			=  {
             '''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 				lowercase_				(  self							)  ->						Tuple:





         """simple docstring"""




         _lowercase:      str   			=  '''cpu'''

         _lowercase:      List[str]   			=  self.get_dummy_components()

         _lowercase:      Optional[int]   			=  self.pipeline_class(**A_							)
         _lowercase:      str   			=  pipe.to(A_							)

         pipe.set_progress_bar_config(disable=A_							)

         _lowercase:      Tuple   			=  pipe(**self.get_dummy_inputs(A_							)							)
         _lowercase:      Optional[int]   			=  output.images

         _lowercase:      Tuple   			=  pipe(
             **self.get_dummy_inputs(A_							)			,					return_dict=A_			,					)[0]

         _lowercase:      List[str]   			=  image[0, -3:, -3:, -1]
         _lowercase:      Optional[int]   			=  image_from_tuple[0, -3:, -3:, -1]

         print(f'''image.shape {image.shape}'''							)

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

         _lowercase:      Union[str, Any]   			=  np.array(
             [0.8_32_69_19, 0.73_79_04_67, 0.20_91_85_81, 0.9_30_96_12, 0.5_51_17_91, 0.43_71_33_28, 0.5_51_33_21, 0.49_92_29_34, 0.59_49_77_86]							)

         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 				lowercase_				(  self							)  ->						List[Any]:





         """simple docstring"""




         super().test_inference_batch_single_identical(expected_max_diff=3E-3							)


@slow
@require_torch_gpu
class       __magic_name__							(							unittest.TestCase			):
       def 				lowercase_				(  self							)  ->						Union[str, Any]:





         """simple docstring"""




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



       def 				lowercase_				(  self							)  ->						List[Any]:





         """simple docstring"""




         _lowercase:      Optional[Any]   			=  load_numpy(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/kandinsky/kandinsky_inpaint_cat_with_hat_fp16.npy'''							)

         _lowercase:      List[Any]   			=  load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/kandinsky/cat.png'''							)
         _lowercase:      Any   			=  np.ones((768, 768)			,					dtype=np.floataa							)
         _lowercase:      int   			=  0

         _lowercase:      str   			=  '''a hat'''

         _lowercase:      List[Any]   			=  KandinskyPriorPipeline.from_pretrained(
             '''kandinsky-community/kandinsky-2-1-prior'''			,					torch_dtype=torch.floataa							)
         pipe_prior.to(A_							)

         _lowercase:      List[str]   			=  KandinskyInpaintPipeline.from_pretrained(
             '''kandinsky-community/kandinsky-2-1-inpaint'''			,					torch_dtype=torch.floataa							)
         _lowercase:      Any   			=  pipeline.to(A_							)
         pipeline.set_progress_bar_config(disable=A_							)

         _lowercase:      Dict   			=  torch.Generator(device='''cpu'''							).manual_seed(0							)
         _lowercase			,					_lowercase:      int   			=  pipe_prior(
             A_			,					generator=A_			,					num_inference_steps=5			,					negative_prompt=''''''			,					).to_tuple()

         _lowercase:      Tuple   			=  pipeline(
             A_			,					image=A_			,					mask_image=A_			,					image_embeds=A_			,					negative_image_embeds=A_			,					generator=A_			,					num_inference_steps=100			,					height=768			,					width=768			,					output_type='''np'''			,					)

         _lowercase:      Optional[int]   			=  output.images[0]

         assert image.shape == (768, 768, 3)

         assert_mean_pixel_difference(A_			,					A_							)


"""simple docstring"""




import itertools
import os
import random
import tempfile
import unittest

import numpy as np
from datasets import load_dataset

from transformers import is_speech_available
from transformers.testing_utils import check_json_file_has_correct_format, require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
  from transformers import WhisperFeatureExtractor

if is_torch_available():
  import torch

A__ :       List[Any]				= random.Random()




def 			_lowerCAmelCase		(		_UpperCamelCase			,					_UpperCamelCase=1.0			,					_UpperCamelCase=None			,					_UpperCamelCase=None    ):






  """simple docstring"""



  if rng is None:
    _lowercase:      Any   			=  global_rng

  _lowercase:      List[str]   			=  []
  for batch_idx in range(shape[0]    ):
    values.append([]    )
    for _ in range(shape[1]    ):
      values[-1].append(rng.random() * scale    )

  return values


@require_torch
@require_torchaudio
class       __magic_name__							(							unittest.TestCase			):
       def __init__(  self			,					A_			,					A_=7			,					A_=400			,					A_=2000			,					A_=10			,					A_=160			,					A_=8			,					A_=0.0			,					A_=4000			,					A_=False			,					A_=True			,					)  ->						List[Any]:





         """simple docstring"""




         _lowercase:      str   			=  parent
         _lowercase:      int   			=  batch_size
         _lowercase:      Tuple   			=  min_seq_length
         _lowercase:      Any   			=  max_seq_length
         _lowercase:      List[str]   			=  (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
         _lowercase:      Any   			=  padding_value
         _lowercase:      List[str]   			=  sampling_rate
         _lowercase:      Union[str, Any]   			=  return_attention_mask
         _lowercase:      Optional[Any]   			=  do_normalize
         _lowercase:      List[str]   			=  feature_size
         _lowercase:      Optional[Any]   			=  chunk_length
         _lowercase:      str   			=  hop_length
       def 				lowercase_				(  self							)  ->						Optional[Any]:





         """simple docstring"""




         return {
             "feature_size": self.feature_size,
             "hop_length": self.hop_length,
             "chunk_length": self.chunk_length,
             "padding_value": self.padding_value,
             "sampling_rate": self.sampling_rate,
             "return_attention_mask": self.return_attention_mask,
             "do_normalize": self.do_normalize,
         }



       def 				lowercase_				(  self			,					A_=False			,					A_=False							)  ->						List[Any]:





         """simple docstring"""




         def _flatten(A_							):
           return list(itertools.chain(*A_							)							)

         if equal_length:
           _lowercase:      List[str]   			=  [floats_list((self.max_seq_length, self.feature_size)							) for _ in range(self.batch_size							)]
         else:
           # make sure that inputs increase in size
           _lowercase:      Union[str, Any]   			=  [
               floats_list((x, self.feature_size)							)
               for x in range(self.min_seq_length			,					self.max_seq_length			,					self.seq_length_diff							)
           ]
         if numpify:
           _lowercase:      str   			=  [np.asarray(A_							) for x in speech_inputs]
         return speech_inputs


@require_torch
@require_torchaudio
class       __magic_name__							(							SCREAMING_SNAKE_CASE__ ,       unittest.TestCase			):
       UpperCamelCase_ =		WhisperFeatureExtractor if is_speech_available() else None
       def 				lowercase_				(  self							)  ->						Dict:





         """simple docstring"""




         _lowercase:      int   			=  WhisperFeatureExtractionTester(self							)
       def 				lowercase_				(  self							)  ->						Tuple:





         """simple docstring"""




         _lowercase:      Union[str, Any]   			=  self.feature_extraction_class(**self.feat_extract_dict							)

         with tempfile.TemporaryDirectory() as tmpdirname:
           _lowercase:      Any   			=  feat_extract_first.save_pretrained(A_							)[0]
           check_json_file_has_correct_format(A_							)
           _lowercase:      str   			=  self.feature_extraction_class.from_pretrained(A_							)

         _lowercase:      Dict   			=  feat_extract_first.to_dict()
         _lowercase:      List[str]   			=  feat_extract_second.to_dict()
         _lowercase:      List[str]   			=  feat_extract_first.mel_filters
         _lowercase:      str   			=  feat_extract_second.mel_filters
         self.assertTrue(np.allclose(A_			,					A_							)							)
         self.assertEqual(A_			,					A_							)
       def 				lowercase_				(  self							)  ->						Optional[int]:





         """simple docstring"""




         _lowercase:      Dict   			=  self.feature_extraction_class(**self.feat_extract_dict							)

         with tempfile.TemporaryDirectory() as tmpdirname:
           _lowercase:      Union[str, Any]   			=  os.path.join(A_			,					'''feat_extract.json'''							)
           feat_extract_first.to_json_file(A_							)
           _lowercase:      Tuple   			=  self.feature_extraction_class.from_json_file(A_							)

         _lowercase:      int   			=  feat_extract_first.to_dict()
         _lowercase:      Optional[int]   			=  feat_extract_second.to_dict()
         _lowercase:      Tuple   			=  feat_extract_first.mel_filters
         _lowercase:      Optional[int]   			=  feat_extract_second.mel_filters
         self.assertTrue(np.allclose(A_			,					A_							)							)
         self.assertEqual(A_			,					A_							)
       def 				lowercase_				(  self							)  ->						Tuple:





         """simple docstring"""




         _lowercase:      Any   			=  self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
         # create three inputs of length 800, 1000, and 1200
         _lowercase:      Tuple   			=  [floats_list((1, x)							)[0] for x in range(800			,					1400			,					200							)]
         _lowercase:      List[str]   			=  [np.asarray(A_							) for speech_input in speech_inputs]

         # Test feature size
         _lowercase:      Dict   			=  feature_extractor(A_			,					padding='''max_length'''			,					return_tensors='''np'''							).input_features
         self.assertTrue(input_features.ndim == 3							)
         self.assertTrue(input_features.shape[-1] == feature_extractor.nb_max_frames							)
         self.assertTrue(input_features.shape[-2] == feature_extractor.feature_size							)

         # Test not batched input
         _lowercase:      Dict   			=  feature_extractor(speech_inputs[0]			,					return_tensors='''np'''							).input_features
         _lowercase:      Dict   			=  feature_extractor(np_speech_inputs[0]			,					return_tensors='''np'''							).input_features
         self.assertTrue(np.allclose(A_			,					A_			,					atol=1E-3							)							)

         # Test batched
         _lowercase:      Union[str, Any]   			=  feature_extractor(A_			,					return_tensors='''np'''							).input_features
         _lowercase:      List[Any]   			=  feature_extractor(A_			,					return_tensors='''np'''							).input_features
         for enc_seq_a, enc_seq_a in zip(A_			,					A_							):
           self.assertTrue(np.allclose(A_			,					A_			,					atol=1E-3							)							)

         # Test 2-D numpy arrays are batched.
         _lowercase:      Optional[Any]   			=  [floats_list((1, x)							)[0] for x in (800, 800, 800)]
         _lowercase:      Any   			=  np.asarray(A_							)
         _lowercase:      Tuple   			=  feature_extractor(A_			,					return_tensors='''np'''							).input_features
         _lowercase:      Dict   			=  feature_extractor(A_			,					return_tensors='''np'''							).input_features
         for enc_seq_a, enc_seq_a in zip(A_			,					A_							):
           self.assertTrue(np.allclose(A_			,					A_			,					atol=1E-3							)							)

         # Test truncation required
         _lowercase:      List[str]   			=  [floats_list((1, x)							)[0] for x in range(200			,					(feature_extractor.n_samples + 500)			,					200							)]
         _lowercase:      List[Any]   			=  [np.asarray(A_							) for speech_input in speech_inputs]

         _lowercase:      Any   			=  [x[: feature_extractor.n_samples] for x in speech_inputs]
         _lowercase:      str   			=  [np.asarray(A_							) for speech_input in speech_inputs_truncated]

         _lowercase:      Optional[Any]   			=  feature_extractor(A_			,					return_tensors='''np'''							).input_features
         _lowercase:      List[str]   			=  feature_extractor(A_			,					return_tensors='''np'''							).input_features
         for enc_seq_a, enc_seq_a in zip(A_			,					A_							):
           self.assertTrue(np.allclose(A_			,					A_			,					atol=1E-3							)							)
       def 				lowercase_				(  self							)  ->						Union[str, Any]:





         """simple docstring"""




         import torch

         _lowercase:      List[str]   			=  self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
         _lowercase:      List[Any]   			=  np.random.rand(100			,					32							).astype(np.floataa							)
         _lowercase:      Tuple   			=  np_speech_inputs.tolist()

         for inputs in [py_speech_inputs, np_speech_inputs]:
           _lowercase:      Optional[Any]   			=  feature_extractor.pad([{'''input_features''': inputs}]			,					return_tensors='''np'''							)
           self.assertTrue(np_processed.input_features.dtype == np.floataa							)
           _lowercase:      int   			=  feature_extractor.pad([{'''input_features''': inputs}]			,					return_tensors='''pt'''							)
           self.assertTrue(pt_processed.input_features.dtype == torch.floataa							)
       def 				lowercase_				(  self			,					A_							)  ->						List[Any]:





         """simple docstring"""




         _lowercase:      Optional[Any]   			=  load_dataset('''hf-internal-testing/librispeech_asr_dummy'''			,					'''clean'''			,					split='''validation'''							)
         # automatic decoding with librispeech
         _lowercase:      List[str]   			=  ds.sort('''id'''							).select(range(A_							)							)[:num_samples]['''audio''']

         return [x["array"] for x in speech_samples]
       def 				lowercase_				(  self							)  ->						Optional[int]:





         """simple docstring"""




         _lowercase:      List[str]   			=  torch.tensor(
             [
                 0.11_93, -0.09_46, -0.10_98, -0.01_96, 0.02_25, -0.06_90, -0.17_36, 0.09_51,
                 0.09_71, -0.08_17, -0.07_02, 0.01_62, 0.02_60, 0.00_17, -0.01_92, -0.16_78,
                 0.07_09, -0.18_67, -0.06_55, -0.02_74, -0.02_34, -0.18_84, -0.05_16, -0.05_54,
                 -0.02_74, -0.14_25, -0.14_23, 0.08_37, 0.03_77, -0.08_54
             ]							)
         # fmt: on

         _lowercase:      Optional[Any]   			=  self._load_datasamples(1							)
         _lowercase:      Optional[int]   			=  WhisperFeatureExtractor()
         _lowercase:      Dict   			=  feature_extractor(A_			,					return_tensors='''pt'''							).input_features
         self.assertEqual(input_features.shape			,					(1, 80, 3000)							)
         self.assertTrue(torch.allclose(input_features[0, 0, :30]			,					A_			,					atol=1E-4							)							)



       def 				lowercase_				(  self							)  ->						Union[str, Any]:





         """simple docstring"""




         _lowercase:      Union[str, Any]   			=  self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
         _lowercase:      List[str]   			=  self._load_datasamples(1							)[0]
         _lowercase:      Union[str, Any]   			=  ((audio - audio.min()) / (audio.max() - audio.min())) * 6_5535  # Rescale to [0, 65535] to show issue
         _lowercase:      Any   			=  feat_extract.zero_mean_unit_var_norm([audio]			,					attention_mask=A_							)[0]

         self.assertTrue(np.all(np.mean(A_							) < 1E-3							)							)
         self.assertTrue(np.all(np.abs(np.var(A_							) - 1							) < 1E-3							)							)
