


import argparse
import json
import os
from collections import OrderedDict

import torch

from transformers import LukeConfig, LukeForMaskedLM, MLukeTokenizer, XLMRobertaTokenizer
from transformers.tokenization_utils_base import AddedToken




@torch.no_grad()
def     __SCREAMING_SNAKE_CASE							(       lowercase_      ,       lowercase_      ,       lowercase_      ,       lowercase_      ,       lowercase_      )		->     Optional[int]:




						'''simple docstring'''







						with open(_SCREAMING_SNAKE_CASE      ) as metadata_file:
												__UpperCAmelCase			:							Union[str, Any]	    =	json.load(_SCREAMING_SNAKE_CASE      )
						__UpperCAmelCase			:							List[str]	    =	LukeConfig(use_entity_aware_attention=_SCREAMING_SNAKE_CASE      ,       **metadata['''model_config''']      )

						# Load in the weights from the checkpoint_path
						__UpperCAmelCase			:							Optional[int]	    =	torch.load(_SCREAMING_SNAKE_CASE      ,       map_location='''cpu'''      )['module']

						# Load the entity vocab file
						__UpperCAmelCase			:							Any	    =	load_original_entity_vocab(_SCREAMING_SNAKE_CASE      )
						# add an entry for [MASK2]
						__UpperCAmelCase			:							Union[str, Any]	    =	max(entity_vocab.values()      ) + 1
						config.entity_vocab_size += 1

						__UpperCAmelCase			:							List[str]	    =	XLMRobertaTokenizer.from_pretrained(metadata['''model_config''']['''bert_model_name''']      )

						# Add special tokens to the token vocabulary for downstream tasks
						__UpperCAmelCase			:							List[Any]	    =	AddedToken('''<ent>'''      ,       lstrip=_SCREAMING_SNAKE_CASE      ,       rstrip=_SCREAMING_SNAKE_CASE      )
						__UpperCAmelCase			:							Optional[Any]	    =	AddedToken('''<ent2>'''      ,       lstrip=_SCREAMING_SNAKE_CASE      ,       rstrip=_SCREAMING_SNAKE_CASE      )
						tokenizer.add_special_tokens({'''additional_special_tokens''': [entity_token_a, entity_token_a]}      )
						config.vocab_size += 2

						print(f"Saving tokenizer to {pytorch_dump_folder_path}"      )
						tokenizer.save_pretrained(_SCREAMING_SNAKE_CASE      )
						with open(os.path.join(_SCREAMING_SNAKE_CASE      ,       '''tokenizer_config.json'''      )      ,       '''r'''      ) as f:
												__UpperCAmelCase			:							Optional[Any]	    =	json.load(_SCREAMING_SNAKE_CASE      )
						__UpperCAmelCase			:							Tuple	    =	'MLukeTokenizer'
						with open(os.path.join(_SCREAMING_SNAKE_CASE      ,       '''tokenizer_config.json'''      )      ,       '''w'''      ) as f:
												json.dump(_SCREAMING_SNAKE_CASE      ,       _SCREAMING_SNAKE_CASE      )

						with open(os.path.join(_SCREAMING_SNAKE_CASE      ,       MLukeTokenizer.vocab_files_names['''entity_vocab_file''']      )      ,       '''w'''      ) as f:
												json.dump(_SCREAMING_SNAKE_CASE      ,       _SCREAMING_SNAKE_CASE      )

						__UpperCAmelCase			:							Union[str, Any]	    =	MLukeTokenizer.from_pretrained(_SCREAMING_SNAKE_CASE      )

						# Initialize the embeddings of the special tokens
						__UpperCAmelCase			:							str	    =	tokenizer.convert_tokens_to_ids(['''@''']      )[0]
						__UpperCAmelCase			:							int	    =	tokenizer.convert_tokens_to_ids(['''#''']      )[0]

						__UpperCAmelCase			:							Dict	    =	state_dict['embeddings.word_embeddings.weight']
						__UpperCAmelCase			:							Union[str, Any]	    =	word_emb[ent_init_index].unsqueeze(0      )
						__UpperCAmelCase			:							str	    =	word_emb[enta_init_index].unsqueeze(0      )
						__UpperCAmelCase			:							Any	    =	torch.cat([word_emb, ent_emb, enta_emb]      )
						# add special tokens for 'entity_predictions.bias'
						for bias_name in ["lm_head.decoder.bias", "lm_head.bias"]:
												__UpperCAmelCase			:							str	    =	state_dict[bias_name]
												__UpperCAmelCase			:							List[str]	    =	decoder_bias[ent_init_index].unsqueeze(0      )
												__UpperCAmelCase			:							str	    =	decoder_bias[enta_init_index].unsqueeze(0      )
												__UpperCAmelCase			:							Union[str, Any]	    =	torch.cat([decoder_bias, ent_decoder_bias, enta_decoder_bias]      )

						# Initialize the query layers of the entity-aware self-attention mechanism
						for layer_index in range(config.num_hidden_layers      ):
												for matrix_name in ["query.weight", "query.bias"]:
																		__UpperCAmelCase			:							Dict	    =	f"encoder.layer.{layer_index}.attention.self."
																		__UpperCAmelCase			:							int	    =	state_dict[prefix + matrix_name]
																		__UpperCAmelCase			:							str	    =	state_dict[prefix + matrix_name]
																		__UpperCAmelCase			:							str	    =	state_dict[prefix + matrix_name]

												# Initialize the embedding of the [MASK2] entity using that of the [MASK] entity for downstream tasks
						__UpperCAmelCase			:							Optional[Any]	    =	state_dict['entity_embeddings.entity_embeddings.weight']
						__UpperCAmelCase			:							Union[str, Any]	    =	entity_emb[entity_vocab['[MASK]']].unsqueeze(0      )
						__UpperCAmelCase			:							Optional[Any]	    =	torch.cat([entity_emb, entity_mask_emb]      )
						# add [MASK2] for 'entity_predictions.bias'
						__UpperCAmelCase			:							List[Any]	    =	state_dict['entity_predictions.bias']
						__UpperCAmelCase			:							Dict	    =	entity_prediction_bias[entity_vocab['[MASK]']].unsqueeze(0      )
						__UpperCAmelCase			:							Any	    =	torch.cat([entity_prediction_bias, entity_mask_bias]      )

						__UpperCAmelCase			:							int	    =	LukeForMaskedLM(config=_SCREAMING_SNAKE_CASE      ).eval()

						state_dict.pop('''entity_predictions.decoder.weight'''      )
						state_dict.pop('''lm_head.decoder.weight'''      )
						state_dict.pop('''lm_head.decoder.bias'''      )
						__UpperCAmelCase			:							Tuple	    =	OrderedDict()
						for key, value in state_dict.items():
												if not (key.startswith('''lm_head'''      ) or key.startswith('''entity_predictions'''      )):
																		__UpperCAmelCase			:							List[str]	    =	state_dict[key]
												else:
																		__UpperCAmelCase			:							Union[str, Any]	    =	state_dict[key]

						__UpperCAmelCase			:							int	    =	model.load_state_dict(_SCREAMING_SNAKE_CASE      ,       strict=_SCREAMING_SNAKE_CASE      )

						if set(_SCREAMING_SNAKE_CASE      ) != {"luke.embeddings.position_ids"}:
												raise ValueError(f"Unexpected unexpected_keys: {unexpected_keys}"      )
						if set(_SCREAMING_SNAKE_CASE      ) != {
						    "lm_head.decoder.weight",
						    "lm_head.decoder.bias",
						    "entity_predictions.decoder.weight",
						}:
												raise ValueError(f"Unexpected missing_keys: {missing_keys}"      )

						model.tie_weights()
						assert (model.luke.embeddings.word_embeddings.weight == model.lm_head.decoder.weight).all()
						assert (model.luke.entity_embeddings.entity_embeddings.weight == model.entity_predictions.decoder.weight).all()

						# Check outputs
						__UpperCAmelCase			:							List[Any]	    =	MLukeTokenizer.from_pretrained(_SCREAMING_SNAKE_CASE      ,       task='''entity_classification'''      )

						__UpperCAmelCase			:							Union[str, Any]	    =	'ISO 639-3 uses the code fas for the dialects spoken across Iran and アフガニスタン (Afghanistan).'
						__UpperCAmelCase			:							Optional[Any]	    =	(0, 9)
						__UpperCAmelCase			:							int	    =	tokenizer(_SCREAMING_SNAKE_CASE      ,       entity_spans=[span]      ,       return_tensors='''pt'''      )

						__UpperCAmelCase			:							List[str]	    =	model(**_SCREAMING_SNAKE_CASE      )

						# Verify word hidden states
						if model_size == "large":
												raise NotImplementedError
						else:  # base
												__UpperCAmelCase			:							Union[str, Any]	    =	torch.Size((1, 33, 768)      )
												__UpperCAmelCase			:							Optional[int]	    =	torch.tensor([[0.0_8_9_2, 0.0_5_9_6, -0.2_8_1_9], [0.0_1_3_4, 0.1_1_9_9, 0.0_5_7_3], [-0.0_1_6_9, 0.0_9_2_7, 0.0_6_4_4]]      )

						if not (outputs.last_hidden_state.shape == expected_shape):
												raise ValueError(
												    f"Outputs.last_hidden_state.shape is {outputs.last_hidden_state.shape}, Expected shape is {expected_shape}"      )
						if not torch.allclose(outputs.last_hidden_state[0, :3, :3]      ,       _SCREAMING_SNAKE_CASE      ,       atol=1e-4      ):
												raise ValueError

						# Verify entity hidden states
						if model_size == "large":
												raise NotImplementedError
						else:  # base
												__UpperCAmelCase			:							Optional[Any]	    =	torch.Size((1, 1, 768)      )
												__UpperCAmelCase			:							Any	    =	torch.tensor([[-0.1_4_8_2, 0.0_6_0_9, 0.0_3_2_2]]      )

						if not (outputs.entity_last_hidden_state.shape == expected_shape):
												raise ValueError(
												    f"Outputs.entity_last_hidden_state.shape is {outputs.entity_last_hidden_state.shape}, Expected shape is"
												    f" {expected_shape}"      )
						if not torch.allclose(outputs.entity_last_hidden_state[0, :3, :3]      ,       _SCREAMING_SNAKE_CASE      ,       atol=1e-4      ):
												raise ValueError

						# Verify masked word/entity prediction
						__UpperCAmelCase			:							str	    =	MLukeTokenizer.from_pretrained(_SCREAMING_SNAKE_CASE      )
						__UpperCAmelCase			:							Tuple	    =	'Tokyo is the capital of <mask>.'
						__UpperCAmelCase			:							Optional[int]	    =	(24, 30)
						__UpperCAmelCase			:							Any	    =	tokenizer(_SCREAMING_SNAKE_CASE      ,       entity_spans=[span]      ,       return_tensors='''pt'''      )

						__UpperCAmelCase			:							int	    =	model(**_SCREAMING_SNAKE_CASE      )

						__UpperCAmelCase			:							Optional[int]	    =	encoding['input_ids'][0].tolist()
						__UpperCAmelCase			:							Optional[int]	    =	input_ids.index(tokenizer.convert_tokens_to_ids('''<mask>'''      )      )
						__UpperCAmelCase			:							str	    =	outputs.logits[0][mask_position_id].argmax(dim=-1      )
						assert "Japan" == tokenizer.decode(_SCREAMING_SNAKE_CASE      )

						__UpperCAmelCase			:							Optional[int]	    =	outputs.entity_logits[0][0].argmax().item()
						__UpperCAmelCase			:							Dict	    =	[
						    entity for entity, entity_id in tokenizer.entity_vocab.items() if entity_id == predicted_entity_id
						]
						assert [e for e in multilingual_predicted_entities if e.startswith('''en:'''      )][0] == "en:Japan"

						# Finally, save our PyTorch model and tokenizer
						print('''Saving PyTorch model to {}'''.format(_SCREAMING_SNAKE_CASE      )      )
						model.save_pretrained(_SCREAMING_SNAKE_CASE      )




