

'''simple docstring'''




import unittest

import numpy as np
import requests

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
  import torch

  from transformers.pytorch_utils import is_torch_greater_or_equal_than_1_11
else:
  _a					:			Any							=    False

if is_vision_available():
  from PIL import Image

  from transformers import PixaStructImageProcessor



class 							lowercase_						(		unittest.TestCase			):

  '''simple docstring'''





  def __init__(  self ,			a_ ,			a_=7 ,			a_=3 ,			a_=1_8 ,			a_=3_0 ,			a_=4_0_0 ,			a_=None ,			a_=True ,			a_=True ,			a_=None ,			)				->					List[Any]:


    """simple docstring"""
    UpperCAmelCase    				=						size if size is not None else {'height': 2_0, 'width': 2_0}
    UpperCAmelCase    				=						parent
    UpperCAmelCase    				=						batch_size
    UpperCAmelCase    				=						num_channels
    UpperCAmelCase    				=						image_size
    UpperCAmelCase    				=						min_resolution
    UpperCAmelCase    				=						max_resolution
    UpperCAmelCase    				=						size
    UpperCAmelCase    				=						do_normalize
    UpperCAmelCase    				=						do_convert_rgb
    UpperCAmelCase    				=						[5_1_2, 1_0_2_4, 2_0_4_8, 4_0_9_6]
    UpperCAmelCase    				=						patch_size if patch_size is not None else {'height': 1_6, 'width': 1_6}
  def     snake_case_      (  self  )				->					Dict:


    """simple docstring"""
    return {"do_normalize": self.do_normalize, "do_convert_rgb": self.do_convert_rgb}





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


    """simple docstring"""
    UpperCAmelCase    				=						'https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/australia.jpg'
    UpperCAmelCase    				=						Image.open(requests.get(a_ ,			stream=a_  ).raw  ).convert('RGB'  )
    return raw_image




@unittest.skipIf(
    not is_torch_greater_or_equal_than_1_11			,							reason="`Pix2StructImageProcessor` requires `torch>=1.11.0`."			,							)
