"""simple docstring"""







import collections
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import (
    is_pt_flax_cross_test,
    require_flax,
    require_torch,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import is_flax_available, is_torch_available, is_vision_available

from ...test_modeling_flax_common import floats_tensor, ids_tensor, random_attention_mask
from ..bert.test_modeling_flax_bert import FlaxBertModelTester
from ..clip.test_modeling_flax_clip import FlaxCLIPVisionModelTester
from ..vit.test_modeling_flax_vit import FlaxViTModelTester


if is_flax_available():
       from transformers import (
           FlaxBertModel,
           FlaxCLIPVisionModel,
           FlaxVisionTextDualEncoderModel,
           FlaxViTModel,
           VisionTextDualEncoderConfig,
           VisionTextDualEncoderProcessor,
       )
       from transformers.modeling_flax_pytorch_utils import (
           convert_pytorch_state_dict_to_flax,
           load_flax_weights_in_pytorch_model,
       )


if is_torch_available():
       import torch

       from transformers import VisionTextDualEncoderModel

if is_vision_available():
       from PIL import Image



def     lowercase__( __SCREAMING_SNAKE_CASE							:  Tuple   ):
     if isinstance(__SCREAMING_SNAKE_CASE  ,    collections.abc.Iterable   ):
          return x
     return (x, x)







@require_flax
class  UpperCamelCase	:
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase    )			->   Any:


        '''simple docstring'''


        pass
   def     _UpperCAmelCase     (    self    )			->   str:


        '''simple docstring'''


        pass
   def     _UpperCAmelCase     (    self    )			->   List[str]:


        '''simple docstring'''


        pass
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase    )			->   Tuple:


        '''simple docstring'''


        lowercase_			:       Tuple									=			np.abs((a - b)    ).max()
        self.assertLessEqual(__UpperCamelCase  ,__UpperCamelCase  ,f'''Difference between torch and flax is {diff} (>= {tol}).'''    )
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase=None  ,**__UpperCamelCase    )			->   Optional[Any]:


        '''simple docstring'''


        lowercase_			:       Optional[Any]									=			VisionTextDualEncoderConfig.from_vision_text_configs(__UpperCamelCase  ,__UpperCamelCase    )

        lowercase_			:       Dict									=			FlaxVisionTextDualEncoderModel(__UpperCamelCase    )

        lowercase_			:       Union[str, Any]									=			model(input_ids=__UpperCamelCase  ,pixel_values=__UpperCamelCase  ,attention_mask=__UpperCamelCase    )

        self.assertEqual(output['text_embeds'].shape  ,(input_ids.shape[0], config.projection_dim)    )
        self.assertEqual(output['image_embeds'].shape  ,(pixel_values.shape[0], config.projection_dim)    )
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase=None  ,**__UpperCamelCase    )			->   List[str]:


        '''simple docstring'''


        lowercase_			:       Tuple									=			self.get_vision_text_model(__UpperCamelCase  ,__UpperCamelCase    )
        lowercase_			:       Union[str, Any]									=			{'vision_model': vision_model, 'text_model': text_model}
        lowercase_			:       List[str]									=			FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**__UpperCamelCase    )

        lowercase_			:       Tuple									=			model(input_ids=__UpperCamelCase  ,pixel_values=__UpperCamelCase  ,attention_mask=__UpperCamelCase    )

        self.assertEqual(output['text_embeds'].shape  ,(input_ids.shape[0], model.config.projection_dim)    )
        self.assertEqual(output['image_embeds'].shape  ,(pixel_values.shape[0], model.config.projection_dim)    )
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase=None  ,**__UpperCamelCase    )			->   int:


        '''simple docstring'''


        lowercase_			:       Any									=			self.get_vision_text_model(__UpperCamelCase  ,__UpperCamelCase    )
        lowercase_			:       Optional[int]									=			{'vision_model': vision_model, 'text_model': text_model}
        lowercase_			:       str									=			FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**__UpperCamelCase    )

        lowercase_			:       Any									=			model(input_ids=__UpperCamelCase  ,pixel_values=__UpperCamelCase  ,attention_mask=__UpperCamelCase    )
        lowercase_			:       List[str]									=			output[0]

        with tempfile.TemporaryDirectory() as tmpdirname:
             model.save_pretrained(__UpperCamelCase    )
             lowercase_			:       List[Any]									=			FlaxVisionTextDualEncoderModel.from_pretrained(__UpperCamelCase    )

             lowercase_			:       str									=			model(input_ids=__UpperCamelCase  ,pixel_values=__UpperCamelCase  ,attention_mask=__UpperCamelCase    )
             lowercase_			:       str									=			after_output[0]
             lowercase_			:       Optional[int]									=			np.amax(np.abs(out_a - out_a    )    )
             self.assertLessEqual(__UpperCamelCase  ,1e-3    )
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase=None  ,**__UpperCamelCase    )			->   int:


        '''simple docstring'''


        lowercase_			:       List[Any]									=			self.get_vision_text_model(__UpperCamelCase  ,__UpperCamelCase    )
        lowercase_			:       Tuple									=			{'vision_model': vision_model, 'text_model': text_model}
        lowercase_			:       List[Any]									=			FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**__UpperCamelCase    )

        lowercase_			:       int									=			model(
            input_ids=__UpperCamelCase  ,pixel_values=__UpperCamelCase  ,attention_mask=__UpperCamelCase  ,output_attentions=__UpperCamelCase    )

        lowercase_			:       str									=			output.vision_model_output.attentions
        self.assertEqual(len(__UpperCamelCase    )  ,vision_config.num_hidden_layers    )

        # in ViT, the seq_len equals the number of patches + 1 (we add 1 for the [CLS] token)
        lowercase_			:       str									=			to_atuple(vision_model.config.image_size    )
        lowercase_			:       List[Any]									=			to_atuple(vision_model.config.patch_size    )
        lowercase_			:       List[Any]									=			(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
        lowercase_			:       List[str]									=			num_patches + 1
        self.assertEqual(vision_attentions[0].shape[-3:]  ,(vision_config.num_attention_heads, seq_len, seq_len)    )

        lowercase_			:       Union[str, Any]									=			output.text_model_output.attentions
        self.assertEqual(len(__UpperCamelCase    )  ,text_config.num_hidden_layers    )

        self.assertEqual(
            text_attentions[0].shape[-3:]  ,(text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])  ,)
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase    )			->   Optional[Any]:


        '''simple docstring'''


        pt_model.to(__UpperCamelCase    )
        pt_model.eval()

        # prepare inputs
        lowercase_			:       List[Any]									=			inputs_dict
        lowercase_			:       Dict									=			{k: torch.tensor(v.tolist()    ) for k, v in flax_inputs.items()}

        with torch.no_grad():
             lowercase_			:       Union[str, Any]									=			pt_model(**__UpperCamelCase    ).to_tuple()

        lowercase_			:       Dict									=			fx_model(**__UpperCamelCase    ).to_tuple()
        self.assertEqual(len(__UpperCamelCase    )  ,len(__UpperCamelCase    )  ,'Output lengths differ between Flax and PyTorch'    )
        for fx_output, pt_output in zip(fx_outputs[:4]  ,pt_outputs[:4]    ):
             self.assert_almost_equals(__UpperCamelCase  ,pt_output.numpy()  ,4e-2    )

        # PT -> Flax
        with tempfile.TemporaryDirectory() as tmpdirname:
             pt_model.save_pretrained(__UpperCamelCase    )
             lowercase_			:       str									=			FlaxVisionTextDualEncoderModel.from_pretrained(__UpperCamelCase  ,from_pt=__UpperCamelCase    )

        lowercase_			:       str									=			fx_model_loaded(**__UpperCamelCase    ).to_tuple()
        self.assertEqual(len(__UpperCamelCase    )  ,len(__UpperCamelCase    )  ,'Output lengths differ between Flax and PyTorch'    )
        for fx_output_loaded, pt_output in zip(fx_outputs_loaded[:4]  ,pt_outputs[:4]    ):
             self.assert_almost_equals(__UpperCamelCase  ,pt_output.numpy()  ,4e-2    )

        # Flax -> PT
        with tempfile.TemporaryDirectory() as tmpdirname:
             fx_model.save_pretrained(__UpperCamelCase    )
             lowercase_			:       int									=			VisionTextDualEncoderModel.from_pretrained(__UpperCamelCase  ,from_flax=__UpperCamelCase    )

        pt_model_loaded.to(__UpperCamelCase    )
        pt_model_loaded.eval()

        with torch.no_grad():
             lowercase_			:       Union[str, Any]									=			pt_model_loaded(**__UpperCamelCase    ).to_tuple()

        self.assertEqual(len(__UpperCamelCase    )  ,len(__UpperCamelCase    )  ,'Output lengths differ between Flax and PyTorch'    )
        for fx_output, pt_output_loaded in zip(fx_outputs[:4]  ,pt_outputs_loaded[:4]    ):
             self.assert_almost_equals(__UpperCamelCase  ,pt_output_loaded.numpy()  ,4e-2    )
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase    )			->   Tuple:


        '''simple docstring'''


        lowercase_			:       str									=			VisionTextDualEncoderConfig.from_vision_text_configs(__UpperCamelCase  ,__UpperCamelCase    )

        lowercase_			:       Dict									=			VisionTextDualEncoderModel(__UpperCamelCase    )
        lowercase_			:       Union[str, Any]									=			FlaxVisionTextDualEncoderModel(__UpperCamelCase    )

        lowercase_			:       List[str]									=			convert_pytorch_state_dict_to_flax(pt_model.state_dict()  ,__UpperCamelCase    )
        lowercase_			:       Union[str, Any]									=			fx_state

        self.check_pt_flax_equivalence(__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase    )
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase    )			->   str:


        '''simple docstring'''


        lowercase_			:       Dict									=			VisionTextDualEncoderConfig.from_vision_text_configs(__UpperCamelCase  ,__UpperCamelCase    )

        lowercase_			:       List[Any]									=			VisionTextDualEncoderModel(__UpperCamelCase    )
        lowercase_			:       List[str]									=			FlaxVisionTextDualEncoderModel(__UpperCamelCase    )

        lowercase_			:       Any									=			load_flax_weights_in_pytorch_model(__UpperCamelCase  ,fx_model.params    )

        self.check_pt_flax_equivalence(__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase    )
   def     _UpperCAmelCase     (    self    )			->   Optional[Any]:


        '''simple docstring'''


        lowercase_			:       Tuple									=			self.prepare_config_and_inputs()
        self.check_model_from_pretrained_configs(**__UpperCamelCase    )
   def     _UpperCAmelCase     (    self    )			->   Tuple:


        '''simple docstring'''


        lowercase_			:       List[str]									=			self.prepare_config_and_inputs()
        self.check_vision_text_dual_encoder_from_pretrained(**__UpperCamelCase    )
   def     _UpperCAmelCase     (    self    )			->   Optional[int]:


        '''simple docstring'''


        lowercase_			:       int									=			self.prepare_config_and_inputs()
        self.check_save_load(**__UpperCamelCase    )
   def     _UpperCAmelCase     (    self    )			->   Any:


        '''simple docstring'''


        lowercase_			:       Dict									=			self.prepare_config_and_inputs()
        self.check_vision_text_output_attention(**__UpperCamelCase    )
   @is_pt_flax_cross_test
   def     _UpperCAmelCase     (    self    )			->   Union[str, Any]:


        '''simple docstring'''


        lowercase_			:       List[str]									=			self.prepare_config_and_inputs()
        lowercase_			:       Optional[int]									=			config_inputs_dict.pop('vision_config'    )
        lowercase_			:       int									=			config_inputs_dict.pop('text_config'    )

        lowercase_			:       Optional[Any]									=			config_inputs_dict

        self.check_equivalence_pt_to_flax(__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase    )
        self.check_equivalence_flax_to_pt(__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase    )







   @slow
   def     _UpperCAmelCase     (    self    )			->   Any:


        '''simple docstring'''


        lowercase_			:       int									=			self.get_pretrained_model_and_inputs()

        lowercase_			:       Union[str, Any]									=			model_a(**__UpperCamelCase    )
        lowercase_			:       Tuple									=			outputs[0]

        with tempfile.TemporaryDirectory() as tmp_dirname:
             model_a.save_pretrained(__UpperCamelCase    )
             lowercase_			:       Any									=			FlaxVisionTextDualEncoderModel.from_pretrained(__UpperCamelCase    )

             lowercase_			:       str									=			model_a(**__UpperCamelCase    )
             lowercase_			:       Optional[int]									=			after_outputs[0]
             lowercase_			:       Union[str, Any]									=			np.amax(np.abs(out_a - out_a    )    )
             self.assertLessEqual(__UpperCamelCase  ,1e-5    )







