


'''simple docstring'''






import argparse
import re

import torch
from CLAP import create_model

from transformers import AutoFeatureExtractor, ClapConfig, ClapModel


__A :  List[str]    							=				{
    "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",
}

__A :  Optional[int]    							=				AutoFeatureExtractor.from_pretrained("laion/clap-htsat-unfused", truncation="rand_trunc")







def     UpperCamelCase_      (   A__ :       str							,   A__ :       Optional[Any]=False ):
  '''simple docstring'''





  lowerCAmelCase_,							lowerCAmelCase_      :				List[Any]      =       create_model(
      """HTSAT-tiny"""							,   """roberta"""							,   A__							,   precision="""fp32"""							,   device="""cuda:0""" if torch.cuda.is_available() else """cpu"""							,   enable_fusion=A__							,   fusion_type="""aff_2d""" if enable_fusion else None							,   )
  return model, model_cfg







def     UpperCamelCase_      (   A__ :       List[str] ):
  '''simple docstring'''





  lowerCAmelCase_      :				Any      =       {}

  lowerCAmelCase_      :				Optional[Any]      =       R""".*sequential.(\d+).*"""
  lowerCAmelCase_      :				Optional[Any]      =       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_      :				Dict      =       key.replace(A__							,   A__ )

    if re.match(A__							,   A__ ):
      # replace sequential layers with list
      lowerCAmelCase_      :				int      =       re.match(A__							,   A__ ).group(1 )

      lowerCAmelCase_      :				Optional[Any]      =       key.replace(f'sequential.{sequential_layer}.'							,   f'layers.{int(A__ )//3}.linear.' )
    elif re.match(A__							,   A__ ):
      lowerCAmelCase_      :				Dict      =       int(re.match(A__							,   A__ ).group(1 ) )

      # Because in CLAP they use `nn.Sequential`...
      lowerCAmelCase_      :				Optional[int]      =       1 if projecton_layer == 0 else 2

      lowerCAmelCase_      :				List[str]      =       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_      :				Optional[int]      =       value
      lowerCAmelCase_      :				int      =       mixed_qkv.size(0 ) // 3

      lowerCAmelCase_      :				str      =       mixed_qkv[:qkv_dim]
      lowerCAmelCase_      :				Union[str, Any]      =       mixed_qkv[qkv_dim : qkv_dim * 2]
      lowerCAmelCase_      :				Optional[Any]      =       mixed_qkv[qkv_dim * 2 :]

      lowerCAmelCase_      :				Optional[Any]      =       query_layer
      lowerCAmelCase_      :				Tuple      =       key_layer
      lowerCAmelCase_      :				str      =       value_layer
    else:
      lowerCAmelCase_      :				Optional[int]      =       value

  return model_state_dict







def     UpperCamelCase_      (   A__ :       Dict							,   A__ :       Any							,   A__ :       str							,   A__ :       Dict=False ):
  '''simple docstring'''





  lowerCAmelCase_,							lowerCAmelCase_      :				Optional[int]      =       init_clap(A__							,   enable_fusion=A__ )

  clap_model.eval()
  lowerCAmelCase_      :				Any      =       clap_model.state_dict()
  lowerCAmelCase_      :				Optional[int]      =       rename_state_dict(A__ )

  lowerCAmelCase_      :				Union[str, Any]      =       ClapConfig()
  lowerCAmelCase_      :				Union[str, Any]      =       enable_fusion
  lowerCAmelCase_      :				Dict      =       ClapModel(A__ )

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

  model.save_pretrained(A__ )
  transformers_config.save_pretrained(A__ )