def     __SCREAMING_SNAKE_CASE							(       lowercase_      )		->     str:




						'''simple docstring'''







						__UpperCAmelCase			:							List[Any]	    =	['[MASK]', '[PAD]', '[UNK]']

						__UpperCAmelCase			:							Any	    =	[json.loads(_SCREAMING_SNAKE_CASE      ) for line in open(_SCREAMING_SNAKE_CASE      )]

						__UpperCAmelCase			:							Optional[Any]	    =	{}
						for entry in data:
												__UpperCAmelCase			:							Optional[int]	    =	entry['id']
												for entity_name, language in entry["entities"]:
																		if entity_name in SPECIAL_TOKENS:
																								__UpperCAmelCase			:							Tuple	    =	entity_id
																								break
																		__UpperCAmelCase			:							int	    =	f"{language}:{entity_name}"
																		__UpperCAmelCase			:							int	    =	entity_id
						return new_mapping


if __name__ == "__main__":
				lowerCAmelCase       		=    argparse.ArgumentParser()
				# Required parameters
				parser.add_argument("""--checkpoint_path""", type=str, help="""Path to a pytorch_model.bin file.""")
				parser.add_argument(
				    """--metadata_path""", default=None, type=str, help="""Path to a metadata.json file, defining the configuration."""
				)
				parser.add_argument(
				    """--entity_vocab_path""",
				    default=None,
				    type=str,
				    help="""Path to an entity_vocab.tsv file, containing the entity vocabulary.""",
				)
				parser.add_argument(
				    """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to where to dump the output PyTorch model."""
				)
				parser.add_argument(
				    """--model_size""", default="""base""", type=str, choices=["""base""", """large"""], help="""Size of the model to be converted."""
				)
				lowerCAmelCase       		=    parser.parse_args()
				convert_luke_checkpoint(
				    args.checkpoint_path,
				    args.metadata_path,
				    args.entity_vocab_path,
				    args.pytorch_dump_folder_path,
				    args.model_size,
				)








