import json
import os
import shutil
import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import AutoConfig, BertConfig, GPTaConfig
from transformers.configuration_utils import PretrainedConfig
from transformers.testing_utils import TOKEN, USER, is_staging_test


sys.path.append(str(Path(__file__).parent.parent / 'utils'))

from test_module.custom_configuration import CustomConfig  # noqa E402


A						=			{
    'return_dict': False,
    'output_hidden_states': True,
    'output_attentions': True,
    'torchscript': True,
    'torch_dtype': 'float16',
    'use_bfloat16': True,
    'tf_legacy_loss': True,
    'pruned_heads': {'a': 1},
    'tie_word_embeddings': False,
    'is_decoder': True,
    'cross_attention_hidden_size': 1_2_8,
    'add_cross_attention': True,
    'tie_encoder_decoder': True,
    'max_length': 5_0,
    'min_length': 3,
    'do_sample': True,
    'early_stopping': True,
    'num_beams': 3,
    'num_beam_groups': 3,
    'diversity_penalty': 0.5,
    'temperature': 2.0,
    'top_k': 1_0,
    'top_p': 0.7,
    'typical_p': 0.2,
    'repetition_penalty': 0.8,
    'length_penalty': 0.8,
    'no_repeat_ngram_size': 5,
    'encoder_no_repeat_ngram_size': 5,
    'bad_words_ids': [1, 2, 3],
    'num_return_sequences': 3,
    'chunk_size_feed_forward': 5,
    'output_scores': True,
    'return_dict_in_generate': True,
    'forced_bos_token_id': 2,
    'forced_eos_token_id': 3,
    'remove_invalid_values': True,
    'architectures': ['BertModel'],
    'finetuning_task': 'translation',
    'id2label': {0: 'label'},
    'label2id': {'label': '0'},
    'tokenizer_class': 'BertTokenizerFast',
    'prefix': 'prefix',
    'bos_token_id': 6,
    'pad_token_id': 7,
    'eos_token_id': 8,
    'sep_token_id': 9,
    'decoder_start_token_id': 1_0,
    'exponential_decay_length_penalty': (5, 1.01),
    'suppress_tokens': [0, 1],
    'begin_suppress_tokens': 2,
    'task_specific_params': {'translation': 'some_params'},
    'problem_type': 'regression',
}





