






'''simple docstring'''







import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModel,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableUnCLIPImgaImgPipeline, UNetaDConditionModel
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    floats_tensor,
    load_image,
    load_numpy,
    require_torch_gpu,
    skip_mps,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()





class 	SCREAMING_SNAKE_CASE					(					SCREAMING_SNAKE_CASE_	,   SCREAMING_SNAKE_CASE_	,   SCREAMING_SNAKE_CASE_	,   unittest.TestCase       ):
			snake_case__  =				StableUnCLIPImgaImgPipeline
			snake_case__  =				TEXT_GUIDED_IMAGE_VARIATION_PARAMS
			snake_case__  =				TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
			snake_case__  =				frozenset(
			    []       )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
			snake_case__  =				frozenset([]       )






			def 			SCREAMING_SNAKE_CASE      ( self      :			List[str]					)			->			Dict:
							a_		:  List[Any]											=      32
							a_		:  Optional[int]											=      embedder_hidden_size

							# image encoding components

							a_		:  str											=      CLIPImageProcessor(crop_size=32  ,						size=32					)

							torch.manual_seed(0					)
							a_		:  Dict											=      CLIPVisionModelWithProjection(
							    CLIPVisionConfig(
							        hidden_size=__SCREAMING_SNAKE_CASE  ,						projection_dim=__SCREAMING_SNAKE_CASE  ,						num_hidden_layers=5  ,						num_attention_heads=4  ,						image_size=32  ,						intermediate_size=37  ,						patch_size=1  ,						)					)

							# regular denoising components

							torch.manual_seed(0					)
							a_		:  str											=      StableUnCLIPImageNormalizer(embedding_dim=__SCREAMING_SNAKE_CASE					)
							a_		:  int											=      DDPMScheduler(beta_schedule='''squaredcos_cap_v2'''					)

							torch.manual_seed(0					)
							a_		:  Optional[int]											=      CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''					)

							torch.manual_seed(0					)
							a_		:  Optional[int]											=      CLIPTextModel(
							    CLIPTextConfig(
							        bos_token_id=0  ,						eos_token_id=2  ,						hidden_size=__SCREAMING_SNAKE_CASE  ,						projection_dim=32  ,						intermediate_size=37  ,						layer_norm_eps=1e-05  ,						num_attention_heads=4  ,						num_hidden_layers=5  ,						pad_token_id=1  ,						vocab_size=1000  ,						)					)

							torch.manual_seed(0					)
							a_		:  List[str]											=      UNetaDConditionModel(
							    sample_size=32  ,						in_channels=4  ,						out_channels=4  ,						down_block_types=('''CrossAttnDownBlock2D''', '''DownBlock2D''')  ,						up_block_types=('''UpBlock2D''', '''CrossAttnUpBlock2D''')  ,						block_out_channels=(32, 64)  ,						attention_head_dim=(2, 4)  ,						class_embed_type='''projection'''  ,						projection_class_embeddings_input_dim=embedder_projection_dim * 2  ,						cross_attention_dim=__SCREAMING_SNAKE_CASE  ,						layers_per_block=1  ,						upcast_attention=__SCREAMING_SNAKE_CASE  ,						use_linear_projection=__SCREAMING_SNAKE_CASE  ,						)

							torch.manual_seed(0					)
							a_		:  Tuple											=      DDIMScheduler(
							    beta_schedule='''scaled_linear'''  ,						beta_start=0.0_0085  ,						beta_end=0.012  ,						prediction_type='''v_prediction'''  ,						set_alpha_to_one=__SCREAMING_SNAKE_CASE  ,						steps_offset=1  ,						)

							torch.manual_seed(0					)
							a_		:  Optional[Any]											=      AutoencoderKL()

							a_		:  Tuple											=      {
							    # image encoding components
							    '''feature_extractor''': feature_extractor,
							    '''image_encoder''': image_encoder.eval(),
							    # image noising components
							    '''image_normalizer''': image_normalizer.eval(),
							    '''image_noising_scheduler''': image_noising_scheduler,
							    # regular denoising components
							    '''tokenizer''': tokenizer,
							    '''text_encoder''': text_encoder.eval(),
							    '''unet''': unet.eval(),
							    '''scheduler''': scheduler,
							    '''vae''': vae.eval(),
							}

							return components






			def 			SCREAMING_SNAKE_CASE      ( self      :			str  ,						__SCREAMING_SNAKE_CASE      :			int  ,						__SCREAMING_SNAKE_CASE      :			Tuple=0  ,						__SCREAMING_SNAKE_CASE      :			Optional[int]=True					)			->			List[str]:
							if str(__SCREAMING_SNAKE_CASE					).startswith('''mps'''					):
											a_		:  List[str]											=      torch.manual_seed(__SCREAMING_SNAKE_CASE					)
							else:
											a_		:  Tuple											=      torch.Generator(device=__SCREAMING_SNAKE_CASE					).manual_seed(__SCREAMING_SNAKE_CASE					)

							a_		:  int											=      floats_tensor((1, 3, 32, 32)  ,						rng=random.Random(__SCREAMING_SNAKE_CASE					)					).to(__SCREAMING_SNAKE_CASE					)

							if pil_image:
											a_		:  List[str]											=      input_image * 0.5 + 0.5
											a_		:  List[Any]											=      input_image.clamp(0  ,						1					)
											a_		:  str											=      input_image.cpu().permute(0  ,						2  ,						3  ,						1					).float().numpy()
											a_		:  Optional[Any]											=      DiffusionPipeline.numpy_to_pil(__SCREAMING_SNAKE_CASE					)[0]

							return {
							    "prompt": "An anime racoon running a marathon",
							    "image": input_image,
							    "generator": generator,
							    "num_inference_steps": 2,
							    "output_type": "np",
							}






			@skip_mps
			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[int]					)			->			Union[str, Any]:
							a_		:  List[str]											=      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							a_		:  List[Any]											=      self.get_dummy_components()
							a_		:  Any											=      StableUnCLIPImgaImgPipeline(**__SCREAMING_SNAKE_CASE					)
							a_		:  int											=      sd_pipe.to(__SCREAMING_SNAKE_CASE					)
							sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE					)

							a_		:  List[Any]											=      self.get_dummy_inputs(__SCREAMING_SNAKE_CASE					)
							inputs.update({'''image_embeds''': None}					)
							a_		:  List[Any]											=      sd_pipe(**__SCREAMING_SNAKE_CASE					).images
							a_		:  Union[str, Any]											=      image[0, -3:, -3:, -1]

							assert image.shape == (1, 32, 32, 3)
							a_		:  List[str]											=      np.array([0.3872, 0.7224, 0.5601, 0.4741, 0.6872, 0.5814, 0.4636, 0.3867, 0.5078]					)

							assert np.abs(image_slice.flatten() - expected_slice					).max() < 1e-3






			def 			SCREAMING_SNAKE_CASE      ( self      :			Any					)			->			Optional[int]:
							a_		:  List[str]											=      torch_device in ['''cpu''', '''mps''']

							self._test_attention_slicing_forward_pass(test_max_difference=__SCREAMING_SNAKE_CASE					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			Union[str, Any]					)			->			Dict:
							a_		:  int											=      torch_device in ['''cpu''', '''mps''']

							self._test_inference_batch_single_identical(test_max_difference=__SCREAMING_SNAKE_CASE					)







			@unittest.skipIf(
			    torch_device != '''cuda''' or not is_xformers_available()  ,						reason='''XFormers attention is only available with CUDA and `xformers` installed'''  ,						)
			def 			SCREAMING_SNAKE_CASE      ( self      :			Tuple					)			->			int:
							self._test_xformers_attention_forwardGenerator_pass(test_max_difference=__SCREAMING_SNAKE_CASE					)