"""simple docstring"""





import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()






class    _lowerCAmelCase  (	a    ,  a    ,  a    ,  unittest.TestCase			):
       """simple docstring"""




       __magic_name__	:int						    =   StableUnCLIPPipeline
       __magic_name__	:int						    =   TEXT_TO_IMAGE_PARAMS
       __magic_name__	:List[str]						    =   TEXT_TO_IMAGE_BATCH_PARAMS
       __magic_name__	:Optional[int]						    =   TEXT_TO_IMAGE_IMAGE_PARAMS
       __magic_name__	:Union[str, Any]						    =   TEXT_TO_IMAGE_IMAGE_PARAMS

       # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
       __magic_name__	:List[str]						    =   False






       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Dict											=				3_2
         lowerCAmelCase__     :List[Any]											=				embedder_hidden_size

         # prior components

         torch.manual_seed(0 )
         lowerCAmelCase__     :Any											=				CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip' )

         torch.manual_seed(0 )
         lowerCAmelCase__     :Any											=				CLIPTextModelWithProjection(
             CLIPTextConfig(
                 bos_token_id=0       ,					eos_token_id=2       ,					hidden_size=__UpperCAmelCase       ,					projection_dim=__UpperCAmelCase       ,					intermediate_size=3_7       ,					layer_norm_eps=1E-05       ,					num_attention_heads=4       ,					num_hidden_layers=5       ,					pad_token_id=1       ,					vocab_size=1_0_0_0       ,					) )

         torch.manual_seed(0 )
         lowerCAmelCase__     :int											=				PriorTransformer(
             num_attention_heads=2       ,					attention_head_dim=1_2       ,					embedding_dim=__UpperCAmelCase       ,					num_layers=1       ,					)

         torch.manual_seed(0 )
         lowerCAmelCase__     :Optional[int]											=				DDPMScheduler(
             variance_type='fixed_small_log'       ,					prediction_type='sample'       ,					num_train_timesteps=1_0_0_0       ,					clip_sample=__UpperCAmelCase       ,					clip_sample_range=5.0       ,					beta_schedule='squaredcos_cap_v2'       ,					)

         # regular denoising components

         torch.manual_seed(0 )
         lowerCAmelCase__     :Dict											=				StableUnCLIPImageNormalizer(embedding_dim=__UpperCAmelCase )
         lowerCAmelCase__     :List[Any]											=				DDPMScheduler(beta_schedule='squaredcos_cap_v2' )

         torch.manual_seed(0 )
         lowerCAmelCase__     :List[Any]											=				CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip' )

         torch.manual_seed(0 )
         lowerCAmelCase__     :str											=				CLIPTextModel(
             CLIPTextConfig(
                 bos_token_id=0       ,					eos_token_id=2       ,					hidden_size=__UpperCAmelCase       ,					projection_dim=3_2       ,					intermediate_size=3_7       ,					layer_norm_eps=1E-05       ,					num_attention_heads=4       ,					num_hidden_layers=5       ,					pad_token_id=1       ,					vocab_size=1_0_0_0       ,					) )

         torch.manual_seed(0 )
         lowerCAmelCase__     :Optional[int]											=				UNetaDConditionModel(
             sample_size=3_2       ,					in_channels=4       ,					out_channels=4       ,					down_block_types=('CrossAttnDownBlock2D', 'DownBlock2D')       ,					up_block_types=('UpBlock2D', 'CrossAttnUpBlock2D')       ,					block_out_channels=(3_2, 6_4)       ,					attention_head_dim=(2, 4)       ,					class_embed_type='projection'       ,					projection_class_embeddings_input_dim=embedder_projection_dim * 2       ,					cross_attention_dim=__UpperCAmelCase       ,					layers_per_block=1       ,					upcast_attention=__UpperCAmelCase       ,					use_linear_projection=__UpperCAmelCase       ,					)

         torch.manual_seed(0 )
         lowerCAmelCase__     :Tuple											=				DDIMScheduler(
             beta_schedule='scaled_linear'       ,					beta_start=0.0_00_85       ,					beta_end=0.0_12       ,					prediction_type='v_prediction'       ,					set_alpha_to_one=__UpperCAmelCase       ,					steps_offset=1       ,					)

         torch.manual_seed(0 )
         lowerCAmelCase__     :Optional[int]											=				AutoencoderKL()

         lowerCAmelCase__     :Optional[int]											=				{
             # prior components
             'prior_tokenizer': prior_tokenizer,
             'prior_text_encoder': prior_text_encoder,
             'prior': prior,
             'prior_scheduler': prior_scheduler,
             # image noising components
             'image_normalizer': image_normalizer,
             'image_noising_scheduler': image_noising_scheduler,
             # regular denoising components
             'tokenizer': tokenizer,
             'text_encoder': text_encoder,
             'unet': unet,
             'scheduler': scheduler,
             'vae': vae,
         }

         return components






       def 	snake_case		(       self       ,					__UpperCAmelCase       ,					__UpperCAmelCase=0 ):



         '''simple docstring'''

         if str(__UpperCAmelCase ).startswith('mps' ):
           lowerCAmelCase__     :str											=				torch.manual_seed(__UpperCAmelCase )
         else:
           lowerCAmelCase__     :Optional[int]											=				torch.Generator(device=__UpperCAmelCase ).manual_seed(__UpperCAmelCase )
         lowerCAmelCase__     :Tuple											=				{
             'prompt': 'A painting of a squirrel eating a burger',
             'generator': generator,
             'num_inference_steps': 2,
             'prior_num_inference_steps': 2,
             'output_type': 'numpy',
         }
         return inputs






       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Optional[int]											=				torch_device == 'cpu'

         self._test_attention_slicing_forward_pass(test_max_difference=__UpperCAmelCase )






       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :str											=				torch_device in ['cpu', 'mps']

         self._test_inference_batch_single_identical(test_max_difference=__UpperCAmelCase )






