



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





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





   @property
   def 	__lowerCAmelCase							(	self			):
         torch.manual_seed(0			)
         _lowerCAmelCase 			=  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 	__lowerCAmelCase							(	self			):
         torch.manual_seed(0			)
         _lowerCAmelCase 			=  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 	__lowerCAmelCase							(	self			):
         torch.manual_seed(0			)
         _lowerCAmelCase 			=  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''')       ,     )
         _lowerCAmelCase 			=  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 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
         _lowerCAmelCase 			=  Mel(
             x_res=self.dummy_unet.config.sample_size[1]       ,     y_res=self.dummy_unet.config.sample_size[0]       ,     )

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

         _lowerCAmelCase 			=  torch.Generator(device=_lowerCAmelCase			).manual_seed(42			)
         _lowerCAmelCase 			=  pipe(generator=_lowerCAmelCase       ,     steps=4			)
         _lowerCAmelCase 			=  output.audios[0]
         _lowerCAmelCase 			=  output.images[0]

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

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

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

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

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

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

         np.random.seed(0			)
         _lowerCAmelCase 			=  torch.rand((1, 1, 10)			)
         _lowerCAmelCase 			=  pipe(generator=_lowerCAmelCase       ,     encoding=_lowerCAmelCase			)
         _lowerCAmelCase 			=  output.images[0]
         _lowerCAmelCase 			=  np.frombuffer(image.tobytes()       ,     dtype='''uint8'''			)[:10]
         _lowerCAmelCase 			=  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 		UpperCAmelCase						(						unittest.TestCase				):





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





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  torch_device

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

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

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




from tempfile import TemporaryDirectory
from unittest import TestCase
from unittest.mock import MagicMock, patch

from transformers import AutoModel, TFAutoModel
from transformers.onnx import FeaturesManager
from transformers.testing_utils import SMALL_MODEL_IDENTIFIER, require_tf, require_torch





@require_torch
@require_tf
class 		UpperCAmelCase						(						snake_case_				):





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  SMALL_MODEL_IDENTIFIER
         _lowerCAmelCase 			=  '''pt'''
         _lowerCAmelCase 			=  '''tf'''





   def 	__lowerCAmelCase							(	self       ,     _lowerCAmelCase			):
         _lowerCAmelCase 			=  AutoModel.from_pretrained(self.test_model			)
         model_pt.save_pretrained(_lowerCAmelCase			)





   def 	__lowerCAmelCase							(	self       ,     _lowerCAmelCase			):
         _lowerCAmelCase 			=  TFAutoModel.from_pretrained(self.test_model       ,     from_pt=_lowerCAmelCase			)
         model_tf.save_pretrained(_lowerCAmelCase			)





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  '''mock_framework'''

         # Framework provided - return whatever the user provides
         _lowerCAmelCase 			=  FeaturesManager.determine_framework(self.test_model       ,     _lowerCAmelCase			)
         self.assertEqual(_lowerCAmelCase       ,     _lowerCAmelCase			)

         # Local checkpoint and framework provided - return provided framework
         # PyTorch checkpoint
         with TemporaryDirectory() as local_pt_ckpt:
               self._setup_pt_ckpt(_lowerCAmelCase			)
               _lowerCAmelCase 			=  FeaturesManager.determine_framework(_lowerCAmelCase       ,     _lowerCAmelCase			)
               self.assertEqual(_lowerCAmelCase       ,     _lowerCAmelCase			)

         # TensorFlow checkpoint
         with TemporaryDirectory() as local_tf_ckpt:
               self._setup_tf_ckpt(_lowerCAmelCase			)
               _lowerCAmelCase 			=  FeaturesManager.determine_framework(_lowerCAmelCase       ,     _lowerCAmelCase			)
               self.assertEqual(_lowerCAmelCase       ,     _lowerCAmelCase			)





   def 	__lowerCAmelCase							(	self			):
         # PyTorch checkpoint
         with TemporaryDirectory() as local_pt_ckpt:
               self._setup_pt_ckpt(_lowerCAmelCase			)
               _lowerCAmelCase 			=  FeaturesManager.determine_framework(_lowerCAmelCase			)
               self.assertEqual(_lowerCAmelCase       ,     self.framework_pt			)

         # TensorFlow checkpoint
         with TemporaryDirectory() as local_tf_ckpt:
               self._setup_tf_ckpt(_lowerCAmelCase			)
               _lowerCAmelCase 			=  FeaturesManager.determine_framework(_lowerCAmelCase			)
               self.assertEqual(_lowerCAmelCase       ,     self.framework_tf			)

         # Invalid local checkpoint
         with TemporaryDirectory() as local_invalid_ckpt:
               with self.assertRaises(_lowerCAmelCase			):
                     _lowerCAmelCase 			=  FeaturesManager.determine_framework(_lowerCAmelCase			)





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  MagicMock(return_value=_lowerCAmelCase			)
         with patch('''transformers.onnx.features.is_tf_available'''       ,     _lowerCAmelCase			):
               _lowerCAmelCase 			=  FeaturesManager.determine_framework(self.test_model			)
               self.assertEqual(_lowerCAmelCase       ,     self.framework_pt			)

         # PyTorch not in environment -> use TensorFlow
         _lowerCAmelCase 			=  MagicMock(return_value=_lowerCAmelCase			)
         with patch('''transformers.onnx.features.is_torch_available'''       ,     _lowerCAmelCase			):
               _lowerCAmelCase 			=  FeaturesManager.determine_framework(self.test_model			)
               self.assertEqual(_lowerCAmelCase       ,     self.framework_tf			)

         # Both in environment -> use PyTorch
         _lowerCAmelCase 			=  MagicMock(return_value=_lowerCAmelCase			)
         _lowerCAmelCase 			=  MagicMock(return_value=_lowerCAmelCase			)
         with patch('''transformers.onnx.features.is_tf_available'''       ,     _lowerCAmelCase			), patch(
             '''transformers.onnx.features.is_torch_available'''       ,     _lowerCAmelCase			):
               _lowerCAmelCase 			=  FeaturesManager.determine_framework(self.test_model			)
               self.assertEqual(_lowerCAmelCase       ,     self.framework_pt			)

         # Both not in environment -> raise error
         _lowerCAmelCase 			=  MagicMock(return_value=_lowerCAmelCase			)
         _lowerCAmelCase 			=  MagicMock(return_value=_lowerCAmelCase			)
         with patch('''transformers.onnx.features.is_tf_available'''       ,     _lowerCAmelCase			), patch(
             '''transformers.onnx.features.is_torch_available'''       ,     _lowerCAmelCase			):
               with self.assertRaises(_lowerCAmelCase			):
                     _lowerCAmelCase 			=  FeaturesManager.determine_framework(self.test_model			)