@is_staging_test
class     lowerCAmelCase__     (				unittest.TestCase			):


		'''simple docstring'''






		@classmethod
		def   _snake_case       (       cls     :     str		)      ->	Optional[int]:
				_lowerCamelCase				   =  TOKEN
				HfFolder.save_token(snake_case__		)
		@classmethod
		def   _snake_case       (       cls     :     str		)      ->	Tuple:
				try:
						delete_repo(token=cls._token     ,  repo_id='test-config'		)
				except HTTPError:
						pass

				try:
						delete_repo(token=cls._token     ,  repo_id='valid_org/test-config-org'		)
				except HTTPError:
						pass

				try:
						delete_repo(token=cls._token     ,  repo_id='test-dynamic-config'		)
				except HTTPError:
						pass
		def   _snake_case       (       self     :     Tuple		)      ->	Dict:
				_lowerCamelCase				   =  BertConfig(
				    vocab_size=9_9     ,  hidden_size=3_2     ,  num_hidden_layers=5     ,  num_attention_heads=4     ,  intermediate_size=3_7		)
				config.push_to_hub('test-config'     ,  use_auth_token=self._token		)

				_lowerCamelCase				   =  BertConfig.from_pretrained(f"""{USER}/test-config"""		)
				for k, v in config.to_dict().items():
						if k != "transformers_version":
								self.assertEqual(snake_case__     ,  getattr(snake_case__     ,  snake_case__		)		)

        # Reset repo
				delete_repo(token=self._token     ,  repo_id='test-config'		)

				# Push to hub via save_pretrained
				with tempfile.TemporaryDirectory() as tmp_dir:
						config.save_pretrained(snake_case__     ,  repo_id='test-config'     ,  push_to_hub=snake_case__     ,  use_auth_token=self._token		)

				_lowerCamelCase				   =  BertConfig.from_pretrained(f"""{USER}/test-config"""		)
				for k, v in config.to_dict().items():
						if k != "transformers_version":
								self.assertEqual(snake_case__     ,  getattr(snake_case__     ,  snake_case__		)		)
		def   _snake_case       (       self     :     Union[str, Any]		)      ->	List[str]:
				_lowerCamelCase				   =  BertConfig(
				    vocab_size=9_9     ,  hidden_size=3_2     ,  num_hidden_layers=5     ,  num_attention_heads=4     ,  intermediate_size=3_7		)
				config.push_to_hub('valid_org/test-config-org'     ,  use_auth_token=self._token		)

				_lowerCamelCase				   =  BertConfig.from_pretrained('valid_org/test-config-org'		)
				for k, v in config.to_dict().items():
						if k != "transformers_version":
								self.assertEqual(snake_case__     ,  getattr(snake_case__     ,  snake_case__		)		)

        # Reset repo
				delete_repo(token=self._token     ,  repo_id='valid_org/test-config-org'		)

				# Push to hub via save_pretrained
				with tempfile.TemporaryDirectory() as tmp_dir:
						config.save_pretrained(
						    snake_case__     ,  repo_id='valid_org/test-config-org'     ,  push_to_hub=snake_case__     ,  use_auth_token=self._token		)

				_lowerCamelCase				   =  BertConfig.from_pretrained('valid_org/test-config-org'		)
				for k, v in config.to_dict().items():
						if k != "transformers_version":
								self.assertEqual(snake_case__     ,  getattr(snake_case__     ,  snake_case__		)		)



		def   _snake_case       (       self     :     Optional[int]		)      ->	Dict:
				CustomConfig.register_for_auto_class()
				_lowerCamelCase				   =  CustomConfig(attribute=4_2		)

				config.push_to_hub('test-dynamic-config'     ,  use_auth_token=self._token		)

				# This has added the proper auto_map field to the config
				self.assertDictEqual(config.auto_map     ,  {'AutoConfig': 'custom_configuration.CustomConfig'}		)

				_lowerCamelCase				   =  AutoConfig.from_pretrained(f"""{USER}/test-dynamic-config"""     ,  trust_remote_code=snake_case__		)
				# Can't make an isinstance check because the new_config is from the FakeConfig class of a dynamic module
				self.assertEqual(new_config.__class__.__name__     ,  'CustomConfig'		)
				self.assertEqual(new_config.attribute     ,  4_2		)