@slow
@require_torch_gpu
class    _lowerCAmelCase  (	unittest.TestCase			):
       """simple docstring"""




       def 	snake_case		(       self ):



         '''simple docstring'''

         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()






       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Tuple											=				load_numpy(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy' )

         lowerCAmelCase__     :List[Any]											=				StableUnCLIPPipeline.from_pretrained('fusing/stable-unclip-2-1-l'       ,					torch_dtype=torch.floataa )
         pipe.to(__UpperCAmelCase )
         pipe.set_progress_bar_config(disable=__UpperCAmelCase )
         # 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()

         lowerCAmelCase__     :List[str]											=				torch.Generator(device='cpu' ).manual_seed(0 )
         lowerCAmelCase__     :List[str]											=				pipe('anime turle'       ,					generator=__UpperCAmelCase       ,					output_type='np' )

         lowerCAmelCase__     :Union[str, Any]											=				output.images[0]

         assert image.shape == (7_6_8, 7_6_8, 3)

         assert_mean_pixel_difference(__UpperCAmelCase       ,					__UpperCAmelCase )






       def 	snake_case		(       self ):



         '''simple docstring'''

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

         lowerCAmelCase__     :int											=				StableUnCLIPPipeline.from_pretrained('fusing/stable-unclip-2-1-l'       ,					torch_dtype=torch.floataa )
         lowerCAmelCase__     :List[str]											=				pipe.to(__UpperCAmelCase )
         pipe.set_progress_bar_config(disable=__UpperCAmelCase )
         pipe.enable_attention_slicing()
         pipe.enable_sequential_cpu_offload()

         lowerCAmelCase__     :Tuple											=				pipe(
             'anime turtle'       ,					prior_num_inference_steps=2       ,					num_inference_steps=2       ,					output_type='np'       ,					)

         lowerCAmelCase__     :str											=				torch.cuda.max_memory_allocated()
         # make sure that less than 7 GB is allocated
         assert mem_bytes < 7 * 1_0**9