@slow
@require_torch_gpu
class 	SCREAMING_SNAKE_CASE					(					unittest.TestCase       ):






			def 			SCREAMING_SNAKE_CASE      ( self      :			str					)			->			Any:
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()






			def 			SCREAMING_SNAKE_CASE      ( self      :			Dict					)			->			Any:
							a_		:  Union[str, Any]											=      load_image(
							    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'''					)

							a_		:  Optional[int]											=      load_numpy(
							    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_img2img_anime_turtle_fp16.npy'''					)

							a_		:  List[Any]											=      StableUnCLIPImgaImgPipeline.from_pretrained(
							    '''fusing/stable-unclip-2-1-l-img2img'''  ,						torch_dtype=torch.floataa					)
							pipe.to(__SCREAMING_SNAKE_CASE					)
							pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE					)
							# stable unclip will oom when integration tests are run on a V100,
							# so turn on memory savings
							pipe.enable_attention_slicing()
							pipe.enable_sequential_cpu_offload()

							a_		:  Tuple											=      torch.Generator(device='''cpu'''					).manual_seed(0					)
							a_		:  Any											=      pipe(__SCREAMING_SNAKE_CASE  ,						'''anime turle'''  ,						generator=__SCREAMING_SNAKE_CASE  ,						output_type='''np'''					)

							a_		:  int											=      output.images[0]

							assert image.shape == (768, 768, 3)

							assert_mean_pixel_difference(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[Any]					)			->			List[str]:
							a_		:  Optional[Any]											=      load_image(
							    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'''					)

							a_		:  int											=      load_numpy(
							    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_h_img2img_anime_turtle_fp16.npy'''					)

							a_		:  Optional[Any]											=      StableUnCLIPImgaImgPipeline.from_pretrained(
							    '''fusing/stable-unclip-2-1-h-img2img'''  ,						torch_dtype=torch.floataa					)
							pipe.to(__SCREAMING_SNAKE_CASE					)
							pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE					)
							# stable unclip will oom when integration tests are run on a V100,
							# so turn on memory savings
							pipe.enable_attention_slicing()
							pipe.enable_sequential_cpu_offload()

							a_		:  str											=      torch.Generator(device='''cpu'''					).manual_seed(0					)
							a_		:  Tuple											=      pipe(__SCREAMING_SNAKE_CASE  ,						'''anime turle'''  ,						generator=__SCREAMING_SNAKE_CASE  ,						output_type='''np'''					)

							a_		:  Tuple											=      output.images[0]

							assert image.shape == (768, 768, 3)

							assert_mean_pixel_difference(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					)







			def 			SCREAMING_SNAKE_CASE      ( self      :			Any					)			->			Any:
							a_		:  int											=      load_image(
							    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'''					)

							torch.cuda.empty_cache()
							torch.cuda.reset_max_memory_allocated()
							torch.cuda.reset_peak_memory_stats()

							a_		:  Any											=      StableUnCLIPImgaImgPipeline.from_pretrained(
							    '''fusing/stable-unclip-2-1-h-img2img'''  ,						torch_dtype=torch.floataa					)
							a_		:  List[str]											=      pipe.to(__SCREAMING_SNAKE_CASE					)
							pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE					)
							pipe.enable_attention_slicing()
							pipe.enable_sequential_cpu_offload()

							a_		:  List[Any]											=      pipe(
							    __SCREAMING_SNAKE_CASE  ,						'''anime turtle'''  ,						num_inference_steps=2  ,						output_type='''np'''  ,						)

							a_		:  Tuple											=      torch.cuda.max_memory_allocated()
							# make sure that less than 7 GB is allocated
							assert mem_bytes < 7 * 10**9












'''simple docstring'''







import argparse

from torch import nn

# transformers_old should correspond to branch `save_old_prophetnet_model_structure` here
# original prophetnet_checkpoints are saved under `patrickvonplaten/..._old` respectively
from transformers_old.modeling_prophetnet import (
    ProphetNetForConditionalGeneration as ProphetNetForConditionalGenerationOld,
)
from transformers_old.modeling_xlm_prophetnet import (
    XLMProphetNetForConditionalGeneration as XLMProphetNetForConditionalGenerationOld,
)

from transformers import ProphetNetForConditionalGeneration, XLMProphetNetForConditionalGeneration, logging


__lowerCAmelCase       		=      logging.get_logger(__name__)
logging.set_verbosity_info()





def 		_UpperCAmelCase   (							__A							:   str						,		__A							:   str      ):
				if "xprophetnet" in prophetnet_checkpoint_path:
								a_		:  Tuple											=      XLMProphetNetForConditionalGenerationOld.from_pretrained(__A      )
								a_				,			a_		:  Optional[Any]											=      XLMProphetNetForConditionalGeneration.from_pretrained(
								    __A						,		output_loading_info=__A      )
				else:
								a_		:  List[Any]											=      ProphetNetForConditionalGenerationOld.from_pretrained(__A      )
								a_				,			a_		:  Any											=      ProphetNetForConditionalGeneration.from_pretrained(
								    __A						,		output_loading_info=__A      )

				a_		:  str											=      ['''key_proj''', '''value_proj''', '''query_proj''']

				a_		:  Tuple											=      {
				    '''self_attn''': '''ngram_self_attn''',
				    '''cross_attn''': '''encoder_attn''',
				    '''cross_attn_layer_norm''': '''encoder_attn_layer_norm''',
				    '''feed_forward_layer_norm''': '''final_layer_norm''',
				    '''feed_forward''': '''''',
				    '''intermediate''': '''fc1''',
				    '''output''': '''fc2''',
				    '''key_proj''': '''k_proj''',
				    '''query_proj''': '''q_proj''',
				    '''value_proj''': '''v_proj''',
				    '''word_embeddings''': '''embed_tokens''',
				    '''embeddings_layer_norm''': '''emb_layer_norm''',
				    '''relative_pos_embeddings''': '''relative_linear''',
				    '''ngram_embeddings''': '''ngram_input_embed''',
				    '''position_embeddings''': '''embed_positions''',
				}

				for key in loading_info["missing_keys"]:
								a_		:  List[str]											=      key.split('''.'''      )

								if attributes[0] == "lm_head":
												a_		:  List[str]											=      prophet
												a_		:  Dict											=      prophet_old
								else:
												a_		:  str											=      prophet.prophetnet
												a_		:  int											=      prophet_old.model

								a_		:  str											=      False
								for attribute in attributes:
												if attribute in mapping:
																a_		:  Dict											=      mapping[attribute]
																if not hasattr(__A						,		__A      ) and len(__A      ) > 0:
																				a_		:  List[str]											=      attribute
												elif hasattr(__A						,		__A      ):
																a_		:  Union[str, Any]											=      attribute

												if attribute == "weight":
																assert old_model.weight.shape == model.weight.shape, "Shapes have to match!"
																a_		:  Tuple											=      old_model.weight
																logger.info(f'{attribute} is initialized.'      )
																a_		:  Union[str, Any]											=      True
																break
												elif attribute == "bias":
																assert old_model.bias.shape == model.bias.shape, "Shapes have to match!"
																a_		:  Union[str, Any]											=      old_model.bias
																logger.info(f'{attribute} is initialized'      )
																a_		:  Dict											=      True
																break
												elif attribute in special_keys and hasattr(__A						,		'''in_proj_weight'''      ):
																a_		:  Tuple											=      old_model.in_proj_weight.shape[0] // 3
																a_		:  Any											=      getattr(__A						,		__A      )
																param.weight.shape == old_model.in_proj_weight[:embed_dim, :].shape, "Shapes have to match"
																param.bias.shape == old_model.in_proj_bias[:embed_dim].shape, "Shapes have to match"
																if attribute == "query_proj":
																				a_		:  Union[str, Any]											=      nn.Parameter(old_model.in_proj_weight[:embed_dim, :]      )
																				a_		:  Optional[Any]											=      nn.Parameter(old_model.in_proj_bias[:embed_dim]      )

																elif attribute == "key_proj":
																				a_		:  List[Any]											=      nn.Parameter(old_model.in_proj_weight[embed_dim : 2 * embed_dim, :]      )
																				a_		:  Optional[int]											=      nn.Parameter(old_model.in_proj_bias[embed_dim : 2 * embed_dim]      )
																elif attribute == "value_proj":
																				a_		:  Tuple											=      nn.Parameter(old_model.in_proj_weight[2 * embed_dim :, :]      )
																				a_		:  Any											=      nn.Parameter(old_model.in_proj_bias[2 * embed_dim :]      )
																a_		:  Dict											=      True
																break
												elif attribute == "position_embeddings":
																assert (
																    model.position_embeddings.weight.shape[-1] == old_model.embed_positions.weight.shape[-1]
																), "Hidden size has to match"
																assert model.position_embeddings.weight.shape[0] == 5_12, "We want 512 position_embeddings."
																a_		:  Union[str, Any]											=      nn.Parameter(old_model.embed_positions.weight[:5_12, :]      )
																a_		:  Optional[Any]											=      True
																break

												if attribute.isdigit():
																a_		:  Union[str, Any]											=      model[int(__A      )]
																a_		:  str											=      old_model[int(__A      )]
												else:
																a_		:  Tuple											=      getattr(__A						,		__A      )

																if old_attribute == "":
																				a_		:  List[str]											=      old_model
																else:
																				if not hasattr(__A						,		__A      ):
																								raise ValueError(f'{old_model} does not have {old_attribute}'      )
																				a_		:  Optional[Any]											=      getattr(__A						,		__A      )

								if not is_key_init:
												raise ValueError(f'{key} was not correctly initialized!'      )

				print(f'Saving model to {pytorch_dump_folder_path}'      )
				prophet.save_pretrained(__A      )


if __name__ == "__main__":
		__lowerCAmelCase       		=      argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    '--prophetnet_checkpoint_path', default=None, type=str, required=True, help='Path the official PyTorch dump.'
		)
		parser.add_argument(
		    '--pytorch_dump_folder_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
		)
		__lowerCAmelCase       		=      parser.parse_args()
		convert_prophetnet_checkpoint_to_pytorch(args.prophetnet_checkpoint_path, args.pytorch_dump_folder_path)