if __name__ == "__main__":
  __A :  List[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")
  __A :  Union[str, Any]    							=				parser.parse_args()

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







'''simple docstring'''






import unittest

from diffusers.pipelines.pipeline_utils import is_safetensors_compatible






class        __snake_case	(       unittest.TestCase):







    """simple docstring"""



    def __lowercase       (       self   :  Tuple		)       ->  Dict:
      lowerCAmelCase_      :				str      =       [
          """safety_checker/pytorch_model.bin""",
          """safety_checker/model.safetensors""",
          """vae/diffusion_pytorch_model.bin""",
          """vae/diffusion_pytorch_model.safetensors""",
          """text_encoder/pytorch_model.bin""",
          """text_encoder/model.safetensors""",
          """unet/diffusion_pytorch_model.bin""",
          """unet/diffusion_pytorch_model.safetensors""",
      ]
      self.assertTrue(is_safetensors_compatible(lowerCamelCase		)		)

    def __lowercase       (       self   :  List[Any]		)       ->  int:
      lowerCAmelCase_      :				Tuple      =       [
          """unet/diffusion_pytorch_model.bin""",
          """unet/diffusion_pytorch_model.safetensors""",
      ]
      self.assertTrue(is_safetensors_compatible(lowerCamelCase		)		)

    def __lowercase       (       self   :  Optional[int]		)       ->  Optional[Any]:
      lowerCAmelCase_      :				int      =       [
          """safety_checker/pytorch_model.bin""",
          """safety_checker/model.safetensors""",
          """vae/diffusion_pytorch_model.bin""",
          """vae/diffusion_pytorch_model.safetensors""",
          """text_encoder/pytorch_model.bin""",
          """text_encoder/model.safetensors""",
          """unet/diffusion_pytorch_model.bin""",
          # Removed: 'unet/diffusion_pytorch_model.safetensors',
      ]
      self.assertFalse(is_safetensors_compatible(lowerCamelCase		)		)

    def __lowercase       (       self   :  int		)       ->  List[Any]:
      lowerCAmelCase_      :				Dict      =       [
          """text_encoder/pytorch_model.bin""",
          """text_encoder/model.safetensors""",
      ]
      self.assertTrue(is_safetensors_compatible(lowerCamelCase		)		)

    def __lowercase       (       self   :  str		)       ->  List[str]:
      lowerCAmelCase_      :				Union[str, Any]      =       [
          """safety_checker/pytorch_model.bin""",
          """safety_checker/model.safetensors""",
          """vae/diffusion_pytorch_model.bin""",
          """vae/diffusion_pytorch_model.safetensors""",
          """text_encoder/pytorch_model.bin""",
          # Removed: 'text_encoder/model.safetensors',
          """unet/diffusion_pytorch_model.bin""",
          """unet/diffusion_pytorch_model.safetensors""",
      ]
      self.assertFalse(is_safetensors_compatible(lowerCamelCase		)		)

    def __lowercase       (       self   :  List[Any]		)       ->  Tuple:
      lowerCAmelCase_      :				Union[str, Any]      =       [
          """safety_checker/pytorch_model.fp16.bin""",
          """safety_checker/model.fp16.safetensors""",
          """vae/diffusion_pytorch_model.fp16.bin""",
          """vae/diffusion_pytorch_model.fp16.safetensors""",
          """text_encoder/pytorch_model.fp16.bin""",
          """text_encoder/model.fp16.safetensors""",
          """unet/diffusion_pytorch_model.fp16.bin""",
          """unet/diffusion_pytorch_model.fp16.safetensors""",
      ]
      lowerCAmelCase_      :				Union[str, Any]      =       """fp16"""
      self.assertTrue(is_safetensors_compatible(lowerCamelCase   ,						variant=lowerCamelCase		)		)

    def __lowercase       (       self   :  Optional[Any]		)       ->  List[str]:
      lowerCAmelCase_      :				str      =       [
          """unet/diffusion_pytorch_model.fp16.bin""",
          """unet/diffusion_pytorch_model.fp16.safetensors""",
      ]
      lowerCAmelCase_      :				Optional[int]      =       """fp16"""
      self.assertTrue(is_safetensors_compatible(lowerCamelCase   ,						variant=lowerCamelCase		)		)

    def __lowercase       (       self   :  Tuple		)       ->  Optional[Any]:
      # pass variant but use the non-variant filenames
      lowerCAmelCase_      :				Dict      =       [
          """unet/diffusion_pytorch_model.bin""",
          """unet/diffusion_pytorch_model.safetensors""",
      ]
      lowerCAmelCase_      :				str      =       """fp16"""
      self.assertTrue(is_safetensors_compatible(lowerCamelCase   ,						variant=lowerCamelCase		)		)

    def __lowercase       (       self   :  Optional[int]		)       ->  List[str]:
      lowerCAmelCase_      :				str      =       [
          """safety_checker/pytorch_model.fp16.bin""",
          """safety_checker/model.fp16.safetensors""",
          """vae/diffusion_pytorch_model.fp16.bin""",
          """vae/diffusion_pytorch_model.fp16.safetensors""",
          """text_encoder/pytorch_model.fp16.bin""",
          """text_encoder/model.fp16.safetensors""",
          """unet/diffusion_pytorch_model.fp16.bin""",
          # Removed: 'unet/diffusion_pytorch_model.fp16.safetensors',
      ]
      lowerCAmelCase_      :				List[str]      =       """fp16"""
      self.assertFalse(is_safetensors_compatible(lowerCamelCase   ,						variant=lowerCamelCase		)		)

    def __lowercase       (       self   :  Union[str, Any]		)       ->  Optional[int]:
      lowerCAmelCase_      :				str      =       [
          """text_encoder/pytorch_model.fp16.bin""",
          """text_encoder/model.fp16.safetensors""",
      ]
      lowerCAmelCase_      :				Optional[Any]      =       """fp16"""
      self.assertTrue(is_safetensors_compatible(lowerCamelCase   ,						variant=lowerCamelCase		)		)

    def __lowercase       (       self   :  List[Any]		)       ->  List[Any]:
      # pass variant but use the non-variant filenames
      lowerCAmelCase_      :				Dict      =       [
          """text_encoder/pytorch_model.bin""",
          """text_encoder/model.safetensors""",
      ]
      lowerCAmelCase_      :				Any      =       """fp16"""
      self.assertTrue(is_safetensors_compatible(lowerCamelCase   ,						variant=lowerCamelCase		)		)





    def __lowercase       (       self   :  Dict		)       ->  Any:
      lowerCAmelCase_      :				Optional[int]      =       [
          """safety_checker/pytorch_model.fp16.bin""",
          """safety_checker/model.fp16.safetensors""",
          """vae/diffusion_pytorch_model.fp16.bin""",
          """vae/diffusion_pytorch_model.fp16.safetensors""",
          """text_encoder/pytorch_model.fp16.bin""",
          # 'text_encoder/model.fp16.safetensors',
          """unet/diffusion_pytorch_model.fp16.bin""",
          """unet/diffusion_pytorch_model.fp16.safetensors""",
      ]
      lowerCAmelCase_      :				int      =       """fp16"""
      self.assertFalse(is_safetensors_compatible(lowerCamelCase   ,						variant=lowerCamelCase		)		)