class     lowerCAmelCase__     (				unittest.TestCase			):


		'''simple docstring'''






		def   _snake_case       (       self     :     List[str]		)      ->	Optional[int]:
				_lowerCamelCase				   =  GPTaConfig()

				# attempt to modify each of int/float/bool/str config records and verify they were updated
				_lowerCamelCase				   =  c.n_embd + 1  # int
				_lowerCamelCase				   =  c.resid_pdrop + 1.0  # float
				_lowerCamelCase				   =  not c.scale_attn_weights  # bool
				_lowerCamelCase				   =  c.summary_type + 'foo'  # str
				c.update_from_string(
				    f"""n_embd={n_embd},resid_pdrop={resid_pdrop},scale_attn_weights={scale_attn_weights},summary_type={summary_type}"""		)
				self.assertEqual(snake_case__     ,  c.n_embd     ,  'mismatch for key: n_embd'		)
				self.assertEqual(snake_case__     ,  c.resid_pdrop     ,  'mismatch for key: resid_pdrop'		)
				self.assertEqual(snake_case__     ,  c.scale_attn_weights     ,  'mismatch for key: scale_attn_weights'		)
				self.assertEqual(snake_case__     ,  c.summary_type     ,  'mismatch for key: summary_type'		)
		def   _snake_case       (       self     :     Dict		)      ->	str:
				_lowerCamelCase				   =  PretrainedConfig()
				_lowerCamelCase				   =  [key for key in base_config.__dict__ if key not in config_common_kwargs]
				# If this part of the test fails, you have arguments to addin config_common_kwargs above.
				self.assertListEqual(
				    snake_case__     ,  ['is_encoder_decoder', '_name_or_path', '_commit_hash', 'transformers_version']		)
				_lowerCamelCase				   =  [key for key, value in config_common_kwargs.items() if value == getattr(snake_case__     ,  snake_case__		)]
				if len(snake_case__		) > 0:
						raise ValueError(
						    'The following keys are set with the default values in'
						    ' `test_configuration_common.config_common_kwargs` pick another value for them:'
						    f""" {", ".join(snake_case__		)}."""		)
		def   _snake_case       (       self     :     Optional[int]		)      ->	Dict:
				with self.assertRaises(snake_case__		):
						# config is in subfolder, the following should not work without specifying the subfolder
						_lowerCamelCase				   =  BertConfig.from_pretrained('hf-internal-testing/tiny-random-bert-subfolder'		)

				_lowerCamelCase				   =  BertConfig.from_pretrained('hf-internal-testing/tiny-random-bert-subfolder'     ,  subfolder='bert'		)

				self.assertIsNotNone(snake_case__		)
		def   _snake_case       (       self     :     Optional[Any]		)      ->	int:
				# A mock response for an HTTP head request to emulate server down
				_lowerCamelCase				   =  mock.Mock()
				_lowerCamelCase				   =  5_0_0
				_lowerCamelCase				   =  {}
				_lowerCamelCase				   =  HTTPError
				_lowerCamelCase				   =  {}

				# Download this model to make sure it's in the cache.
				_lowerCamelCase				   =  BertConfig.from_pretrained('hf-internal-testing/tiny-random-bert'		)

				# Under the mock environment we get a 500 error when trying to reach the model.
				with mock.patch('requests.Session.request'     ,  return_value=snake_case__		) as mock_head:
						_lowerCamelCase				   =  BertConfig.from_pretrained('hf-internal-testing/tiny-random-bert'		)
						# This check we did call the fake head request
						mock_head.assert_called()
		def   _snake_case       (       self     :     Tuple		)      ->	List[str]:
				# This test is for deprecated behavior and can be removed in v5
				_lowerCamelCase				   =  BertConfig.from_pretrained(
				    'https://huggingface.co/hf-internal-testing/tiny-random-bert/resolve/main/config.json'		)
		def   _snake_case       (       self     :     Optional[Any]		)      ->	Dict:
				_lowerCamelCase				   =  AutoConfig.from_pretrained('bert-base-cased'		)
				_lowerCamelCase				   =  ['config.4.0.0.json']

				with tempfile.TemporaryDirectory() as tmp_dir:
						configuration.save_pretrained(snake_case__		)
						_lowerCamelCase				   =  2
						json.dump(configuration.to_dict()     ,  open(os.path.join(snake_case__     ,  'config.4.0.0.json'		)     ,  'w'		)		)

						# This should pick the new configuration file as the version of Transformers is > 4.0.0
						_lowerCamelCase				   =  AutoConfig.from_pretrained(snake_case__		)
						self.assertEqual(new_configuration.hidden_size     ,  2		)

						# Will need to be adjusted if we reach v42 and this test is still here.
						# Should pick the old configuration file as the version of Transformers is < 4.42.0
						_lowerCamelCase				   =  ['config.42.0.0.json']
						_lowerCamelCase				   =  7_6_8
						configuration.save_pretrained(snake_case__		)
						shutil.move(os.path.join(snake_case__     ,  'config.4.0.0.json'		)     ,  os.path.join(snake_case__     ,  'config.42.0.0.json'		)		)
						_lowerCamelCase				   =  AutoConfig.from_pretrained(snake_case__		)
						self.assertEqual(new_configuration.hidden_size     ,  7_6_8		)



		def   _snake_case       (       self     :     Dict		)      ->	List[Any]:
				# This repo has two configuration files, one for v4.0.0 and above with a different hidden size.
				_lowerCamelCase				   =  'hf-internal-testing/test-two-configs'

				import transformers as new_transformers

				_lowerCamelCase				   =  'v4.0.0'
				_lowerCamelCase   ,							_lowerCamelCase				   =  new_transformers.models.auto.AutoConfig.from_pretrained(
				    snake_case__     ,  return_unused_kwargs=snake_case__		)
				self.assertEqual(new_configuration.hidden_size     ,  2		)
				# This checks `_configuration_file` ia not kept in the kwargs by mistake.
				self.assertDictEqual(snake_case__     ,  {}		)

				# Testing an older version by monkey-patching the version in the module it's used.
				import transformers as old_transformers

				_lowerCamelCase				   =  'v3.0.0'
				_lowerCamelCase				   =  old_transformers.models.auto.AutoConfig.from_pretrained(snake_case__		)
				self.assertEqual(old_configuration.hidden_size     ,  7_6_8		)
