




"""simple docstring"""




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      __UpperCAmelCase     ( __UpperCamelCase	,				__UpperCamelCase	):
			return torch.atana(__UpperCamelCase	,				__UpperCamelCase	) / math.pi * 2







def      __UpperCAmelCase     ( __UpperCamelCase	):
			__lowercase			:  str								=       torch.sin(t * math.pi / 2	) ** 2
			__lowercase			:  Dict								=       (1 - sigma**2) ** 0.5
			return alpha_sigma_to_t(__UpperCamelCase	,				__UpperCamelCase	)

class   UpperCAmelCase_    (			snake_case		):
		pass


class   UpperCAmelCase_    (			nn.Module		):
		def __init__(						self				,						UpperCamelCase_       )						->       Tuple:
					super().__init__()

					__lowercase			:  str								=       DiffusionAttnUnetaD(UpperCamelCase_				,						n_attn_layers=4       )
					__lowercase			:  int								=       deepcopy(self.diffusion       )
					__lowercase			:  Union[str, Any]								=       torch.quasirandom.SobolEngine(1				,						scramble=UpperCamelCase_       )







def      __UpperCAmelCase     ( __UpperCamelCase	):
			__lowercase			:  Any								=       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      __UpperCAmelCase     ( __UpperCamelCase	):
			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      __UpperCAmelCase     ( __UpperCamelCase	):
			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      __UpperCAmelCase     ( __UpperCamelCase	,				__UpperCamelCase=13	):
			__lowercase			:  Dict								=       input_string

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

			__lowercase			:  Any								=       0
			if string.startswith('''net.3.'''	):
						depth += 1
						__lowercase			:  List[str]								=       string[6:]
			elif string.startswith('''net.'''	):
						__lowercase			:  Union[str, Any]								=       string[4:]

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

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

			# mid block
			if string[:2].isdigit():
						__lowercase			:  str								=       string[:2]
						__lowercase			:  List[str]								=       string[2:]
			else:
						__lowercase			:  Dict								=       string[0]
						__lowercase			:  Tuple								=       string[1:]

			if depth == max_depth:
						__lowercase			:  List[Any]								=       MID_NUM_TO_LAYER[layer_num]
						__lowercase			:  Union[str, Any]								=       '''mid_block'''
			elif depth > 0 and int(__UpperCamelCase	) < 7:
						__lowercase			:  str								=       DOWN_NUM_TO_LAYER[layer_num]
						__lowercase			:  int								=       f"""down_blocks.{depth}"""
			elif depth > 0 and int(__UpperCamelCase	) > 7:
						__lowercase			:  List[Any]								=       UP_NUM_TO_LAYER[layer_num]
						__lowercase			:  List[str]								=       f"""up_blocks.{max_depth - depth - 1}"""
			elif depth == 0:
						__lowercase			:  str								=       DEPTH_0_TO_LAYER[layer_num]
						__lowercase			:  Any								=       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}."""	)

			__lowercase			:  str								=       string_left[1:]

			if "resnets" in new_layer:
						__lowercase			:  Dict								=       convert_resconv_naming(__UpperCamelCase	)
			elif "attentions" in new_layer:
						__lowercase			:  str								=       convert_attn_naming(__UpperCamelCase	)
						__lowercase			:  Optional[int]								=       new_string_left

			if not isinstance(__UpperCamelCase	,				__UpperCamelCase	):
						__lowercase			:  int								=       prefix + '''.''' + new_layer + '''.''' + string_left
			else:
						__lowercase			:  List[Any]								=       [prefix + '''.''' + new_layer + '''.''' + s for s in string_left]
			return new_string







def      __UpperCAmelCase     ( __UpperCamelCase	):
			__lowercase			:  List[str]								=       {}
			for k, v in state_dict.items():
						if k.endswith('''kernel'''	):
									# up- and downsample layers, don't have trainable weights
									continue

						__lowercase			:  List[str]								=       rename(__UpperCamelCase	)

						# check if we need to transform from Conv => Linear for attention
						if isinstance(__UpperCamelCase	,				__UpperCamelCase	):
									__lowercase			:  List[Any]								=       transform_conv_attns(__UpperCamelCase	,				__UpperCamelCase	,				__UpperCamelCase	)
						else:
									__lowercase			:  Dict								=       v

			return new_state_dict







def      __UpperCAmelCase     ( __UpperCamelCase	,				__UpperCamelCase	,				__UpperCamelCase	):
			if len(__UpperCamelCase	) == 1:
						if len(v.shape	) == 3:
									# weight
									__lowercase			:  Optional[Any]								=       v[:, :, 0]
						else:
									# bias
									__lowercase			:  Dict								=       v
			else:
						# qkv matrices
						__lowercase			:  Any								=       v.shape[0]
						__lowercase			:  int								=       trippled_shape // 3
						for i in range(3	):
									if len(v.shape	) == 3:
												__lowercase			:  int								=       v[i * single_shape : (i + 1) * single_shape, :, 0]
									else:
												__lowercase			:  Dict								=       v[i * single_shape : (i + 1) * single_shape]
			return new_state_dict







def      __UpperCAmelCase     ( __UpperCamelCase	):
			__lowercase			:  int								=       torch.device('''cuda''' if torch.cuda.is_available() else '''cpu'''	)

			__lowercase			:  List[str]								=       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()}"""
						__lowercase			:  List[Any]								=       download(__UpperCamelCase	)

			__lowercase			:  Union[str, Any]								=       MODELS_MAP[model_name]['''sample_rate''']
			__lowercase			:  List[Any]								=       MODELS_MAP[model_name]['''sample_size''']

			__lowercase			:  int								=       Object()
			__lowercase			:  List[str]								=       sample_size
			__lowercase			:  Tuple								=       sample_rate
			__lowercase			:  Optional[int]								=       0

			__lowercase			:  List[str]								=       UNetaDModel(sample_size=__UpperCamelCase	,				sample_rate=__UpperCamelCase	)
			__lowercase			:  Dict								=       diffusers_model.state_dict()

			__lowercase			:  int								=       DiffusionUncond(__UpperCamelCase	)
			orig_model.load_state_dict(torch.load(args.model_path	,				map_location=__UpperCamelCase	)['''state_dict''']	)
			__lowercase			:  Dict								=       orig_model.diffusion_ema.eval()
			__lowercase			:  Optional[int]								=       orig_model.state_dict()
			__lowercase			:  Any								=       rename_orig_weights(__UpperCamelCase	)

			__lowercase			:  int								=       set(renamed_state_dict.keys()	) - set(diffusers_state_dict.keys()	)
			__lowercase			:  Tuple								=       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":
									__lowercase			:  Optional[Any]								=       value.squeeze()

						__lowercase			:  str								=       value

			diffusers_model.load_state_dict(__UpperCamelCase	)

			__lowercase			:  Union[str, Any]								=       1_00
			__lowercase			:  Tuple								=       33

			__lowercase			:  Optional[Any]								=       IPNDMScheduler(num_train_timesteps=__UpperCamelCase	)

			__lowercase			:  Dict								=       torch.manual_seed(__UpperCamelCase	)
			__lowercase			:  List[str]								=       torch.randn([1, 2, config.sample_size]	,				generator=__UpperCamelCase	).to(__UpperCamelCase	)

			__lowercase			:  List[Any]								=       torch.linspace(1	,				0	,				steps + 1	,				device=__UpperCamelCase	)[:-1]
			__lowercase			:  Any								=       get_crash_schedule(__UpperCamelCase	)

			__lowercase			:  int								=       DanceDiffusionPipeline(unet=__UpperCamelCase	,				scheduler=__UpperCamelCase	)

			__lowercase			:  List[str]								=       torch.manual_seed(33	)
			__lowercase			:  Optional[int]								=       pipe(num_inference_steps=__UpperCamelCase	,				generator=__UpperCamelCase	).audios

			__lowercase			:  Dict								=       sampling.iplms_sample(__UpperCamelCase	,				__UpperCamelCase	,				__UpperCamelCase	,				{}	)
			__lowercase			:  Tuple								=       generated.clamp(-1	,				1	)

			__lowercase			:  Union[str, Any]								=       (generated - audio).abs().sum()
			__lowercase			:  Dict								=       (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)








