



"""simple docstring"""





import argparse

import requests
import torch

# pip3 install salesforce-lavis
# I'm actually installing a slightly modified version: pip3 install git+https://github.com/nielsrogge/LAVIS.git@fix_lavis
from lavis.models import load_model_and_preprocess
from PIL import Image

from transformers import (
    AutoTokenizer,
    BlipaConfig,
    BlipaForConditionalGeneration,
    BlipaProcessor,
    BlipaVisionConfig,
    BlipImageProcessor,
    OPTConfig,
    TaConfig,
)
from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD
def        a__     (				):
      UpperCAmelCase_      			=							'https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'
      UpperCAmelCase_      			=							Image.open(requests.get(snake_case__     ,      stream=snake_case__				).raw				).convert("RGB"				)

      return image
def        a__     (				lowerCAmelCase__				):
      UpperCAmelCase_      			=							[]
      # fmt: off

      # vision encoder
      rename_keys.append(("visual_encoder.cls_token", "vision_model.embeddings.class_embedding")				)
      rename_keys.append(("visual_encoder.pos_embed", "vision_model.embeddings.position_embedding")				)
      rename_keys.append(("visual_encoder.patch_embed.proj.weight", "vision_model.embeddings.patch_embedding.weight")				)
      rename_keys.append(("visual_encoder.patch_embed.proj.bias", "vision_model.embeddings.patch_embedding.bias")				)
      rename_keys.append(("ln_vision.weight", "vision_model.post_layernorm.weight")				)
      rename_keys.append(("ln_vision.bias", "vision_model.post_layernorm.bias")				)

      for i in range(config.vision_config.num_hidden_layers				):
            rename_keys.append((f"""visual_encoder.blocks.{i}.norm1.weight""", f"""vision_model.encoder.layers.{i}.layer_norm1.weight""")				)
            rename_keys.append((f"""visual_encoder.blocks.{i}.norm1.bias""", f"""vision_model.encoder.layers.{i}.layer_norm1.bias""")				)
            rename_keys.append((f"""visual_encoder.blocks.{i}.norm2.weight""", f"""vision_model.encoder.layers.{i}.layer_norm2.weight""")				)
            rename_keys.append((f"""visual_encoder.blocks.{i}.norm2.bias""", f"""vision_model.encoder.layers.{i}.layer_norm2.bias""")				)
            rename_keys.append((f"""visual_encoder.blocks.{i}.attn.qkv.weight""", f"""vision_model.encoder.layers.{i}.self_attn.qkv.weight""")				)
            rename_keys.append((f"""visual_encoder.blocks.{i}.attn.proj.weight""", f"""vision_model.encoder.layers.{i}.self_attn.projection.weight""",)				)
            rename_keys.append((f"""visual_encoder.blocks.{i}.attn.proj.bias""", f"""vision_model.encoder.layers.{i}.self_attn.projection.bias""")				)
            rename_keys.append((f"""visual_encoder.blocks.{i}.mlp.fc1.weight""", f"""vision_model.encoder.layers.{i}.mlp.fc1.weight""")				)
            rename_keys.append((f"""visual_encoder.blocks.{i}.mlp.fc1.bias""", f"""vision_model.encoder.layers.{i}.mlp.fc1.bias""")				)
            rename_keys.append((f"""visual_encoder.blocks.{i}.mlp.fc2.weight""", f"""vision_model.encoder.layers.{i}.mlp.fc2.weight""")				)
            rename_keys.append((f"""visual_encoder.blocks.{i}.mlp.fc2.bias""", f"""vision_model.encoder.layers.{i}.mlp.fc2.bias""")				)

      # QFormer
      rename_keys.append(("Qformer.bert.embeddings.LayerNorm.weight", "qformer.layernorm.weight")				)
      rename_keys.append(("Qformer.bert.embeddings.LayerNorm.bias", "qformer.layernorm.bias")				)

      # fmt: on
      return rename_keys
def        a__     (				lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__				):
      UpperCAmelCase_      			=							dct.pop(snake_case__				)
      UpperCAmelCase_      			=							val
