

'''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 			lowercase_	(unittest.TestCase      ):




  """simple docstring"""




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






  @property
  def 					SCREAMING_SNAKE_CASE      (      self   :     Any			):
   torch.manual_seed(0			)
   __lowercase							  = UNetaDModel(
       sample_size=(3_2, 6_4)					,in_channels=1					,out_channels=1					,layers_per_block=2					,block_out_channels=(1_2_8, 1_2_8)					,down_block_types=('''AttnDownBlock2D''', '''DownBlock2D''')					,up_block_types=('''UpBlock2D''', '''AttnUpBlock2D''')					,)
   return model






  @property
  def 					SCREAMING_SNAKE_CASE      (      self   :     Union[str, Any]			):
   torch.manual_seed(0			)
   __lowercase							  = UNetaDConditionModel(
       sample_size=(6_4, 3_2)					,in_channels=1					,out_channels=1					,layers_per_block=2					,block_out_channels=(1_2_8, 1_2_8)					,down_block_types=('''CrossAttnDownBlock2D''', '''DownBlock2D''')					,up_block_types=('''UpBlock2D''', '''CrossAttnUpBlock2D''')					,cross_attention_dim=1_0					,)
   return model






  @property
  def 					SCREAMING_SNAKE_CASE      (      self   :     int			):
   torch.manual_seed(0			)
   __lowercase							  = AutoencoderKL(
       sample_size=(1_2_8, 6_4)					,in_channels=1					,out_channels=1					,latent_channels=1					,layers_per_block=2					,block_out_channels=(1_2_8, 1_2_8)					,down_block_types=('''DownEncoderBlock2D''', '''DownEncoderBlock2D''')					,up_block_types=('''UpDecoderBlock2D''', '''UpDecoderBlock2D''')					,)
   __lowercase							  = UNetaDModel(
       sample_size=(6_4, 3_2)					,in_channels=1					,out_channels=1					,layers_per_block=2					,block_out_channels=(1_2_8, 1_2_8)					,down_block_types=('''AttnDownBlock2D''', '''DownBlock2D''')					,up_block_types=('''UpBlock2D''', '''AttnUpBlock2D''')					,)
   return vqvae, unet






  @slow
  def 					SCREAMING_SNAKE_CASE      (      self   :     List[str]			):
   __lowercase							  = '''cpu'''  # ensure determinism for the device-dependent torch.Generator
   __lowercase							  = Mel(
       x_res=self.dummy_unet.config.sample_size[1]					,y_res=self.dummy_unet.config.sample_size[0]					,)

   __lowercase							  = DDPMScheduler()
   __lowercase							  = AudioDiffusionPipeline(vqvae=lowercase__					,unet=self.dummy_unet					,mel=lowercase__					,scheduler=lowercase__			)
   __lowercase							  = pipe.to(lowercase__			)
   pipe.set_progress_bar_config(disable=lowercase__			)

   __lowercase							  = torch.Generator(device=lowercase__			).manual_seed(4_2			)
   __lowercase							  = pipe(generator=lowercase__					,steps=4			)
   __lowercase							  = output.audios[0]
   __lowercase							  = output.images[0]

   __lowercase							  = torch.Generator(device=lowercase__			).manual_seed(4_2			)
   __lowercase							  = pipe(generator=lowercase__					,steps=4					,return_dict=lowercase__			)
   __lowercase							  = 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]
   )
   __lowercase							  = np.frombuffer(image.tobytes()					,dtype='''uint8'''			)[:1_0]
   __lowercase							  = np.frombuffer(image_from_tuple.tobytes()					,dtype='''uint8'''			)[:1_0]
   __lowercase							  = np.array([6_9, 2_5_5, 2_5_5, 2_5_5, 0, 0, 7_7, 1_8_1, 1_2, 1_2_7]			)

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

   __lowercase							  = 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]					,)

   __lowercase							  = DDIMScheduler()
   __lowercase							  = self.dummy_vqvae_and_unet
   __lowercase							  = AudioDiffusionPipeline(
       vqvae=self.dummy_vqvae_and_unet[0]					,unet=dummy_vqvae_and_unet[1]					,mel=lowercase__					,scheduler=lowercase__			)
   __lowercase							  = pipe.to(lowercase__			)
   pipe.set_progress_bar_config(disable=lowercase__			)

   np.random.seed(0			)
   __lowercase							  = np.random.uniform(-1					,1					,((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)			)
   __lowercase							  = torch.Generator(device=lowercase__			).manual_seed(4_2			)
   __lowercase							  = pipe(raw_audio=lowercase__					,generator=lowercase__					,start_step=5					,steps=1_0			)
   __lowercase							  = 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]
   )
   __lowercase							  = np.frombuffer(image.tobytes()					,dtype='''uint8'''			)[:1_0]
   __lowercase							  = np.array([1_2_0, 1_1_7, 1_1_0, 1_0_9, 1_3_8, 1_6_7, 1_3_8, 1_4_8, 1_3_2, 1_2_1]			)

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

   __lowercase							  = self.dummy_unet_condition
   __lowercase							  = AudioDiffusionPipeline(
       vqvae=self.dummy_vqvae_and_unet[0]					,unet=lowercase__					,mel=lowercase__					,scheduler=lowercase__			)
   __lowercase							  = pipe.to(lowercase__			)
   pipe.set_progress_bar_config(disable=lowercase__			)

   np.random.seed(0			)
   __lowercase							  = torch.rand((1, 1, 1_0)			)
   __lowercase							  = pipe(generator=lowercase__					,encoding=lowercase__			)
   __lowercase							  = output.images[0]
   __lowercase							  = np.frombuffer(image.tobytes()					,dtype='''uint8'''			)[:1_0]
   __lowercase							  = np.array([1_0_7, 1_0_3, 1_2_0, 1_2_7, 1_4_2, 1_2_2, 1_1_3, 1_2_2, 9_7, 1_1_1]			)

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






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




  """simple docstring"""




  def 					SCREAMING_SNAKE_CASE      (      self   :     Dict			):
   # clean up the VRAM after each test
   super().tearDown()
   gc.collect()
   torch.cuda.empty_cache()






  def 					SCREAMING_SNAKE_CASE      (      self   :     List[str]			):
   __lowercase							  = torch_device

   __lowercase							  = DiffusionPipeline.from_pretrained('''teticio/audio-diffusion-ddim-256'''			)
   __lowercase							  = pipe.to(lowercase__			)
   pipe.set_progress_bar_config(disable=lowercase__			)

   __lowercase							  = torch.Generator(device=lowercase__			).manual_seed(4_2			)
   __lowercase							  = pipe(generator=lowercase__			)
   __lowercase							  = output.audios[0]
   __lowercase							  = 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]
   __lowercase							  = np.frombuffer(image.tobytes()					,dtype='''uint8'''			)[:1_0]
   __lowercase							  = np.array([1_5_1, 1_6_7, 1_5_4, 1_4_4, 1_2_2, 1_3_4, 1_2_1, 1_0_5, 7_0, 2_6]			)

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






"""simple docstring"""







import json
import logging
import os
import sys
from time import time
from unittest.mock import patch

from transformers.testing_utils import TestCasePlus, require_torch_tpu


logging.basicConfig(level=logging.DEBUG)

lowercase__		:				List[str]          =	logging.getLogger()

def 				UpperCamelCase_    (    lowerCAmelCase__      :				str						)						->   Dict:

  """simple docstring"""




  lowerCAmelCase_   :				Dict								= {}
  lowerCAmelCase_   :				Union[str, Any]								= os.path.join(lowerCAmelCase__    ,   'all_results.json'						)
  if os.path.exists(lowerCAmelCase__						):
    with open(lowerCAmelCase__    ,   'r'						) as f:
      lowerCAmelCase_   :				Any								= json.load(lowerCAmelCase__						)
  else:
    raise ValueError(f"can't find {path}"						)
  return results


lowercase__		:				Optional[int]          =	logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)







@require_torch_tpu
class       UpperCamelCase__   (	lowercase_  ):




      """simple docstring"""




      def      SCREAMING_SNAKE_CASE__ (      self				:      Any       ):
        import xla_spawn

        lowerCAmelCase_   :				Union[str, Any]								= self.get_auto_remove_tmp_dir()
        lowerCAmelCase_   :				Optional[Any]								= F"\n            ./examples/pytorch/text-classification/run_glue.py\n            --num_cores=8\n            ./examples/pytorch/text-classification/run_glue.py\n            --model_name_or_path distilbert-base-uncased\n            --output_dir {tmp_dir}\n            --overwrite_output_dir\n            --train_file ./tests/fixtures/tests_samples/MRPC/train.csv\n            --validation_file ./tests/fixtures/tests_samples/MRPC/dev.csv\n            --do_train\n            --do_eval\n            --debug tpu_metrics_debug\n            --per_device_train_batch_size=2\n            --per_device_eval_batch_size=1\n            --learning_rate=1e-4\n            --max_steps=10\n            --warmup_steps=2\n            --seed=42\n            --max_seq_length=128\n            ".split()

        with patch.object(SCREAMING_SNAKE_CASE_    ,			'argv'    ,			SCREAMING_SNAKE_CASE_       ):
          lowerCAmelCase_   :				Optional[Any]								= time()
          xla_spawn.main()
          lowerCAmelCase_   :				Any								= time()

          lowerCAmelCase_   :				Dict								= get_results(SCREAMING_SNAKE_CASE_       )
          self.assertGreaterEqual(result['eval_accuracy']    ,			0.75       )

          # Assert that the script takes less than 500 seconds to make sure it doesn't hang.
          self.assertLess(end - start    ,			5_0_0       )





      def      SCREAMING_SNAKE_CASE__ (      self				:      List[str]       ):
        import xla_spawn

        lowerCAmelCase_   :				Dict								= '\n            ./tests/test_trainer_tpu.py\n            --num_cores=8\n            ./tests/test_trainer_tpu.py\n            '.split()
        with patch.object(SCREAMING_SNAKE_CASE_    ,			'argv'    ,			SCREAMING_SNAKE_CASE_       ):
          xla_spawn.main()





