
import gc
import unittest

import numpy as np
import torch

from diffusers import DanceDiffusionPipeline, IPNDMScheduler, UNetaDModel
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS, UNCONDITIONAL_AUDIO_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()







class 			__A( A__  ,   unittest.TestCase			):
				snake_case_						   =    DanceDiffusionPipeline
				snake_case_						   =    UNCONDITIONAL_AUDIO_GENERATION_PARAMS
				snake_case_						   =    PipelineTesterMixin.required_optional_params - {
				    '''callback''',
				    '''latents''',
				    '''callback_steps''',
				    '''output_type''',
				    '''num_images_per_prompt''',
				}
				snake_case_						   =    UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
				snake_case_						   =    False
				snake_case_						   =    False




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      str:


								'''simple docstring'''




								torch.manual_seed(0				)
								__a										=    UNetaDModel(
								    block_out_channels=(32, 32, 64)		,     extra_in_channels=16		,     sample_size=512		,     sample_rate=16_000		,     in_channels=2		,     out_channels=2		,     flip_sin_to_cos=_snake_case		,     use_timestep_embedding=_snake_case		,     time_embedding_type='''fourier'''		,     mid_block_type='''UNetMidBlock1D'''		,     down_block_types=('''DownBlock1DNoSkip''', '''DownBlock1D''', '''AttnDownBlock1D''')		,     up_block_types=('''AttnUpBlock1D''', '''UpBlock1D''', '''UpBlock1DNoSkip''')		,     )
								__a										=    IPNDMScheduler()

								__a										=    {
								    '''unet''': unet,
								    '''scheduler''': scheduler,
								}
								return components




				def 				SCREAMING_SNAKE_CASE_		(     self		,     _snake_case		,     _snake_case=0				)    ->      int:


								'''simple docstring'''




								if str(_snake_case				).startswith('''mps'''				):
												__a										=    torch.manual_seed(_snake_case				)
								else:
												__a										=    torch.Generator(device=_snake_case				).manual_seed(_snake_case				)
								__a										=    {
								    '''batch_size''': 1,
								    '''generator''': generator,
								    '''num_inference_steps''': 4,
								}
								return inputs




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


								'''simple docstring'''




								__a										=    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
								__a										=    self.get_dummy_components()
								__a										=    DanceDiffusionPipeline(**_snake_case				)
								__a										=    pipe.to(_snake_case				)
								pipe.set_progress_bar_config(disable=_snake_case				)

								__a										=    self.get_dummy_inputs(_snake_case				)
								__a										=    pipe(**_snake_case				)
								__a										=    output.audios

								__a										=    audio[0, -3:, -3:]

								assert audio.shape == (1, 2, components["unet"].sample_size)
								__a										=    np.array([-0.7265, 1.0000, -0.8388, 0.1175, 0.9498, -1.0000]				)
								assert np.abs(audio_slice.flatten() - expected_slice				).max() < 1E-2




				@skip_mps
				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      str:


								'''simple docstring'''




								return super().test_save_load_local()




				@skip_mps
				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      List[Any]:


								'''simple docstring'''




								return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3				)




				@skip_mps
				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      List[str]:


								'''simple docstring'''




								return super().test_save_load_optional_components()




				@skip_mps
				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      str:


								'''simple docstring'''




								return super().test_attention_slicing_forward_pass()




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Tuple:


								'''simple docstring'''




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







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




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Any:


								'''simple docstring'''




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




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Optional[Any]:


								'''simple docstring'''




								__a										=    torch_device

								__a										=    DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k'''				)
								__a										=    pipe.to(_snake_case				)
								pipe.set_progress_bar_config(disable=_snake_case				)

								__a										=    torch.manual_seed(0				)
								__a										=    pipe(generator=_snake_case		,     num_inference_steps=100		,     audio_length_in_s=4.096				)
								__a										=    output.audios

								__a										=    audio[0, -3:, -3:]

								assert audio.shape == (1, 2, pipe.unet.sample_size)
								__a										=    np.array([-0.0192, -0.0231, -0.0318, -0.0059, 0.0002, -0.0020]				)

								assert np.abs(audio_slice.flatten() - expected_slice				).max() < 1E-2




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Optional[Any]:


								'''simple docstring'''




								__a										=    torch_device

								__a										=    DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k'''		,     torch_dtype=torch.floataa				)
								__a										=    pipe.to(_snake_case				)
								pipe.set_progress_bar_config(disable=_snake_case				)

								__a										=    torch.manual_seed(0				)
								__a										=    pipe(generator=_snake_case		,     num_inference_steps=100		,     audio_length_in_s=4.096				)
								__a										=    output.audios

								__a										=    audio[0, -3:, -3:]

								assert audio.shape == (1, 2, pipe.unet.sample_size)
								__a										=    np.array([-0.0367, -0.0488, -0.0771, -0.0525, -0.0444, -0.0341]				)

								assert np.abs(audio_slice.flatten() - expected_slice				).max() < 1E-2







