




import gc
import random
import unittest

import numpy as np
import torch

from diffusers import DDIMScheduler, KandinskyVaaPipeline, KandinskyVaaPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.utils import floats_tensor, 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   __a(					_lowerCAmelCase    ,						unittest.TestCase       ):






    """simple docstring"""





    lowerCAmelCase			      =				KandinskyVaaPipeline
    lowerCAmelCase			      =				[
        "image_embeds",
        "negative_image_embeds",
    ]
    lowerCAmelCase			      =				["image_embeds", "negative_image_embeds"]
    lowerCAmelCase			      =				[
        "generator",
        "height",
        "width",
        "latents",
        "guidance_scale",
        "num_inference_steps",
        "return_dict",
        "guidance_scale",
        "num_images_per_prompt",
        "output_type",
        "return_dict",
    ]
    lowerCAmelCase			      =				False





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





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





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





    @property
    def 							a__				(				self  )			->		int:
         return self.time_input_dim * 4





    @property
    def 							a__				(				self  )			->		Any:
         return 100





    @property
    def 							a__				(				self  )			->		List[str]:
         torch.manual_seed(0  )

         UpperCAmelCase_      :							str       =      {
             '''in_channels''': 4,
             # Out channels is double in channels because predicts mean and variance
             '''out_channels''': 8,
             '''addition_embed_type''': '''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''': '''image_proj''',
             '''cross_attention_dim''': self.cross_attention_dim,
             '''attention_head_dim''': 4,
             '''resnet_time_scale_shift''': '''scale_shift''',
             '''class_embed_type''': None,
         }

         UpperCAmelCase_      :							List[Any]       =      UNetaDConditionModel(**_lowercase  )
         return model





    @property
    def 							a__				(				self  )			->		int:
         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  )			->		Tuple:
         torch.manual_seed(0  )
         UpperCAmelCase_      :							Any       =      VQModel(**self.dummy_movq_kwargs  )
         return model





    def 							a__				(				self  )			->		List[Any]:
         UpperCAmelCase_      :							Optional[Any]       =      self.dummy_unet
         UpperCAmelCase_      :							Tuple       =      self.dummy_movq

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

         UpperCAmelCase_      :							Union[str, Any]       =      {
             '''unet''': unet,
             '''scheduler''': scheduler,
             '''movq''': movq,
         }
         return components





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE=0  )			->		List[str]:
         UpperCAmelCase_      :							Optional[int]       =      floats_tensor((1, self.text_embedder_hidden_size)       ,rng=random.Random(_lowercase  )  ).to(_lowercase  )
         UpperCAmelCase_      :							Any       =      floats_tensor((1, self.text_embedder_hidden_size)       ,rng=random.Random(seed + 1  )  ).to(
             _lowercase  )
         if str(_lowercase  ).startswith('''mps'''  ):
              UpperCAmelCase_      :							int       =      torch.manual_seed(_lowercase  )
         else:
              UpperCAmelCase_      :							Any       =      torch.Generator(device=_lowercase  ).manual_seed(_lowercase  )
         UpperCAmelCase_      :							Any       =      {
             '''image_embeds''': image_embeds,
             '''negative_image_embeds''': negative_image_embeds,
             '''generator''': generator,
             '''height''': 64,
             '''width''': 64,
             '''guidance_scale''': 4.0,
             '''num_inference_steps''': 2,
             '''output_type''': '''np''',
         }
         return inputs





    def 							a__				(				self  )			->		Optional[Any]:
         UpperCAmelCase_      :							Optional[int]       =      '''cpu'''

         UpperCAmelCase_      :							Union[str, Any]       =      self.get_dummy_components()

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

         pipe.set_progress_bar_config(disable=_lowercase  )

         UpperCAmelCase_      :							int       =      pipe(**self.get_dummy_inputs(_lowercase  )  )
         UpperCAmelCase_      :							Tuple       =      output.images

         UpperCAmelCase_      :							str       =      pipe(
             **self.get_dummy_inputs(_lowercase  )       ,return_dict=_lowercase       ,)[0]

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

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

         UpperCAmelCase_      :							Optional[Any]       =      np.array(
             [0.6_23_79_76, 1.0, 0.36_44_13_32, 1.0, 0.70_63_96_34, 0.29_87_71_86, 0.85_65_21_25, 0.5_21_68_43, 0.54_45_40_46]  )

         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()}'''


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






    """simple docstring"""





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





    def 							a__				(				self  )			->		Optional[Any]:
         UpperCAmelCase_      :							List[str]       =      load_numpy(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/kandinskyv22/kandinskyv22_text2img_cat_fp16.npy'''  )

         UpperCAmelCase_      :							int       =      KandinskyVaaPriorPipeline.from_pretrained(
             '''kandinsky-community/kandinsky-2-2-prior'''       ,torch_dtype=torch.floataa  )
         pipe_prior.to(_lowercase  )

         UpperCAmelCase_      :							List[str]       =      KandinskyVaaPipeline.from_pretrained(
             '''kandinsky-community/kandinsky-2-2-decoder'''       ,torch_dtype=torch.floataa  )
         UpperCAmelCase_      :							Tuple       =      pipeline.to(_lowercase  )
         pipeline.set_progress_bar_config(disable=_lowercase  )

         UpperCAmelCase_      :							str       =      '''red cat, 4k photo'''

         UpperCAmelCase_      :							Dict       =      torch.Generator(device='''cuda'''  ).manual_seed(0  )
         UpperCAmelCase_,   UpperCAmelCase_      :							Union[str, Any]       =      pipe_prior(
             _lowercase       ,generator=_lowercase       ,num_inference_steps=5       ,negative_prompt=''''''       ,).to_tuple()

         UpperCAmelCase_      :							List[Any]       =      torch.Generator(device='''cuda'''  ).manual_seed(0  )
         UpperCAmelCase_      :							Union[str, Any]       =      pipeline(
             image_embeds=_lowercase       ,negative_image_embeds=_lowercase       ,generator=_lowercase       ,num_inference_steps=100       ,output_type='''np'''       ,)

         UpperCAmelCase_      :							Optional[Any]       =      output.images[0]

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

         assert_mean_pixel_difference(_lowercase       ,_lowercase  )