@require_torch
@require_vision
class 							lowercase_						(		a			,							unittest.TestCase			):

  '''simple docstring'''





  __lowerCAmelCase       :			List[Any]     	= PixaStructImageProcessor if is_vision_available() else None
  def     snake_case_      (  self  )				->					Dict:


    """simple docstring"""
    UpperCAmelCase    				=						PixaStructImageProcessingTester(self  )
  @property
  def     snake_case_      (  self  )				->					Tuple:


    """simple docstring"""
    return self.image_processor_tester.prepare_image_processor_dict()
  def     snake_case_      (  self  )				->					Optional[Any]:


    """simple docstring"""
    UpperCAmelCase    				=						self.image_processing_class(**self.image_processor_dict  )
    self.assertTrue(hasattr(a_ ,			'do_normalize'  )  )
    self.assertTrue(hasattr(a_ ,			'do_convert_rgb'  )  )
  def     snake_case_      (  self  )				->					Union[str, Any]:


    """simple docstring"""
    UpperCAmelCase    				=						self.image_processor_tester.prepare_dummy_image()

    UpperCAmelCase    				=						self.image_processing_class(**self.image_processor_dict  )
    UpperCAmelCase    				=						2_0_4_8

    UpperCAmelCase    				=						image_processor(a_ ,			return_tensors='pt' ,			max_patches=a_  )
    self.assertTrue(torch.allclose(inputs.flattened_patches.mean() ,			torch.tensor(0.0606  ) ,			atol=1E-3 ,			rtol=1E-3  )  )
  def     snake_case_      (  self  )				->					str:


    """simple docstring"""
    # Initialize image_processor
    UpperCAmelCase    				=						self.image_processing_class(**self.image_processor_dict  )
    # create random PIL images
    UpperCAmelCase    				=						prepare_image_inputs(self.image_processor_tester ,			equal_resolution=a_  )
    for image in image_inputs:
      self.assertIsInstance(a_ ,			Image.Image  )

    # Test not batched input
    UpperCAmelCase    				=						(
        (self.image_processor_tester.patch_size['height'] * self.image_processor_tester.patch_size['width'])
        * self.image_processor_tester.num_channels
    ) + 2

    for max_patch in self.image_processor_tester.max_patches:
      # Test not batched input
      UpperCAmelCase    				=						image_processor(
          image_inputs[0] ,			return_tensors='pt' ,			max_patches=a_  ).flattened_patches
      self.assertEqual(
          encoded_images.shape ,			(1, max_patch, expected_hidden_dim) ,			)

      # Test batched
      UpperCAmelCase    				=						image_processor(
          a_ ,			return_tensors='pt' ,			max_patches=a_  ).flattened_patches
      self.assertEqual(
          encoded_images.shape ,			(self.image_processor_tester.batch_size, max_patch, expected_hidden_dim) ,			)
  def     snake_case_      (  self  )				->					Union[str, Any]:


    """simple docstring"""
    # Initialize image_processor
    UpperCAmelCase    				=						self.image_processing_class(**self.image_processor_dict  )
    # create random PIL images
    UpperCAmelCase    				=						prepare_image_inputs(self.image_processor_tester ,			equal_resolution=a_  )
    for image in image_inputs:
      self.assertIsInstance(a_ ,			Image.Image  )

    # Test not batched input
    UpperCAmelCase    				=						(
        (self.image_processor_tester.patch_size['height'] * self.image_processor_tester.patch_size['width'])
        * self.image_processor_tester.num_channels
    ) + 2

    UpperCAmelCase    				=						True

    for max_patch in self.image_processor_tester.max_patches:
      # Test not batched input
      with self.assertRaises(a_  ):
        UpperCAmelCase    				=						image_processor(
            image_inputs[0] ,			return_tensors='pt' ,			max_patches=a_  ).flattened_patches

      UpperCAmelCase    				=						'Hello'

      UpperCAmelCase    				=						image_processor(
          image_inputs[0] ,			return_tensors='pt' ,			max_patches=a_ ,			header_text=a_  ).flattened_patches
      self.assertEqual(
          encoded_images.shape ,			(1, max_patch, expected_hidden_dim) ,			)

      # Test batched
      UpperCAmelCase    				=						image_processor(
          a_ ,			return_tensors='pt' ,			max_patches=a_ ,			header_text=a_  ).flattened_patches
      self.assertEqual(
          encoded_images.shape ,			(self.image_processor_tester.batch_size, max_patch, expected_hidden_dim) ,			)
  def     snake_case_      (  self  )				->					Dict:


    """simple docstring"""
    # Initialize image_processor
    UpperCAmelCase    				=						self.image_processing_class(**self.image_processor_dict  )
    # create random numpy tensors
    UpperCAmelCase    				=						prepare_image_inputs(self.image_processor_tester ,			equal_resolution=a_ ,			numpify=a_  )
    for image in image_inputs:
      self.assertIsInstance(a_ ,			np.ndarray  )

    UpperCAmelCase    				=						(
        (self.image_processor_tester.patch_size['height'] * self.image_processor_tester.patch_size['width'])
        * self.image_processor_tester.num_channels
    ) + 2

    for max_patch in self.image_processor_tester.max_patches:
      # Test not batched input
      UpperCAmelCase    				=						image_processor(
          image_inputs[0] ,			return_tensors='pt' ,			max_patches=a_  ).flattened_patches
      self.assertEqual(
          encoded_images.shape ,			(1, max_patch, expected_hidden_dim) ,			)

      # Test batched
      UpperCAmelCase    				=						image_processor(
          a_ ,			return_tensors='pt' ,			max_patches=a_  ).flattened_patches
      self.assertEqual(
          encoded_images.shape ,			(self.image_processor_tester.batch_size, max_patch, expected_hidden_dim) ,			)





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


    """simple docstring"""
    # Initialize image_processor
    UpperCAmelCase    				=						self.image_processing_class(**self.image_processor_dict  )
    # create random PyTorch tensors
    UpperCAmelCase    				=						prepare_image_inputs(self.image_processor_tester ,			equal_resolution=a_ ,			torchify=a_  )
    for image in image_inputs:
      self.assertIsInstance(a_ ,			torch.Tensor  )

    # Test not batched input
    UpperCAmelCase    				=						(
        (self.image_processor_tester.patch_size['height'] * self.image_processor_tester.patch_size['width'])
        * self.image_processor_tester.num_channels
    ) + 2

    for max_patch in self.image_processor_tester.max_patches:
      # Test not batched input
      UpperCAmelCase    				=						image_processor(
          image_inputs[0] ,			return_tensors='pt' ,			max_patches=a_  ).flattened_patches
      self.assertEqual(
          encoded_images.shape ,			(1, max_patch, expected_hidden_dim) ,			)

      # Test batched
      UpperCAmelCase    				=						image_processor(
          a_ ,			return_tensors='pt' ,			max_patches=a_  ).flattened_patches
      self.assertEqual(
          encoded_images.shape ,			(self.image_processor_tester.batch_size, max_patch, expected_hidden_dim) ,			)