def        a__     (				lowerCAmelCase__     ,      lowerCAmelCase__				):
      for i in range(config.vision_config.num_hidden_layers				):
            # read in original q and v biases
            UpperCAmelCase_      			=							state_dict.pop(f"""visual_encoder.blocks.{i}.attn.q_bias"""				)
            UpperCAmelCase_      			=							state_dict.pop(f"""visual_encoder.blocks.{i}.attn.v_bias"""				)

            # next, set bias in the state dict
            UpperCAmelCase_      			=							torch.cat((q_bias, torch.zeros_like(snake_case__     ,      requires_grad=snake_case__				), v_bias)				)
            UpperCAmelCase_      			=							qkv_bias
def        a__     (				lowerCAmelCase__     ,      lowerCAmelCase__				):
      UpperCAmelCase_      			=							364 if 'coco' in model_name else 224
      UpperCAmelCase_      			=							BlipaVisionConfig(image_size=snake_case__				).to_dict()

      # make sure the models have proper bos_token_id and eos_token_id set (important for generation)
      # seems like flan-T5 models don't have bos_token_id properly set?
      if "opt-2.7b" in model_name:
            UpperCAmelCase_      			=							OPTConfig.from_pretrained("facebook/opt-2.7b"     ,      eos_token_id=snake_case__				).to_dict()
      elif "opt-6.7b" in model_name:
            UpperCAmelCase_      			=							OPTConfig.from_pretrained("facebook/opt-6.7b"     ,      eos_token_id=snake_case__				).to_dict()
      elif "t5-xl" in model_name:
            UpperCAmelCase_      			=							TaConfig.from_pretrained("google/flan-t5-xl"     ,      dense_act_fn="gelu"     ,      bos_token_id=1				).to_dict()
      elif "t5-xxl" in model_name:
            UpperCAmelCase_      			=							TaConfig.from_pretrained("google/flan-t5-xxl"     ,      dense_act_fn="gelu"     ,      bos_token_id=1				).to_dict()

      UpperCAmelCase_      			=							BlipaConfig(vision_config=snake_case__     ,      text_config=snake_case__				)

      return config, image_size