import argparse
import os

import numpy as np
import tensorflow as tf
import torch

from transformers import BertModel






def   a			(						lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_				):






				'''simple docstring'''
				lowercase__    =			('''dense.weight''', '''attention.self.query''', '''attention.self.key''', '''attention.self.value''')

				lowercase__    =			(
				    ('''layer.''', '''layer_'''),
				    ('''word_embeddings.weight''', '''word_embeddings'''),
				    ('''position_embeddings.weight''', '''position_embeddings'''),
				    ('''token_type_embeddings.weight''', '''token_type_embeddings'''),
				    ('''.''', '''/'''),
				    ('''LayerNorm/weight''', '''LayerNorm/gamma'''),
				    ('''LayerNorm/bias''', '''LayerNorm/beta'''),
				    ('''weight''', '''kernel'''),
				)

				if not os.path.isdir(lowerCamelCase_				):
								os.makedirs(lowerCamelCase_				)

				lowercase__    =			model.state_dict()

				def to_tf_var_name(lowerCamelCase_				):
								for patt, repl in iter(lowerCamelCase_				):
												lowercase__    =			name.replace(lowerCamelCase_		,   lowerCamelCase_				)
								return F"""bert/{name}"""

				def create_tf_var(lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_				):
								lowercase__    =			tf.dtypes.as_dtype(tensor.dtype				)
								lowercase__    =			tf.get_variable(dtype=lowerCamelCase_		,   shape=tensor.shape		,   name=lowerCamelCase_		,   initializer=tf.zeros_initializer()				)
								session.run(tf.variables_initializer([tf_var]				)				)
								session.run(lowerCamelCase_				)
								return tf_var

				tf.reset_default_graph()
				with tf.Session() as session:
								for var_name in state_dict:
												lowercase__    =			to_tf_var_name(lowerCamelCase_				)
												lowercase__    =			state_dict[var_name].numpy()
												if any(x in var_name for x in tensors_to_transpose				):
																lowercase__    =			torch_tensor.T
												lowercase__    =			create_tf_var(tensor=lowerCamelCase_		,   name=lowerCamelCase_		,   session=lowerCamelCase_				)
												tf.keras.backend.set_value(lowerCamelCase_		,   lowerCamelCase_				)
												lowercase__    =			session.run(lowerCamelCase_				)
												print(F"""Successfully created {tf_name}: {np.allclose(lowerCamelCase_		,   lowerCamelCase_				)}"""				)

								lowercase__    =			tf.train.Saver(tf.trainable_variables()				)
								saver.save(lowerCamelCase_		,   os.path.join(lowerCamelCase_		,   model_name.replace('''-'''		,   '''_'''				) + '''.ckpt'''				)				)






def   a			(						lowerCamelCase_=None				):






				'''simple docstring'''
				lowercase__    =			argparse.ArgumentParser()
				parser.add_argument('''--model_name'''		,   type=lowerCamelCase_		,   required=lowerCamelCase_		,   help='''model name e.g. bert-base-uncased'''				)
				parser.add_argument(
				    '''--cache_dir'''		,   type=lowerCamelCase_		,   default=lowerCamelCase_		,   required=lowerCamelCase_		,   help='''Directory containing pytorch model'''				)
				parser.add_argument('''--pytorch_model_path'''		,   type=lowerCamelCase_		,   required=lowerCamelCase_		,   help='''/path/to/<pytorch-model-name>.bin'''				)
				parser.add_argument('''--tf_cache_dir'''		,   type=lowerCamelCase_		,   required=lowerCamelCase_		,   help='''Directory in which to save tensorflow model'''				)
				lowercase__    =			parser.parse_args(lowerCamelCase_				)

				lowercase__    =			BertModel.from_pretrained(
				    pretrained_model_name_or_path=args.model_name		,   state_dict=torch.load(args.pytorch_model_path				)		,   cache_dir=args.cache_dir		,   )

				convert_pytorch_checkpoint_to_tf(model=lowerCamelCase_		,   ckpt_dir=args.tf_cache_dir		,   model_name=args.model_name				)


if __name__ == "__main__":
							main()




