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     UpperCamelCase						(	__lowercase						:			Tuple     ):
		'''simple docstring'''




		if isinstance(__lowercase						,collections.abc.Iterable     ):
				return x
		return (x, x)





@require_flax
class UpperCAmelCase						:


				'''simple docstring'''






				def   lowerCAmelCase_       (						self		,     lowercase		,     lowercase ):



						"""simple docstring"""





						pass
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						pass
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						pass
				def   lowerCAmelCase_       (						self		,     lowercase		,     lowercase		,     lowercase ):



						"""simple docstring"""





						A_		:    Optional[Any]         =				np.abs((a - b) ).max()
						self.assertLessEqual(lowercase		,     lowercase		,     F'''Difference between torch and flax is {diff} (>= {tol}).''' )
				def   lowerCAmelCase_       (						self		,     lowercase		,     lowercase		,     lowercase		,     lowercase		,     lowercase=None		,     **lowercase ):



						"""simple docstring"""





						A_		:    Optional[int]         =				VisionTextDualEncoderConfig.from_vision_text_configs(lowercase		,     lowercase )

						A_		:    Dict         =				FlaxVisionTextDualEncoderModel(lowercase )

						A_		:    Optional[int]         =				model(input_ids=lowercase		,     pixel_values=lowercase		,     attention_mask=lowercase )

						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   lowerCAmelCase_       (						self		,     lowercase		,     lowercase		,     lowercase		,     lowercase		,     lowercase=None		,     **lowercase ):



						"""simple docstring"""





						A_       ,      A_		:    Dict         =				self.get_vision_text_model(lowercase		,     lowercase )
						A_		:    Tuple         =				{'vision_model': vision_model, 'text_model': text_model}
						A_		:    Union[str, Any]         =				FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**lowercase )

						A_		:    Union[str, Any]         =				model(input_ids=lowercase		,     pixel_values=lowercase		,     attention_mask=lowercase )

						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   lowerCAmelCase_       (						self		,     lowercase		,     lowercase		,     lowercase		,     lowercase		,     lowercase=None		,     **lowercase ):



						"""simple docstring"""





						A_       ,      A_		:    Union[str, Any]         =				self.get_vision_text_model(lowercase		,     lowercase )
						A_		:    str         =				{'vision_model': vision_model, 'text_model': text_model}
						A_		:    Union[str, Any]         =				FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**lowercase )

						A_		:    str         =				model(input_ids=lowercase		,     pixel_values=lowercase		,     attention_mask=lowercase )
						A_		:    List[str]         =				output[0]

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

								A_		:    Optional[int]         =				model(input_ids=lowercase		,     pixel_values=lowercase		,     attention_mask=lowercase )
								A_		:    Tuple         =				after_output[0]
								A_		:    Union[str, Any]         =				np.amax(np.abs(out_a - out_a ) )
								self.assertLessEqual(lowercase		,     1E-3 )
				def   lowerCAmelCase_       (						self		,     lowercase		,     lowercase		,     lowercase		,     lowercase		,     lowercase=None		,     **lowercase ):



						"""simple docstring"""





						A_       ,      A_		:    str         =				self.get_vision_text_model(lowercase		,     lowercase )
						A_		:    Dict         =				{'vision_model': vision_model, 'text_model': text_model}
						A_		:    str         =				FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**lowercase )

						A_		:    Optional[int]         =				model(
						    input_ids=lowercase		,     pixel_values=lowercase		,     attention_mask=lowercase		,     output_attentions=lowercase )

						A_		:    int         =				output.vision_model_output.attentions
						self.assertEqual(len(lowercase )		,     vision_config.num_hidden_layers )

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

						A_		:    Optional[int]         =				output.text_model_output.attentions
						self.assertEqual(len(lowercase )		,     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   lowerCAmelCase_       (						self		,     lowercase		,     lowercase		,     lowercase ):



						"""simple docstring"""





						pt_model.to(lowercase )
						pt_model.eval()

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

						with torch.no_grad():
								A_		:    int         =				pt_model(**lowercase ).to_tuple()

						A_		:    Optional[int]         =				fx_model(**lowercase ).to_tuple()
						self.assertEqual(len(lowercase )		,     len(lowercase )		,     'Output lengths differ between Flax and PyTorch' )
						for fx_output, pt_output in zip(fx_outputs[:4]		,     pt_outputs[:4] ):
								self.assert_almost_equals(lowercase		,     pt_output.numpy()		,     4E-2 )

						# PT -> Flax
						with tempfile.TemporaryDirectory() as tmpdirname:
								pt_model.save_pretrained(lowercase )
								A_		:    Union[str, Any]         =				FlaxVisionTextDualEncoderModel.from_pretrained(lowercase		,     from_pt=lowercase )

						A_		:    Union[str, Any]         =				fx_model_loaded(**lowercase ).to_tuple()
						self.assertEqual(len(lowercase )		,     len(lowercase )		,     '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(lowercase		,     pt_output.numpy()		,     4E-2 )

						# Flax -> PT
						with tempfile.TemporaryDirectory() as tmpdirname:
								fx_model.save_pretrained(lowercase )
								A_		:    List[Any]         =				VisionTextDualEncoderModel.from_pretrained(lowercase		,     from_flax=lowercase )

						pt_model_loaded.to(lowercase )
						pt_model_loaded.eval()

						with torch.no_grad():
								A_		:    Tuple         =				pt_model_loaded(**lowercase ).to_tuple()

						self.assertEqual(len(lowercase )		,     len(lowercase )		,     '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(lowercase		,     pt_output_loaded.numpy()		,     4E-2 )
				def   lowerCAmelCase_       (						self		,     lowercase		,     lowercase		,     lowercase ):



						"""simple docstring"""





						A_		:    Tuple         =				VisionTextDualEncoderConfig.from_vision_text_configs(lowercase		,     lowercase )

						A_		:    Any         =				VisionTextDualEncoderModel(lowercase )
						A_		:    Tuple         =				FlaxVisionTextDualEncoderModel(lowercase )

						A_		:    Tuple         =				convert_pytorch_state_dict_to_flax(pt_model.state_dict()		,     lowercase )
						A_		:    Union[str, Any]         =				fx_state

						self.check_pt_flax_equivalence(lowercase		,     lowercase		,     lowercase )
				def   lowerCAmelCase_       (						self		,     lowercase		,     lowercase		,     lowercase ):



						"""simple docstring"""





						A_		:    List[Any]         =				VisionTextDualEncoderConfig.from_vision_text_configs(lowercase		,     lowercase )

						A_		:    Dict         =				VisionTextDualEncoderModel(lowercase )
						A_		:    Optional[Any]         =				FlaxVisionTextDualEncoderModel(lowercase )

						A_		:    Any         =				load_flax_weights_in_pytorch_model(lowercase		,     fx_model.params )

						self.check_pt_flax_equivalence(lowercase		,     lowercase		,     lowercase )
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						A_		:    Union[str, Any]         =				self.prepare_config_and_inputs()
						self.check_model_from_pretrained_configs(**lowercase )
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						A_		:    Union[str, Any]         =				self.prepare_config_and_inputs()
						self.check_vision_text_dual_encoder_from_pretrained(**lowercase )
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						A_		:    List[Any]         =				self.prepare_config_and_inputs()
						self.check_save_load(**lowercase )
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						A_		:    Optional[Any]         =				self.prepare_config_and_inputs()
						self.check_vision_text_output_attention(**lowercase )
				@is_pt_flax_cross_test
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						A_		:    str         =				self.prepare_config_and_inputs()
						A_		:    Union[str, Any]         =				config_inputs_dict.pop('vision_config' )
						A_		:    List[Any]         =				config_inputs_dict.pop('text_config' )

						A_		:    int         =				config_inputs_dict

						self.check_equivalence_pt_to_flax(lowercase		,     lowercase		,     lowercase )
						self.check_equivalence_flax_to_pt(lowercase		,     lowercase		,     lowercase )







				@slow
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						A_       ,      A_		:    Tuple         =				self.get_pretrained_model_and_inputs()

						A_		:    Union[str, Any]         =				model_a(**lowercase )
						A_		:    List[str]         =				outputs[0]

						with tempfile.TemporaryDirectory() as tmp_dirname:
								model_a.save_pretrained(lowercase )
								A_		:    List[Any]         =				FlaxVisionTextDualEncoderModel.from_pretrained(lowercase )

								A_		:    Any         =				model_a(**lowercase )
								A_		:    int         =				after_outputs[0]
								A_		:    Dict         =				np.amax(np.abs(out_a - out_a ) )
								self.assertLessEqual(lowercase		,     1E-5 )





@require_flax
class UpperCAmelCase						( __A ,	unittest.TestCase  ):


				'''simple docstring'''






				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





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

						return model, inputs
				def   lowerCAmelCase_       (						self		,     lowercase		,     lowercase ):



						"""simple docstring"""





						A_		:    Union[str, Any]         =				FlaxViTModel(lowercase )
						A_		:    str         =				FlaxBertModel(lowercase )
						return vision_model, text_model







				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						A_		:    Dict         =				FlaxViTModelTester(self )
						A_		:    Optional[int]         =				FlaxBertModelTester(self )
						A_		:    Optional[Any]         =				vit_model_tester.prepare_config_and_inputs()
						A_		:    Dict         =				bert_model_tester.prepare_config_and_inputs()

						A_       ,      A_		:    Dict         =				vision_config_and_inputs

						A_       ,      A_       ,      A_       ,      A_		:    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_torch
class UpperCAmelCase						( __A ,	unittest.TestCase  ):


				'''simple docstring'''






				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





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

						return model, inputs
				def   lowerCAmelCase_       (						self		,     lowercase		,     lowercase ):



						"""simple docstring"""





						A_		:    Optional[int]         =				FlaxCLIPVisionModel(lowercase )
						A_		:    int         =				FlaxBertModel(lowercase )
						return vision_model, text_model







				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





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

						A_       ,      A_		:    Union[str, Any]         =				vision_config_and_inputs

						A_       ,      A_       ,      A_       ,      A_		:    List[str]         =				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  ):


				'''simple docstring'''






				@slow
				def   lowerCAmelCase_       (						self ):



						"""simple docstring"""





						A_		:    str         =				FlaxVisionTextDualEncoderModel.from_pretrained('clip-italian/clip-italian'		,     logit_scale_init_value=1.0 )
						A_		:    Optional[Any]         =				VisionTextDualEncoderProcessor.from_pretrained('clip-italian/clip-italian' )

						A_		:    Tuple         =				Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png' )
						A_		:    int         =				processor(
						    text=['una foto di un gatto', 'una foto di un cane']		,     images=lowercase		,     padding=lowercase		,     return_tensors='np' )

						A_		:    Optional[int]         =				model(**lowercase )

						# 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])		,     )

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

						self.assertTrue(np.allclose(outputs.logits_per_image		,     lowercase		,     atol=1E-3 ) )







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_UpperCAmelCase            =		{
    """configuration_instructblip""": [
        """INSTRUCTBLIP_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """InstructBlipConfig""",
        """InstructBlipQFormerConfig""",
        """InstructBlipVisionConfig""",
    ],
    """processing_instructblip""": ["""InstructBlipProcessor"""],
}

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_UpperCAmelCase            =		[
	    """INSTRUCTBLIP_PRETRAINED_MODEL_ARCHIVE_LIST""",
	    """InstructBlipQFormerModel""",
	    """InstructBlipPreTrainedModel""",
	    """InstructBlipForConditionalGeneration""",
	    """InstructBlipVisionModel""",
	]

if TYPE_CHECKING:
	from .configuration_instructblip import (
	    INSTRUCTBLIP_PRETRAINED_CONFIG_ARCHIVE_MAP,
	    InstructBlipConfig,
	    InstructBlipQFormerConfig,
	    InstructBlipVisionConfig,
	)
	from .processing_instructblip import InstructBlipProcessor

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_instructblip import (
		    INSTRUCTBLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
		    InstructBlipForConditionalGeneration,
		    InstructBlipPreTrainedModel,
		    InstructBlipQFormerModel,
		    InstructBlipVisionModel,
		)

else:
	import sys

	_UpperCAmelCase            =		_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)