@torch.no_grad()
def        a__     (				lowerCAmelCase__     ,      lowerCAmelCase__=None     ,      lowerCAmelCase__=False				):
      UpperCAmelCase_      			=							(
          AutoTokenizer.from_pretrained("facebook/opt-2.7b"				)
          if 'opt' in model_name
          else AutoTokenizer.from_pretrained("google/flan-t5-xl"				)
      )
      UpperCAmelCase_      			=							tokenizer("\n"     ,      add_special_tokens=snake_case__				).input_ids[0]
      UpperCAmelCase_      			=							get_blipa_config(snake_case__     ,      eos_token_id=snake_case__				)

      UpperCAmelCase_      			=							BlipaForConditionalGeneration(snake_case__				).eval()

      UpperCAmelCase_      			=							{
          'blip2-opt-2.7b': ('blip2_opt', 'pretrain_opt2.7b'),
          'blip2-opt-6.7b': ('blip2_opt', 'pretrain_opt6.7b'),
          'blip2-opt-2.7b-coco': ('blip2_opt', 'caption_coco_opt2.7b'),
          'blip2-opt-6.7b-coco': ('blip2_opt', 'caption_coco_opt6.7b'),
          'blip2-flan-t5-xl': ('blip2_t5', 'pretrain_flant5xl'),
          'blip2-flan-t5-xl-coco': ('blip2_t5', 'caption_coco_flant5xl'),
          'blip2-flan-t5-xxl': ('blip2_t5', 'pretrain_flant5xxl'),
      }

      UpperCAmelCase_      			=							model_name_to_original[model_name]

      # load original model
      print("Loading original model..."				)
      UpperCAmelCase_      			=							'cuda' if torch.cuda.is_available() else 'cpu'
      UpperCAmelCase_      			=							load_model_and_preprocess(
          name=snake_case__     ,      model_type=snake_case__     ,      is_eval=snake_case__     ,      device=snake_case__				)
      original_model.eval()
      print("Done!"				)

      # update state dict keys
      UpperCAmelCase_      			=							original_model.state_dict()
      UpperCAmelCase_      			=							create_rename_keys(snake_case__				)
      for src, dest in rename_keys:
            rename_key(snake_case__     ,      snake_case__     ,      snake_case__				)

      # some keys can be renamed efficiently
      for key, val in state_dict.copy().items():
            UpperCAmelCase_      			=							state_dict.pop(snake_case__				)
            if key.startswith("Qformer.bert"				):
                  UpperCAmelCase_      			=							key.replace("Qformer.bert"     ,      "qformer"				)
            if "attention.self" in key:
                  UpperCAmelCase_      			=							key.replace("self"     ,      "attention"				)
            if "opt_proj" in key:
                  UpperCAmelCase_      			=							key.replace("opt_proj"     ,      "language_projection"				)
            if "t5_proj" in key:
                  UpperCAmelCase_      			=							key.replace("t5_proj"     ,      "language_projection"				)
            if key.startswith("opt"				):
                  UpperCAmelCase_      			=							key.replace("opt"     ,      "language"				)
            if key.startswith("t5"				):
                  UpperCAmelCase_      			=							key.replace("t5"     ,      "language"				)
            UpperCAmelCase_      			=							val

      # read in qv biases
      read_in_q_v_bias(snake_case__     ,      snake_case__				)

      UpperCAmelCase_      			=							hf_model.load_state_dict(snake_case__     ,      strict=snake_case__				)
      assert len(snake_case__				) == 0
      assert unexpected_keys == ["qformer.embeddings.position_ids"]

      UpperCAmelCase_      			=							load_demo_image()
      UpperCAmelCase_      			=							vis_processors['eval'](snake_case__				).unsqueeze(0				).to(snake_case__				)
      UpperCAmelCase_      			=							tokenizer(["\n"]     ,      return_tensors="pt"				).input_ids.to(snake_case__				)

      # create processor
      UpperCAmelCase_      			=							BlipImageProcessor(
          size={"height": image_size, "width": image_size}     ,      image_mean=snake_case__     ,      image_std=snake_case__				)
      UpperCAmelCase_      			=							BlipaProcessor(image_processor=snake_case__     ,      tokenizer=snake_case__				)
      UpperCAmelCase_      			=							processor(images=snake_case__     ,      return_tensors="pt"				).pixel_values.to(snake_case__				)

      # make sure processor creates exact same pixel values
      assert torch.allclose(snake_case__     ,      snake_case__				)

      original_model.to(snake_case__				)
      hf_model.to(snake_case__				)
      with torch.no_grad():
            if "opt" in model_name:
                  UpperCAmelCase_      			=							original_model({"image": original_pixel_values, "text_input": [""]}				).logits
                  UpperCAmelCase_      			=							hf_model(snake_case__     ,      snake_case__				).logits
            else:
                  UpperCAmelCase_      			=							original_model(
                      {"image": original_pixel_values, "text_input": ["\n"], "text_output": ["\n"]}				).logits
                  UpperCAmelCase_      			=							input_ids.masked_fill(input_ids == tokenizer.pad_token_id     ,      -100				)
                  UpperCAmelCase_      			=							hf_model(snake_case__     ,      snake_case__     ,      labels=snake_case__				).logits

      assert original_logits.shape == logits.shape
      print("First values of original logits:"     ,      original_logits[0, :3, :3]				)
      print("First values of HF logits:"     ,      logits[0, :3, :3]				)

      # assert values
      if model_name == "blip2-flan-t5-xl":
            UpperCAmelCase_      			=							torch.tensor(
                [[-41.5850, -4.4440, -8.9922], [-47.4322, -5.9143, -1.7340]]     ,      device=snake_case__				)
            assert torch.allclose(logits[0, :3, :3]     ,      snake_case__     ,      atol=1e-4				)
      elif model_name == "blip2-flan-t5-xl-coco":
            UpperCAmelCase_      			=							torch.tensor(
                [[-57.0109, -9.8967, -12.6280], [-68.6578, -12.7191, -10.5065]]     ,      device=snake_case__				)
      else:
            # cast to same type
            UpperCAmelCase_      			=							logits.dtype
            assert torch.allclose(original_logits.to(snake_case__				)     ,      snake_case__     ,      atol=1e-2				)
      print("Looks ok!"				)

      print("Generating a caption..."				)
      UpperCAmelCase_      			=							''
      UpperCAmelCase_      			=							tokenizer(snake_case__     ,      return_tensors="pt"				).input_ids.to(snake_case__				)

      UpperCAmelCase_      			=							original_model.generate({"image": original_pixel_values}				)
      UpperCAmelCase_      			=							hf_model.generate(
          snake_case__     ,      snake_case__     ,      do_sample=snake_case__     ,      num_beams=5     ,      max_length=30     ,      min_length=1     ,      top_p=0.9     ,      repetition_penalty=1.0     ,      length_penalty=1.0     ,      temperature=1     ,      )
      print("Original generation:"     ,      snake_case__				)
      UpperCAmelCase_      			=							input_ids.shape[1]
      UpperCAmelCase_      			=							processor.batch_decode(outputs[:, prompt_length:]     ,      skip_special_tokens=snake_case__				)
      UpperCAmelCase_      			=							[text.strip() for text in output_text]
      print("HF generation:"     ,      snake_case__				)

      if pytorch_dump_folder_path is not None:
            processor.save_pretrained(snake_case__				)
            hf_model.save_pretrained(snake_case__				)

      if push_to_hub:
            processor.push_to_hub(f"""nielsr/{model_name}"""				)
            hf_model.push_to_hub(f"""nielsr/{model_name}"""				)


