

'''simple docstring'''







import argparse
from collections import OrderedDict
from pathlib import Path

import torch

from transformers import (
    VisualBertConfig,
    VisualBertForMultipleChoice,
    VisualBertForPreTraining,
    VisualBertForQuestionAnswering,
    VisualBertForVisualReasoning,
)
from transformers.utils import logging


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

__lowerCamelCase		       =    [
    ('''bert.bert''', '''visual_bert'''),
    ('''bert.cls''', '''cls'''),
    ('''bert.classifier''', '''cls'''),
    ('''token_type_embeddings_visual''', '''visual_token_type_embeddings'''),
    ('''position_embeddings_visual''', '''visual_position_embeddings'''),
    ('''projection''', '''visual_projection'''),
]

__lowerCamelCase		       =    [
    '''nlvr2_coco_pre_trained.th''',
    '''nlvr2_fine_tuned.th''',
    '''nlvr2_pre_trained.th''',
    '''vcr_coco_pre_train.th''',
    '''vcr_fine_tune.th''',
    '''vcr_pre_train.th''',
    '''vqa_coco_pre_trained.th''',
    '''vqa_fine_tuned.th''',
    '''vqa_pre_trained.th''',
]
def       UpperCAmelCase__    (      UpperCAmelCase__				)			->			str:
       A_						=      torch.load(UpperCAmelCase__,       map_location="""cpu"""				)
       return sd
def       UpperCAmelCase__    (      UpperCAmelCase__,       UpperCAmelCase__,       UpperCAmelCase__=rename_keys_prefix				)			->			Optional[int]:
       A_						=      OrderedDict()
       A_						=      torch.arange(config.max_position_embeddings				).expand((1, -1)				)
       # detector_d = OrderedDict()
       for key in d:
              if "detector" in key:
                     # detector_d[key.replace('detector.','')] = d[key]
                     continue
              A_						=      key
              for name_pair in rename_keys_prefix:
                     A_						=      new_key.replace(name_pair[0],       name_pair[1]				)
              A_						=      d[key]
              if key == "bert.cls.predictions.decoder.weight":
                     # Old bert code didn't have `decoder.bias`, but was added separately
                     A_						=      new_d["""cls.predictions.bias"""]
       return new_d
@torch.no_grad()
def       UpperCAmelCase__    (      UpperCAmelCase__,       UpperCAmelCase__				)			->			Dict:
       assert (
           checkpoint_path.split("""/"""				)[-1] in ACCEPTABLE_CHECKPOINTS
       ), F'''The checkpoint provided must be in {ACCEPTABLE_CHECKPOINTS}.'''

       # Get Config
       if "pre" in checkpoint_path:
              A_						=      """pretraining"""
              if "vcr" in checkpoint_path:
                     A_						=      {"""visual_embedding_dim""": 5_12}
              elif "vqa_advanced" in checkpoint_path:
                     A_						=      {"""visual_embedding_dim""": 20_48}
              elif "vqa" in checkpoint_path:
                     A_						=      {"""visual_embedding_dim""": 20_48}
              elif "nlvr" in checkpoint_path:
                     A_						=      {"""visual_embedding_dim""": 10_24}
              else:
                     raise NotImplementedError(F'''No implementation found for `{checkpoint_path}`.'''				)
       else:
              if "vcr" in checkpoint_path:
                     A_						=      {"""visual_embedding_dim""": 5_12}
                     A_						=      """multichoice"""
              elif "vqa_advanced" in checkpoint_path:
                     A_						=      {"""visual_embedding_dim""": 20_48}
                     A_						=      """vqa_advanced"""
              elif "vqa" in checkpoint_path:
                     A_						=      {"""visual_embedding_dim""": 20_48, """num_labels""": 31_29}
                     A_						=      """vqa"""
              elif "nlvr" in checkpoint_path:
                     A_						=      {
                         """visual_embedding_dim""": 10_24,
                         """num_labels""": 2,
                     }
                     A_						=      """nlvr"""

       A_						=      VisualBertConfig(**UpperCAmelCase__				)

       # Load State Dict
       A_						=      load_state_dict(UpperCAmelCase__				)

       A_						=      get_new_dict(UpperCAmelCase__,       UpperCAmelCase__				)

       if model_type == "pretraining":
              A_						=      VisualBertForPreTraining(UpperCAmelCase__				)
       elif model_type == "vqa":
              A_						=      VisualBertForQuestionAnswering(UpperCAmelCase__				)
       elif model_type == "nlvr":
              A_						=      VisualBertForVisualReasoning(UpperCAmelCase__				)
       elif model_type == "multichoice":
              A_						=      VisualBertForMultipleChoice(UpperCAmelCase__				)

       model.load_state_dict(UpperCAmelCase__				)
       # Save Checkpoints
       Path(UpperCAmelCase__				).mkdir(exist_ok=UpperCAmelCase__				)
       model.save_pretrained(UpperCAmelCase__				)