@require_flax
class  UpperCamelCase	(    lowercase_		,   unittest.TestCase						):
   def     _UpperCAmelCase     (    self    )			->   Optional[Any]:


        '''simple docstring'''


        lowercase_			:       Optional[Any]									=			FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
            'hf-internal-testing/tiny-random-vit'  ,'hf-internal-testing/tiny-bert'  ,vision_from_pt=__UpperCamelCase  ,text_from_pt=__UpperCamelCase  ,)
        lowercase_			:       List[str]									=			13
        lowercase_			:       Tuple									=			floats_tensor(
            [
                batch_size,
                model.config.vision_config.num_channels,
                model.config.vision_config.image_size,
                model.config.vision_config.image_size,
            ]    )
        lowercase_			:       Tuple									=			ids_tensor([batch_size, 4]  ,model.config.text_config.vocab_size    )
        lowercase_			:       List[str]									=			random_attention_mask([batch_size, 4]    )
        lowercase_			:       str									=			{'pixel_values': pixel_values, 'input_ids': input_ids, 'attention_mask': attention_mask}

        return model, inputs
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase    )			->   Union[str, Any]:


        '''simple docstring'''


        lowercase_			:       List[Any]									=			FlaxViTModel(__UpperCamelCase    )
        lowercase_			:       Optional[Any]									=			FlaxBertModel(__UpperCamelCase    )
        return vision_model, text_model







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


        '''simple docstring'''


        lowercase_			:       Any									=			FlaxViTModelTester(self    )
        lowercase_			:       Dict									=			FlaxBertModelTester(self    )
        lowercase_			:       Union[str, Any]									=			vit_model_tester.prepare_config_and_inputs()
        lowercase_			:       Optional[Any]									=			bert_model_tester.prepare_config_and_inputs()

        lowercase_			:       int									=			vision_config_and_inputs

        lowercase_			:       Optional[Any]									=			text_config_and_inputs

        # make sure that cross attention layers are added
        return {
            "text_config": text_config,
            "vision_config": vision_config,
            "pixel_values": pixel_values,
            "attention_mask": attention_mask,
            "input_ids": input_ids,
            "token_type_ids": token_type_ids,
        }