@unittest.skipIf(
    not is_torch_greater_or_equal_than_1_11			,							reason="`Pix2StructImageProcessor` requires `torch>=1.11.0`."			,							)
@require_torch
@require_vision
class 							lowercase_						(		a			,							unittest.TestCase			):

  '''simple docstring'''





  __lowerCAmelCase       :			Optional[Any]     	= PixaStructImageProcessor if is_vision_available() else None
  def     snake_case_      (  self  )				->					Any:


    """simple docstring"""
    UpperCAmelCase    				=						PixaStructImageProcessingTester(self ,			num_channels=4  )
    UpperCAmelCase    				=						3
  @property
  def     snake_case_      (  self  )				->					str:


    """simple docstring"""
    return self.image_processor_tester.prepare_image_processor_dict()
  def     snake_case_      (  self  )				->					Optional[Any]:


    """simple docstring"""
    UpperCAmelCase    				=						self.image_processing_class(**self.image_processor_dict  )
    self.assertTrue(hasattr(a_ ,			'do_normalize'  )  )
    self.assertTrue(hasattr(a_ ,			'do_convert_rgb'  )  )





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


    """simple docstring"""
    # Initialize image_processor
    UpperCAmelCase    				=						self.image_processing_class(**self.image_processor_dict  )
    # create random PIL images
    UpperCAmelCase    				=						prepare_image_inputs(self.image_processor_tester ,			equal_resolution=a_  )
    for image in image_inputs:
      self.assertIsInstance(a_ ,			Image.Image  )

    # Test not batched input
    UpperCAmelCase    				=						(
        (self.image_processor_tester.patch_size['height'] * self.image_processor_tester.patch_size['width'])
        * (self.image_processor_tester.num_channels - 1)
    ) + 2

    for max_patch in self.image_processor_tester.max_patches:
      # Test not batched input
      UpperCAmelCase    				=						image_processor(
          image_inputs[0] ,			return_tensors='pt' ,			max_patches=a_  ).flattened_patches
      self.assertEqual(
          encoded_images.shape ,			(1, max_patch, expected_hidden_dim) ,			)

      # Test batched
      UpperCAmelCase    				=						image_processor(
          a_ ,			return_tensors='pt' ,			max_patches=a_  ).flattened_patches
      self.assertEqual(
          encoded_images.shape ,			(self.image_processor_tester.batch_size, max_patch, expected_hidden_dim) ,			)





'''simple docstring'''




import argparse
import os

import torch

from transformers import FlavaConfig, FlavaForPreTraining
from transformers.models.flava.convert_dalle_to_flava_codebook import convert_dalle_checkpoint


def    lowerCamelCase__		(    SCREAMING_SNAKE_CASE   : str    ):
  # encoder.embeddings are double copied in original FLAVA
  return sum(param.float().sum() if 'encoder.embeddings' not in key else 0 for key, param in state_dict.items()    )


