





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


__magic_name__       :					Any						=			{
    """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							(SCREAMING_SNAKE_CASE							,     SCREAMING_SNAKE_CASE     ):
						return torch.atana(SCREAMING_SNAKE_CASE							,     SCREAMING_SNAKE_CASE     ) / math.pi * 2

def      UpperCamelCase							(SCREAMING_SNAKE_CASE     ):
						UpperCamelCase				:		Optional[Any]				   =			torch.sin(t * math.pi / 2     ) ** 2
						UpperCamelCase				:		Dict				   =			(1 - sigma**2) ** 0.5
						return alpha_sigma_to_t(SCREAMING_SNAKE_CASE							,     SCREAMING_SNAKE_CASE     )






class 							lowercase__			(   __SCREAMING_SNAKE_CASE	):

			"""simple docstring"""
			pass






class 							lowercase__			(   nn.Module	):

			"""simple docstring"""







			def __init__(      self						,   _A    ):






									'''simple docstring'''







									super().__init__()

									UpperCamelCase				:		List[str]				   =			DiffusionAttnUnetaD(_A						,   n_attn_layers=4    )
									UpperCamelCase				:		Tuple				   =			deepcopy(self.diffusion    )
									UpperCamelCase				:		Dict				   =			torch.quasirandom.SobolEngine(1						,   scramble=_A    )


def      UpperCamelCase							(SCREAMING_SNAKE_CASE     ):
						UpperCamelCase				:		Optional[int]				   =			MODELS_MAP[model_name]["""url"""]
						os.system(f"""wget {url} ./"""     )

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


__magic_name__       :					str						=			{
    """1""": """resnets.0""",
    """2""": """attentions.0""",
    """3""": """resnets.1""",
    """4""": """attentions.1""",
    """5""": """resnets.2""",
    """6""": """attentions.2""",
}
__magic_name__       :					Dict						=			{
    """8""": """resnets.0""",
    """9""": """attentions.0""",
    """10""": """resnets.1""",
    """11""": """attentions.1""",
    """12""": """resnets.2""",
    """13""": """attentions.2""",
}
__magic_name__       :					Any						=			{
    """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""",
}
__magic_name__       :					Dict						=			{
    """0""": """resnets.0""",
    """1""": """resnets.1""",
    """2""": """resnets.2""",
    """4""": """resnets.0""",
    """5""": """resnets.1""",
    """6""": """resnets.2""",
}

__magic_name__       :					Any						=			{
    """skip""": """conv_skip""",
    """main.0""": """conv_1""",
    """main.1""": """group_norm_1""",
    """main.3""": """conv_2""",
    """main.4""": """group_norm_2""",
}

__magic_name__       :					Any						=			{
    """norm""": """group_norm""",
    """qkv_proj""": ["""query""", """key""", """value"""],
    """out_proj""": ["""proj_attn"""],
}

def      UpperCamelCase							(SCREAMING_SNAKE_CASE     ):
						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							(SCREAMING_SNAKE_CASE     ):
						for key, value in ATTN_MAP.items():
												if name.startswith(SCREAMING_SNAKE_CASE     ) and not isinstance(SCREAMING_SNAKE_CASE							,     SCREAMING_SNAKE_CASE     ):
																		return name.replace(SCREAMING_SNAKE_CASE							,     SCREAMING_SNAKE_CASE     )
												elif name.startswith(SCREAMING_SNAKE_CASE     ):
																		return [name.replace(SCREAMING_SNAKE_CASE							,     SCREAMING_SNAKE_CASE     ) for v in value]
						raise ValueError(f"""Attn error with {name}"""     )

def      UpperCamelCase							(SCREAMING_SNAKE_CASE							,     SCREAMING_SNAKE_CASE=13     ):
						UpperCamelCase				:		Optional[Any]				   =			input_string

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

						UpperCamelCase				:		int				   =			0
						if string.startswith("""net.3."""     ):
												depth += 1
												UpperCamelCase				:		List[str]				   =			string[6:]
						elif string.startswith("""net."""     ):
												UpperCamelCase				:		List[Any]				   =			string[4:]

						while string.startswith("""main.7."""     ):
												depth += 1
												UpperCamelCase				:		List[Any]				   =			string[7:]

						if string.startswith("""main."""     ):
												UpperCamelCase				:		Union[str, Any]				   =			string[5:]

						# mid block
						if string[:2].isdigit():
												UpperCamelCase				:		List[Any]				   =			string[:2]
												UpperCamelCase				:		Optional[Any]				   =			string[2:]
						else:
												UpperCamelCase				:		str				   =			string[0]
												UpperCamelCase				:		int				   =			string[1:]

						if depth == max_depth:
												UpperCamelCase				:		List[str]				   =			MID_NUM_TO_LAYER[layer_num]
												UpperCamelCase				:		Optional[Any]				   =			"""mid_block"""
						elif depth > 0 and int(SCREAMING_SNAKE_CASE     ) < 7:
												UpperCamelCase				:		List[str]				   =			DOWN_NUM_TO_LAYER[layer_num]
												UpperCamelCase				:		Dict				   =			f"""down_blocks.{depth}"""
						elif depth > 0 and int(SCREAMING_SNAKE_CASE     ) > 7:
												UpperCamelCase				:		Any				   =			UP_NUM_TO_LAYER[layer_num]
												UpperCamelCase				:		Optional[int]				   =			f"""up_blocks.{max_depth - depth - 1}"""
						elif depth == 0:
												UpperCamelCase				:		str				   =			DEPTH_0_TO_LAYER[layer_num]
												UpperCamelCase				:		int				   =			f"""up_blocks.{max_depth - 1}""" if int(SCREAMING_SNAKE_CASE     ) > 3 else """down_blocks.0"""

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

						UpperCamelCase				:		Dict				   =			string_left[1:]

						if "resnets" in new_layer:
												UpperCamelCase				:		Optional[Any]				   =			convert_resconv_naming(SCREAMING_SNAKE_CASE     )
						elif "attentions" in new_layer:
												UpperCamelCase				:		Tuple				   =			convert_attn_naming(SCREAMING_SNAKE_CASE     )
												UpperCamelCase				:		Union[str, Any]				   =			new_string_left

						if not isinstance(SCREAMING_SNAKE_CASE							,     SCREAMING_SNAKE_CASE     ):
												UpperCamelCase				:		Dict				   =			prefix + """.""" + new_layer + """.""" + string_left
						else:
												UpperCamelCase				:		Tuple				   =			[prefix + """.""" + new_layer + """.""" + s for s in string_left]
						return new_string

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

												UpperCamelCase				:		List[Any]				   =			rename(SCREAMING_SNAKE_CASE     )

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

						return new_state_dict

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

def      UpperCamelCase							(SCREAMING_SNAKE_CASE     ):
						UpperCamelCase				:		List[str]				   =			torch.device("""cuda""" if torch.cuda.is_available() else """cpu"""     )

						UpperCamelCase				:		Any				   =			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()}"""
												UpperCamelCase				:		Union[str, Any]				   =			download(SCREAMING_SNAKE_CASE     )

						UpperCamelCase				:		Dict				   =			MODELS_MAP[model_name]["""sample_rate"""]
						UpperCamelCase				:		Tuple				   =			MODELS_MAP[model_name]["""sample_size"""]

						UpperCamelCase				:		List[str]				   =			Object()
						UpperCamelCase				:		Union[str, Any]				   =			sample_size
						UpperCamelCase				:		Dict				   =			sample_rate
						UpperCamelCase				:		Any				   =			0

						UpperCamelCase				:		int				   =			UNetaDModel(sample_size=SCREAMING_SNAKE_CASE							,     sample_rate=SCREAMING_SNAKE_CASE     )
						UpperCamelCase				:		Any				   =			diffusers_model.state_dict()

						UpperCamelCase				:		Optional[Any]				   =			DiffusionUncond(SCREAMING_SNAKE_CASE     )
						orig_model.load_state_dict(torch.load(args.model_path							,     map_location=SCREAMING_SNAKE_CASE     )["""state_dict"""]     )
						UpperCamelCase				:		Optional[Any]				   =			orig_model.diffusion_ema.eval()
						UpperCamelCase				:		Optional[int]				   =			orig_model.state_dict()
						UpperCamelCase				:		str				   =			rename_orig_weights(SCREAMING_SNAKE_CASE     )

						UpperCamelCase				:		Optional[Any]				   =			set(renamed_state_dict.keys()     ) - set(diffusers_state_dict.keys()     )
						UpperCamelCase				:		Optional[Any]				   =			set(diffusers_state_dict.keys()     ) - set(renamed_state_dict.keys()     )

						assert len(SCREAMING_SNAKE_CASE     ) == 0, f"""Problem with {renamed_minus_diffusers}"""
						assert all(k.endswith("""kernel"""     ) for k in list(SCREAMING_SNAKE_CASE     )     ), 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":
																		UpperCamelCase				:		int				   =			value.squeeze()

												UpperCamelCase				:		str				   =			value

						diffusers_model.load_state_dict(SCREAMING_SNAKE_CASE     )

						UpperCamelCase				:		Optional[int]				   =			100
						UpperCamelCase				:		List[Any]				   =			33

						UpperCamelCase				:		Tuple				   =			IPNDMScheduler(num_train_timesteps=SCREAMING_SNAKE_CASE     )

						UpperCamelCase				:		Optional[int]				   =			torch.manual_seed(SCREAMING_SNAKE_CASE     )
						UpperCamelCase				:		List[str]				   =			torch.randn([1, 2, config.sample_size]							,     generator=SCREAMING_SNAKE_CASE     ).to(SCREAMING_SNAKE_CASE     )

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

						UpperCamelCase				:		Union[str, Any]				   =			DanceDiffusionPipeline(unet=SCREAMING_SNAKE_CASE							,     scheduler=SCREAMING_SNAKE_CASE     )

						UpperCamelCase				:		Optional[Any]				   =			torch.manual_seed(33     )
						UpperCamelCase				:		Dict				   =			pipe(num_inference_steps=SCREAMING_SNAKE_CASE							,     generator=SCREAMING_SNAKE_CASE     ).audios

						UpperCamelCase				:		Optional[Any]				   =			sampling.iplms_sample(SCREAMING_SNAKE_CASE							,     SCREAMING_SNAKE_CASE							,     SCREAMING_SNAKE_CASE							,     {}     )
						UpperCamelCase				:		Any				   =			generated.clamp(-1							,     1     )

						UpperCamelCase				:		List[Any]				   =			(generated - audio).abs().sum()
						UpperCamelCase				:		Optional[int]				   =			(generated - audio).abs().max()

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

						print("""Diff sum"""							,     SCREAMING_SNAKE_CASE     )
						print("""Diff max"""							,     SCREAMING_SNAKE_CASE     )

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

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


if __name__ == "__main__":
		__magic_name__       :					Tuple						=			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.""")
		__magic_name__       :					List[str]						=			parser.parse_args()

		main(args)









from math import atan, cos, radians, sin, tan

from .haversine_distance import haversine_distance

_lowerCamelCase      :						Dict							=  6_378_137.0
_lowerCamelCase      :						Union[str, Any]							=  6_356_752.314_245
_lowerCamelCase      :						List[Any]							=  6378137







def SCREAMING_SNAKE_CASE							(			lowercase_					,							lowercase_					,							lowercase_					,							lowercase_		)  ->			float:



       """simple docstring"""

       A__ 					=       (AXIS_A - AXIS_B) / AXIS_A
       # Parametric latitudes
       # https://en.wikipedia.org/wiki/Latitude#Parametric_(or_reduced)_latitude
       A__ 					=       atan((1 - flattening) * tan(radians(lowercase_		)		)		)
       A__ 					=       atan((1 - flattening) * tan(radians(lowercase_		)		)		)

       # Compute central angle between two points
       # using haversine theta. sigma =  haversine_distance / equatorial radius
       A__ 					=       haversine_distance(lowercase_					,							lowercase_					,							lowercase_					,							lowercase_		) / EQUATORIAL_RADIUS

       # Intermediate P and Q values
       A__ 					=       (b_lata + b_lata) / 2
       A__ 					=       (b_lata - b_lata) / 2

       # Intermediate X value
       # X = (sigma - sin(sigma)) * sin^2Pcos^2Q / cos^2(sigma/2)
       A__ 					=       (sin(lowercase_		) ** 2) * (cos(lowercase_		) ** 2)
       A__ 					=       cos(sigma / 2		) ** 2
       A__ 					=       (sigma - sin(lowercase_		)) * (x_numerator / x_demonimator)

       # Intermediate Y value
       # Y = (sigma + sin(sigma)) * cos^2Psin^2Q / sin^2(sigma/2)
       A__ 					=       (cos(lowercase_		) ** 2) * (sin(lowercase_		) ** 2)
       A__ 					=       sin(sigma / 2		) ** 2
       A__ 					=       (sigma + sin(lowercase_		)) * (y_numerator / y_denominator)

       return EQUATORIAL_RADIUS * (sigma - ((flattening / 2) * (x_value + y_value)))


if __name__ == "__main__":
  import doctest

  doctest.testmod()