"""simple docstring"""




import random

import torch
from huggingface_hub import HfApi

from diffusers import UNetaDModel


_lowerCAmelCase							:   Any				=				HfApi()

_lowerCAmelCase							:   int				=				{}
# fmt: off
_lowerCAmelCase							:   List[Any]				=				torch.tensor([
    -0.7_5_1_5, -1.6_8_8_3, 0.2_4_2_0, 0.0_3_0_0, 0.6_3_4_7, 1.3_4_3_3, -1.1_7_4_3, -3.7_4_6_7,
    1.2_3_4_2, -2.2_4_8_5, 0.4_6_3_6, 0.8_0_7_6, -0.7_9_9_1, 0.3_9_6_9, 0.8_4_9_8, 0.9_1_8_9,
    -1.8_8_8_7, -3.3_5_2_2, 0.7_6_3_9, 0.2_0_4_0, 0.6_2_7_1, -2.7_1_4_8, -1.6_3_1_6, 3.0_8_3_9,
    0.3_1_8_6, 0.2_7_2_1, -0.9_7_5_9, -1.2_4_6_1, 2.6_2_5_7, 1.3_5_5_7
])
_lowerCAmelCase							:   Any				=				torch.tensor([
    -2.3_6_3_9, -2.5_3_4_4, 0.0_0_5_4, -0.6_6_7_4, 1.5_9_9_0, 1.0_1_5_8, 0.3_1_2_4, -2.1_4_3_6,
    1.8_7_9_5, -2.5_4_2_9, -0.1_5_6_6, -0.3_9_7_3, 1.2_4_9_0, 2.6_4_4_7, 1.2_2_8_3, -0.5_2_0_8,
    -2.8_1_5_4, -3.5_1_1_9, 2.3_8_3_8, 1.2_0_3_3, 1.7_2_0_1, -2.1_2_5_6, -1.4_5_7_6, 2.7_9_4_8,
    2.4_2_0_4, -0.9_7_5_2, -1.2_5_4_6, 0.8_0_2_7, 3.2_7_5_8, 3.1_3_6_5
])
_lowerCAmelCase							:   List[str]				=				torch.tensor([
    -0.6_5_3_1, -0.6_8_9_1, -0.3_1_7_2, -0.5_3_7_5, -0.9_1_4_0, -0.5_3_6_7, -0.1_1_7_5, -0.7_8_6_9,
    -0.3_8_0_8, -0.4_5_1_3, -0.2_0_9_8, -0.0_0_8_3, 0.3_1_8_3, 0.5_1_4_0, 0.2_2_4_7, -0.1_3_0_4,
    -0.1_3_0_2, -0.2_8_0_2, -0.2_0_8_4, -0.2_0_2_5, -0.4_9_6_7, -0.4_8_7_3, -0.0_8_6_1, 0.6_9_2_5,
    0.0_2_5_0, 0.1_2_9_0, -0.1_5_4_3, 0.6_3_1_6, 1.0_4_6_0, 1.4_9_4_3
])
_lowerCAmelCase							:   Optional[int]				=				torch.tensor([
    0.0_9_1_1, 0.1_1_0_7, 0.0_1_8_2, 0.0_4_3_5, -0.0_8_0_5, -0.0_6_0_8, 0.0_3_8_1, 0.2_1_7_2,
    -0.0_2_8_0, 0.1_3_2_7, -0.0_2_9_9, -0.0_2_5_5, -0.0_0_5_0, -0.1_1_7_0, -0.1_0_4_6, 0.0_3_0_9,
    0.1_3_6_7, 0.1_7_2_8, -0.0_5_3_3, -0.0_7_4_8, -0.0_5_3_4, 0.1_6_2_4, 0.0_3_8_4, -0.1_8_0_5,
    -0.0_7_0_7, 0.0_6_4_2, 0.0_2_2_0, -0.0_1_3_4, -0.1_3_3_3, -0.1_5_0_5
])
_lowerCAmelCase							:   List[Any]				=				torch.tensor([
    0.1_3_2_1, 0.1_3_3_7, 0.0_4_4_0, 0.0_6_2_2, -0.0_5_9_1, -0.0_3_7_0, 0.0_5_0_3, 0.2_1_3_3,
    -0.0_1_7_7, 0.1_4_1_5, -0.0_1_1_6, -0.0_1_1_2, 0.0_0_4_4, -0.0_9_8_0, -0.0_7_8_9, 0.0_3_9_5,
    0.1_5_0_2, 0.1_7_8_5, -0.0_4_8_8, -0.0_5_1_4, -0.0_4_0_4, 0.1_5_3_9, 0.0_4_5_4, -0.1_5_5_9,
    -0.0_6_6_5, 0.0_6_5_9, 0.0_3_8_3, -0.0_0_0_5, -0.1_2_6_6, -0.1_3_8_6
])
_lowerCAmelCase							:   Optional[int]				=				torch.tensor([
    0.1_1_5_4, 0.1_2_1_8, 0.0_3_0_7, 0.0_5_2_6, -0.0_7_1_1, -0.0_5_4_1, 0.0_3_6_6, 0.2_0_7_8,
    -0.0_2_6_7, 0.1_3_1_7, -0.0_2_2_6, -0.0_1_9_3, -0.0_0_1_4, -0.1_0_5_5, -0.0_9_0_2, 0.0_3_3_0,
    0.1_3_9_1, 0.1_7_0_9, -0.0_5_6_2, -0.0_6_9_3, -0.0_5_6_0, 0.1_4_8_2, 0.0_3_8_1, -0.1_6_8_3,
    -0.0_6_8_1, 0.0_6_6_1, 0.0_3_3_1, -0.0_0_4_6, -0.1_2_6_8, -0.1_4_3_1
])
_lowerCAmelCase							:   Optional[Any]				=				torch.tensor([
    0.1_1_9_2, 0.1_2_4_0, 0.0_4_1_4, 0.0_6_0_6, -0.0_5_5_7, -0.0_4_1_2, 0.0_4_3_0, 0.2_0_4_2,
    -0.0_2_0_0, 0.1_3_8_5, -0.0_1_1_5, -0.0_1_3_2, 0.0_0_1_7, -0.0_9_6_5, -0.0_8_0_2, 0.0_3_9_8,
    0.1_4_3_3, 0.1_7_4_7, -0.0_4_5_8, -0.0_5_3_3, -0.0_4_0_7, 0.1_5_4_5, 0.0_4_1_9, -0.1_5_7_4,
    -0.0_6_4_5, 0.0_6_2_6, 0.0_3_4_1, -0.0_0_1_0, -0.1_1_9_9, -0.1_3_9_0
])
_lowerCAmelCase							:   List[str]				=				torch.tensor([
    0.1_0_7_5, 0.1_0_7_4, 0.0_2_0_5, 0.0_4_3_1, -0.0_7_7_4, -0.0_6_0_7, 0.0_2_9_8, 0.2_0_4_2,
    -0.0_3_2_0, 0.1_2_6_7, -0.0_2_8_1, -0.0_2_5_0, -0.0_0_6_4, -0.1_0_9_1, -0.0_9_4_6, 0.0_2_9_0,
    0.1_3_2_8, 0.1_6_5_0, -0.0_5_8_0, -0.0_7_3_8, -0.0_5_8_6, 0.1_4_4_0, 0.0_3_3_7, -0.1_7_4_6,
    -0.0_7_1_2, 0.0_6_0_5, 0.0_2_5_0, -0.0_0_9_9, -0.1_3_1_6, -0.1_4_7_3
])
_lowerCAmelCase							:   Any				=				torch.tensor([
    -1.4_5_7_2, -2.0_4_8_1, -0.0_4_1_4, -0.6_0_0_5, 1.4_1_3_6, 0.5_8_4_8, 0.4_0_2_8, -2.7_3_3_0,
    1.2_2_1_2, -2.1_2_2_8, 0.2_1_5_5, 0.4_0_3_9, 0.7_6_6_2, 2.0_5_3_5, 0.7_4_7_7, -0.3_2_4_3,
    -2.1_7_5_8, -2.7_6_4_8, 1.6_9_4_7, 0.7_0_2_6, 1.2_3_3_8, -1.6_0_7_8, -0.8_6_8_2, 2.2_8_1_0,
    1.8_5_7_4, -0.5_7_1_8, -0.5_5_8_6, -0.0_1_8_6, 2.3_4_1_5, 2.1_2_5_1])
