





import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionPanoramaPipeline,
    UNetaDConditionModel,
)
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 TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()



@skip_mps
class    __lowerCAmelCase			(     SCREAMING_SNAKE_CASE_,	SCREAMING_SNAKE_CASE_,	unittest.TestCase	):
			lowerCamelCase_	:	int			     = StableDiffusionPanoramaPipeline
			lowerCamelCase_	:	List[str]			     = TEXT_TO_IMAGE_PARAMS
			lowerCamelCase_	:	int			     = TEXT_TO_IMAGE_BATCH_PARAMS
			lowerCamelCase_	:	List[str]			     = TEXT_TO_IMAGE_IMAGE_PARAMS
			lowerCamelCase_	:	Optional[Any]			     = TEXT_TO_IMAGE_IMAGE_PARAMS






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


							'''simple docstring'''



							torch.manual_seed(0       )
							snake_case_		:       List[str]           =    UNetaDConditionModel(
							    block_out_channels=(32, 64)   ,				layers_per_block=1   ,				sample_size=32   ,				in_channels=4   ,				out_channels=4   ,				down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')   ,				up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')   ,				cross_attention_dim=32   ,				)
							snake_case_		:       Tuple           =    DDIMScheduler()
							torch.manual_seed(0       )
							snake_case_		:       List[str]           =    AutoencoderKL(
							    block_out_channels=[32, 64]   ,				in_channels=3   ,				out_channels=3   ,				down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']   ,				up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']   ,				latent_channels=4   ,				)
							torch.manual_seed(0       )
							snake_case_		:       Tuple           =    CLIPTextConfig(
							    bos_token_id=0   ,				eos_token_id=2   ,				hidden_size=32   ,				intermediate_size=37   ,				layer_norm_eps=1e-05   ,				num_attention_heads=4   ,				num_hidden_layers=5   ,				pad_token_id=1   ,				vocab_size=1000   ,				)
							snake_case_		:       Optional[int]           =    CLIPTextModel(snake_case__       )
							snake_case_		:       Union[str, Any]           =    CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''       )

							snake_case_		:       Tuple           =    {
							    '''unet''': unet,
							    '''scheduler''': scheduler,
							    '''vae''': vae,
							    '''text_encoder''': text_encoder,
							    '''tokenizer''': tokenizer,
							    '''safety_checker''': None,
							    '''feature_extractor''': None,
							}
							return components






			def      lowerCamelCase    (self   ,				__magic_name__   ,				__magic_name__=0       )	->  str:


							'''simple docstring'''



							snake_case_		:       Tuple           =    torch.manual_seed(snake_case__       )
							snake_case_		:       List[Any]           =    {
							    '''prompt''': '''a photo of the dolomites''',
							    '''generator''': generator,
							    # Setting height and width to None to prevent OOMs on CPU.
							    '''height''': None,
							    '''width''': None,
							    '''num_inference_steps''': 1,
							    '''guidance_scale''': 6.0,
							    '''output_type''': '''numpy''',
							}
							return inputs






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


							'''simple docstring'''



							snake_case_		:       int           =    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							snake_case_		:       Dict           =    self.get_dummy_components()
							snake_case_		:       Tuple           =    StableDiffusionPanoramaPipeline(**snake_case__       )
							snake_case_		:       int           =    sd_pipe.to(snake_case__       )
							sd_pipe.set_progress_bar_config(disable=snake_case__       )

							snake_case_		:       Union[str, Any]           =    self.get_dummy_inputs(snake_case__       )
							snake_case_		:       int           =    sd_pipe(**snake_case__       ).images
							snake_case_		:       List[Any]           =    image[0, -3:, -3:, -1]
							assert image.shape == (1, 64, 64, 3)

							snake_case_		:       List[str]           =    np.array([0.6_186, 0.5_374, 0.4_915, 0.4_135, 0.4_114, 0.4_563, 0.5_128, 0.4_977, 0.4_757]       )

							assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-2






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


							'''simple docstring'''



							super().test_inference_batch_consistent(batch_sizes=[1, 2]       )






			def      lowerCamelCase    (self       )	->  int:


							'''simple docstring'''



							super().test_inference_batch_single_identical(batch_size=2   ,				expected_max_diff=3.25e-3       )






			def      lowerCamelCase    (self       )	->  Any:


							'''simple docstring'''



							snake_case_		:       List[Any]           =    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							snake_case_		:       int           =    self.get_dummy_components()
							snake_case_		:       Optional[Any]           =    StableDiffusionPanoramaPipeline(**snake_case__       )
							snake_case_		:       Union[str, Any]           =    sd_pipe.to(snake_case__       )
							sd_pipe.set_progress_bar_config(disable=snake_case__       )

							snake_case_		:       List[Any]           =    self.get_dummy_inputs(snake_case__       )
							snake_case_		:       List[Any]           =    '''french fries'''
							snake_case_		:       Union[str, Any]           =    sd_pipe(**snake_case__   ,				negative_prompt=snake_case__       )
							snake_case_		:       Optional[Any]           =    output.images
							snake_case_		:       Tuple           =    image[0, -3:, -3:, -1]

							assert image.shape == (1, 64, 64, 3)

							snake_case_		:       Dict           =    np.array([0.6_187, 0.5_375, 0.4_915, 0.4_136, 0.4_114, 0.4_563, 0.5_128, 0.4_976, 0.4_757]       )

							assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-2






			def      lowerCamelCase    (self       )	->  int:


							'''simple docstring'''



							snake_case_		:       int           =    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							snake_case_		:       Union[str, Any]           =    self.get_dummy_components()
							snake_case_		:       Optional[Any]           =    StableDiffusionPanoramaPipeline(**snake_case__       )
							snake_case_		:       Dict           =    sd_pipe.to(snake_case__       )
							sd_pipe.set_progress_bar_config(disable=snake_case__       )

							snake_case_		:       Optional[Any]           =    self.get_dummy_inputs(snake_case__       )
							snake_case_		:       Optional[Any]           =    sd_pipe(**snake_case__   ,				view_batch_size=2       )
							snake_case_		:       str           =    output.images
							snake_case_		:       Tuple           =    image[0, -3:, -3:, -1]

							assert image.shape == (1, 64, 64, 3)

							snake_case_		:       Tuple           =    np.array([0.6_187, 0.5_375, 0.4_915, 0.4_136, 0.4_114, 0.4_563, 0.5_128, 0.4_976, 0.4_757]       )

							assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-2






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


							'''simple docstring'''



							snake_case_		:       Tuple           =    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							snake_case_		:       List[str]           =    self.get_dummy_components()
							snake_case_		:       List[Any]           =    EulerAncestralDiscreteScheduler(
							    beta_start=0.00_085   ,				beta_end=0.012   ,				beta_schedule='''scaled_linear'''       )
							snake_case_		:       Tuple           =    StableDiffusionPanoramaPipeline(**snake_case__       )
							snake_case_		:       Tuple           =    sd_pipe.to(snake_case__       )
							sd_pipe.set_progress_bar_config(disable=snake_case__       )

							snake_case_		:       Any           =    self.get_dummy_inputs(snake_case__       )
							snake_case_		:       List[Any]           =    sd_pipe(**snake_case__       ).images
							snake_case_		:       Optional[Any]           =    image[0, -3:, -3:, -1]

							assert image.shape == (1, 64, 64, 3)

							snake_case_		:       List[str]           =    np.array([0.4_024, 0.6_510, 0.4_901, 0.5_378, 0.5_813, 0.5_622, 0.4_795, 0.4_467, 0.4_952]       )

							assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-2






			def      lowerCamelCase    (self       )	->  str:


							'''simple docstring'''



							snake_case_		:       List[Any]           =    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							snake_case_		:       List[Any]           =    self.get_dummy_components()
							snake_case_		:       Optional[int]           =    PNDMScheduler(
							    beta_start=0.00_085   ,				beta_end=0.012   ,				beta_schedule='''scaled_linear'''   ,				skip_prk_steps=snake_case__       )
							snake_case_		:       Union[str, Any]           =    StableDiffusionPanoramaPipeline(**snake_case__       )
							snake_case_		:       Union[str, Any]           =    sd_pipe.to(snake_case__       )
							sd_pipe.set_progress_bar_config(disable=snake_case__       )

							snake_case_		:       str           =    self.get_dummy_inputs(snake_case__       )
							snake_case_		:       str           =    sd_pipe(**snake_case__       ).images
							snake_case_		:       Dict           =    image[0, -3:, -3:, -1]

							assert image.shape == (1, 64, 64, 3)

							snake_case_		:       Any           =    np.array([0.6_391, 0.6_291, 0.4_861, 0.5_134, 0.5_552, 0.4_578, 0.5_032, 0.5_023, 0.4_539]       )

							assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-2




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






			def      lowerCamelCase    (self       )	->  Optional[int]:


							'''simple docstring'''



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






			def      lowerCamelCase    (self   ,				__magic_name__=0       )	->  Tuple:


							'''simple docstring'''



							snake_case_		:       Tuple           =    torch.manual_seed(snake_case__       )
							snake_case_		:       List[str]           =    {
							    '''prompt''': '''a photo of the dolomites''',
							    '''generator''': generator,
							    '''num_inference_steps''': 3,
							    '''guidance_scale''': 7.5,
							    '''output_type''': '''numpy''',
							}
							return inputs






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


							'''simple docstring'''



							snake_case_		:       List[str]           =    '''stabilityai/stable-diffusion-2-base'''
							snake_case_		:       str           =    DDIMScheduler.from_pretrained(snake_case__   ,				subfolder='''scheduler'''       )
							snake_case_		:       Tuple           =    StableDiffusionPanoramaPipeline.from_pretrained(snake_case__   ,				scheduler=snake_case__   ,				safety_checker=snake_case__       )
							pipe.to(snake_case__       )
							pipe.set_progress_bar_config(disable=snake_case__       )
							pipe.enable_attention_slicing()

							snake_case_		:       Optional[Any]           =    self.get_inputs()
							snake_case_		:       Dict           =    pipe(**snake_case__       ).images
							snake_case_		:       Any           =    image[0, -3:, -3:, -1].flatten()

							assert image.shape == (1, 512, 2048, 3)

							snake_case_		:       Any           =    np.array(
							    [
							        0.36_968_392,
							        0.27_025_372,
							        0.32_446_766,
							        0.28_379_387,
							        0.36_363_274,
							        0.30_733_347,
							        0.27_100_027,
							        0.27_054_125,
							        0.25_536_096,
							    ]       )

							assert np.abs(expected_slice - image_slice       ).max() < 1e-2






			def      lowerCamelCase    (self       )	->  Tuple:


							'''simple docstring'''



							snake_case_		:       Any           =    StableDiffusionPanoramaPipeline.from_pretrained(
							    '''stabilityai/stable-diffusion-2-base'''   ,				safety_checker=snake_case__       )
							snake_case_		:       Optional[Any]           =    LMSDiscreteScheduler.from_config(pipe.scheduler.config       )
							pipe.to(snake_case__       )
							pipe.set_progress_bar_config(disable=snake_case__       )
							pipe.enable_attention_slicing()

							snake_case_		:       Any           =    self.get_inputs()
							snake_case_		:       Any           =    pipe(**snake_case__       ).images
							snake_case_		:       Tuple           =    image[0, -3:, -3:, -1].flatten()

							assert image.shape == (1, 512, 2048, 3)

							snake_case_		:       Any           =    np.array(
							    [
							        [
							            0.0,
							            0.0,
							            0.0,
							            0.0,
							            0.0,
							            0.0,
							            0.0,
							            0.0,
							            0.0,
							        ]
							    ]       )

							assert np.abs(expected_slice - image_slice       ).max() < 1e-3






			def      lowerCamelCase    (self       )	->  str:


							'''simple docstring'''



							snake_case_		:       str           =    0

							def callback_fn(__magic_name__   ,				__magic_name__   ,				__magic_name__       ) -> None:
											snake_case_		:       Optional[int]           =    True
											nonlocal number_of_steps
											number_of_steps += 1
											if step == 1:
															snake_case_		:       Union[str, Any]           =    latents.detach().cpu().numpy()
															assert latents.shape == (1, 4, 64, 256)
															snake_case_		:       Dict           =    latents[0, -3:, -3:, -1]

															snake_case_		:       int           =    np.array(
															    [
															        0.18_681_869,
															        0.33_907_816,
															        0.5_361_276,
															        0.14_432_865,
															        -0.02_856_611,
															        -0.73_941_123,
															        0.23_397_987,
															        0.47_322_682,
															        -0.37_823_164,
															    ]       )
															assert np.abs(latents_slice.flatten() - expected_slice       ).max() < 5e-2
											elif step == 2:
															snake_case_		:       Union[str, Any]           =    latents.detach().cpu().numpy()
															assert latents.shape == (1, 4, 64, 256)
															snake_case_		:       Optional[Any]           =    latents[0, -3:, -3:, -1]

															snake_case_		:       int           =    np.array(
															    [
															        0.18_539_645,
															        0.33_987_248,
															        0.5_378_559,
															        0.14_437_142,
															        -0.02_455_261,
															        -0.7_338_317,
															        0.23_990_755,
															        0.47_356_272,
															        -0.3_786_505,
															    ]       )

															assert np.abs(latents_slice.flatten() - expected_slice       ).max() < 5e-2

							snake_case_		:       str           =    False

							snake_case_		:       int           =    '''stabilityai/stable-diffusion-2-base'''
							snake_case_		:       Optional[int]           =    DDIMScheduler.from_pretrained(snake_case__   ,				subfolder='''scheduler'''       )
							snake_case_		:       str           =    StableDiffusionPanoramaPipeline.from_pretrained(snake_case__   ,				scheduler=snake_case__   ,				safety_checker=snake_case__       )
							snake_case_		:       Dict           =    pipe.to(snake_case__       )
							pipe.set_progress_bar_config(disable=snake_case__       )
							pipe.enable_attention_slicing()

							snake_case_		:       str           =    self.get_inputs()
							pipe(**snake_case__   ,				callback=snake_case__   ,				callback_steps=1       )
							assert callback_fn.has_been_called
							assert number_of_steps == 3






			def      lowerCamelCase    (self       )	->  Any:


							'''simple docstring'''



							torch.cuda.empty_cache()
							torch.cuda.reset_max_memory_allocated()
							torch.cuda.reset_peak_memory_stats()

							snake_case_		:       Optional[Any]           =    '''stabilityai/stable-diffusion-2-base'''
							snake_case_		:       Tuple           =    DDIMScheduler.from_pretrained(snake_case__   ,				subfolder='''scheduler'''       )
							snake_case_		:       List[Any]           =    StableDiffusionPanoramaPipeline.from_pretrained(snake_case__   ,				scheduler=snake_case__   ,				safety_checker=snake_case__       )
							snake_case_		:       List[str]           =    pipe.to(snake_case__       )
							pipe.set_progress_bar_config(disable=snake_case__       )
							pipe.enable_attention_slicing(1       )
							pipe.enable_sequential_cpu_offload()

							snake_case_		:       List[str]           =    self.get_inputs()
							snake_case_		:       Optional[Any]           =    pipe(**snake_case__       )

							snake_case_		:       Any           =    torch.cuda.max_memory_allocated()
							# make sure that less than 5.2 GB is allocated
							assert mem_bytes < 5.5 * 10**9

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)