if __name__ == "__main__":
      __lowerCamelCase		       =    argparse.ArgumentParser()
      # Required parameters
      parser.add_argument('''orig_checkpoint_path''', type=str, help='''A path to .th on local filesystem.''')
      parser.add_argument('''pytorch_dump_folder_path''', type=str, help='''Path to the output PyTorch model.''')
      __lowerCamelCase		       =    parser.parse_args()
      convert_visual_bert_checkpoint(args.orig_checkpoint_path, args.pytorch_dump_folder_path)



'''simple docstring'''







import unittest

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

from diffusers import KandinskyVaaPriorPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


class 					A__      (      _snake_case						,						unittest.TestCase			):
    lowercase						=		KandinskyVaaPriorPipeline
    lowercase						=		["prompt"]
    lowercase						=		["prompt", "negative_prompt"]
    lowercase						=		[
        "num_images_per_prompt",
        "generator",
        "num_inference_steps",
        "latents",
        "negative_prompt",
        "guidance_scale",
        "output_type",
        "return_dict",
    ]
    lowercase						=		False



    @property
    def 					snake_case_			(					self       )						->      Any:





           '''simple docstring'''
           return 32



    @property
    def 					snake_case_			(					self       )						->      Optional[int]:





           '''simple docstring'''
           return 32



    @property
    def 					snake_case_			(					self       )						->      Optional[Any]:





           '''simple docstring'''
           return self.time_input_dim



    @property
    def 					snake_case_			(					self       )						->      str:





           '''simple docstring'''
           return self.time_input_dim * 4



    @property
    def 					snake_case_			(					self       )						->      int:





           '''simple docstring'''
           return 100



    @property
    def 					snake_case_			(					self       )						->      Optional[Any]:





           '''simple docstring'''
           A_						=      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""       )
           return tokenizer



    @property
    def 					snake_case_			(					self       )						->      Union[str, Any]:





           '''simple docstring'''
           torch.manual_seed(0       )
           A_						=      CLIPTextConfig(
               bos_token_id=0		,				eos_token_id=2		,				hidden_size=self.text_embedder_hidden_size		,				projection_dim=self.text_embedder_hidden_size		,				intermediate_size=37		,				layer_norm_eps=1e-0_5		,				num_attention_heads=4		,				num_hidden_layers=5		,				pad_token_id=1		,				vocab_size=1000		,				)
           return CLIPTextModelWithProjection(UpperCamelCase__       )



    @property
    def 					snake_case_			(					self       )						->      Tuple:





           '''simple docstring'''
           torch.manual_seed(0       )

           A_						=      {
               """num_attention_heads""": 2,
               """attention_head_dim""": 12,
               """embedding_dim""": self.text_embedder_hidden_size,
               """num_layers""": 1,
           }

           A_						=      PriorTransformer(**UpperCamelCase__       )
           # clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
           A_						=      nn.Parameter(torch.ones(model.clip_std.shape       )       )
           return model



    @property
    def 					snake_case_			(					self       )						->      str:





           '''simple docstring'''
           torch.manual_seed(0       )
           A_						=      CLIPVisionConfig(
               hidden_size=self.text_embedder_hidden_size		,				image_size=224		,				projection_dim=self.text_embedder_hidden_size		,				intermediate_size=37		,				num_attention_heads=4		,				num_channels=3		,				num_hidden_layers=5		,				patch_size=14		,				)

           A_						=      CLIPVisionModelWithProjection(UpperCamelCase__       )
           return model



    @property
    def 					snake_case_			(					self       )						->      str:





           '''simple docstring'''
           A_						=      CLIPImageProcessor(
               crop_size=224		,				do_center_crop=UpperCamelCase__		,				do_normalize=UpperCamelCase__		,				do_resize=UpperCamelCase__		,				image_mean=[0.48145466, 0.4578275, 0.40821073]		,				image_std=[0.26862954, 0.26130258, 0.27577711]		,				resample=3		,				size=224		,				)

           return image_processor



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





           '''simple docstring'''
           A_						=      self.dummy_prior
           A_						=      self.dummy_image_encoder
           A_						=      self.dummy_text_encoder
           A_						=      self.dummy_tokenizer
           A_						=      self.dummy_image_processor

           A_						=      UnCLIPScheduler(
               variance_type="""fixed_small_log"""		,				prediction_type="""sample"""		,				num_train_timesteps=1000		,				clip_sample=UpperCamelCase__		,				clip_sample_range=10.0		,				)

           A_						=      {
               """prior""": prior,
               """image_encoder""": image_encoder,
               """text_encoder""": text_encoder,
               """tokenizer""": tokenizer,
               """scheduler""": scheduler,
               """image_processor""": image_processor,
           }

           return components



    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__=0       )						->      Optional[int]:





           '''simple docstring'''
           if str(UpperCamelCase__       ).startswith("""mps"""       ):
                  A_						=      torch.manual_seed(UpperCamelCase__       )
           else:
                  A_						=      torch.Generator(device=UpperCamelCase__       ).manual_seed(UpperCamelCase__       )
           A_						=      {
               """prompt""": """horse""",
               """generator""": generator,
               """guidance_scale""": 4.0,
               """num_inference_steps""": 2,
               """output_type""": """np""",
           }
           return inputs



    def 					snake_case_			(					self       )						->      List[str]:





           '''simple docstring'''
           A_						=      """cpu"""

           A_						=      self.get_dummy_components()

           A_						=      self.pipeline_class(**UpperCamelCase__       )
           A_						=      pipe.to(UpperCamelCase__       )

           pipe.set_progress_bar_config(disable=UpperCamelCase__       )

           A_						=      pipe(**self.get_dummy_inputs(UpperCamelCase__       )       )
           A_						=      output.image_embeds

           A_						=      pipe(
               **self.get_dummy_inputs(UpperCamelCase__       )		,				return_dict=UpperCamelCase__		,				)[0]

           A_						=      image[0, -10:]
           A_						=      image_from_tuple[0, -10:]

           assert image.shape == (1, 32)

           A_						=      np.array(
               [-0.0532, 1.7120, 0.3656, -1.0852, -0.8946, -1.1756, 0.4348, 0.2482, 0.5146, -0.1156]       )

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



    @skip_mps
    def 					snake_case_			(					self       )						->      int:





           '''simple docstring'''
           A_						=      torch_device == """cpu"""
           A_						=      True
           A_						=      False

           self._test_inference_batch_single_identical(
               test_max_difference=UpperCamelCase__		,				relax_max_difference=UpperCamelCase__		,				test_mean_pixel_difference=UpperCamelCase__		,				)







    @skip_mps
    def 					snake_case_			(					self       )						->      Union[str, Any]:





           '''simple docstring'''
           A_						=      torch_device == """cpu"""
           A_						=      False

           self._test_attention_slicing_forward_pass(
               test_max_difference=UpperCamelCase__		,				test_mean_pixel_difference=UpperCamelCase__		,				)
