




"""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)





import torch
import torch.nn as nn
from transformers import CLIPConfig, CLIPVisionModel, PreTrainedModel

from ...utils import logging


UpperCamelCase_						:     int	       =							logging.get_logger(__name__)



def 		UpperCamelCase     (       _UpperCAmelCase  :   Optional[int] ,			_UpperCAmelCase  :   Dict  )   ->			Any:






						'''simple docstring'''

						_lowercase       :	List[Any]           =   nn.functional.normalize(_UpperCAmelCase  )
						_lowercase       :	List[str]           =   nn.functional.normalize(_UpperCAmelCase  )
						return torch.mm(_UpperCAmelCase ,			normalized_text_embeds.t()  )




class   __lowercase						(						__snake_case       ):
	_A       =   CLIPConfig

	_A       =   ["CLIPEncoderLayer"]






	def __init__(self				: List[str]      ,	snake_case				: CLIPConfig    ) ->			Optional[int]:
							super().__init__(snake_case    )

							_lowercase       :	Tuple           =   CLIPVisionModel(config.vision_config    )
							_lowercase       :	Any           =   nn.Linear(config.vision_config.hidden_size      ,	config.projection_dim      ,	bias=snake_case    )

							_lowercase       :	int           =   nn.Parameter(torch.ones(17      ,	config.projection_dim    )      ,	requires_grad=snake_case    )
							_lowercase       :	str           =   nn.Parameter(torch.ones(3      ,	config.projection_dim    )      ,	requires_grad=snake_case    )

							_lowercase       :	Any           =   nn.Parameter(torch.ones(17    )      ,	requires_grad=snake_case    )
							_lowercase       :	Dict           =   nn.Parameter(torch.ones(3    )      ,	requires_grad=snake_case    )






	@torch.no_grad()
	def       _a(self				: int      ,	snake_case				: Tuple      ,	snake_case				: Optional[Any]    ) ->			Union[str, Any]:
							_lowercase       :	str           =   self.vision_model(snake_case    )[1]  # pooled_output
							_lowercase       :	List[Any]           =   self.visual_projection(snake_case    )

							# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
							_lowercase       :	str           =   cosine_distance(snake_case      ,	self.special_care_embeds    ).cpu().float().numpy()
							_lowercase       :	Optional[int]           =   cosine_distance(snake_case      ,	self.concept_embeds    ).cpu().float().numpy()

							_lowercase       :	List[str]           =   []
							_lowercase       :	int           =   image_embeds.shape[0]
							for i in range(snake_case    ):
													_lowercase       :	int           =   {"special_scores": {}, "special_care": [], "concept_scores": {}, "bad_concepts": []}

													# increase this value to create a stronger `nfsw` filter
													# at the cost of increasing the possibility of filtering benign images
													_lowercase       :	Optional[int]           =   0.0

													for concept_idx in range(len(special_cos_dist[0]    )    ):
																			_lowercase       :	Union[str, Any]           =   special_cos_dist[i][concept_idx]
																			_lowercase       :	Optional[int]           =   self.special_care_embeds_weights[concept_idx].item()
																			_lowercase       :	Optional[int]           =   round(concept_cos - concept_threshold + adjustment      ,	3    )
																			if result_img["special_scores"][concept_idx] > 0:
																									result_img["special_care"].append({concept_idx, result_img["special_scores"][concept_idx]}    )
																									_lowercase       :	Union[str, Any]           =   0.01

													for concept_idx in range(len(cos_dist[0]    )    ):
																			_lowercase       :	Tuple           =   cos_dist[i][concept_idx]
																			_lowercase       :	Union[str, Any]           =   self.concept_embeds_weights[concept_idx].item()
																			_lowercase       :	str           =   round(concept_cos - concept_threshold + adjustment      ,	3    )
																			if result_img["concept_scores"][concept_idx] > 0:
																									result_img["bad_concepts"].append(snake_case    )

													result.append(snake_case    )

							_lowercase       :	Dict           =   [len(res["bad_concepts"]    ) > 0 for res in result]

							return images, has_nsfw_concepts






	@torch.no_grad()
	def       _a(self				: str      ,	snake_case				: torch.FloatTensor      ,	snake_case				: torch.FloatTensor    ) ->			List[str]:
							_lowercase       :	Dict           =   self.vision_model(snake_case    )[1]  # pooled_output
							_lowercase       :	Any           =   self.visual_projection(snake_case    )

							_lowercase       :	Optional[Any]           =   cosine_distance(snake_case      ,	self.special_care_embeds    )
							_lowercase       :	Optional[Any]           =   cosine_distance(snake_case      ,	self.concept_embeds    )

							# increase this value to create a stronger `nsfw` filter
							# at the cost of increasing the possibility of filtering benign images
							_lowercase       :	Optional[int]           =   0.0

							_lowercase       :	Optional[int]           =   special_cos_dist - self.special_care_embeds_weights + adjustment
							# special_scores = special_scores.round(decimals=3)
							_lowercase       :	List[str]           =   torch.any(special_scores > 0      ,	dim=1    )
							_lowercase       :	List[Any]           =   special_care * 0.01
							_lowercase       :	Optional[Any]           =   special_adjustment.unsqueeze(1    ).expand(-1      ,	cos_dist.shape[1]    )

							_lowercase       :	Any           =   (cos_dist - self.concept_embeds_weights) + special_adjustment
							# concept_scores = concept_scores.round(decimals=3)
							_lowercase       :	List[Any]           =   torch.any(concept_scores > 0      ,	dim=1    )

							return images, has_nsfw_concepts




