





"""simple docstring"""






import argparse
import io

import requests
import torch
from omegaconf import OmegaConf

from diffusers import AutoencoderKL
from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (
    assign_to_checkpoint,
    conv_attn_to_linear,
    create_vae_diffusers_config,
    renew_vae_attention_paths,
    renew_vae_resnet_paths,
)






def      lowercase__ (					lowerCAmelCase__   :			Union[str, Any]							,			lowerCAmelCase__   :			Any			)    ->			Tuple:



			'''simple docstring'''
			a__  :	str		 =			checkpoint

			a__  :	List[Any]		 =			{}

			a__  :	List[str]		 =			vae_state_dict["encoder.conv_in.weight"]
			a__  :	List[str]		 =			vae_state_dict["encoder.conv_in.bias"]
			a__  :	Union[str, Any]		 =			vae_state_dict["encoder.conv_out.weight"]
			a__  :	str		 =			vae_state_dict["encoder.conv_out.bias"]
			a__  :	Dict		 =			vae_state_dict["encoder.norm_out.weight"]
			a__  :	Optional[int]		 =			vae_state_dict["encoder.norm_out.bias"]

			a__  :	Union[str, Any]		 =			vae_state_dict["decoder.conv_in.weight"]
			a__  :	Optional[int]		 =			vae_state_dict["decoder.conv_in.bias"]
			a__  :	Any		 =			vae_state_dict["decoder.conv_out.weight"]
			a__  :	Optional[Any]		 =			vae_state_dict["decoder.conv_out.bias"]
			a__  :	Dict		 =			vae_state_dict["decoder.norm_out.weight"]
			a__  :	int		 =			vae_state_dict["decoder.norm_out.bias"]

			a__  :	Optional[Any]		 =			vae_state_dict["quant_conv.weight"]
			a__  :	List[Any]		 =			vae_state_dict["quant_conv.bias"]
			a__  :	Optional[Any]		 =			vae_state_dict["post_quant_conv.weight"]
			a__  :	Any		 =			vae_state_dict["post_quant_conv.bias"]

			# Retrieves the keys for the encoder down blocks only
			a__  :	int		 =			len({".".join(layer.split("."			)[:3]			) for layer in vae_state_dict if "encoder.down" in layer}			)
			a__  :	Tuple		 =			{
			    layer_id: [key for key in vae_state_dict if F"down.{layer_id}" in key] for layer_id in range(lowerCAmelCase__			)
			}

			# Retrieves the keys for the decoder up blocks only
			a__  :	Dict		 =			len({".".join(layer.split("."			)[:3]			) for layer in vae_state_dict if "decoder.up" in layer}			)
			a__  :	Any		 =			{
			    layer_id: [key for key in vae_state_dict if F"up.{layer_id}" in key] for layer_id in range(lowerCAmelCase__			)
			}

			for i in range(lowerCAmelCase__			):
						a__  :	Optional[int]		 =			[key for key in down_blocks[i] if F"down.{i}" in key and F"down.{i}.downsample" not in key]

						if F"encoder.down.{i}.downsample.conv.weight" in vae_state_dict:
									a__  :	Optional[Any]		 =			vae_state_dict.pop(
									    F"encoder.down.{i}.downsample.conv.weight"			)
									a__  :	str		 =			vae_state_dict.pop(
									    F"encoder.down.{i}.downsample.conv.bias"			)

						a__  :	Optional[Any]		 =			renew_vae_resnet_paths(lowerCAmelCase__			)
						a__  :	Union[str, Any]		 =			{"old": F"down.{i}.block", "new": F"down_blocks.{i}.resnets"}
						assign_to_checkpoint(lowerCAmelCase__							,			lowerCAmelCase__							,			lowerCAmelCase__							,			additional_replacements=[meta_path]							,			config=lowerCAmelCase__			)

			a__  :	List[str]		 =			[key for key in vae_state_dict if "encoder.mid.block" in key]
			a__  :	Any		 =			2
			for i in range(1							,			num_mid_res_blocks + 1			):
						a__  :	Union[str, Any]		 =			[key for key in mid_resnets if F"encoder.mid.block_{i}" in key]

						a__  :	str		 =			renew_vae_resnet_paths(lowerCAmelCase__			)
						a__  :	List[str]		 =			{"old": F"mid.block_{i}", "new": F"mid_block.resnets.{i - 1}"}
						assign_to_checkpoint(lowerCAmelCase__							,			lowerCAmelCase__							,			lowerCAmelCase__							,			additional_replacements=[meta_path]							,			config=lowerCAmelCase__			)

			a__  :	Union[str, Any]		 =			[key for key in vae_state_dict if "encoder.mid.attn" in key]
			a__  :	str		 =			renew_vae_attention_paths(lowerCAmelCase__			)
			a__  :	Optional[int]		 =			{"old": "mid.attn_1", "new": "mid_block.attentions.0"}
			assign_to_checkpoint(lowerCAmelCase__							,			lowerCAmelCase__							,			lowerCAmelCase__							,			additional_replacements=[meta_path]							,			config=lowerCAmelCase__			)
			conv_attn_to_linear(lowerCAmelCase__			)

			for i in range(lowerCAmelCase__			):
						a__  :	List[Any]		 =			num_up_blocks - 1 - i
						a__  :	str		 =			[
						    key for key in up_blocks[block_id] if F"up.{block_id}" in key and F"up.{block_id}.upsample" not in key
						]

						if F"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict:
									a__  :	Dict		 =			vae_state_dict[
									    F"decoder.up.{block_id}.upsample.conv.weight"
									]
									a__  :	Any		 =			vae_state_dict[
									    F"decoder.up.{block_id}.upsample.conv.bias"
									]

						a__  :	Optional[int]		 =			renew_vae_resnet_paths(lowerCAmelCase__			)
						a__  :	List[Any]		 =			{"old": F"up.{block_id}.block", "new": F"up_blocks.{i}.resnets"}
						assign_to_checkpoint(lowerCAmelCase__							,			lowerCAmelCase__							,			lowerCAmelCase__							,			additional_replacements=[meta_path]							,			config=lowerCAmelCase__			)

			a__  :	str		 =			[key for key in vae_state_dict if "decoder.mid.block" in key]
			a__  :	int		 =			2
			for i in range(1							,			num_mid_res_blocks + 1			):
						a__  :	List[str]		 =			[key for key in mid_resnets if F"decoder.mid.block_{i}" in key]

						a__  :	Optional[int]		 =			renew_vae_resnet_paths(lowerCAmelCase__			)
						a__  :	Any		 =			{"old": F"mid.block_{i}", "new": F"mid_block.resnets.{i - 1}"}
						assign_to_checkpoint(lowerCAmelCase__							,			lowerCAmelCase__							,			lowerCAmelCase__							,			additional_replacements=[meta_path]							,			config=lowerCAmelCase__			)

			a__  :	List[Any]		 =			[key for key in vae_state_dict if "decoder.mid.attn" in key]
			a__  :	Optional[int]		 =			renew_vae_attention_paths(lowerCAmelCase__			)
			a__  :	Optional[int]		 =			{"old": "mid.attn_1", "new": "mid_block.attentions.0"}
			assign_to_checkpoint(lowerCAmelCase__							,			lowerCAmelCase__							,			lowerCAmelCase__							,			additional_replacements=[meta_path]							,			config=lowerCAmelCase__			)
			conv_attn_to_linear(lowerCAmelCase__			)
			return new_checkpoint