def    lowerCamelCase__		(    SCREAMING_SNAKE_CASE   : int							, SCREAMING_SNAKE_CASE   : Optional[int]    ):
  UpperCAmelCase    				=						{}

  for key, value in state_dict.items():
    if "text_encoder.embeddings" in key or "image_encoder.embeddings" in key:
      continue

    UpperCAmelCase    				=						key.replace('heads.cmd.mim_head.cls.predictions'							, 'mmm_image_head'    )
    UpperCAmelCase    				=						key.replace('heads.cmd.mlm_head.cls.predictions'							, 'mmm_text_head'    )
    UpperCAmelCase    				=						key.replace('heads.cmd.itm_head.cls'							, 'itm_head'    )
    UpperCAmelCase    				=						key.replace('heads.cmd.itm_head.pooler'							, 'itm_head.pooler'    )
    UpperCAmelCase    				=						key.replace('heads.cmd.clip_head.logit_scale'							, 'flava.logit_scale'    )
    UpperCAmelCase    				=						key.replace('heads.fairseq_mlm.cls.predictions'							, 'mlm_head'    )
    UpperCAmelCase    				=						key.replace('heads.imagenet.mim_head.cls.predictions'							, 'mim_head'    )
    UpperCAmelCase    				=						key.replace('mm_text_projection'							, 'flava.text_to_mm_projection'    )
    UpperCAmelCase    				=						key.replace('mm_image_projection'							, 'flava.image_to_mm_projection'    )
    UpperCAmelCase    				=						key.replace('image_encoder.module'							, 'flava.image_model'    )
    UpperCAmelCase    				=						key.replace('text_encoder.module'							, 'flava.text_model'    )
    UpperCAmelCase    				=						key.replace('mm_encoder.module.encoder.cls_token'							, 'flava.multimodal_model.cls_token'    )
    UpperCAmelCase    				=						key.replace('mm_encoder.module'							, 'flava.multimodal_model'    )
    UpperCAmelCase    				=						key.replace('text_projection'							, 'flava.text_projection'    )
    UpperCAmelCase    				=						key.replace('image_projection'							, 'flava.image_projection'    )

    UpperCAmelCase    				=						value.float()

  for key, value in codebook_state_dict.items():
    UpperCAmelCase    				=						value

  return upgrade


@torch.no_grad()
def    lowerCamelCase__		(    SCREAMING_SNAKE_CASE   : List[str]							, SCREAMING_SNAKE_CASE   : Optional[Any]							, SCREAMING_SNAKE_CASE   : Optional[int]							, SCREAMING_SNAKE_CASE   : Any=None    ):
  if config_path is not None:
    UpperCAmelCase    				=						FlavaConfig.from_pretrained(SCREAMING_SNAKE_CASE    )
  else:
    UpperCAmelCase    				=						FlavaConfig()

  UpperCAmelCase    				=						FlavaForPreTraining(SCREAMING_SNAKE_CASE    ).eval()

  UpperCAmelCase    				=						convert_dalle_checkpoint(SCREAMING_SNAKE_CASE							, SCREAMING_SNAKE_CASE							, save_checkpoint=SCREAMING_SNAKE_CASE    )

  if os.path.exists(SCREAMING_SNAKE_CASE    ):
    UpperCAmelCase    				=						torch.load(SCREAMING_SNAKE_CASE							, map_location='cpu'    )
  else:
    UpperCAmelCase    				=						torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE							, map_location='cpu'    )

  UpperCAmelCase    				=						upgrade_state_dict(SCREAMING_SNAKE_CASE							, SCREAMING_SNAKE_CASE    )
  hf_model.load_state_dict(SCREAMING_SNAKE_CASE    )
  UpperCAmelCase    				=						hf_model.state_dict()
  UpperCAmelCase    				=						count_parameters(SCREAMING_SNAKE_CASE    )
  UpperCAmelCase    				=						count_parameters(SCREAMING_SNAKE_CASE    ) + count_parameters(SCREAMING_SNAKE_CASE    )

  assert torch.allclose(SCREAMING_SNAKE_CASE							, SCREAMING_SNAKE_CASE							, atol=1E-3    )

  hf_model.save_pretrained(SCREAMING_SNAKE_CASE    )


if __name__ == "__main__":
  _a					:			str							=    argparse.ArgumentParser()
  parser.add_argument('--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
  parser.add_argument('--checkpoint_path', default=None, type=str, help='Path to flava checkpoint')
  parser.add_argument('--codebook_path', default=None, type=str, help='Path to flava codebook checkpoint')
  parser.add_argument('--config_path', default=None, type=str, help='Path to hf config.json of model to convert')
  _a					:			str							=    parser.parse_args()

  convert_flava_checkpoint(args.checkpoint_path, args.codebook_path, args.pytorch_dump_folder_path, args.config_path)