from ...processing_utils import ProcessorMixin

class   __a(					_a       ):






    """simple docstring"""





    lowerCAmelCase        =				'''SpeechT5FeatureExtractor'''
    lowerCAmelCase        =				'''SpeechT5Tokenizer'''





    def __init__(				self       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )			->		int:
         super().__init__(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )





    def __call__(				self       ,*_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  )			->		int:
         UpperCAmelCase_      :							List[str]       =      kwargs.pop('''audio'''       ,_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							str       =      kwargs.pop('''text'''       ,_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							List[str]       =      kwargs.pop('''text_target'''       ,_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							int       =      kwargs.pop('''audio_target'''       ,_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Tuple       =      kwargs.pop('''sampling_rate'''       ,_SCREAMING_SNAKE_CASE  )

         if audio is not None and text is not None:
              raise ValueError(
                  '''Cannot process both `audio` and `text` inputs. Did you mean `audio_target` or `text_target`?'''  )
         if audio_target is not None and text_target is not None:
              raise ValueError(
                  '''Cannot process both `audio_target` and `text_target` inputs. Did you mean `audio` or `text`?'''  )
         if audio is None and audio_target is None and text is None and text_target is None:
              raise ValueError(
                  '''You need to specify either an `audio`, `audio_target`, `text`, or `text_target` input to process.'''  )

         if audio is not None:
              UpperCAmelCase_      :							Optional[Any]       =      self.feature_extractor(_SCREAMING_SNAKE_CASE       ,*_SCREAMING_SNAKE_CASE       ,sampling_rate=_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  )
         elif text is not None:
              UpperCAmelCase_      :							Dict       =      self.tokenizer(_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  )
         else:
              UpperCAmelCase_      :							List[str]       =      None

         if audio_target is not None:
              UpperCAmelCase_      :							List[Any]       =      self.feature_extractor(audio_target=_SCREAMING_SNAKE_CASE       ,*_SCREAMING_SNAKE_CASE       ,sampling_rate=_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  )
              UpperCAmelCase_      :							int       =      targets['''input_values''']
         elif text_target is not None:
              UpperCAmelCase_      :							Optional[int]       =      self.tokenizer(_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  )
              UpperCAmelCase_      :							Optional[Any]       =      targets['''input_ids''']
         else:
              UpperCAmelCase_      :							Tuple       =      None

         if inputs is None:
              return targets

         if targets is not None:
              UpperCAmelCase_      :							Dict       =      labels

              UpperCAmelCase_      :							Optional[int]       =      targets.get('''attention_mask'''  )
              if decoder_attention_mask is not None:
                   UpperCAmelCase_      :							List[Any]       =      decoder_attention_mask

         return inputs





    def 							a__				(				self       ,*_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  )			->		Tuple:
         UpperCAmelCase_      :							Dict       =      kwargs.pop('''input_values'''       ,_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Union[str, Any]       =      kwargs.pop('''input_ids'''       ,_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							List[str]       =      kwargs.pop('''labels'''       ,_SCREAMING_SNAKE_CASE  )

         if input_values is not None and input_ids is not None:
              raise ValueError('''Cannot process both `input_values` and `input_ids` inputs.'''  )
         if input_values is None and input_ids is None and labels is None:
              raise ValueError(
                  '''You need to specify either an `input_values`, `input_ids`, or `labels` input to be padded.'''  )

         if input_values is not None:
              UpperCAmelCase_      :							Tuple       =      self.feature_extractor.pad(_SCREAMING_SNAKE_CASE       ,*_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  )
         elif input_ids is not None:
              UpperCAmelCase_      :							Tuple       =      self.tokenizer.pad(_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  )
         else:
              UpperCAmelCase_      :							List[str]       =      None

         if labels is not None:
              if "input_ids" in labels or (isinstance(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  ) and "input_ids" in labels[0]):
                   UpperCAmelCase_      :							Tuple       =      self.tokenizer.pad(_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  )
                   UpperCAmelCase_      :							Any       =      targets['''input_ids''']
              else:
                   UpperCAmelCase_      :							int       =      self.feature_extractor.feature_size
                   UpperCAmelCase_      :							List[str]       =      self.feature_extractor.num_mel_bins
                   UpperCAmelCase_      :							Optional[Any]       =      self.feature_extractor.pad(_SCREAMING_SNAKE_CASE       ,*_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  )
                   UpperCAmelCase_      :							List[Any]       =      feature_size_hack
                   UpperCAmelCase_      :							List[Any]       =      targets['''input_values''']
         else:
              UpperCAmelCase_      :							Optional[int]       =      None

         if inputs is None:
              return targets

         if targets is not None:
              UpperCAmelCase_      :							Optional[int]       =      labels

              UpperCAmelCase_      :							Union[str, Any]       =      targets.get('''attention_mask'''  )
              if decoder_attention_mask is not None:
                   UpperCAmelCase_      :							List[Any]       =      decoder_attention_mask

         return inputs





    def 							a__				(				self       ,*_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  )			->		Optional[Any]:
         return self.tokenizer.batch_decode(*_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  )





    def 							a__				(				self       ,*_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  )			->		Union[str, Any]:
         return self.tokenizer.decode(*_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  )