_lowerCAmelCase							:   int				=				torch.tensor([
    -1.3_6_9_0, -1.9_7_2_0, -0.4_0_9_0, -0.6_9_6_6, 1.4_6_6_0, 0.9_9_3_8, -0.1_3_8_5, -2.7_3_2_4,
    0.7_7_3_6, -1.8_9_1_7, 0.2_9_2_3, 0.4_2_9_3, 0.1_6_9_3, 1.4_1_1_2, 1.1_8_8_7, -0.3_1_8_1,
    -2.2_1_6_0, -2.6_3_8_1, 1.3_1_7_0, 0.8_1_6_3, 0.9_2_4_0, -1.6_5_4_4, -0.6_0_9_9, 2.5_2_5_9,
    1.6_4_3_0, -0.9_0_9_0, -0.9_3_9_2, -0.0_1_2_6, 2.4_2_6_8, 2.3_2_6_6
])
_lowerCAmelCase							:   Union[str, Any]				=				torch.tensor([
    -1.3_5_2_5, -1.9_6_2_8, -0.3_9_5_6, -0.6_8_6_0, 1.4_6_6_4, 1.0_0_1_4, -0.1_2_5_9, -2.7_2_1_2,
    0.7_7_7_2, -1.8_8_1_1, 0.2_9_9_6, 0.4_3_8_8, 0.1_7_0_4, 1.4_0_2_9, 1.1_7_0_1, -0.3_0_2_7,
    -2.2_0_5_3, -2.6_2_8_7, 1.3_3_5_0, 0.8_1_3_1, 0.9_2_7_4, -1.6_2_9_2, -0.6_0_9_8, 2.5_1_3_1,
    1.6_5_0_5, -0.8_9_5_8, -0.9_2_9_8, -0.0_1_5_1, 2.4_2_5_7, 2.3_3_5_5
])
_lowerCAmelCase							:   str				=				torch.tensor([
    -2.0_5_8_5, -2.7_8_9_7, -0.2_8_5_0, -0.8_9_4_0, 1.9_0_5_2, 0.5_7_0_2, 0.6_3_4_5, -3.8_9_5_9,
    1.5_9_3_2, -3.2_3_1_9, 0.1_9_7_4, 0.0_2_8_7, 1.7_5_6_6, 2.6_5_4_3, 0.8_3_8_7, -0.5_3_5_1,
    -3.2_7_3_6, -4.3_3_7_5, 2.9_0_2_9, 1.6_3_9_0, 1.4_6_4_0, -2.1_7_0_1, -1.9_0_1_3, 2.9_3_4_1,
    3.4_9_8_1, -0.6_2_5_5, -1.1_6_4_4, -0.1_5_9_1, 3.7_0_9_7, 3.2_0_6_6
])
_lowerCAmelCase							:   Optional[int]				=				torch.tensor([
    -2.3_1_3_9, -2.5_5_9_4, -0.0_1_9_7, -0.6_7_8_5, 1.7_0_0_1, 1.1_6_0_6, 0.3_0_7_5, -2.1_7_4_0,
    1.8_0_7_1, -2.5_6_3_0, -0.0_9_2_6, -0.3_8_1_1, 1.2_1_1_6, 2.6_2_4_6, 1.2_7_3_1, -0.5_3_9_8,
    -2.8_1_5_3, -3.6_1_4_0, 2.3_8_9_3, 1.3_2_6_2, 1.6_2_5_8, -2.1_8_5_6, -1.3_2_6_7, 2.8_3_9_5,
    2.3_7_7_9, -1.0_6_2_3, -1.2_4_6_8, 0.8_9_5_9, 3.3_3_6_7, 3.2_2_4_3
])
_lowerCAmelCase							:   Union[str, Any]				=				torch.tensor([
    -2.0_6_2_8, -2.7_6_6_7, -0.2_0_8_9, -0.8_2_6_3, 2.0_5_3_9, 0.5_9_9_2, 0.6_4_9_5, -3.8_3_3_6,
    1.6_0_2_5, -3.2_8_1_7, 0.1_7_2_1, -0.0_6_3_3, 1.7_5_1_6, 2.7_0_3_9, 0.8_1_0_0, -0.5_9_0_8,
    -3.2_1_1_3, -4.4_3_4_3, 2.9_2_5_7, 1.3_6_3_2, 1.5_5_6_2, -2.1_4_8_9, -1.9_8_9_4, 3.0_5_6_0,
    3.3_3_9_6, -0.7_3_2_8, -1.0_4_1_7, 0.0_3_8_3, 3.7_0_9_3, 3.2_3_4_3
])
_lowerCAmelCase							:   Any				=				torch.tensor([
    -1.4_5_7_4, -2.0_5_6_9, -0.0_4_7_3, -0.6_1_1_7, 1.4_0_1_8, 0.5_7_6_9, 0.4_1_2_9, -2.7_3_4_4,
    1.2_2_4_1, -2.1_3_9_7, 0.2_0_0_0, 0.3_9_3_7, 0.7_6_1_6, 2.0_4_5_3, 0.7_3_2_4, -0.3_3_9_1,
    -2.1_7_4_6, -2.7_7_4_4, 1.6_9_6_3, 0.6_9_2_1, 1.2_1_8_7, -1.6_1_7_2, -0.8_8_7_7, 2.2_4_3_9,
    1.8_4_7_1, -0.5_8_3_9, -0.5_6_0_5, -0.0_4_6_4, 2.3_2_5_0, 2.1_2_1_9
])
# fmt: on

