





from typing import Optional

import torch
import torch.utils.checkpoint
from torch import Tensor, nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_outputs import (
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import logging
from .configuration_regnet import RegNetConfig


lowerCAmelCase_        =							logging.get_logger(__name__)

# General docstring
lowerCAmelCase_        =							'''RegNetConfig'''

# Base docstring
lowerCAmelCase_        =							'''facebook/regnet-y-040'''
lowerCAmelCase_        =							[1, 1_0_8_8, 7, 7]

# Image classification docstring
lowerCAmelCase_        =							'''facebook/regnet-y-040'''
lowerCAmelCase_        =							'''tabby, tabby cat'''

lowerCAmelCase_        =							[
    '''facebook/regnet-y-040''',
    # See all regnet models at https://huggingface.co/models?filter=regnet
]



class    __lowerCAmelCase			(     nn.Module	):






			def __init__(self   ,				__magic_name__   ,				__magic_name__   ,				__magic_name__ = 3   ,				__magic_name__ = 1   ,				__magic_name__ = 1   ,				__magic_name__ = "relu"   ,				)	->  int:


							'''simple docstring'''



							super().__init__()
							snake_case_		:       Dict           =    nn.Convad(
							    snake_case__   ,				snake_case__   ,				kernel_size=snake_case__   ,				stride=snake_case__   ,				padding=kernel_size // 2   ,				groups=snake_case__   ,				bias=snake_case__   ,				)
							snake_case_		:       Union[str, Any]           =    nn.BatchNormad(snake_case__       )
							snake_case_		:       List[Any]           =    ACTaFN[activation] if activation is not None else nn.Identity()






			def      lowerCamelCase    (self   ,				__magic_name__       )	->  List[str]:


							'''simple docstring'''



							snake_case_		:       str           =    self.convolution(snake_case__       )
							snake_case_		:       int           =    self.normalization(snake_case__       )
							snake_case_		:       int           =    self.activation(snake_case__       )
							return hidden_state




class    __lowerCAmelCase			(     nn.Module	):






			def __init__(self   ,				__magic_name__       )	->  Optional[int]:


							'''simple docstring'''



							super().__init__()
							snake_case_		:       int           =    RegNetConvLayer(
							    config.num_channels   ,				config.embedding_size   ,				kernel_size=3   ,				stride=2   ,				activation=config.hidden_act       )
							snake_case_		:       Union[str, Any]           =    config.num_channels






			def      lowerCamelCase    (self   ,				__magic_name__       )	->  Union[str, Any]:


							'''simple docstring'''



							snake_case_		:       List[str]           =    pixel_values.shape[1]
							if num_channels != self.num_channels:
											raise ValueError(
											    '''Make sure that the channel dimension of the pixel values match with the one set in the configuration.'''       )
							snake_case_		:       List[str]           =    self.embedder(snake_case__       )
							return hidden_state




class    __lowerCAmelCase			(     nn.Module	):






			def __init__(self   ,				__magic_name__   ,				__magic_name__   ,				__magic_name__ = 2       )	->  List[Any]:


							'''simple docstring'''



							super().__init__()
							snake_case_		:       Union[str, Any]           =    nn.Convad(snake_case__   ,				snake_case__   ,				kernel_size=1   ,				stride=snake_case__   ,				bias=snake_case__       )
							snake_case_		:       List[Any]           =    nn.BatchNormad(snake_case__       )






			def      lowerCamelCase    (self   ,				__magic_name__       )	->  Tensor:


							'''simple docstring'''



							snake_case_		:       List[Any]           =    self.convolution(snake_case__       )
							snake_case_		:       List[str]           =    self.normalization(snake_case__       )
							return hidden_state




class    __lowerCAmelCase			(     nn.Module	):






			def __init__(self   ,				__magic_name__   ,				__magic_name__       )	->  int:


							'''simple docstring'''



							super().__init__()

							snake_case_		:       Union[str, Any]           =    nn.AdaptiveAvgPoolad((1, 1)       )
							snake_case_		:       List[str]           =    nn.Sequential(
							    nn.Convad(snake_case__   ,				snake_case__   ,				kernel_size=1       )   ,				nn.ReLU()   ,				nn.Convad(snake_case__   ,				snake_case__   ,				kernel_size=1       )   ,				nn.Sigmoid()   ,				)






			def      lowerCamelCase    (self   ,				__magic_name__       )	->  Tuple:


							'''simple docstring'''



							snake_case_		:       Union[str, Any]           =    self.pooler(snake_case__       )
							snake_case_		:       Optional[int]           =    self.attention(snake_case__       )
							snake_case_		:       List[str]           =    hidden_state * attention
							return hidden_state




class    __lowerCAmelCase			(     nn.Module	):






			def __init__(self   ,				__magic_name__   ,				__magic_name__   ,				__magic_name__   ,				__magic_name__ = 1       )	->  Optional[int]:


							'''simple docstring'''



							super().__init__()
							snake_case_		:       Optional[Any]           =    in_channels != out_channels or stride != 1
							snake_case_		:       int           =    max(1   ,				out_channels // config.groups_width       )
							snake_case_		:       List[str]           =    (
							    RegNetShortCut(snake_case__   ,				snake_case__   ,				stride=snake_case__       ) if should_apply_shortcut else nn.Identity()
							)
							snake_case_		:       str           =    nn.Sequential(
							    RegNetConvLayer(snake_case__   ,				snake_case__   ,				kernel_size=1   ,				activation=config.hidden_act       )   ,				RegNetConvLayer(snake_case__   ,				snake_case__   ,				stride=snake_case__   ,				groups=snake_case__   ,				activation=config.hidden_act       )   ,				RegNetConvLayer(snake_case__   ,				snake_case__   ,				kernel_size=1   ,				activation=snake_case__       )   ,				)
							snake_case_		:       List[str]           =    ACTaFN[config.hidden_act]






			def      lowerCamelCase    (self   ,				__magic_name__       )	->  Any:


							'''simple docstring'''



							snake_case_		:       Any           =    hidden_state
							snake_case_		:       Optional[int]           =    self.layer(snake_case__       )
							snake_case_		:       Optional[Any]           =    self.shortcut(snake_case__       )
							hidden_state += residual
							snake_case_		:       List[Any]           =    self.activation(snake_case__       )
							return hidden_state




class    __lowerCAmelCase			(     nn.Module	):






			def __init__(self   ,				__magic_name__   ,				__magic_name__   ,				__magic_name__   ,				__magic_name__ = 1       )	->  Tuple:


							'''simple docstring'''



							super().__init__()
							snake_case_		:       int           =    in_channels != out_channels or stride != 1
							snake_case_		:       Dict           =    max(1   ,				out_channels // config.groups_width       )
							snake_case_		:       List[str]           =    (
							    RegNetShortCut(snake_case__   ,				snake_case__   ,				stride=snake_case__       ) if should_apply_shortcut else nn.Identity()
							)
							snake_case_		:       Union[str, Any]           =    nn.Sequential(
							    RegNetConvLayer(snake_case__   ,				snake_case__   ,				kernel_size=1   ,				activation=config.hidden_act       )   ,				RegNetConvLayer(snake_case__   ,				snake_case__   ,				stride=snake_case__   ,				groups=snake_case__   ,				activation=config.hidden_act       )   ,				RegNetSELayer(snake_case__   ,				reduced_channels=int(round(in_channels / 4       )       )       )   ,				RegNetConvLayer(snake_case__   ,				snake_case__   ,				kernel_size=1   ,				activation=snake_case__       )   ,				)
							snake_case_		:       List[str]           =    ACTaFN[config.hidden_act]






			def      lowerCamelCase    (self   ,				__magic_name__       )	->  Optional[Any]:


							'''simple docstring'''



							snake_case_		:       Tuple           =    hidden_state
							snake_case_		:       List[Any]           =    self.layer(snake_case__       )
							snake_case_		:       Any           =    self.shortcut(snake_case__       )
							hidden_state += residual
							snake_case_		:       Any           =    self.activation(snake_case__       )
							return hidden_state




class    __lowerCAmelCase			(     nn.Module	):






			def __init__(self   ,				__magic_name__   ,				__magic_name__   ,				__magic_name__   ,				__magic_name__ = 2   ,				__magic_name__ = 2   ,				)	->  List[Any]:


							'''simple docstring'''



							super().__init__()

							snake_case_		:       Any           =    RegNetXLayer if config.layer_type == '''x''' else RegNetYLayer

							snake_case_		:       int           =    nn.Sequential(
							    # downsampling is done in the first layer with stride of 2
							    layer(
							        snake_case__   ,				snake_case__   ,				snake_case__   ,				stride=snake_case__   ,				)   ,				*[layer(snake_case__   ,				snake_case__   ,				snake_case__       ) for _ in range(depth - 1       )]   ,				)






			def      lowerCamelCase    (self   ,				__magic_name__       )	->  str:


							'''simple docstring'''



							snake_case_		:       List[Any]           =    self.layers(snake_case__       )
							return hidden_state




class    __lowerCAmelCase			(     nn.Module	):






			def __init__(self   ,				__magic_name__       )	->  Any:


							'''simple docstring'''



							super().__init__()
							snake_case_		:       Any           =    nn.ModuleList([]       )
							# based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
							self.stages.append(
							    RegNetStage(
							        snake_case__   ,				config.embedding_size   ,				config.hidden_sizes[0]   ,				stride=2 if config.downsample_in_first_stage else 1   ,				depth=config.depths[0]   ,				)       )
							snake_case_		:       Dict           =    zip(config.hidden_sizes   ,				config.hidden_sizes[1:]       )
							for (in_channels, out_channels), depth in zip(snake_case__   ,				config.depths[1:]       ):
											self.stages.append(RegNetStage(snake_case__   ,				snake_case__   ,				snake_case__   ,				depth=snake_case__       )       )






			def      lowerCamelCase    (self   ,				__magic_name__   ,				__magic_name__ = False   ,				__magic_name__ = True       )	->  BaseModelOutputWithNoAttention:


							'''simple docstring'''



							snake_case_		:       Optional[int]           =    () if output_hidden_states else None

							for stage_module in self.stages:
											if output_hidden_states:
															snake_case_		:       Tuple           =    hidden_states + (hidden_state,)

											snake_case_		:       Optional[Any]           =    stage_module(snake_case__       )

							if output_hidden_states:
											snake_case_		:       List[str]           =    hidden_states + (hidden_state,)

							if not return_dict:
											return tuple(v for v in [hidden_state, hidden_states] if v is not None       )

							return BaseModelOutputWithNoAttention(last_hidden_state=snake_case__   ,				hidden_states=snake_case__       )




class    __lowerCAmelCase			(     SCREAMING_SNAKE_CASE_	):
			lowerCamelCase_	:	Union[str, Any]			     = RegNetConfig
			lowerCamelCase_	:	str			     = '''regnet'''
			lowerCamelCase_	:	List[str]			     = '''pixel_values'''
			lowerCamelCase_	:	Union[str, Any]			     = True






			def      lowerCamelCase    (self   ,				__magic_name__       )	->  Any:


							'''simple docstring'''



							if isinstance(snake_case__   ,				nn.Convad       ):
											nn.init.kaiming_normal_(module.weight   ,				mode='''fan_out'''   ,				nonlinearity='''relu'''       )
							elif isinstance(snake_case__   ,				(nn.BatchNormad, nn.GroupNorm)       ):
											nn.init.constant_(module.weight   ,				1       )
											nn.init.constant_(module.bias   ,				0       )






			def      lowerCamelCase    (self   ,				__magic_name__   ,				__magic_name__=False       )	->  Any:


							'''simple docstring'''



							if isinstance(snake_case__   ,				snake_case__       ):
											snake_case_		:       Any           =    value




lowerCAmelCase_        =							r'''\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it\n    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\n        config ([`RegNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n'''

lowerCAmelCase_        =							r'''\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConvNextImageProcessor.__call__`] for details.\n\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.\n'''



@add_start_docstrings(
    '''The bare RegNet model outputting raw features without any specific head on top.''',	SCREAMING_SNAKE_CASE_,	)
# Copied from transformers.models.resnet.modeling_resnet.ResNetModel with RESNET->REGNET,ResNet->RegNet
class    __lowerCAmelCase			(     SCREAMING_SNAKE_CASE_	):






			def __init__(self   ,				__magic_name__       )	->  Union[str, Any]:


							'''simple docstring'''



							super().__init__(snake_case__       )
							snake_case_		:       Optional[int]           =    config
							snake_case_		:       Tuple           =    RegNetEmbeddings(snake_case__       )
							snake_case_		:       Optional[Any]           =    RegNetEncoder(snake_case__       )
							snake_case_		:       Dict           =    nn.AdaptiveAvgPoolad((1, 1)       )
							# Initialize weights and apply final processing
							self.post_init()






			@add_start_docstrings_to_model_forward(snake_case__       )
			@add_code_sample_docstrings(
			    checkpoint=_CHECKPOINT_FOR_DOC   ,				output_type=snake_case__   ,				config_class=_CONFIG_FOR_DOC   ,				modality='''vision'''   ,				expected_output=_EXPECTED_OUTPUT_SHAPE   ,				)
			def      lowerCamelCase    (self   ,				__magic_name__   ,				__magic_name__ = None   ,				__magic_name__ = None       )	->  BaseModelOutputWithPoolingAndNoAttention:


							'''simple docstring'''



							snake_case_		:       Dict           =    (
							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
							)
							snake_case_		:       str           =    return_dict if return_dict is not None else self.config.use_return_dict

							snake_case_		:       int           =    self.embedder(snake_case__       )

							snake_case_		:       Optional[Any]           =    self.encoder(
							    snake_case__   ,				output_hidden_states=snake_case__   ,				return_dict=snake_case__       )

							snake_case_		:       List[Any]           =    encoder_outputs[0]

							snake_case_		:       Union[str, Any]           =    self.pooler(snake_case__       )

							if not return_dict:
											return (last_hidden_state, pooled_output) + encoder_outputs[1:]

							return BaseModelOutputWithPoolingAndNoAttention(
							    last_hidden_state=snake_case__   ,				pooler_output=snake_case__   ,				hidden_states=encoder_outputs.hidden_states   ,				)




@add_start_docstrings(
    '''\n    RegNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    ''',	SCREAMING_SNAKE_CASE_,	)
# Copied from transformers.models.resnet.modeling_resnet.ResNetForImageClassification with RESNET->REGNET,ResNet->RegNet,resnet->regnet
class    __lowerCAmelCase			(     SCREAMING_SNAKE_CASE_	):






			def __init__(self   ,				__magic_name__       )	->  Any:


							'''simple docstring'''



							super().__init__(snake_case__       )
							snake_case_		:       Union[str, Any]           =    config.num_labels
							snake_case_		:       Any           =    RegNetModel(snake_case__       )
							# classification head
							snake_case_		:       Tuple           =    nn.Sequential(
							    nn.Flatten()   ,				nn.Linear(config.hidden_sizes[-1]   ,				config.num_labels       ) if config.num_labels > 0 else nn.Identity()   ,				)
							# initialize weights and apply final processing
							self.post_init()






			@add_start_docstrings_to_model_forward(snake_case__       )
			@add_code_sample_docstrings(
			    checkpoint=_IMAGE_CLASS_CHECKPOINT   ,				output_type=snake_case__   ,				config_class=_CONFIG_FOR_DOC   ,				expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT   ,				)
			def      lowerCamelCase    (self   ,				__magic_name__ = None   ,				__magic_name__ = None   ,				__magic_name__ = None   ,				__magic_name__ = None   ,				)	->  ImageClassifierOutputWithNoAttention:


							'''simple docstring'''



							snake_case_		:       Tuple           =    return_dict if return_dict is not None else self.config.use_return_dict

							snake_case_		:       int           =    self.regnet(snake_case__   ,				output_hidden_states=snake_case__   ,				return_dict=snake_case__       )

							snake_case_		:       Dict           =    outputs.pooler_output if return_dict else outputs[1]

							snake_case_		:       int           =    self.classifier(snake_case__       )

							snake_case_		:       Dict           =    None

							if labels is not None:
											if self.config.problem_type is None:
															if self.num_labels == 1:
																			snake_case_		:       List[Any]           =    '''regression'''
															elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
																			snake_case_		:       List[Any]           =    '''single_label_classification'''
															else:
																			snake_case_		:       Union[str, Any]           =    '''multi_label_classification'''
											if self.config.problem_type == "regression":
															snake_case_		:       Any           =    MSELoss()
															if self.num_labels == 1:
																			snake_case_		:       List[Any]           =    loss_fct(logits.squeeze()   ,				labels.squeeze()       )
															else:
																			snake_case_		:       Tuple           =    loss_fct(snake_case__   ,				snake_case__       )
											elif self.config.problem_type == "single_label_classification":
															snake_case_		:       List[Any]           =    CrossEntropyLoss()
															snake_case_		:       Union[str, Any]           =    loss_fct(logits.view(-1   ,				self.num_labels       )   ,				labels.view(-1       )       )
											elif self.config.problem_type == "multi_label_classification":
															snake_case_		:       Optional[Any]           =    BCEWithLogitsLoss()
															snake_case_		:       int           =    loss_fct(snake_case__   ,				snake_case__       )

							if not return_dict:
											snake_case_		:       Optional[int]           =    (logits,) + outputs[2:]
											return (loss,) + output if loss is not None else output

							return ImageClassifierOutputWithNoAttention(loss=snake_case__   ,				logits=snake_case__   ,				hidden_states=outputs.hidden_states       )

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)