@require_torch
class  UpperCamelCase	(    lowercase_		,   unittest.TestCase						):
   def     _UpperCAmelCase     (    self    )			->   Dict:


        '''simple docstring'''


        lowercase_			:       int									=			FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
            'hf-internal-testing/tiny-random-clip'  ,'hf-internal-testing/tiny-bert'  ,vision_from_pt=__UpperCamelCase  ,text_from_pt=__UpperCamelCase  ,)
        lowercase_			:       List[str]									=			13
        lowercase_			:       Any									=			floats_tensor(
            [
                batch_size,
                model.config.vision_config.num_channels,
                model.config.vision_config.image_size,
                model.config.vision_config.image_size,
            ]    )
        lowercase_			:       Optional[Any]									=			ids_tensor([batch_size, 4]  ,model.config.text_config.vocab_size    )
        lowercase_			:       Any									=			random_attention_mask([batch_size, 4]    )
        lowercase_			:       Union[str, Any]									=			{'pixel_values': pixel_values, 'input_ids': input_ids, 'attention_mask': attention_mask}

        return model, inputs
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase    )			->   Any:


        '''simple docstring'''


        lowercase_			:       Optional[Any]									=			FlaxCLIPVisionModel(__UpperCamelCase    )
        lowercase_			:       Tuple									=			FlaxBertModel(__UpperCamelCase    )
        return vision_model, text_model







   def     _UpperCAmelCase     (    self    )			->   Any:


        '''simple docstring'''


        lowercase_			:       Any									=			FlaxCLIPVisionModelTester(self    )
        lowercase_			:       Dict									=			FlaxBertModelTester(self    )
        lowercase_			:       List[str]									=			clip_model_tester.prepare_config_and_inputs()
        lowercase_			:       Optional[int]									=			bert_model_tester.prepare_config_and_inputs()

        lowercase_			:       Any									=			vision_config_and_inputs

        lowercase_			:       Optional[int]									=			text_config_and_inputs

        # make sure that cross attention layers are added
        return {
            "text_config": text_config,
            "vision_config": vision_config,
            "pixel_values": pixel_values,
            "attention_mask": attention_mask,
            "input_ids": input_ids,
            "token_type_ids": token_type_ids,
        }







