
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        _UpperCamelCase				(					_A     ,       unittest.TestCase  ):






			'''simple docstring'''

			__UpperCamelCase      :						str					      = DanceDiffusionPipeline
			__UpperCamelCase      :						Optional[Any]					      = UNCONDITIONAL_AUDIO_GENERATION_PARAMS
			__UpperCamelCase      :						int					      = PipelineTesterMixin.required_optional_params - {
			    """callback""",
			    """latents""",
			    """callback_steps""",
			    """output_type""",
			    """num_images_per_prompt""",
			}
			__UpperCamelCase      :						Optional[int]					      = UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
			__UpperCamelCase      :						Union[str, Any]					      = False
			__UpperCamelCase      :						Optional[Any]					      = False

			def       lowerCAmelCase__	(       self :       List[Any]      ):
							torch.manual_seed(0      )
							UpperCamelCase_:							Optional[int]	=						UNetaDModel(
							    block_out_channels=(32, 32, 64)						,	extra_in_channels=16						,	sample_size=512						,	sample_rate=1_6000						,	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""")						,	)
							UpperCamelCase_:							Optional[int]	=						IPNDMScheduler()

							UpperCamelCase_:							Optional[Any]	=						{
							    """unet""": unet,
							    """scheduler""": scheduler,
							}
							return components

			def       lowerCAmelCase__	(       self :       List[Any]						,	snake_case_ :       List[str]						,	snake_case_ :       int=0      ):
							if str(snake_case_      ).startswith("""mps"""      ):
											UpperCamelCase_:							List[Any]	=						torch.manual_seed(snake_case_      )
							else:
											UpperCamelCase_:							List[Any]	=						torch.Generator(device=snake_case_      ).manual_seed(snake_case_      )
							UpperCamelCase_:							Tuple	=						{
							    """batch_size""": 1,
							    """generator""": generator,
							    """num_inference_steps""": 4,
							}
							return inputs

			def       lowerCAmelCase__	(       self :       List[str]      ):
							UpperCamelCase_:							List[Any]	=						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
							UpperCamelCase_:							Any	=						self.get_dummy_components()
							UpperCamelCase_:							Optional[Any]	=						DanceDiffusionPipeline(**snake_case_      )
							UpperCamelCase_:							str	=						pipe.to(snake_case_      )
							pipe.set_progress_bar_config(disable=snake_case_      )

							UpperCamelCase_:							Union[str, Any]	=						self.get_dummy_inputs(snake_case_      )
							UpperCamelCase_:							Union[str, Any]	=						pipe(**snake_case_      )
							UpperCamelCase_:							Optional[int]	=						output.audios

							UpperCamelCase_:							str	=						audio[0, -3:, -3:]

							assert audio.shape == (1, 2, components["unet"].sample_size)
							UpperCamelCase_:							Tuple	=						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       lowerCAmelCase__	(       self :       Any      ):
							return super().test_save_load_local()

			@skip_mps
			def       lowerCAmelCase__	(       self :       Any      ):
							return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3      )

			@skip_mps
			def       lowerCAmelCase__	(       self :       str      ):
							return super().test_save_load_optional_components()

			@skip_mps
			def       lowerCAmelCase__	(       self :       str      ):
							return super().test_attention_slicing_forward_pass()




			def       lowerCAmelCase__	(       self :       Optional[int]      ):
							super().test_inference_batch_single_identical(expected_max_diff=3e-3      )




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






			'''simple docstring'''

			def       lowerCAmelCase__	(       self :       Optional[int]      ):
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()

			def       lowerCAmelCase__	(       self :       List[Any]      ):
							UpperCamelCase_:							Optional[int]	=						torch_device

							UpperCamelCase_:							Tuple	=						DanceDiffusionPipeline.from_pretrained("""harmonai/maestro-150k"""      )
							UpperCamelCase_:							List[Any]	=						pipe.to(snake_case_      )
							pipe.set_progress_bar_config(disable=snake_case_      )

							UpperCamelCase_:							List[Any]	=						torch.manual_seed(0      )
							UpperCamelCase_:							int	=						pipe(generator=snake_case_						,	num_inference_steps=100						,	audio_length_in_s=4.096      )
							UpperCamelCase_:							Optional[int]	=						output.audios

							UpperCamelCase_:							Tuple	=						audio[0, -3:, -3:]

							assert audio.shape == (1, 2, pipe.unet.sample_size)
							UpperCamelCase_:							List[Any]	=						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       lowerCAmelCase__	(       self :       int      ):
							UpperCamelCase_:							Tuple	=						torch_device

							UpperCamelCase_:							Union[str, Any]	=						DanceDiffusionPipeline.from_pretrained("""harmonai/maestro-150k"""						,	torch_dtype=torch.floataa      )
							UpperCamelCase_:							List[Any]	=						pipe.to(snake_case_      )
							pipe.set_progress_bar_config(disable=snake_case_      )

							UpperCamelCase_:							Optional[int]	=						torch.manual_seed(0      )
							UpperCamelCase_:							Union[str, Any]	=						pipe(generator=snake_case_						,	num_inference_steps=100						,	audio_length_in_s=4.096      )
							UpperCamelCase_:							Optional[Any]	=						output.audios

							UpperCamelCase_:							Union[str, Any]	=						audio[0, -3:, -3:]

							assert audio.shape == (1, 2, pipe.unet.sample_size)
							UpperCamelCase_:							List[Any]	=						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 numpy as np
import torch

from transformers import SpeechTaHifiGan, SpeechTaHifiGanConfig, logging


logging.set_verbosity_info()
lowerCamelCase_						: Optional[int]         =  logging.get_logger("""transformers.models.speecht5""")







def        A__					(						lowerCamelCase     ,    lowerCamelCase     ,    lowerCamelCase						)  ->	Optional[Any]:
				hf_model.apply_weight_norm()

				UpperCamelCase_:							Union[str, Any]	=						checkpoint["""input_conv.weight_g"""]
				UpperCamelCase_:							Optional[int]	=						checkpoint["""input_conv.weight_v"""]
				UpperCamelCase_:							List[Any]	=						checkpoint["""input_conv.bias"""]

				for i in range(len(config.upsample_rates						)						):
								UpperCamelCase_:							List[str]	=						checkpoint[F'''upsamples.{i}.1.weight_g''']
								UpperCamelCase_:							Dict	=						checkpoint[F'''upsamples.{i}.1.weight_v''']
								UpperCamelCase_:							List[str]	=						checkpoint[F'''upsamples.{i}.1.bias''']

				for i in range(len(config.upsample_rates						) * len(config.resblock_kernel_sizes						)						):
								for j in range(len(config.resblock_dilation_sizes						)						):
												UpperCamelCase_:							Tuple	=						checkpoint[F'''blocks.{i}.convs1.{j}.1.weight_g''']
												UpperCamelCase_:							Any	=						checkpoint[F'''blocks.{i}.convs1.{j}.1.weight_v''']
												UpperCamelCase_:							Tuple	=						checkpoint[F'''blocks.{i}.convs1.{j}.1.bias''']

												UpperCamelCase_:							Union[str, Any]	=						checkpoint[F'''blocks.{i}.convs2.{j}.1.weight_g''']
												UpperCamelCase_:							Any	=						checkpoint[F'''blocks.{i}.convs2.{j}.1.weight_v''']
												UpperCamelCase_:							int	=						checkpoint[F'''blocks.{i}.convs2.{j}.1.bias''']

				UpperCamelCase_:							int	=						checkpoint["""output_conv.1.weight_g"""]
				UpperCamelCase_:							Tuple	=						checkpoint["""output_conv.1.weight_v"""]
				UpperCamelCase_:							List[str]	=						checkpoint["""output_conv.1.bias"""]

				hf_model.remove_weight_norm()







@torch.no_grad()
def        A__					(						lowerCamelCase     ,    lowerCamelCase     ,    lowerCamelCase     ,    lowerCamelCase=None     ,    lowerCamelCase=None     ,    )  ->	Optional[int]:
				if config_path is not None:
								UpperCamelCase_:							Union[str, Any]	=						SpeechTaHifiGanConfig.from_pretrained(lowerCamelCase						)
				else:
								UpperCamelCase_:							str	=						SpeechTaHifiGanConfig()

				UpperCamelCase_:							Union[str, Any]	=						SpeechTaHifiGan(lowerCamelCase						)

				UpperCamelCase_:							str	=						torch.load(lowerCamelCase						)
				load_weights(orig_checkpoint["""model"""]["""generator"""]     ,    lowerCamelCase     ,    lowerCamelCase						)

				UpperCamelCase_:							Union[str, Any]	=						np.load(lowerCamelCase						)
				UpperCamelCase_:							int	=						stats[0].reshape(-1						)
				UpperCamelCase_:							Union[str, Any]	=						stats[1].reshape(-1						)
				UpperCamelCase_:							Dict	=						torch.from_numpy(lowerCamelCase						).float()
				UpperCamelCase_:							Optional[Any]	=						torch.from_numpy(lowerCamelCase						).float()

				model.save_pretrained(lowerCamelCase						)

				if repo_id:
								print("""Pushing to the hub..."""						)
								model.push_to_hub(lowerCamelCase						)


if __name__ == "__main__":
							lowerCamelCase_						: Union[str, Any]         =  argparse.ArgumentParser()
							parser.add_argument("""--checkpoint_path""", required=True, default=None, type=str, help="""Path to original checkpoint""")
							parser.add_argument("""--stats_path""", required=True, default=None, type=str, help="""Path to stats.npy file""")
							parser.add_argument("""--config_path""", default=None, type=str, help="""Path to hf config.json of model to convert""")
							parser.add_argument(
							    """--pytorch_dump_folder_path""", required=True, default=None, type=str, help="""Path to the output PyTorch model."""
							)
							parser.add_argument(
							    """--push_to_hub""", default=None, type=str, help="""Where to upload the converted model on the 🤗 hub."""
							)

							lowerCamelCase_						: Optional[int]         =  parser.parse_args()
							convert_hifigan_checkpoint(
							    args.checkpoint_path,
							    args.stats_path,
							    args.pytorch_dump_folder_path,
							    args.config_path,
							    args.push_to_hub,
							)