import argparse
import math
import os
from copy import deepcopy

import torch
from audio_diffusion.models import DiffusionAttnUnetaD
from diffusion import sampling
from torch import nn

from diffusers import DanceDiffusionPipeline, IPNDMScheduler, UNetaDModel


A						=			{
    'gwf-440k': {
        'url': 'https://model-server.zqevans2.workers.dev/gwf-440k.ckpt',
        'sample_rate': 4_8_0_0_0,
        'sample_size': 6_5_5_3_6,
    },
    'jmann-small-190k': {
        'url': 'https://model-server.zqevans2.workers.dev/jmann-small-190k.ckpt',
        'sample_rate': 4_8_0_0_0,
        'sample_size': 6_5_5_3_6,
    },
    'jmann-large-580k': {
        'url': 'https://model-server.zqevans2.workers.dev/jmann-large-580k.ckpt',
        'sample_rate': 4_8_0_0_0,
        'sample_size': 1_3_1_0_7_2,
    },
    'maestro-uncond-150k': {
        'url': 'https://model-server.zqevans2.workers.dev/maestro-uncond-150k.ckpt',
        'sample_rate': 1_6_0_0_0,
        'sample_size': 6_5_5_3_6,
    },
    'unlocked-uncond-250k': {
        'url': 'https://model-server.zqevans2.workers.dev/unlocked-uncond-250k.ckpt',
        'sample_rate': 1_6_0_0_0,
        'sample_size': 6_5_5_3_6,
    },
    'honk-140k': {
        'url': 'https://model-server.zqevans2.workers.dev/honk-140k.ckpt',
        'sample_rate': 1_6_0_0_0,
        'sample_size': 6_5_5_3_6,
    },
}


def 	lowerCamelCase					(	UpperCamelCase     :				Optional[int]			,       UpperCamelCase     :				Union[str, Any]							)   ->    Optional[Any]:
		return torch.atana(UpperCamelCase			,       UpperCamelCase							) / math.pi * 2


def 	lowerCamelCase					(	UpperCamelCase     :				str							)   ->    Union[str, Any]:
		_lowerCamelCase				   =  torch.sin(t * math.pi / 2							) ** 2
		_lowerCamelCase				   =  (1 - sigma**2) ** 0.5
		return alpha_sigma_to_t(UpperCamelCase			,       UpperCamelCase							)





class     lowerCAmelCase__     (				SCREAMING_SNAKE_CASE_			):


		'''simple docstring'''






		pass





class     lowerCAmelCase__     (				nn.Module			):


		'''simple docstring'''






		def __init__(       self     :     Dict     ,  snake_case__     :     Any		)      ->	Optional[Any]:
				super().__init__()

				_lowerCamelCase				   =  DiffusionAttnUnetaD(snake_case__     ,  n_attn_layers=4		)
				_lowerCamelCase				   =  deepcopy(self.diffusion		)
				_lowerCamelCase				   =  torch.quasirandom.SobolEngine(1     ,  scramble=snake_case__		)



