
import random
import unittest

import torch

from diffusers import IFInpaintingSuperResolutionPipeline
from diffusers.utils import floats_tensor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import skip_mps, torch_device

from ..pipeline_params import (
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin






@skip_mps
class     UpperCAmelCase						(		A_			,A_			,unittest.TestCase					):
   A__   :						Optional[int]			=	IFInpaintingSuperResolutionPipeline
   A__   :						Union[str, Any]			=	TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {"width", "height"}
   A__   :						int			=	TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS.union({"original_image"}					)
   A__   :						List[str]			=	PipelineTesterMixin.required_optional_params - {"latents"}



   def 							_SCREAMING_SNAKE_CASE  (self						:   Any					)						->      Optional[int]:


          '''simple docstring'''


          return self._get_superresolution_dummy_components()



   def 							_SCREAMING_SNAKE_CASE  (self						:   Any       ,		snake_case__						:   Dict       ,		snake_case__						:   Optional[Any]=0					)						->      int:


          '''simple docstring'''


          if str(snake_case__					).startswith("mps"					):
                 snake_case			:  Any     							= torch.manual_seed(snake_case__					)
          else:
                 snake_case			:  int     							= torch.Generator(device=snake_case__					).manual_seed(snake_case__					)

          snake_case			:  Any     							= floats_tensor((1, 3, 16, 16)       ,		rng=random.Random(snake_case__					)					).to(snake_case__					)
          snake_case			:  Optional[int]     							= floats_tensor((1, 3, 32, 32)       ,		rng=random.Random(snake_case__					)					).to(snake_case__					)
          snake_case			:  Any     							= floats_tensor((1, 3, 32, 32)       ,		rng=random.Random(snake_case__					)					).to(snake_case__					)

          snake_case			:  int     							= {
              "prompt": "A painting of a squirrel eating a burger",
              "image": image,
              "original_image": original_image,
              "mask_image": mask_image,
              "generator": generator,
              "num_inference_steps": 2,
              "output_type": "numpy",
          }

          return inputs



   @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						:   Dict					)						->      Dict:


          '''simple docstring'''


          self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3					)



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


          '''simple docstring'''


          self._test_save_load_optional_components()



   @unittest.skipIf(torch_device != "cuda"       ,		reason="float16 requires CUDA"					)
   def 							_SCREAMING_SNAKE_CASE  (self						:   Optional[int]					)						->      Any:


          '''simple docstring'''


          super().test_save_load_floataa(expected_max_diff=1e-1					)



   def 							_SCREAMING_SNAKE_CASE  (self						:   int					)						->      Any:


          '''simple docstring'''


          self._test_attention_slicing_forward_pass(expected_max_diff=1e-2					)



   def 							_SCREAMING_SNAKE_CASE  (self						:   List[Any]					)						->      List[Any]:


          '''simple docstring'''


          self._test_save_load_local()



   def 							_SCREAMING_SNAKE_CASE  (self						:   Any					)						->      str:


          '''simple docstring'''


          self._test_inference_batch_single_identical(
              expected_max_diff=1e-2       ,		)







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   (       __lowerCamelCase	:					int       ):
       snake_case			:  List[Any]     							= torch.load(__lowerCamelCase							,   map_location="cpu"       )
       return sd


def      UpperCamelCase   (       __lowerCamelCase	:					List[Any]							,   __lowerCamelCase	:					Optional[int]							,   __lowerCamelCase	:					Optional[Any]=rename_keys_prefix       ):
       snake_case			:  Optional[Any]     							= OrderedDict()
       snake_case			:  Optional[int]     							= 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
              snake_case			:  Any     							= key
              for name_pair in rename_keys_prefix:
                     snake_case			:  List[Any]     							= new_key.replace(name_pair[0]							,   name_pair[1]       )
              snake_case			:  List[str]     							= d[key]
              if key == "bert.cls.predictions.decoder.weight":
                     # Old bert code didn't have `decoder.bias`, but was added separately
                     snake_case			:  int     							= new_d["cls.predictions.bias"]
       return new_d


@torch.no_grad()
def      UpperCamelCase   (       __lowerCamelCase	:					List[Any]							,   __lowerCamelCase	:					Optional[Any]       ):
       assert (
           checkpoint_path.split("/"       )[-1] in ACCEPTABLE_CHECKPOINTS
       ), f"""The checkpoint provided must be in {ACCEPTABLE_CHECKPOINTS}."""

       # Get Config
       if "pre" in checkpoint_path:
              snake_case			:  Optional[Any]     							= "pretraining"
              if "vcr" in checkpoint_path:
                     snake_case			:  Optional[int]     							= {"visual_embedding_dim": 512}
              elif "vqa_advanced" in checkpoint_path:
                     snake_case			:  int     							= {"visual_embedding_dim": 2048}
              elif "vqa" in checkpoint_path:
                     snake_case			:  Dict     							= {"visual_embedding_dim": 2048}
              elif "nlvr" in checkpoint_path:
                     snake_case			:  Tuple     							= {"visual_embedding_dim": 1024}
              else:
                     raise NotImplementedError(f"""No implementation found for `{checkpoint_path}`."""       )
       else:
              if "vcr" in checkpoint_path:
                     snake_case			:  Optional[Any]     							= {"visual_embedding_dim": 512}
                     snake_case			:  List[str]     							= "multichoice"
              elif "vqa_advanced" in checkpoint_path:
                     snake_case			:  Dict     							= {"visual_embedding_dim": 2048}
                     snake_case			:  Tuple     							= "vqa_advanced"
              elif "vqa" in checkpoint_path:
                     snake_case			:  Optional[Any]     							= {"visual_embedding_dim": 2048, "num_labels": 3129}
                     snake_case			:  Tuple     							= "vqa"
              elif "nlvr" in checkpoint_path:
                     snake_case			:  str     							= {
                         "visual_embedding_dim": 1024,
                         "num_labels": 2,
                     }
                     snake_case			:  str     							= "nlvr"

       snake_case			:  int     							= VisualBertConfig(**__lowerCamelCase       )

       # Load State Dict
       snake_case			:  Any     							= load_state_dict(__lowerCamelCase       )

       snake_case			:  List[Any]     							= get_new_dict(__lowerCamelCase							,   __lowerCamelCase       )

       if model_type == "pretraining":
              snake_case			:  str     							= VisualBertForPreTraining(__lowerCamelCase       )
       elif model_type == "vqa":
              snake_case			:  Any     							= VisualBertForQuestionAnswering(__lowerCamelCase       )
       elif model_type == "nlvr":
              snake_case			:  Optional[Any]     							= VisualBertForVisualReasoning(__lowerCamelCase       )
       elif model_type == "multichoice":
              snake_case			:  Optional[int]     							= VisualBertForMultipleChoice(__lowerCamelCase       )

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


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)