if __name__ == "__main__":
  lowerCamelCase       	=				argparse.ArgumentParser()
  lowerCamelCase       	=				[
      """blip2-opt-2.7b""",
      """blip2-opt-6.7b""",
      """blip2-opt-2.7b-coco""",
      """blip2-opt-6.7b-coco""",
      """blip2-flan-t5-xl""",
      """blip2-flan-t5-xl-coco""",
      """blip2-flan-t5-xxl""",
  ]
  parser.add_argument(
      """--model_name""",
      default="""blip2-opt-2.7b""",
      choices=choices,
      type=str,
      help="""Path to hf config.json of model to convert""",
  )
  parser.add_argument("""--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model.""")
  parser.add_argument(
      """--push_to_hub""",
      action="""store_true""",
      help="""Whether to push the model and processor to the hub after converting""",
  )

  lowerCamelCase       	=				parser.parse_args()

  convert_blipa_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)









"""simple docstring"""





import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import MaskaFormerConfig, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import MaskaFormerForUniversalSegmentation, MaskaFormerModel

  if is_vision_available():
    from transformers import MaskaFormerImageProcessor

if is_vision_available():
  from PIL import Image
class 				lowercase__       :







       '''simple docstring'''


       def __init__(							self      :  Any		,    _UpperCAmelCase      :  str		,    _UpperCAmelCase      :  Union[str, Any]=2		,    _UpperCAmelCase      :  Any=True		,    _UpperCAmelCase      :  int=False		,    _UpperCAmelCase      :  str=10		,    _UpperCAmelCase      :  int=3		,    _UpperCAmelCase      :  Optional[int]=32 * 8		,    _UpperCAmelCase      :  str=32 * 8		,    _UpperCAmelCase      :  Any=4		,    _UpperCAmelCase      :  List[Any]=64		,    )			->						str:


             '''simple docstring'''




             UpperCAmelCase_      			=							parent
             UpperCAmelCase_      			=							batch_size
             UpperCAmelCase_      			=							is_training
             UpperCAmelCase_      			=							use_auxiliary_loss
             UpperCAmelCase_      			=							num_queries
             UpperCAmelCase_      			=							num_channels
             UpperCAmelCase_      			=							min_size
             UpperCAmelCase_      			=							max_size
             UpperCAmelCase_      			=							num_labels
             UpperCAmelCase_      			=							hidden_dim
             UpperCAmelCase_      			=							hidden_dim



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


             '''simple docstring'''




             UpperCAmelCase_      			=							floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]   ).to(
                 _UpperCAmelCase   )

             UpperCAmelCase_      			=							torch.ones([self.batch_size, self.min_size, self.max_size]		,    device=_UpperCAmelCase   )

             UpperCAmelCase_      			=							(
                 torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]		,    device=_UpperCAmelCase   ) > 0.5
             ).float()
             UpperCAmelCase_      			=							(torch.rand((self.batch_size, self.num_labels)		,    device=_UpperCAmelCase   ) > 0.5).long()

             UpperCAmelCase_      			=							self.get_config()
             return config, pixel_values, pixel_mask, mask_labels, class_labels



       def 					lowercase__    (							self      :  Any   )			->						Tuple:


             '''simple docstring'''




             UpperCAmelCase_      			=							MaskaFormerConfig(
                 hidden_size=self.hidden_dim		,    )
             UpperCAmelCase_      			=							self.num_queries
             UpperCAmelCase_      			=							self.num_labels

             UpperCAmelCase_      			=							[1, 1, 1, 1]
             UpperCAmelCase_      			=							self.num_channels

             UpperCAmelCase_      			=							64
             UpperCAmelCase_      			=							128
             UpperCAmelCase_      			=							self.hidden_dim
             UpperCAmelCase_      			=							self.hidden_dim
             UpperCAmelCase_      			=							self.hidden_dim
             return config



       def 					lowercase__    (							self      :  Dict   )			->						Dict:


             '''simple docstring'''




             UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_      			=							self.prepare_config_and_inputs()
             UpperCAmelCase_      			=							{"pixel_values": pixel_values, "pixel_mask": pixel_mask}
             return config, inputs_dict



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


             '''simple docstring'''




             UpperCAmelCase_      			=							output.encoder_hidden_states
             UpperCAmelCase_      			=							output.pixel_decoder_hidden_states
             UpperCAmelCase_      			=							output.transformer_decoder_hidden_states

             self.parent.assertTrue(len(_UpperCAmelCase   )		,    len(config.backbone_config.depths   )   )
             self.parent.assertTrue(len(_UpperCAmelCase   )		,    len(config.backbone_config.depths   )   )
             self.parent.assertTrue(len(_UpperCAmelCase   )		,    config.decoder_layers   )



       def 					lowercase__    (							self      :  Tuple		,    _UpperCAmelCase      :  List[Any]		,    _UpperCAmelCase      :  Any		,    _UpperCAmelCase      :  Union[str, Any]		,    _UpperCAmelCase      :  Optional[int]=False   )			->						str:


             '''simple docstring'''




             with torch.no_grad():
                   UpperCAmelCase_      			=							MaskaFormerModel(config=_UpperCAmelCase   )
                   model.to(_UpperCAmelCase   )
                   model.eval()

                   UpperCAmelCase_      			=							model(pixel_values=_UpperCAmelCase		,    pixel_mask=_UpperCAmelCase   )
                   UpperCAmelCase_      			=							model(_UpperCAmelCase		,    output_hidden_states=_UpperCAmelCase   )

             self.parent.assertEqual(
                 output.transformer_decoder_last_hidden_state.shape		,    (self.batch_size, self.num_queries, self.hidden_dim)		,    )
             # let's ensure the other two hidden state exists
             self.parent.assertTrue(output.pixel_decoder_last_hidden_state is not None   )
             self.parent.assertTrue(output.encoder_last_hidden_state is not None   )

             if output_hidden_states:
                   self.check_output_hidden_state(_UpperCAmelCase		,    _UpperCAmelCase   )







       def 					lowercase__    (							self      :  Optional[Any]		,    _UpperCAmelCase      :  Optional[int]		,    _UpperCAmelCase      :  Dict		,    _UpperCAmelCase      :  Optional[int]		,    _UpperCAmelCase      :  Union[str, Any]		,    _UpperCAmelCase      :  int   )			->						Optional[Any]:


             '''simple docstring'''




             UpperCAmelCase_      			=							MaskaFormerForUniversalSegmentation(config=_UpperCAmelCase   )
             model.to(_UpperCAmelCase   )
             model.eval()

             def comm_check_on_output(_UpperCAmelCase      :  List[Any]   ):
                   # let's still check that all the required stuff is there
                   self.parent.assertTrue(result.transformer_decoder_last_hidden_state is not None   )
                   self.parent.assertTrue(result.pixel_decoder_last_hidden_state is not None   )
                   self.parent.assertTrue(result.encoder_last_hidden_state is not None   )
                   # okay, now we need to check the logits shape
                   # due to the encoder compression, masks have a //4 spatial size
                   self.parent.assertEqual(
                       result.masks_queries_logits.shape		,    (self.batch_size, self.num_queries, self.min_size // 4, self.max_size // 4)		,    )
                   # + 1 for null class
                   self.parent.assertEqual(
                       result.class_queries_logits.shape		,    (self.batch_size, self.num_queries, self.num_labels + 1)   )

             with torch.no_grad():
                   UpperCAmelCase_      			=							model(pixel_values=_UpperCAmelCase		,    pixel_mask=_UpperCAmelCase   )
                   UpperCAmelCase_      			=							model(_UpperCAmelCase   )

                   comm_check_on_output(_UpperCAmelCase   )

                   UpperCAmelCase_      			=							model(
                       pixel_values=_UpperCAmelCase		,    pixel_mask=_UpperCAmelCase		,    mask_labels=_UpperCAmelCase		,    class_labels=_UpperCAmelCase   )

             comm_check_on_output(_UpperCAmelCase   )

             self.parent.assertTrue(result.loss is not None   )
             self.parent.assertEqual(result.loss.shape		,    torch.Size([1]   )   )



@require_torch
class 				lowercase__       (   SCREAMING_SNAKE_CASE    , SCREAMING_SNAKE_CASE    , unittest.TestCase				):







       '''simple docstring'''


       UpperCamelCase						=					(MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
       UpperCamelCase						=					{'''feature-extraction''': MaskaFormerModel} if is_torch_available() else {}

       UpperCamelCase						=					False
       UpperCamelCase						=					False
       UpperCamelCase						=					False
       UpperCamelCase						=					False



       def 					lowercase__    (							self      :  Tuple   )			->						Tuple:


             '''simple docstring'''




             UpperCAmelCase_      			=							MaskaFormerModelTester(self   )
             UpperCAmelCase_      			=							ConfigTester(self		,    config_class=_UpperCAmelCase		,    has_text_modality=_UpperCAmelCase   )



       def 					lowercase__    (							self      :  List[str]   )			->						Dict:


             '''simple docstring'''




             self.config_tester.run_common_tests()



       def 					lowercase__    (							self      :  Dict   )			->						str:


             '''simple docstring'''




             UpperCAmelCase_			,   UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs_for_common()
             self.model_tester.create_and_check_maskaformer_model(_UpperCAmelCase		,    **_UpperCAmelCase		,    output_hidden_states=_UpperCAmelCase   )



       def 					lowercase__    (							self      :  Optional[int]   )			->						Dict:


             '''simple docstring'''




             UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*_UpperCAmelCase   )



       @unittest.skip(reason="Mask2Former does not use inputs_embeds"   )
       def 					lowercase__    (							self      :  Tuple   )			->						Any:


             '''simple docstring'''




             pass



       @unittest.skip(reason="Mask2Former does not have a get_input_embeddings method"   )
       def 					lowercase__    (							self      :  Tuple   )			->						Optional[int]:


             '''simple docstring'''




             pass



       @unittest.skip(reason="Mask2Former is not a generative model"   )
       def 					lowercase__    (							self      :  str   )			->						Dict:


             '''simple docstring'''




             pass



       @unittest.skip(reason="Mask2Former does not use token embeddings"   )
       def 					lowercase__    (							self      :  Union[str, Any]   )			->						Union[str, Any]:


             '''simple docstring'''




             pass



       @require_torch_multi_gpu
       @unittest.skip(
           reason="Mask2Former has some layers using `add_module` which doesn't work well with `nn.DataParallel`"   )
       def 					lowercase__    (							self      :  List[Any]   )			->						Dict:


             '''simple docstring'''




             pass



       @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."   )
       def 					lowercase__    (							self      :  List[str]   )			->						List[str]:


             '''simple docstring'''




             pass



       def 					lowercase__    (							self      :  Tuple   )			->						str:


             '''simple docstring'''




             UpperCAmelCase_			,   UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                   UpperCAmelCase_      			=							model_class(_UpperCAmelCase   )
                   UpperCAmelCase_      			=							inspect.signature(model.forward   )
                   # signature.parameters is an OrderedDict => so arg_names order is deterministic
                   UpperCAmelCase_      			=							[*signature.parameters.keys()]

                   UpperCAmelCase_      			=							["pixel_values"]
                   self.assertListEqual(arg_names[:1]		,    _UpperCAmelCase   )



       @slow
       def 					lowercase__    (							self      :  List[Any]   )			->						Dict:


             '''simple docstring'''




             for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
                   UpperCAmelCase_      			=							MaskaFormerModel.from_pretrained(_UpperCAmelCase   )
                   self.assertIsNotNone(_UpperCAmelCase   )



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


             '''simple docstring'''




             UpperCAmelCase_      			=							(self.model_tester.min_size,) * 2
             UpperCAmelCase_      			=							{
                 "pixel_values": torch.randn((2, 3, *size)		,    device=_UpperCAmelCase   ),
                 "mask_labels": torch.randn((2, 10, *size)		,    device=_UpperCAmelCase   ),
                 "class_labels": torch.zeros(2		,    10		,    device=_UpperCAmelCase   ).long(),
             }
             UpperCAmelCase_      			=							self.model_tester.get_config()

             UpperCAmelCase_      			=							MaskaFormerForUniversalSegmentation(_UpperCAmelCase   ).to(_UpperCAmelCase   )
             UpperCAmelCase_      			=							model(**_UpperCAmelCase   )
             self.assertTrue(outputs.loss is not None   )



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


             '''simple docstring'''




             UpperCAmelCase_			,   UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs_for_common()
             self.model_tester.create_and_check_maskaformer_model(_UpperCAmelCase		,    **_UpperCAmelCase		,    output_hidden_states=_UpperCAmelCase   )



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


             '''simple docstring'''




             UpperCAmelCase_			,   UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                   UpperCAmelCase_      			=							model_class(_UpperCAmelCase   ).to(_UpperCAmelCase   )
                   UpperCAmelCase_      			=							model(**_UpperCAmelCase		,    output_attentions=_UpperCAmelCase   )
                   self.assertTrue(outputs.attentions is not None   )



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


             '''simple docstring'''




             if not self.model_tester.is_training:
                   return

             UpperCAmelCase_      			=							self.all_model_classes[1]
             UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs()

             UpperCAmelCase_      			=							model_class(_UpperCAmelCase   )
             model.to(_UpperCAmelCase   )
             model.train()

             UpperCAmelCase_      			=							model(_UpperCAmelCase		,    mask_labels=_UpperCAmelCase		,    class_labels=_UpperCAmelCase   ).loss
             loss.backward()







       def 					lowercase__    (							self      :  str   )			->						Optional[int]:


             '''simple docstring'''




             UpperCAmelCase_      			=							self.all_model_classes[1]
             UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs()
             UpperCAmelCase_      			=							True
             UpperCAmelCase_      			=							True

             UpperCAmelCase_      			=							model_class(_UpperCAmelCase   ).to(_UpperCAmelCase   )
             model.train()

             UpperCAmelCase_      			=							model(_UpperCAmelCase		,    mask_labels=_UpperCAmelCase		,    class_labels=_UpperCAmelCase   )

             UpperCAmelCase_      			=							outputs.encoder_hidden_states[0]
             encoder_hidden_states.retain_grad()

             UpperCAmelCase_      			=							outputs.pixel_decoder_hidden_states[0]
             pixel_decoder_hidden_states.retain_grad()

             UpperCAmelCase_      			=							outputs.transformer_decoder_hidden_states[0]
             transformer_decoder_hidden_states.retain_grad()

             UpperCAmelCase_      			=							outputs.attentions[0]
             attentions.retain_grad()

             outputs.loss.backward(retain_graph=_UpperCAmelCase   )

             self.assertIsNotNone(encoder_hidden_states.grad   )
             self.assertIsNotNone(pixel_decoder_hidden_states.grad   )
             self.assertIsNotNone(transformer_decoder_hidden_states.grad   )
             self.assertIsNotNone(attentions.grad   )



lowerCamelCase       	=						1e-4
def        a__     (				):
      UpperCAmelCase_      			=							Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"				)
      return image


@require_vision
@slow
class 				lowercase__       (   unittest.TestCase				):







       '''simple docstring'''


       @cached_property
       def 					lowercase__    (							self      :  Optional[Any]   )			->						int:


             '''simple docstring'''




             return "facebook/mask2former-swin-small-coco-instance"



       @cached_property
       def 					lowercase__    (							self      :  Dict   )			->						Optional[int]:


             '''simple docstring'''




             return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints   ) if is_vision_available() else None



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


             '''simple docstring'''




             UpperCAmelCase_      			=							MaskaFormerModel.from_pretrained(self.model_checkpoints   ).to(_UpperCAmelCase   )
             UpperCAmelCase_      			=							self.default_image_processor
             UpperCAmelCase_      			=							prepare_img()
             UpperCAmelCase_      			=							image_processor(_UpperCAmelCase		,    return_tensors="pt"   ).to(_UpperCAmelCase   )
             UpperCAmelCase_      			=							inputs["pixel_values"].shape
             # check size is divisible by 32
             self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0   )
             # check size
             self.assertEqual(_UpperCAmelCase		,    (1, 3, 384, 384)   )

             with torch.no_grad():
                   UpperCAmelCase_      			=							model(**_UpperCAmelCase   )

             UpperCAmelCase_      			=							torch.tensor(
                 [[-0.2790, -1.0717, -1.1668], [-0.5128, -0.3128, -0.4987], [-0.5832, 0.1971, -0.0197]]   ).to(_UpperCAmelCase   )
             self.assertTrue(
                 torch.allclose(
                     outputs.encoder_last_hidden_state[0, 0, :3, :3]		,    _UpperCAmelCase		,    atol=_UpperCAmelCase   )   )

             UpperCAmelCase_      			=							torch.tensor(
                 [[0.8973, 1.1847, 1.1776], [1.1934, 1.5040, 1.5128], [1.1153, 1.4486, 1.4951]]   ).to(_UpperCAmelCase   )
             self.assertTrue(
                 torch.allclose(
                     outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]		,    _UpperCAmelCase		,    atol=_UpperCAmelCase   )   )

             UpperCAmelCase_      			=							torch.tensor(
                 [[2.1152, 1.7000, -0.8603], [1.5808, 1.8004, -0.9353], [1.6043, 1.7495, -0.5999]]   ).to(_UpperCAmelCase   )
             self.assertTrue(
                 torch.allclose(
                     outputs.transformer_decoder_last_hidden_state[0, :3, :3]		,    _UpperCAmelCase		,    atol=_UpperCAmelCase   )   )



       def 					lowercase__    (							self      :  int   )			->						List[str]:


             '''simple docstring'''




             UpperCAmelCase_      			=							MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints   ).to(_UpperCAmelCase   ).eval()
             UpperCAmelCase_      			=							self.default_image_processor
             UpperCAmelCase_      			=							prepare_img()
             UpperCAmelCase_      			=							image_processor(_UpperCAmelCase		,    return_tensors="pt"   ).to(_UpperCAmelCase   )
             UpperCAmelCase_      			=							inputs["pixel_values"].shape
             # check size is divisible by 32
             self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0   )
             # check size
             self.assertEqual(_UpperCAmelCase		,    (1, 3, 384, 384)   )

             with torch.no_grad():
                   UpperCAmelCase_      			=							model(**_UpperCAmelCase   )
             # masks_queries_logits
             UpperCAmelCase_      			=							outputs.masks_queries_logits
             self.assertEqual(
                 masks_queries_logits.shape		,    (1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)   )
             UpperCAmelCase_      			=							[
                 [-8.7839, -9.0056, -8.8121],
                 [-7.4104, -7.0313, -6.5401],
                 [-6.6105, -6.3427, -6.4675],
             ]
             UpperCAmelCase_      			=							torch.tensor(_UpperCAmelCase   ).to(_UpperCAmelCase   )
             self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]		,    _UpperCAmelCase		,    atol=_UpperCAmelCase   )   )
             # class_queries_logits
             UpperCAmelCase_      			=							outputs.class_queries_logits
             self.assertEqual(class_queries_logits.shape		,    (1, model.config.num_queries, model.config.num_labels + 1)   )
             UpperCAmelCase_      			=							torch.tensor(
                 [
                     [1.8324, -8.0835, -4.1922],
                     [0.8450, -9.0050, -3.6053],
                     [0.3045, -7.7293, -3.0275],
                 ]   ).to(_UpperCAmelCase   )
             self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]		,    _UpperCAmelCase		,    atol=_UpperCAmelCase   )   )







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


             '''simple docstring'''




             UpperCAmelCase_      			=							MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints   ).to(_UpperCAmelCase   ).eval()
             UpperCAmelCase_      			=							self.default_image_processor

             UpperCAmelCase_      			=							image_processor(
                 [np.zeros((3, 800, 1333)   ), np.zeros((3, 800, 1333)   )]		,    segmentation_maps=[np.zeros((384, 384)   ).astype(np.floataa   ), np.zeros((384, 384)   ).astype(np.floataa   )]		,    return_tensors="pt"		,    )

             UpperCAmelCase_      			=							inputs["pixel_values"].to(_UpperCAmelCase   )
             UpperCAmelCase_      			=							[el.to(_UpperCAmelCase   ) for el in inputs["mask_labels"]]
             UpperCAmelCase_      			=							[el.to(_UpperCAmelCase   ) for el in inputs["class_labels"]]

             with torch.no_grad():
                   UpperCAmelCase_      			=							model(**_UpperCAmelCase   )

             self.assertTrue(outputs.loss is not None   )