def 	lowerCamelCase					(	UpperCamelCase     :				List[Any]							)   ->    List[str]:
		_lowerCamelCase				   =  MODELS_MAP[model_name]['url']
		os.system(F"""wget {url} ./"""							)

		return F"""./{model_name}.ckpt"""


A						=			{
    '1': 'resnets.0',
    '2': 'attentions.0',
    '3': 'resnets.1',
    '4': 'attentions.1',
    '5': 'resnets.2',
    '6': 'attentions.2',
}
A						=			{
    '8': 'resnets.0',
    '9': 'attentions.0',
    '10': 'resnets.1',
    '11': 'attentions.1',
    '12': 'resnets.2',
    '13': 'attentions.2',
}
A						=			{
    '1': 'resnets.0',
    '2': 'attentions.0',
    '3': 'resnets.1',
    '4': 'attentions.1',
    '5': 'resnets.2',
    '6': 'attentions.2',
    '8': 'resnets.3',
    '9': 'attentions.3',
    '10': 'resnets.4',
    '11': 'attentions.4',
    '12': 'resnets.5',
    '13': 'attentions.5',
}
A						=			{
    '0': 'resnets.0',
    '1': 'resnets.1',
    '2': 'resnets.2',
    '4': 'resnets.0',
    '5': 'resnets.1',
    '6': 'resnets.2',
}

A						=			{
    'skip': 'conv_skip',
    'main.0': 'conv_1',
    'main.1': 'group_norm_1',
    'main.3': 'conv_2',
    'main.4': 'group_norm_2',
}

A						=			{
    'norm': 'group_norm',
    'qkv_proj': ['query', 'key', 'value'],
    'out_proj': ['proj_attn'],
}


def 	lowerCamelCase					(	UpperCamelCase     :				Tuple							)   ->    int:
		if name.startswith('skip'							):
				return name.replace('skip'			,       RES_CONV_MAP['skip']							)

		# name has to be of format main.{digit}
		if not name.startswith('main.'							):
				raise ValueError(F"""ResConvBlock error with {name}"""							)

		return name.replace(name[:6]			,       RES_CONV_MAP[name[:6]]							)


def 	lowerCamelCase					(	UpperCamelCase     :				Optional[Any]							)   ->    Tuple:
		for key, value in ATTN_MAP.items():
				if name.startswith(UpperCamelCase							) and not isinstance(UpperCamelCase			,       UpperCamelCase							):
						return name.replace(UpperCamelCase			,       UpperCamelCase							)
				elif name.startswith(UpperCamelCase							):
						return [name.replace(UpperCamelCase			,       UpperCamelCase							) for v in value]
		raise ValueError(F"""Attn error with {name}"""							)