def      lowercase__ (					lowerCAmelCase__   :			str							,			lowerCAmelCase__   :			str							,			)    ->			Tuple:



			'''simple docstring'''
			# Only support V1
			a__  :	Optional[Any]		 =			requests.get(
			    " https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml"			)
			a__  :	Dict		 =			io.BytesIO(r.content			)

			a__  :	Union[str, Any]		 =			OmegaConf.load(lowerCAmelCase__			)
			a__  :	List[Any]		 =			5_1_2
			a__  :	List[Any]		 =			"cuda" if torch.cuda.is_available() else "cpu"
			if checkpoint_path.endswith("safetensors"			):
						from safetensors import safe_open

						a__  :	List[str]		 =			{}
						with safe_open(lowerCAmelCase__							,			framework="pt"							,			device="cpu"			) as f:
									for key in f.keys():
												a__  :	Union[str, Any]		 =			f.get_tensor(lowerCAmelCase__			)
			else:
						a__  :	Dict		 =			torch.load(lowerCAmelCase__							,			map_location=lowerCAmelCase__			)["state_dict"]

			# Convert the VAE model.
			a__  :	Dict		 =			create_vae_diffusers_config(lowerCAmelCase__							,			image_size=lowerCAmelCase__			)
			a__  :	str		 =			custom_convert_ldm_vae_checkpoint(lowerCAmelCase__							,			lowerCAmelCase__			)

			a__  :	Union[str, Any]		 =			AutoencoderKL(**lowerCAmelCase__			)
			vae.load_state_dict(lowerCAmelCase__			)
			vae.save_pretrained(lowerCAmelCase__			)


if __name__ == "__main__":
		__UpperCAmelCase						=			argparse.ArgumentParser()

		parser.add_argument('''--vae_pt_path''', default=None, type=str, required=True, help='''Path to the VAE.pt to convert.''')
		parser.add_argument('''--dump_path''', default=None, type=str, required=True, help='''Path to the VAE.pt to convert.''')

		__UpperCAmelCase						=			parser.parse_args()

		vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)






"""simple docstring"""






import argparse
import re

import torch
from CLAP import create_model

from transformers import AutoFeatureExtractor, ClapConfig, ClapModel


