



from __future__ import annotations

import copy
import inspect
import json
import math
import os
import tempfile
import unittest
from importlib import import_module

import numpy as np

from transformers import ViTMAEConfig
from transformers.file_utils import cached_property, is_tf_available, is_vision_available
from transformers.testing_utils import require_tf, require_vision, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
    import tensorflow as tf

    from transformers import TFViTMAEForPreTraining, TFViTMAEModel


if is_vision_available():
    from PIL import Image

    from transformers import ViTImageProcessor



class 							lowerCAmelCase_						:
       def __init__(     self,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_=13,     SCREAMING_SNAKE_CASE_=30,     SCREAMING_SNAKE_CASE_=2,     SCREAMING_SNAKE_CASE_=3,     SCREAMING_SNAKE_CASE_=True,     SCREAMING_SNAKE_CASE_=True,     SCREAMING_SNAKE_CASE_=32,     SCREAMING_SNAKE_CASE_=2,     SCREAMING_SNAKE_CASE_=4,     SCREAMING_SNAKE_CASE_=37,     SCREAMING_SNAKE_CASE_="gelu",     SCREAMING_SNAKE_CASE_=0.1,     SCREAMING_SNAKE_CASE_=0.1,     SCREAMING_SNAKE_CASE_=10,     SCREAMING_SNAKE_CASE_=0.02,     SCREAMING_SNAKE_CASE_=3,     SCREAMING_SNAKE_CASE_=0.6,     SCREAMING_SNAKE_CASE_=None,     )	->						List[str]:
          UpperCamelCase       :    Union[str, Any]		      =      parent
          UpperCamelCase       :    Optional[Any]		      =      batch_size
          UpperCamelCase       :    Optional[int]		      =      image_size
          UpperCamelCase       :    Union[str, Any]		      =      patch_size
          UpperCamelCase       :    List[Any]		      =      num_channels
          UpperCamelCase       :    Any		      =      is_training
          UpperCamelCase       :    Tuple		      =      use_labels
          UpperCamelCase       :    List[Any]		      =      hidden_size
          UpperCamelCase       :    Union[str, Any]		      =      num_hidden_layers
          UpperCamelCase       :    Optional[int]		      =      num_attention_heads
          UpperCamelCase       :    Tuple		      =      intermediate_size
          UpperCamelCase       :    Optional[int]		      =      hidden_act
          UpperCamelCase       :    Tuple		      =      hidden_dropout_prob
          UpperCamelCase       :    List[str]		      =      attention_probs_dropout_prob
          UpperCamelCase       :    str		      =      type_sequence_label_size
          UpperCamelCase       :    Optional[int]		      =      initializer_range
          UpperCamelCase       :    Tuple		      =      mask_ratio
          UpperCamelCase       :    Optional[int]		      =      scope

          # in ViTMAE, the expected sequence length = (num_patches + 1) * (1 - config.mask_ratio), rounded above
          # (we add 1 for the [CLS] token)
          UpperCamelCase       :    Any		      =      (image_size // patch_size) ** 2
          UpperCamelCase       :    Union[str, Any]		      =      int(math.ceil((1 - mask_ratio) * (num_patches + 1)			)			)
       def    snake_case_			(     self			)	->						str:
          UpperCamelCase       :    List[str]		      =      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]			)

          UpperCamelCase       :    Optional[Any]		      =      None
          if self.use_labels:
             UpperCamelCase       :    Any		      =      ids_tensor([self.batch_size],     self.type_sequence_label_size			)

          UpperCamelCase       :    Any		      =      self.get_config()

          return config, pixel_values, labels
       def    snake_case_			(     self			)	->						Optional[Any]:
          return ViTMAEConfig(
              image_size=self.image_size,     patch_size=self.patch_size,     num_channels=self.num_channels,     hidden_size=self.hidden_size,     num_hidden_layers=self.num_hidden_layers,     num_attention_heads=self.num_attention_heads,     intermediate_size=self.intermediate_size,     decoder_hidden_size=self.hidden_size,     decoder_num_hidden_layers=self.num_hidden_layers,     decoder_num_attention_heads=self.num_attention_heads,     decoder_intermediate_size=self.intermediate_size,     hidden_act=self.hidden_act,     hidden_dropout_prob=self.hidden_dropout_prob,     attention_probs_dropout_prob=self.attention_probs_dropout_prob,     is_decoder=__lowerCamelCase,     initializer_range=self.initializer_range,     mask_ratio=self.mask_ratio,     )
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)	->						Union[str, Any]:
          UpperCamelCase       :    str		      =      TFViTMAEModel(config=__lowerCamelCase			)
          UpperCamelCase       :    List[str]		      =      model(__lowerCamelCase,     training=__lowerCamelCase			)
          self.parent.assertEqual(result.last_hidden_state.shape,     (self.batch_size, self.seq_length, self.hidden_size)			)
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)	->						Tuple:
          UpperCamelCase       :    List[Any]		      =      TFViTMAEForPreTraining(__lowerCamelCase			)
          UpperCamelCase       :    Tuple		      =      model(__lowerCamelCase,     training=__lowerCamelCase			)
          # expected sequence length = num_patches
          UpperCamelCase       :    List[str]		      =      (self.image_size // self.patch_size) ** 2
          UpperCamelCase       :    str		      =      self.patch_size**2 * self.num_channels
          self.parent.assertEqual(result.logits.shape,     (self.batch_size, num_patches, expected_num_channels)			)

          # test greyscale images
          UpperCamelCase       :    Tuple		      =      1
          UpperCamelCase       :    Tuple		      =      TFViTMAEForPreTraining(__lowerCamelCase			)

          UpperCamelCase       :    Any		      =      floats_tensor([self.batch_size, 1, self.image_size, self.image_size]			)
          UpperCamelCase       :    Union[str, Any]		      =      model(__lowerCamelCase,     training=__lowerCamelCase			)
          UpperCamelCase       :    Any		      =      self.patch_size**2
          self.parent.assertEqual(result.logits.shape,     (self.batch_size, num_patches, expected_num_channels)			)






       def    snake_case_			(     self			)	->						Union[str, Any]:
          UpperCamelCase       :    Tuple		      =      self.prepare_config_and_inputs()
          (UpperCamelCase)       :    List[Any]		      =      config_and_inputs
          UpperCamelCase       :    List[Any]		      =      {'''pixel_values''': pixel_values}
          return config, inputs_dict






@require_tf
class 							lowerCAmelCase_						(    SCREAMING_SNAKE_CASE__			,       SCREAMING_SNAKE_CASE__			,       unittest.TestCase				):
       UpperCAmelCase__							:   Any  	=    (TFViTMAEModel, TFViTMAEForPreTraining) if is_tf_available() else ()
       UpperCAmelCase__							:   Optional[Any]  	=    {"feature-extraction": TFViTMAEModel} if is_tf_available() else {}

       UpperCAmelCase__							:   List[str]  	=    False
       UpperCAmelCase__							:   Optional[Any]  	=    False
       UpperCAmelCase__							:   Optional[int]  	=    False
       UpperCAmelCase__							:   Optional[int]  	=    False
       def    snake_case_			(     self			)	->						List[str]:
          UpperCamelCase       :    Union[str, Any]		      =      TFViTMAEModelTester(self			)
          UpperCamelCase       :    Any		      =      ConfigTester(self,     config_class=__lowerCamelCase,     has_text_modality=__lowerCamelCase,     hidden_size=37			)
       def    snake_case_			(     self			)	->						Tuple:
          self.config_tester.run_common_tests()
       @unittest.skip(reason='ViTMAE does not use inputs_embeds'			)
       def    snake_case_			(     self			)	->						Any:
          pass
       def    snake_case_			(     self			)	->						Union[str, Any]:
          UpperCamelCase       :    Optional[int]		      =      self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
             UpperCamelCase       :    List[Any]		      =      model_class(__lowerCamelCase			)
             self.assertIsInstance(model.get_input_embeddings(),     (tf.keras.layers.Layer)			)
             UpperCamelCase       :    List[Any]		      =      model.get_output_embeddings()
             self.assertTrue(x is None or isinstance(__lowerCamelCase,     tf.keras.layers.Layer			)			)
       def    snake_case_			(     self			)	->						str:
          UpperCamelCase       :    int		      =      self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
             UpperCamelCase       :    Tuple		      =      model_class(__lowerCamelCase			)
             UpperCamelCase       :    List[str]		      =      inspect.signature(model.call			)
             # signature.parameters is an OrderedDict => so arg_names order is deterministic
             UpperCamelCase       :    Optional[int]		      =      [*signature.parameters.keys()]

             UpperCamelCase       :    Tuple		      =      ['''pixel_values''']
             self.assertListEqual(arg_names[:1],     __lowerCamelCase			)
       def    snake_case_			(     self			)	->						Optional[int]:
          UpperCamelCase       :    Any		      =      self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*__lowerCamelCase			)
       def    snake_case_			(     self			)	->						str:
          UpperCamelCase       :    Union[str, Any]		      =      self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_pretraining(*__lowerCamelCase			)
       def    snake_case_			(     self			)	->						Any:
          np.random.seed(2			)

          UpperCamelCase       :    Dict		      =      self.model_tester.prepare_config_and_inputs_for_common()

          UpperCamelCase       :    int		      =      int((config.image_size // config.patch_size) ** 2			)
          UpperCamelCase       :    Optional[int]		      =      np.random.uniform(size=(self.model_tester.batch_size, num_patches)			)

          for model_class in self.all_model_classes:
             UpperCamelCase       :    Optional[Any]		      =      model_class(__lowerCamelCase			)
             UpperCamelCase       :    Any		      =      self._prepare_for_class(__lowerCamelCase,     __lowerCamelCase			)

             UpperCamelCase       :    Tuple		      =      model(__lowerCamelCase,     noise=__lowerCamelCase			)

             UpperCamelCase       :    str		      =      copy.deepcopy(self._prepare_for_class(__lowerCamelCase,     __lowerCamelCase			)			)
             UpperCamelCase       :    int		      =      model(**__lowerCamelCase,     noise=__lowerCamelCase			)
             UpperCamelCase       :    Optional[int]		      =      outputs_dict[0].numpy()
             UpperCamelCase       :    List[str]		      =      outputs_keywords[0].numpy()

             self.assertLess(np.sum(np.abs(output_dict - output_keywords			)			),     1e-6			)
       def    snake_case_			(     self			)	->						int:
          np.random.seed(2			)

          UpperCamelCase       :    Optional[Any]		      =      self.model_tester.prepare_config_and_inputs_for_common()

          UpperCamelCase       :    int		      =      int((config.image_size // config.patch_size) ** 2			)
          UpperCamelCase       :    Tuple		      =      np.random.uniform(size=(self.model_tester.batch_size, num_patches)			)

          def prepare_numpy_arrays(SCREAMING_SNAKE_CASE_			):
             UpperCamelCase       :    Optional[Any]		      =      {}
             for k, v in inputs_dict.items():
                if tf.is_tensor(__lowerCamelCase			):
                   UpperCamelCase       :    str		      =      v.numpy()
                else:
                   UpperCamelCase       :    str		      =      np.array(__lowerCamelCase			)

             return inputs_np_dict

          for model_class in self.all_model_classes:
             UpperCamelCase       :    Any		      =      model_class(__lowerCamelCase			)

             UpperCamelCase       :    List[Any]		      =      self._prepare_for_class(__lowerCamelCase,     __lowerCamelCase			)
             UpperCamelCase       :    List[Any]		      =      prepare_numpy_arrays(__lowerCamelCase			)

             UpperCamelCase       :    Dict		      =      model(__lowerCamelCase,     noise=__lowerCamelCase			)
             UpperCamelCase       :    Union[str, Any]		      =      model(**__lowerCamelCase,     noise=__lowerCamelCase			)
             self.assert_outputs_same(__lowerCamelCase,     __lowerCamelCase			)
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)	->						Dict:
          np.random.seed(2			)

          UpperCamelCase       :    Union[str, Any]		      =      int((tf_model.config.image_size // tf_model.config.patch_size) ** 2			)
          UpperCamelCase       :    Optional[int]		      =      np.random.uniform(size=(self.model_tester.batch_size, num_patches)			)
          UpperCamelCase       :    Union[str, Any]		      =      tf.constant(__lowerCamelCase			)

          # Add `noise` argument.
          # PT inputs will be prepared in `super().check_pt_tf_models()` with this added `noise` argument
          UpperCamelCase       :    Optional[int]		      =      tf_noise

          super().check_pt_tf_models(__lowerCamelCase,     __lowerCamelCase,     __lowerCamelCase			)
       def    snake_case_			(     self			)	->						str:
          np.random.seed(2			)

          UpperCamelCase       :    Optional[Any]		      =      self.model_tester.prepare_config_and_inputs_for_common()

          UpperCamelCase       :    str		      =      {
              module_member
              for model_class in self.all_model_classes
              for module in (import_module(model_class.__module__			),)
              for module_member_name in dir(__lowerCamelCase			)
              if module_member_name.endswith('MainLayer'			)
              # This condition is required, since `modeling_tf_clip.py` has 3 classes whose names end with `MainLayer`.
              and module_member_name[: -len('MainLayer'			)] == model_class.__name__[: -len('Model'			)]
              for module_member in (getattr(__lowerCamelCase,     __lowerCamelCase			),)
              if isinstance(__lowerCamelCase,     __lowerCamelCase			)
              and tf.keras.layers.Layer in module_member.__bases__
              and getattr(__lowerCamelCase,     '_keras_serializable',     __lowerCamelCase			)
          }

          UpperCamelCase       :    Optional[int]		      =      int((config.image_size // config.patch_size) ** 2			)
          UpperCamelCase       :    str		      =      np.random.uniform(size=(self.model_tester.batch_size, num_patches)			)
          UpperCamelCase       :    Optional[int]		      =      tf.convert_to_tensor(__lowerCamelCase			)
          inputs_dict.update({'noise': noise}			)

          for main_layer_class in tf_main_layer_classes:
             UpperCamelCase       :    Optional[int]		      =      main_layer_class(__lowerCamelCase			)

             UpperCamelCase       :    List[str]		      =      {
                 name: tf.keras.Input(tensor.shape[1:],     dtype=tensor.dtype			) for name, tensor in inputs_dict.items()
             }

             UpperCamelCase       :    Dict		      =      tf.keras.Model(__lowerCamelCase,     outputs=main_layer(__lowerCamelCase			)			)
             UpperCamelCase       :    List[Any]		      =      model(__lowerCamelCase			)

             with tempfile.TemporaryDirectory() as tmpdirname:
                UpperCamelCase       :    Any		      =      os.path.join(__lowerCamelCase,     'keras_model.h5'			)
                model.save(__lowerCamelCase			)
                UpperCamelCase       :    List[str]		      =      tf.keras.models.load_model(
                    __lowerCamelCase,     custom_objects={main_layer_class.__name__: main_layer_class}			)
                assert isinstance(__lowerCamelCase,     tf.keras.Model			)
                UpperCamelCase       :    int		      =      model(__lowerCamelCase			)
                self.assert_outputs_same(__lowerCamelCase,     __lowerCamelCase			)
       @slow
       def    snake_case_			(     self			)	->						List[Any]:
          np.random.seed(2			)

          UpperCamelCase       :    List[str]		      =      self.model_tester.prepare_config_and_inputs_for_common()

          UpperCamelCase       :    Any		      =      int((config.image_size // config.patch_size) ** 2			)
          UpperCamelCase       :    Union[str, Any]		      =      np.random.uniform(size=(self.model_tester.batch_size, num_patches)			)

          for model_class in self.all_model_classes:
             UpperCamelCase       :    Union[str, Any]		      =      model_class(__lowerCamelCase			)
             UpperCamelCase       :    List[str]		      =      self._prepare_for_class(__lowerCamelCase,     __lowerCamelCase			)
             UpperCamelCase       :    int		      =      model(__lowerCamelCase,     noise=__lowerCamelCase			)

             if model_class.__name__ == "TFViTMAEModel":
                UpperCamelCase       :    Dict		      =      outputs.last_hidden_state.numpy()
                UpperCamelCase       :    Union[str, Any]		      =      0
             else:
                UpperCamelCase       :    List[Any]		      =      outputs.logits.numpy()
                UpperCamelCase       :    Optional[Any]		      =      0

             with tempfile.TemporaryDirectory() as tmpdirname:
                model.save_pretrained(__lowerCamelCase,     saved_model=__lowerCamelCase			)
                UpperCamelCase       :    List[str]		      =      model_class.from_pretrained(__lowerCamelCase			)
                UpperCamelCase       :    Any		      =      model(__lowerCamelCase,     noise=__lowerCamelCase			)

                if model_class.__name__ == "TFViTMAEModel":
                   UpperCamelCase       :    List[str]		      =      after_outputs['''last_hidden_state'''].numpy()
                   UpperCamelCase       :    int		      =      0
                else:
                   UpperCamelCase       :    List[str]		      =      after_outputs['''logits'''].numpy()
                   UpperCamelCase       :    Dict		      =      0

                UpperCamelCase       :    List[str]		      =      np.amax(np.abs(out_a - out_a			)			)
                self.assertLessEqual(__lowerCamelCase,     1e-5			)
       def    snake_case_			(     self			)	->						Tuple:
          np.random.seed(2			)

          UpperCamelCase       :    int		      =      self.model_tester.prepare_config_and_inputs_for_common()

          UpperCamelCase       :    Optional[Any]		      =      int((config.image_size // config.patch_size) ** 2			)
          UpperCamelCase       :    Tuple		      =      np.random.uniform(size=(self.model_tester.batch_size, num_patches)			)

          for model_class in self.all_model_classes:
             UpperCamelCase       :    List[Any]		      =      model_class(__lowerCamelCase			)
             UpperCamelCase       :    Optional[Any]		      =      self._prepare_for_class(__lowerCamelCase,     __lowerCamelCase			)

             UpperCamelCase       :    Optional[Any]		      =      model(__lowerCamelCase,     noise=__lowerCamelCase			)
             UpperCamelCase       :    Any		      =      model.get_config()
             # make sure that returned config is jsonifiable, which is required by keras
             json.dumps(__lowerCamelCase			)
             UpperCamelCase       :    Dict		      =      model_class.from_config(model.get_config()			)
             # make sure it also accepts a normal config
             UpperCamelCase       :    Union[str, Any]		      =      model_class.from_config(model.config			)
             UpperCamelCase       :    int		      =      new_model(__lowerCamelCase			)  # Build model
             new_model.set_weights(model.get_weights()			)
             UpperCamelCase       :    Optional[int]		      =      new_model(__lowerCamelCase,     noise=__lowerCamelCase			)

             self.assert_outputs_same(__lowerCamelCase,     __lowerCamelCase			)
       @unittest.skip(
           reason='ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results.'			)
       def    snake_case_			(     self			)	->						Optional[int]:
          pass
       @unittest.skip(reason='ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load'			)
       def    snake_case_			(     self			)	->						Optional[Any]:
          pass






       @slow
       def    snake_case_			(     self			)	->						Union[str, Any]:
          UpperCamelCase       :    str		      =      TFViTMAEModel.from_pretrained('google/vit-base-patch16-224'			)
          self.assertIsNotNone(__lowerCamelCase			)






def 						UpperCamelCase      (  )      ->	Dict:
   UpperCamelCase       :    Tuple		      =      Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'      )
   return image



@require_tf
@require_vision
class 							lowerCAmelCase_						(    unittest.TestCase				):
       @cached_property
       def    snake_case_			(     self			)	->						Optional[int]:
          return ViTImageProcessor.from_pretrained('facebook/vit-mae-base'			) if is_vision_available() else None






       @slow
       def    snake_case_			(     self			)	->						Any:
          np.random.seed(2			)

          UpperCamelCase       :    List[str]		      =      TFViTMAEForPreTraining.from_pretrained('facebook/vit-mae-base'			)

          UpperCamelCase       :    Union[str, Any]		      =      self.default_image_processor
          UpperCamelCase       :    Any		      =      prepare_img()
          UpperCamelCase       :    Optional[int]		      =      image_processor(images=__lowerCamelCase,     return_tensors='tf'			)

          # prepare a noise vector that will be also used for testing the TF model
          # (this way we can ensure that the PT and TF models operate on the same inputs)
          UpperCamelCase       :    Optional[Any]		      =      ViTMAEConfig()
          UpperCamelCase       :    Union[str, Any]		      =      int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2			)
          UpperCamelCase       :    List[str]		      =      np.random.uniform(size=(1, num_patches)			)

          # forward pass
          UpperCamelCase       :    Optional[Any]		      =      model(**__lowerCamelCase,     noise=__lowerCamelCase			)

          # verify the logits
          UpperCamelCase       :    Union[str, Any]		      =      tf.convert_to_tensor([1, 196, 768]			)
          self.assertEqual(outputs.logits.shape,     __lowerCamelCase			)

          UpperCamelCase       :    List[Any]		      =      tf.convert_to_tensor(
              [[-0.05_48, -1.70_23, -0.93_25], [0.37_21, -0.56_70, -0.22_33], [0.82_35, -1.38_78, -0.35_24]]			)

          tf.debugging.assert_near(outputs.logits[0, :3, :3],     __lowerCamelCase,     atol=1e-4			)










"""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 						__snake_case   (						unittest.TestCase	):





 """simple docstring"""


 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




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

 @property
 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     torch.manual_seed(0    )
     __A       :		List[Any]							=	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    UpperCamelCase__(       self    ):



     '''simple docstring'''




     torch.manual_seed(0    )
     __A       :		str							=	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    UpperCamelCase__(       self    ):



     '''simple docstring'''




     torch.manual_seed(0    )
     __A       :		int							=	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''')   ,		)
     __A       :		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    UpperCamelCase__(       self    ):



     '''simple docstring'''




     __A       :		str							=	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
     __A       :		Union[str, Any]							=	Mel(
         x_res=self.dummy_unet.config.sample_size[1]   ,		y_res=self.dummy_unet.config.sample_size[0]   ,		)

     __A       :		List[Any]							=	DDPMScheduler()
     __A       :		Optional[Any]							=	AudioDiffusionPipeline(vqvae=__lowerCamelCase   ,		unet=self.dummy_unet   ,		mel=__lowerCamelCase   ,		scheduler=__lowerCamelCase    )
     __A       :		Tuple							=	pipe.to(__lowerCamelCase    )
     pipe.set_progress_bar_config(disable=__lowerCamelCase    )

     __A       :		Tuple							=	torch.Generator(device=__lowerCamelCase    ).manual_seed(42    )
     __A       :		List[str]							=	pipe(generator=__lowerCamelCase   ,		steps=4    )
     __A       :		Union[str, Any]							=	output.audios[0]
     __A       :		str							=	output.images[0]

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

     __A       :		int							=	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]   ,		)

     __A       :		Any							=	DDIMScheduler()
     __A       :		Optional[Any]							=	self.dummy_vqvae_and_unet
     __A       :		int							=	AudioDiffusionPipeline(
         vqvae=self.dummy_vqvae_and_unet[0]   ,		unet=dummy_vqvae_and_unet[1]   ,		mel=__lowerCamelCase   ,		scheduler=__lowerCamelCase    )
     __A       :		Any							=	pipe.to(__lowerCamelCase    )
     pipe.set_progress_bar_config(disable=__lowerCamelCase    )

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

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

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

     np.random.seed(0    )
     __A       :		int							=	torch.rand((1, 1, 10)    )
     __A       :		Optional[Any]							=	pipe(generator=__lowerCamelCase   ,		encoding=__lowerCamelCase    )
     __A       :		str							=	output.images[0]
     __A       :		Optional[Any]							=	np.frombuffer(image.tobytes()   ,		dtype='''uint8'''    )[:10]
     __A       :		Tuple							=	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 						__snake_case   (						unittest.TestCase	):





 """simple docstring"""


 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




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







 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     __A       :		int							=	torch_device

     __A       :		List[Any]							=	DiffusionPipeline.from_pretrained('''teticio/audio-diffusion-ddim-256'''    )
     __A       :		List[Any]							=	pipe.to(__lowerCamelCase    )
     pipe.set_progress_bar_config(disable=__lowerCamelCase    )

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

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