def 	lowerCamelCase					(	UpperCamelCase     :				Any			,       UpperCamelCase     :				int=13							)   ->    Optional[int]:
		_lowerCamelCase				   =  input_string

		if string.split('.'							)[0] == "timestep_embed":
				return string.replace('timestep_embed'			,       'time_proj'							)

		_lowerCamelCase				   =  0
		if string.startswith('net.3.'							):
				depth += 1
				_lowerCamelCase				   =  string[6:]
		elif string.startswith('net.'							):
				_lowerCamelCase				   =  string[4:]

		while string.startswith('main.7.'							):
				depth += 1
				_lowerCamelCase				   =  string[7:]

		if string.startswith('main.'							):
				_lowerCamelCase				   =  string[5:]

		# mid block
		if string[:2].isdigit():
				_lowerCamelCase				   =  string[:2]
				_lowerCamelCase				   =  string[2:]
		else:
				_lowerCamelCase				   =  string[0]
				_lowerCamelCase				   =  string[1:]

		if depth == max_depth:
				_lowerCamelCase				   =  MID_NUM_TO_LAYER[layer_num]
				_lowerCamelCase				   =  'mid_block'
		elif depth > 0 and int(UpperCamelCase							) < 7:
				_lowerCamelCase				   =  DOWN_NUM_TO_LAYER[layer_num]
				_lowerCamelCase				   =  F"""down_blocks.{depth}"""
		elif depth > 0 and int(UpperCamelCase							) > 7:
				_lowerCamelCase				   =  UP_NUM_TO_LAYER[layer_num]
				_lowerCamelCase				   =  F"""up_blocks.{max_depth - depth - 1}"""
		elif depth == 0:
				_lowerCamelCase				   =  DEPTH_0_TO_LAYER[layer_num]
				_lowerCamelCase				   =  F"""up_blocks.{max_depth - 1}""" if int(UpperCamelCase							) > 3 else 'down_blocks.0'

		if not string_left.startswith('.'							):
				raise ValueError(F"""Naming error with {input_string} and string_left: {string_left}."""							)

		_lowerCamelCase				   =  string_left[1:]

		if "resnets" in new_layer:
				_lowerCamelCase				   =  convert_resconv_naming(UpperCamelCase							)
		elif "attentions" in new_layer:
				_lowerCamelCase				   =  convert_attn_naming(UpperCamelCase							)
				_lowerCamelCase				   =  new_string_left

		if not isinstance(UpperCamelCase			,       UpperCamelCase							):
				_lowerCamelCase				   =  prefix + '.' + new_layer + '.' + string_left
		else:
				_lowerCamelCase				   =  [prefix + '.' + new_layer + '.' + s for s in string_left]
		return new_string


def 	lowerCamelCase					(	UpperCamelCase     :				List[Any]							)   ->    int:
		_lowerCamelCase				   =  {}
		for k, v in state_dict.items():
				if k.endswith('kernel'							):
						# up- and downsample layers, don't have trainable weights
						continue

				_lowerCamelCase				   =  rename(UpperCamelCase							)

				# check if we need to transform from Conv => Linear for attention
				if isinstance(UpperCamelCase			,       UpperCamelCase							):
						_lowerCamelCase				   =  transform_conv_attns(UpperCamelCase			,       UpperCamelCase			,       UpperCamelCase							)
				else:
						_lowerCamelCase				   =  v

		return new_state_dict


def 	lowerCamelCase					(	UpperCamelCase     :				List[str]			,       UpperCamelCase     :				Tuple			,       UpperCamelCase     :				Dict							)   ->    Optional[Any]:
		if len(UpperCamelCase							) == 1:
				if len(v.shape							) == 3:
						# weight
						_lowerCamelCase				   =  v[:, :, 0]
				else:
						# bias
						_lowerCamelCase				   =  v
		else:
				# qkv matrices
				_lowerCamelCase				   =  v.shape[0]
				_lowerCamelCase				   =  trippled_shape // 3
				for i in range(3							):
						if len(v.shape							) == 3:
								_lowerCamelCase				   =  v[i * single_shape : (i + 1) * single_shape, :, 0]
						else:
								_lowerCamelCase				   =  v[i * single_shape : (i + 1) * single_shape]
		return new_state_dict