_lowerCAmelCase							:   Optional[Any]				=				api.list_models(filter="diffusers")
for mod in models:
 if "google" in mod.author or mod.modelId == "CompVis/ldm-celebahq-256":
  _lowerCAmelCase							:   List[Any]				=				"/home/patrick/google_checkpoints/" + mod.modelId.split("/")[-1]

  print(F"Started running {mod.modelId}!!!")

  if mod.modelId.startswith("CompVis"):
   _lowerCAmelCase							:   Tuple				=				UNetaDModel.from_pretrained(local_checkpoint, subfolder="unet")
  else:
   _lowerCAmelCase							:   List[Any]				=				UNetaDModel.from_pretrained(local_checkpoint)

  torch.manual_seed(0)
  random.seed(0)

  _lowerCAmelCase							:   Optional[Any]				=				torch.randn(1, model.config.in_channels, model.config.sample_size, model.config.sample_size)
  _lowerCAmelCase							:   Tuple				=				torch.tensor([10] * noise.shape[0])
  with torch.no_grad():
   _lowerCAmelCase							:   Optional[Any]				=				model(noise, time_step).sample

  assert torch.allclose(
      logits[0, 0, 0, :30], results["_".join("_".join(mod.modelId.split("/")).split("-"))], atol=1E-3
  )
  print(F"{mod.modelId} has passed successfully!!!")
