





"""simple docstring"""






import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()



class 					__UpperCAmelCase							(						_UpperCamelCase						,  _UpperCamelCase						,  unittest.TestCase  ):
				__lowerCamelCase				:	Optional[Any]				     =				StableDiffusionSAGPipeline
				__lowerCamelCase				:	Optional[int]				     =				TEXT_TO_IMAGE_PARAMS
				__lowerCamelCase				:	Optional[Any]				     =				TEXT_TO_IMAGE_BATCH_PARAMS
				__lowerCamelCase				:	List[str]				     =				TEXT_TO_IMAGE_IMAGE_PARAMS
				__lowerCamelCase				:	List[Any]				     =				TEXT_TO_IMAGE_IMAGE_PARAMS
				__lowerCamelCase				:	List[str]				     =				False
				def       UpperCAmelCase      (		self    :			Tuple    )	->	List[Any]:



							'''simple docstring'''



							torch.manual_seed(0    )
							a__  :	Dict		 =			UNetaDConditionModel(
							    block_out_channels=(32, 64)	, layers_per_block=2	, sample_size=32	, in_channels=4	, out_channels=4	, down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")	, up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")	, cross_attention_dim=32	, )
							a__  :	int		 =			DDIMScheduler(
							    beta_start=0.0_0085	, beta_end=0.012	, beta_schedule="scaled_linear"	, clip_sample=a_	, set_alpha_to_one=a_	, )
							torch.manual_seed(0    )
							a__  :	Optional[Any]		 =			AutoencoderKL(
							    block_out_channels=[32, 64]	, in_channels=3	, out_channels=3	, down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]	, up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]	, latent_channels=4	, )
							torch.manual_seed(0    )
							a__  :	Dict		 =			CLIPTextConfig(
							    bos_token_id=0	, eos_token_id=2	, hidden_size=32	, intermediate_size=37	, layer_norm_eps=1E-05	, num_attention_heads=4	, num_hidden_layers=5	, pad_token_id=1	, vocab_size=10_00	, )
							a__  :	int		 =			CLIPTextModel(a_    )
							a__  :	str		 =			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"    )

							a__  :	List[str]		 =			{
							    "unet": unet,
							    "scheduler": scheduler,
							    "vae": vae,
							    "text_encoder": text_encoder,
							    "tokenizer": tokenizer,
							    "safety_checker": None,
							    "feature_extractor": None,
							}
							return components
				def       UpperCAmelCase      (		self    :			List[str]	, a_    :			Tuple	, a_    :			List[str]=0    )	->	Tuple:



							'''simple docstring'''



							if str(a_    ).startswith("mps"    ):
										a__  :	List[Any]		 =			torch.manual_seed(a_    )
							else:
										a__  :	Tuple		 =			torch.Generator(device=a_    ).manual_seed(a_    )
							a__  :	Optional[Any]		 =			{
							    "prompt": ".",
							    "generator": generator,
							    "num_inference_steps": 2,
							    "guidance_scale": 1.0,
							    "sag_scale": 1.0,
							    "output_type": "numpy",
							}
							return inputs






				def       UpperCAmelCase      (		self    :			Any    )	->	int:



							'''simple docstring'''



							super().test_inference_batch_single_identical(expected_max_diff=3E-3    )



@slow
@require_torch_gpu
class 					__UpperCAmelCase							(						unittest.TestCase  ):
				def       UpperCAmelCase      (		self    :			Optional[int]    )	->	Dict:



							'''simple docstring'''



							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()
				def       UpperCAmelCase      (		self    :			List[Any]    )	->	Optional[int]:



							'''simple docstring'''



							a__  :	Any		 =			StableDiffusionSAGPipeline.from_pretrained("CompVis/stable-diffusion-v1-4"    )
							a__  :	Any		 =			sag_pipe.to(a_    )
							sag_pipe.set_progress_bar_config(disable=a_    )

							a__  :	List[Any]		 =			"."
							a__  :	int		 =			torch.manual_seed(0    )
							a__  :	Dict		 =			sag_pipe(
							    [prompt]	, generator=a_	, guidance_scale=7.5	, sag_scale=1.0	, num_inference_steps=20	, output_type="np"    )

							a__  :	int		 =			output.images

							a__  :	Tuple		 =			image[0, -3:, -3:, -1]

							assert image.shape == (1, 5_12, 5_12, 3)
							a__  :	Dict		 =			np.array([0.1568, 0.1738, 0.1695, 0.1693, 0.1507, 0.1705, 0.1547, 0.1751, 0.1949]    )

							assert np.abs(image_slice.flatten() - expected_slice    ).max() < 5E-2
				def       UpperCAmelCase      (		self    :			List[Any]    )	->	List[str]:



							'''simple docstring'''



							a__  :	int		 =			StableDiffusionSAGPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base"    )
							a__  :	Tuple		 =			sag_pipe.to(a_    )
							sag_pipe.set_progress_bar_config(disable=a_    )

							a__  :	Dict		 =			"."
							a__  :	Optional[Any]		 =			torch.manual_seed(0    )
							a__  :	Union[str, Any]		 =			sag_pipe(
							    [prompt]	, generator=a_	, guidance_scale=7.5	, sag_scale=1.0	, num_inference_steps=20	, output_type="np"    )

							a__  :	Optional[int]		 =			output.images

							a__  :	str		 =			image[0, -3:, -3:, -1]

							assert image.shape == (1, 5_12, 5_12, 3)
							a__  :	List[str]		 =			np.array([0.3459, 0.2876, 0.2537, 0.3002, 0.2671, 0.2160, 0.3026, 0.2262, 0.2371]    )

							assert np.abs(image_slice.flatten() - expected_slice    ).max() < 5E-2






				def       UpperCAmelCase      (		self    :			Union[str, Any]    )	->	Optional[Any]:



							'''simple docstring'''



							a__  :	str		 =			StableDiffusionSAGPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base"    )
							a__  :	List[str]		 =			sag_pipe.to(a_    )
							sag_pipe.set_progress_bar_config(disable=a_    )

							a__  :	Tuple		 =			"."
							a__  :	str		 =			torch.manual_seed(0    )
							a__  :	List[str]		 =			sag_pipe(
							    [prompt]	, width=7_68	, height=5_12	, generator=a_	, guidance_scale=7.5	, sag_scale=1.0	, num_inference_steps=20	, output_type="np"	, )

							a__  :	Optional[Any]		 =			output.images

							assert image.shape == (1, 5_12, 7_68, 3)






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