def 	lowerCamelCase					(	UpperCamelCase     :				Any							)   ->    Optional[Any]:
		_lowerCamelCase				   =  torch.device('cuda' if torch.cuda.is_available() else 'cpu'							)

		_lowerCamelCase				   =  args.model_path.split('/'							)[-1].split('.'							)[0]
		if not os.path.isfile(args.model_path							):
				assert (
				    model_name == args.model_path
				), F"""Make sure to provide one of the official model names {MODELS_MAP.keys()}"""
				_lowerCamelCase				   =  download(UpperCamelCase							)

		_lowerCamelCase				   =  MODELS_MAP[model_name]['sample_rate']
		_lowerCamelCase				   =  MODELS_MAP[model_name]['sample_size']

		_lowerCamelCase				   =  Object()
		_lowerCamelCase				   =  sample_size
		_lowerCamelCase				   =  sample_rate
		_lowerCamelCase				   =  0

		_lowerCamelCase				   =  UNetaDModel(sample_size=UpperCamelCase			,       sample_rate=UpperCamelCase							)
		_lowerCamelCase				   =  diffusers_model.state_dict()

		_lowerCamelCase				   =  DiffusionUncond(UpperCamelCase							)
		orig_model.load_state_dict(torch.load(args.model_path			,       map_location=UpperCamelCase							)['state_dict']							)
		_lowerCamelCase				   =  orig_model.diffusion_ema.eval()
		_lowerCamelCase				   =  orig_model.state_dict()
		_lowerCamelCase				   =  rename_orig_weights(UpperCamelCase							)

		_lowerCamelCase				   =  set(renamed_state_dict.keys()							) - set(diffusers_state_dict.keys()							)
		_lowerCamelCase				   =  set(diffusers_state_dict.keys()							) - set(renamed_state_dict.keys()							)

		assert len(UpperCamelCase							) == 0, F"""Problem with {renamed_minus_diffusers}"""
		assert all(k.endswith('kernel'							) for k in list(UpperCamelCase							)							), F"""Problem with {diffusers_minus_renamed}"""

		for key, value in renamed_state_dict.items():
				assert (
				    diffusers_state_dict[key].squeeze().shape == value.squeeze().shape
				), F"""Shape for {key} doesn't match. Diffusers: {diffusers_state_dict[key].shape} vs. {value.shape}"""
				if key == "time_proj.weight":
						_lowerCamelCase				   =  value.squeeze()

				_lowerCamelCase				   =  value

		diffusers_model.load_state_dict(UpperCamelCase							)

		_lowerCamelCase				   =  1_00
		_lowerCamelCase				   =  33

		_lowerCamelCase				   =  IPNDMScheduler(num_train_timesteps=UpperCamelCase							)

		_lowerCamelCase				   =  torch.manual_seed(UpperCamelCase							)
		_lowerCamelCase				   =  torch.randn([1, 2, config.sample_size]			,       generator=UpperCamelCase							).to(UpperCamelCase							)

		_lowerCamelCase				   =  torch.linspace(1			,       0			,       steps + 1			,       device=UpperCamelCase							)[:-1]
		_lowerCamelCase				   =  get_crash_schedule(UpperCamelCase							)

		_lowerCamelCase				   =  DanceDiffusionPipeline(unet=UpperCamelCase			,       scheduler=UpperCamelCase							)

		_lowerCamelCase				   =  torch.manual_seed(33							)
		_lowerCamelCase				   =  pipe(num_inference_steps=UpperCamelCase			,       generator=UpperCamelCase							).audios

		_lowerCamelCase				   =  sampling.iplms_sample(UpperCamelCase			,       UpperCamelCase			,       UpperCamelCase			,       {}							)
		_lowerCamelCase				   =  generated.clamp(-1			,       1							)

		_lowerCamelCase				   =  (generated - audio).abs().sum()
		_lowerCamelCase				   =  (generated - audio).abs().max()

		if args.save:
				pipe.save_pretrained(args.checkpoint_path							)

		print('Diff sum'			,       UpperCamelCase							)
		print('Diff max'			,       UpperCamelCase							)

		assert diff_max < 1e-3, F"""Diff max: {diff_max} is too much :-/"""

		print(F"""Conversion for {model_name} successful!"""							)


if __name__ == "__main__":
				A						=			argparse.ArgumentParser()

				parser.add_argument('--model_path', default=None, type=str, required=True, help='Path to the model to convert.')
				parser.add_argument(
				    '--save', default=True, type=bool, required=False, help='Whether to save the converted model or not.'
				)
				parser.add_argument('--checkpoint_path', default=None, type=str, required=True, help='Path to the output model.')
				A						=			parser.parse_args()

				main(args)