@require_flax
@require_vision
class  UpperCamelCase	(    unittest.TestCase						):
   @slow
   def     _UpperCAmelCase     (    self    )			->   Any:


        '''simple docstring'''


        lowercase_			:       Any									=			FlaxVisionTextDualEncoderModel.from_pretrained('clip-italian/clip-italian'  ,logit_scale_init_value=1.0    )
        lowercase_			:       Any									=			VisionTextDualEncoderProcessor.from_pretrained('clip-italian/clip-italian'    )

        lowercase_			:       Union[str, Any]									=			Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'    )
        lowercase_			:       Dict									=			processor(
            text=['una foto di un gatto', 'una foto di un cane']  ,images=__UpperCamelCase  ,padding=__UpperCamelCase  ,return_tensors='np'    )

        lowercase_			:       Union[str, Any]									=			model(**__UpperCamelCase    )

        # verify the logits
        self.assertEqual(outputs.logits_per_image.shape  ,(inputs.pixel_values.shape[0], inputs.input_ids.shape[0])    )
        self.assertEqual(
            outputs.logits_per_text.shape  ,(inputs.input_ids.shape[0], inputs.pixel_values.shape[0])  ,)

        lowercase_			:       Any									=			np.array([[1.228_4727, 0.310_4122]]    )

        self.assertTrue(np.allclose(outputs.logits_per_image  ,__UpperCamelCase  ,atol=1e-3    )    )


"""simple docstring"""







from diffusers.utils.testing_utils import require_onnxruntime







@require_onnxruntime
class  UpperCamelCase	:
   pass

