

'''simple docstring'''




import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers import is_speech_available, is_vision_available
from transformers.testing_utils import require_torch


if is_vision_available():
  from transformers import TvltImageProcessor

if is_speech_available():
  from transformers import TvltFeatureExtractor

from transformers import TvltProcessor



@require_torch
class 							lowercase_						(		unittest.TestCase			):

  '''simple docstring'''





  def     snake_case_      (  self  )				->					Tuple:


    """simple docstring"""
    UpperCAmelCase    				=						'ZinengTang/tvlt-base'
    UpperCAmelCase    				=						tempfile.mkdtemp()
  def     snake_case_      (  self ,			**a_  )				->					List[Any]:


    """simple docstring"""
    return TvltImageProcessor.from_pretrained(self.checkpoint ,			**a_  )
  def     snake_case_      (  self ,			**a_  )				->					List[Any]:


    """simple docstring"""
    return TvltFeatureExtractor.from_pretrained(self.checkpoint ,			**a_  )
  def     snake_case_      (  self  )				->					List[str]:


    """simple docstring"""
    shutil.rmtree(self.tmpdirname  )
  def     snake_case_      (  self  )				->					Any:


    """simple docstring"""
    UpperCAmelCase    				=						self.get_image_processor()
    UpperCAmelCase    				=						self.get_feature_extractor()

    UpperCAmelCase    				=						TvltProcessor(image_processor=a_ ,			feature_extractor=a_  )
    processor.save_pretrained(self.tmpdirname  )
    UpperCAmelCase    				=						TvltProcessor.from_pretrained(self.tmpdirname  )

    self.assertIsInstance(processor.feature_extractor ,			a_  )
    self.assertIsInstance(processor.image_processor ,			a_  )
  def     snake_case_      (  self  )				->					Any:


    """simple docstring"""
    UpperCAmelCase    				=						self.get_image_processor()
    UpperCAmelCase    				=						self.get_feature_extractor()

    UpperCAmelCase    				=						TvltProcessor(image_processor=a_ ,			feature_extractor=a_  )

    UpperCAmelCase    				=						np.ones([1_2_0_0_0]  )

    UpperCAmelCase    				=						feature_extractor(a_ ,			return_tensors='np'  )
    UpperCAmelCase    				=						processor(audio=a_ ,			return_tensors='np'  )

    for key in audio_dict.keys():
      self.assertAlmostEqual(audio_dict[key].sum() ,			input_processor[key].sum() ,			delta=1E-2  )
  def     snake_case_      (  self  )				->					Union[str, Any]:


    """simple docstring"""
    UpperCAmelCase    				=						self.get_image_processor()
    UpperCAmelCase    				=						self.get_feature_extractor()

    UpperCAmelCase    				=						TvltProcessor(image_processor=a_ ,			feature_extractor=a_  )

    UpperCAmelCase    				=						np.ones([3, 2_2_4, 2_2_4]  )

    UpperCAmelCase    				=						image_processor(a_ ,			return_tensors='np'  )
    UpperCAmelCase    				=						processor(images=a_ ,			return_tensors='np'  )

    for key in image_dict.keys():
      self.assertAlmostEqual(image_dict[key].sum() ,			input_processor[key].sum() ,			delta=1E-2  )
  def     snake_case_      (  self  )				->					Dict:


    """simple docstring"""
    UpperCAmelCase    				=						self.get_image_processor()
    UpperCAmelCase    				=						self.get_feature_extractor()

    UpperCAmelCase    				=						TvltProcessor(image_processor=a_ ,			feature_extractor=a_  )

    UpperCAmelCase    				=						np.ones([1_2_0_0_0]  )
    UpperCAmelCase    				=						np.ones([3, 2_2_4, 2_2_4]  )

    UpperCAmelCase    				=						processor(audio=a_ ,			images=a_  )

    self.assertListEqual(list(inputs.keys()  ) ,			['audio_values', 'audio_mask', 'pixel_values', 'pixel_mask']  )

    # test if it raises when no input is passed
    with pytest.raises(a_  ):
      processor()





  def     snake_case_      (  self  )				->					int:


    """simple docstring"""
    UpperCAmelCase    				=						self.get_image_processor()
    UpperCAmelCase    				=						self.get_feature_extractor()

    UpperCAmelCase    				=						TvltProcessor(image_processor=a_ ,			feature_extractor=a_  )

    self.assertListEqual(
        processor.model_input_names ,			image_processor.model_input_names + feature_extractor.model_input_names ,			msg='`processor` and `image_processor`+`feature_extractor` model input names do not match' ,			)





'''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)