__UpperCAmelCase						=			{
    '''text_branch''': '''text_model''',
    '''audio_branch''': '''audio_model.audio_encoder''',
    '''attn''': '''attention.self''',
    '''self.proj''': '''output.dense''',
    '''attention.self_mask''': '''attn_mask''',
    '''mlp.fc1''': '''intermediate.dense''',
    '''mlp.fc2''': '''output.dense''',
    '''norm1''': '''layernorm_before''',
    '''norm2''': '''layernorm_after''',
    '''bn0''': '''batch_norm''',
}

__UpperCAmelCase						=			AutoFeatureExtractor.from_pretrained('''laion/clap-htsat-unfused''', truncation='''rand_trunc''')






def      lowercase__ (					lowerCAmelCase__   :			Optional[int]							,			lowerCAmelCase__   :			Optional[int]=False			)    ->			List[str]:



			'''simple docstring'''
			a__ ,	a__  :	Optional[int]		 =			create_model(
			    "HTSAT-tiny"							,			"roberta"							,			lowerCAmelCase__							,			precision="fp32"							,			device="cuda:0" if torch.cuda.is_available() else "cpu"							,			enable_fusion=lowerCAmelCase__							,			fusion_type="aff_2d" if enable_fusion else None							,			)
			return model, model_cfg






def      lowercase__ (					lowerCAmelCase__   :			Optional[int]			)    ->			Dict:



			'''simple docstring'''
			a__  :	Any		 =			{}

			a__  :	Tuple		 =			R".*sequential.(\d+).*"
			a__  :	Tuple		 =			R".*_projection.(\d+).*"

			for key, value in state_dict.items():
						# check if any key needs to be modified
						for key_to_modify, new_key in KEYS_TO_MODIFY_MAPPING.items():
									if key_to_modify in key:
												a__  :	str		 =			key.replace(lowerCAmelCase__							,			lowerCAmelCase__			)

						if re.match(lowerCAmelCase__							,			lowerCAmelCase__			):
									# replace sequential layers with list
									a__  :	Optional[Any]		 =			re.match(lowerCAmelCase__							,			lowerCAmelCase__			).group(1			)

									a__  :	int		 =			key.replace(F"sequential.{sequential_layer}."							,			F"layers.{int(lowerCAmelCase__			)//3}.linear."			)
						elif re.match(lowerCAmelCase__							,			lowerCAmelCase__			):
									a__  :	Tuple		 =			int(re.match(lowerCAmelCase__							,			lowerCAmelCase__			).group(1			)			)

									# Because in CLAP they use `nn.Sequential`...
									a__  :	List[Any]		 =			1 if projecton_layer == 0 else 2

									a__  :	Dict		 =			key.replace(F"_projection.{projecton_layer}."							,			F"_projection.linear{transformers_projection_layer}."			)

						if "audio" and "qkv" in key:
									# split qkv into query key and value
									a__  :	List[Any]		 =			value
									a__  :	List[Any]		 =			mixed_qkv.size(0			) // 3

									a__  :	Optional[int]		 =			mixed_qkv[:qkv_dim]
									a__  :	List[str]		 =			mixed_qkv[qkv_dim : qkv_dim * 2]
									a__  :	Optional[Any]		 =			mixed_qkv[qkv_dim * 2 :]

									a__  :	Tuple		 =			query_layer
									a__  :	int		 =			key_layer
									a__  :	Optional[int]		 =			value_layer
						else:
									a__  :	List[str]		 =			value

			return model_state_dict






def      lowercase__ (					lowerCAmelCase__   :			List[str]							,			lowerCAmelCase__   :			Any							,			lowerCAmelCase__   :			Optional[Any]							,			lowerCAmelCase__   :			Tuple=False			)    ->			Tuple:



			'''simple docstring'''
			a__ ,	a__  :	Tuple		 =			init_clap(lowerCAmelCase__							,			enable_fusion=lowerCAmelCase__			)

			clap_model.eval()
			a__  :	Optional[int]		 =			clap_model.state_dict()
			a__  :	Optional[Any]		 =			rename_state_dict(lowerCAmelCase__			)

			a__  :	Union[str, Any]		 =			ClapConfig()
			a__  :	Dict		 =			enable_fusion
			a__  :	Any		 =			ClapModel(lowerCAmelCase__			)

			# ignore the spectrogram embedding layer
			model.load_state_dict(lowerCAmelCase__							,			strict=lowerCAmelCase__			)

			model.save_pretrained(lowerCAmelCase__			)
			transformers_config.save_pretrained(lowerCAmelCase__			)


if __name__ == "__main__":
		__UpperCAmelCase						=			argparse.ArgumentParser()
		parser.add_argument('''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model.''')
		parser.add_argument('''--checkpoint_path''', default=None, type=str, help='''Path to fairseq checkpoint''')
		parser.add_argument('''--config_path''', default=None, type=str, help='''Path to hf config.json of model to convert''')
		parser.add_argument('''--enable_fusion''', action='''store_true''', help='''Whether to enable fusion or not''')
		__UpperCAmelCase						=			parser.parse_args()

		convert_clap_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.enable_fusion)