


import gc
import random
import unittest

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

from diffusers import AutoencoderKL, CycleDiffusionPipeline, DDIMScheduler, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

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


enable_full_determinism()

class 							lowercase  (						_UpperCAmelCase    ,       _UpperCAmelCase    ,       unittest.TestCase			):
     lowerCamelCase			:			Optional[int]      = CycleDiffusionPipeline
     lowerCamelCase			:			int      = TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
         '''negative_prompt''',
         '''height''',
         '''width''',
         '''negative_prompt_embeds''',
     }
     lowerCamelCase			:			Tuple      = PipelineTesterMixin.required_optional_params - {'''latents'''}
     lowerCamelCase			:			str      = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({'''source_prompt'''}			)
     lowerCamelCase			:			Optional[Any]      = IMAGE_TO_IMAGE_IMAGE_PARAMS
     lowerCamelCase			:			Tuple      = IMAGE_TO_IMAGE_IMAGE_PARAMS







     def 							lowercase__       (						self :						Optional[int]      ):
           torch.manual_seed(0      )
           SCREAMING_SNAKE_CASE__       :  List[Any]						     =		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   ,   )
           SCREAMING_SNAKE_CASE__       :  int						     =		DDIMScheduler(
               beta_start=0.00085   ,   beta_end=0.012   ,   beta_schedule='''scaled_linear'''   ,   num_train_timesteps=10_00   ,   clip_sample=_lowercase   ,   set_alpha_to_one=_lowercase   ,   )
           torch.manual_seed(0      )
           SCREAMING_SNAKE_CASE__       :  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      )
           SCREAMING_SNAKE_CASE__       :  List[Any]						     =		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   ,   )
           SCREAMING_SNAKE_CASE__       :  int						     =		CLIPTextModel(_lowercase      )
           SCREAMING_SNAKE_CASE__       :  Optional[int]						     =		CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''      )

           SCREAMING_SNAKE_CASE__       :  Optional[int]						     =		{
               '''unet''': unet,
               '''scheduler''': scheduler,
               '''vae''': vae,
               '''text_encoder''': text_encoder,
               '''tokenizer''': tokenizer,
               '''safety_checker''': None,
               '''feature_extractor''': None,
           }
           return components







     def 							lowercase__       (						self :						int   ,   _lowercase :						List[str]   ,   _lowercase :						int=0      ):
           SCREAMING_SNAKE_CASE__       :  List[str]						     =		floats_tensor((1, 3, 32, 32)   ,   rng=random.Random(_lowercase      )      ).to(_lowercase      )
           SCREAMING_SNAKE_CASE__       :  int						     =		image / 2 + 0.5
           if str(_lowercase      ).startswith('''mps'''      ):
                 SCREAMING_SNAKE_CASE__       :  List[Any]						     =		torch.manual_seed(_lowercase      )
           else:
                 SCREAMING_SNAKE_CASE__       :  Optional[Any]						     =		torch.Generator(device=_lowercase      ).manual_seed(_lowercase      )
           SCREAMING_SNAKE_CASE__       :  List[str]						     =		{
               '''prompt''': '''An astronaut riding an elephant''',
               '''source_prompt''': '''An astronaut riding a horse''',
               '''image''': image,
               '''generator''': generator,
               '''num_inference_steps''': 2,
               '''eta''': 0.1,
               '''strength''': 0.8,
               '''guidance_scale''': 3,
               '''source_guidance_scale''': 1,
               '''output_type''': '''numpy''',
           }
           return inputs







     def 							lowercase__       (						self :						Union[str, Any]      ):
           SCREAMING_SNAKE_CASE__       :  Any						     =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator

           SCREAMING_SNAKE_CASE__       :  Any						     =		self.get_dummy_components()
           SCREAMING_SNAKE_CASE__       :  str						     =		CycleDiffusionPipeline(**_lowercase      )
           SCREAMING_SNAKE_CASE__       :  List[str]						     =		pipe.to(_lowercase      )
           pipe.set_progress_bar_config(disable=_lowercase      )

           SCREAMING_SNAKE_CASE__       :  List[Any]						     =		self.get_dummy_inputs(_lowercase      )
           SCREAMING_SNAKE_CASE__       :  Union[str, Any]						     =		pipe(**_lowercase      )
           SCREAMING_SNAKE_CASE__       :  Tuple						     =		output.images

           SCREAMING_SNAKE_CASE__       :  int						     =		images[0, -3:, -3:, -1]

           assert images.shape == (1, 32, 32, 3)
           SCREAMING_SNAKE_CASE__       :  str						     =		np.array([0.4459, 0.4943, 0.4544, 0.6643, 0.5474, 0.4327, 0.5701, 0.5959, 0.5179]      )

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







     @unittest.skipIf(torch_device != '''cuda'''   ,   '''This test requires a GPU'''      )
     def 							lowercase__       (						self :						int      ):
           SCREAMING_SNAKE_CASE__       :  Union[str, Any]						     =		self.get_dummy_components()
           for name, module in components.items():
                 if hasattr(_lowercase   ,   '''half'''      ):
                       SCREAMING_SNAKE_CASE__       :  str						     =		module.half()
           SCREAMING_SNAKE_CASE__       :  str						     =		CycleDiffusionPipeline(**_lowercase      )
           SCREAMING_SNAKE_CASE__       :  Optional[Any]						     =		pipe.to(_lowercase      )
           pipe.set_progress_bar_config(disable=_lowercase      )

           SCREAMING_SNAKE_CASE__       :  Any						     =		self.get_dummy_inputs(_lowercase      )
           SCREAMING_SNAKE_CASE__       :  Dict						     =		pipe(**_lowercase      )
           SCREAMING_SNAKE_CASE__       :  str						     =		output.images

           SCREAMING_SNAKE_CASE__       :  List[str]						     =		images[0, -3:, -3:, -1]

           assert images.shape == (1, 32, 32, 3)
           SCREAMING_SNAKE_CASE__       :  Optional[Any]						     =		np.array([0.3506, 0.4543, 0.446, 0.4575, 0.5195, 0.4155, 0.5273, 0.518, 0.4116]      )

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







     @skip_mps
     def 							lowercase__       (						self :						List[Any]      ):
           return super().test_save_load_local()







     @unittest.skip('''non-deterministic pipeline'''      )
     def 							lowercase__       (						self :						str      ):
           return super().test_inference_batch_single_identical()







     @skip_mps
     def 							lowercase__       (						self :						Union[str, Any]      ):
           return super().test_dict_tuple_outputs_equivalent()







     @skip_mps
     def 							lowercase__       (						self :						Optional[Any]      ):
           return super().test_save_load_optional_components()







     @skip_mps
     def 							lowercase__       (						self :						Optional[Any]      ):
           return super().test_attention_slicing_forward_pass()




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







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







     def 							lowercase__       (						self :						Union[str, Any]      ):
           SCREAMING_SNAKE_CASE__       :  List[Any]						     =		load_image(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
               '''/cycle-diffusion/black_colored_car.png'''      )
           SCREAMING_SNAKE_CASE__       :  List[str]						     =		load_numpy(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/cycle-diffusion/blue_colored_car_fp16.npy'''      )
           SCREAMING_SNAKE_CASE__       :  Optional[Any]						     =		init_image.resize((5_12, 5_12)      )

           SCREAMING_SNAKE_CASE__       :  List[Any]						     =		'''CompVis/stable-diffusion-v1-4'''
           SCREAMING_SNAKE_CASE__       :  List[str]						     =		DDIMScheduler.from_pretrained(_lowercase   ,   subfolder='''scheduler'''      )
           SCREAMING_SNAKE_CASE__       :  Tuple						     =		CycleDiffusionPipeline.from_pretrained(
               _lowercase   ,   scheduler=_lowercase   ,   safety_checker=_lowercase   ,   torch_dtype=torch.floataa   ,   revision='''fp16'''      )

           pipe.to(_lowercase      )
           pipe.set_progress_bar_config(disable=_lowercase      )
           pipe.enable_attention_slicing()

           SCREAMING_SNAKE_CASE__       :  List[str]						     =		'''A black colored car'''
           SCREAMING_SNAKE_CASE__       :  Optional[int]						     =		'''A blue colored car'''

           SCREAMING_SNAKE_CASE__       :  Optional[Any]						     =		torch.manual_seed(0      )
           SCREAMING_SNAKE_CASE__       :  Optional[Any]						     =		pipe(
               prompt=_lowercase   ,   source_prompt=_lowercase   ,   image=_lowercase   ,   num_inference_steps=1_00   ,   eta=0.1   ,   strength=0.85   ,   guidance_scale=3   ,   source_guidance_scale=1   ,   generator=_lowercase   ,   output_type='''np'''   ,   )
           SCREAMING_SNAKE_CASE__       :  List[str]						     =		output.images

           # the values aren't exactly equal, but the images look the same visually
           assert np.abs(image - expected_image      ).max() < 5E-1







     def 							lowercase__       (						self :						str      ):
           SCREAMING_SNAKE_CASE__       :  Any						     =		load_image(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
               '''/cycle-diffusion/black_colored_car.png'''      )
           SCREAMING_SNAKE_CASE__       :  str						     =		load_numpy(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/cycle-diffusion/blue_colored_car.npy'''      )
           SCREAMING_SNAKE_CASE__       :  Union[str, Any]						     =		init_image.resize((5_12, 5_12)      )

           SCREAMING_SNAKE_CASE__       :  int						     =		'''CompVis/stable-diffusion-v1-4'''
           SCREAMING_SNAKE_CASE__       :  Dict						     =		DDIMScheduler.from_pretrained(_lowercase   ,   subfolder='''scheduler'''      )
           SCREAMING_SNAKE_CASE__       :  str						     =		CycleDiffusionPipeline.from_pretrained(_lowercase   ,   scheduler=_lowercase   ,   safety_checker=_lowercase      )

           pipe.to(_lowercase      )
           pipe.set_progress_bar_config(disable=_lowercase      )
           pipe.enable_attention_slicing()

           SCREAMING_SNAKE_CASE__       :  Optional[int]						     =		'''A black colored car'''
           SCREAMING_SNAKE_CASE__       :  int						     =		'''A blue colored car'''

           SCREAMING_SNAKE_CASE__       :  int						     =		torch.manual_seed(0      )
           SCREAMING_SNAKE_CASE__       :  List[Any]						     =		pipe(
               prompt=_lowercase   ,   source_prompt=_lowercase   ,   image=_lowercase   ,   num_inference_steps=1_00   ,   eta=0.1   ,   strength=0.85   ,   guidance_scale=3   ,   source_guidance_scale=1   ,   generator=_lowercase   ,   output_type='''np'''   ,   )
           SCREAMING_SNAKE_CASE__       :  int						     =		output.images

           assert np.abs(image - expected_image      ).max() < 2E-2













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


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



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				str     ,    SCREAMING_SNAKE_CASE							:				str ):


							'''simple docstring'''







							if "xprophetnet" in prophetnet_checkpoint_path:
														__lowerCamelCase   :   Optional[Any]       =XLMProphetNetForConditionalGenerationOld.from_pretrained(SCREAMING_SNAKE_CASE )
														__lowerCamelCase     ,    __lowerCamelCase   :   List[Any]       =XLMProphetNetForConditionalGeneration.from_pretrained(
														    SCREAMING_SNAKE_CASE     ,    output_loading_info=SCREAMING_SNAKE_CASE )
							else:
														__lowerCamelCase   :   int       =ProphetNetForConditionalGenerationOld.from_pretrained(SCREAMING_SNAKE_CASE )
														__lowerCamelCase     ,    __lowerCamelCase   :   int       =ProphetNetForConditionalGeneration.from_pretrained(
														    SCREAMING_SNAKE_CASE     ,    output_loading_info=SCREAMING_SNAKE_CASE )

							__lowerCamelCase   :   Optional[Any]       =['''key_proj''', '''value_proj''', '''query_proj''']

							__lowerCamelCase   :   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"]:
														__lowerCamelCase   :   int       =key.split('''.''' )

														if attributes[0] == "lm_head":
																					__lowerCamelCase   :   int       =prophet
																					__lowerCamelCase   :   Optional[int]       =prophet_old
														else:
																					__lowerCamelCase   :   Any       =prophet.prophetnet
																					__lowerCamelCase   :   Union[str, Any]       =prophet_old.model

														__lowerCamelCase   :   Optional[Any]       =False
														for attribute in attributes:
																					if attribute in mapping:
																												__lowerCamelCase   :   Optional[Any]       =mapping[attribute]
																												if not hasattr(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE ) and len(SCREAMING_SNAKE_CASE ) > 0:
																																			__lowerCamelCase   :   Any       =attribute
																					elif hasattr(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE ):
																												__lowerCamelCase   :   Any       =attribute

																					if attribute == "weight":
																												assert old_model.weight.shape == model.weight.shape, "Shapes have to match!"
																												__lowerCamelCase   :   str       =old_model.weight
																												logger.info(F'{attribute} is initialized.' )
																												__lowerCamelCase   :   Any       =True
																												break
																					elif attribute == "bias":
																												assert old_model.bias.shape == model.bias.shape, "Shapes have to match!"
																												__lowerCamelCase   :   Union[str, Any]       =old_model.bias
																												logger.info(F'{attribute} is initialized' )
																												__lowerCamelCase   :   str       =True
																												break
																					elif attribute in special_keys and hasattr(SCREAMING_SNAKE_CASE     ,    '''in_proj_weight''' ):
																												__lowerCamelCase   :   int       =old_model.in_proj_weight.shape[0] // 3
																												__lowerCamelCase   :   Union[str, Any]       =getattr(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE )
																												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":
																																			__lowerCamelCase   :   List[str]       =nn.Parameter(old_model.in_proj_weight[:embed_dim, :] )
																																			__lowerCamelCase   :   str       =nn.Parameter(old_model.in_proj_bias[:embed_dim] )

																												elif attribute == "key_proj":
																																			__lowerCamelCase   :   List[str]       =nn.Parameter(old_model.in_proj_weight[embed_dim : 2 * embed_dim, :] )
																																			__lowerCamelCase   :   Tuple       =nn.Parameter(old_model.in_proj_bias[embed_dim : 2 * embed_dim] )
																												elif attribute == "value_proj":
																																			__lowerCamelCase   :   Optional[Any]       =nn.Parameter(old_model.in_proj_weight[2 * embed_dim :, :] )
																																			__lowerCamelCase   :   int       =nn.Parameter(old_model.in_proj_bias[2 * embed_dim :] )
																												__lowerCamelCase   :   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] == 512, "We want 512 position_embeddings."
																												__lowerCamelCase   :   str       =nn.Parameter(old_model.embed_positions.weight[:512, :] )
																												__lowerCamelCase   :   Dict       =True
																												break

																					if attribute.isdigit():
																												__lowerCamelCase   :   List[str]       =model[int(SCREAMING_SNAKE_CASE )]
																												__lowerCamelCase   :   Optional[Any]       =old_model[int(SCREAMING_SNAKE_CASE )]
																					else:
																												__lowerCamelCase   :   int       =getattr(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE )

																												if old_attribute == "":
																																			__lowerCamelCase   :   Dict       =old_model
																												else:
																																			if not hasattr(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE ):
																																										raise ValueError(F'{old_model} does not have {old_attribute}' )
																																			__lowerCamelCase   :   Tuple       =getattr(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE )

														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(SCREAMING_SNAKE_CASE )


if __name__ == "__main__":
							_UpperCamelCase					      =   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.'
							)
							_UpperCamelCase					      =   parser.parse_args()
							convert_prophetnet_checkpoint_to_pytorch(args.prophetnet_checkpoint_path, args.pytorch_dump_folder_path)




