






import argparse
import re

import torch
from CLAP import create_model

from transformers import AutoFeatureExtractor, ClapConfig, ClapModel


__snake_case    :    Dict         =     {
    """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""",
}

__snake_case    :    List[Any]         =     AutoFeatureExtractor.from_pretrained("""laion/clap-htsat-unfused""", truncation="""rand_trunc""")







def 	_UpperCamelCase   (    UpperCamelCase_       :					int    ,       UpperCamelCase_       :					Tuple=False		)  ->   str:

						"""simple docstring"""



						lowerCAmelCase__						,       lowerCAmelCase__    						=						create_model(
						    'HTSAT-tiny'    ,       'roberta'    ,       _lowercase    ,       precision='fp32'    ,       device='cuda:0' if torch.cuda.is_available() else 'cpu'    ,       enable_fusion=_lowercase    ,       fusion_type='aff_2d' if enable_fusion else None    ,       )
						return model, model_cfg







def 	_UpperCamelCase   (    UpperCamelCase_       :					List[str]		)  ->   int:

						"""simple docstring"""



						lowerCAmelCase__    						=						{}

						lowerCAmelCase__    						=						r'.*sequential.(\d+).*'
						lowerCAmelCase__    						=						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:
																								lowerCAmelCase__    						=						key.replace(_lowercase    ,       _lowercase		)

												if re.match(_lowercase    ,       _lowercase		):
																		# replace sequential layers with list
																		lowerCAmelCase__    						=						re.match(_lowercase    ,       _lowercase		).group(1		)

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

																		# Because in CLAP they use `nn.Sequential`...
																		lowerCAmelCase__    						=						1 if projecton_layer == 0 else 2

																		lowerCAmelCase__    						=						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
																		lowerCAmelCase__    						=						value
																		lowerCAmelCase__    						=						mixed_qkv.size(0		) // 3

																		lowerCAmelCase__    						=						mixed_qkv[:qkv_dim]
																		lowerCAmelCase__    						=						mixed_qkv[qkv_dim : qkv_dim * 2]
																		lowerCAmelCase__    						=						mixed_qkv[qkv_dim * 2 :]

																		lowerCAmelCase__    						=						query_layer
																		lowerCAmelCase__    						=						key_layer
																		lowerCAmelCase__    						=						value_layer
												else:
																		lowerCAmelCase__    						=						value

						return model_state_dict







def 	_UpperCamelCase   (    UpperCamelCase_       :					str    ,       UpperCamelCase_       :					str    ,       UpperCamelCase_       :					Dict    ,       UpperCamelCase_       :					str=False		)  ->   int:

						"""simple docstring"""



						lowerCAmelCase__						,       lowerCAmelCase__    						=						init_clap(_lowercase    ,       enable_fusion=_lowercase		)

						clap_model.eval()
						lowerCAmelCase__    						=						clap_model.state_dict()
						lowerCAmelCase__    						=						rename_state_dict(_lowercase		)

						lowerCAmelCase__    						=						ClapConfig()
						lowerCAmelCase__    						=						enable_fusion
						lowerCAmelCase__    						=						ClapModel(_lowercase		)

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

						model.save_pretrained(_lowercase		)
						transformers_config.save_pretrained(_lowercase		)


if __name__ == "__main__":
					__snake_case    :    Any         =     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""")
					__snake_case    :    Optional[int]         =     parser.parse_args()

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













import numpy as np
from PIL import Image







def 	_UpperCamelCase   (    UpperCamelCase_       :					np.ndarray    ,       UpperCamelCase_       :					int    ,       UpperCamelCase_       :					int		)  ->   np.ndarray:

						"""simple docstring"""



						lowerCAmelCase__    						=						np.array(UpperCamelCase_		)
						if arr.shape[0] != arr.shape[1]:
												raise ValueError('The input array is not a square matrix'		)
						lowerCAmelCase__    						=						0
						lowerCAmelCase__    						=						0
						lowerCAmelCase__    						=						0
						lowerCAmelCase__    						=						0

						# compute the shape of the output matrix
						lowerCAmelCase__    						=						(arr.shape[0] - size) // stride + 1
						# initialize the output matrix with zeros of shape maxpool_shape
						lowerCAmelCase__    						=						np.zeros((maxpool_shape, maxpool_shape)		)

						while i < arr.shape[0]:
												if i + size > arr.shape[0]:
																		# if the end of the matrix is reached, break
																		break
												while j < arr.shape[1]:
																		# if the end of the matrix is reached, break
																		if j + size > arr.shape[1]:
																								break
																		# compute the maximum of the pooling matrix
																		lowerCAmelCase__    						=						np.max(arr[i : i + size, j : j + size]		)
																		# shift the pooling matrix by stride of column pixels
																		j += stride
																		mat_j += 1

												# shift the pooling matrix by stride of row pixels
												i += stride
												mat_i += 1

												# reset the column index to 0
												lowerCAmelCase__    						=						0
												lowerCAmelCase__    						=						0

						return updated_arr







def 	_UpperCamelCase   (    UpperCamelCase_       :					np.ndarray    ,       UpperCamelCase_       :					int    ,       UpperCamelCase_       :					int		)  ->   np.ndarray:

						"""simple docstring"""



						lowerCAmelCase__    						=						np.array(UpperCamelCase_		)
						if arr.shape[0] != arr.shape[1]:
												raise ValueError('The input array is not a square matrix'		)
						lowerCAmelCase__    						=						0
						lowerCAmelCase__    						=						0
						lowerCAmelCase__    						=						0
						lowerCAmelCase__    						=						0

						# compute the shape of the output matrix
						lowerCAmelCase__    						=						(arr.shape[0] - size) // stride + 1
						# initialize the output matrix with zeros of shape avgpool_shape
						lowerCAmelCase__    						=						np.zeros((avgpool_shape, avgpool_shape)		)

						while i < arr.shape[0]:
												# if the end of the matrix is reached, break
												if i + size > arr.shape[0]:
																		break
												while j < arr.shape[1]:
																		# if the end of the matrix is reached, break
																		if j + size > arr.shape[1]:
																								break
																		# compute the average of the pooling matrix
																		lowerCAmelCase__    						=						int(np.average(arr[i : i + size, j : j + size]		)		)
																		# shift the pooling matrix by stride of column pixels
																		j += stride
																		mat_j += 1

												# shift the pooling matrix by stride of row pixels
												i += stride
												mat_i += 1
												# reset the column index to 0
												lowerCAmelCase__    						=						0
												lowerCAmelCase__    						=						0

						return updated_arr


# Main Function
if __name__ == "__main__":
					from doctest import testmod

					testmod(name="""avgpooling""", verbose=True)

					# Loading the image
					__snake_case    :    List[str]         =     Image.open("""path_to_image""")

					# Converting the image to numpy array and maxpooling, displaying the result
					# Ensure that the image is a square matrix

					Image.fromarray(maxpooling(np.array(image), size=3, stride=2)).show()

					# Converting the image to numpy array and averagepooling, displaying the result
					# Ensure that the image is a square matrix

					Image.fromarray(avgpooling(np.array(image), size=3, stride=2)).show()





