



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      __snake_case	(			_UpperCAmelCase				):
    if isinstance(_UpperCAmelCase   , collections.abc.Iterable				):
        return x
    return (x, x)

@require_flax
class     _A       :



       def 		_lowerCamelCase			(		self							:     List[Any]     ,		__SCREAMING_SNAKE_CASE							:     Optional[Any]     ,		__SCREAMING_SNAKE_CASE							:     int):


           '''simple docstring'''

           pass



       def 		_lowerCamelCase			(		self							:     Tuple):


           '''simple docstring'''

           pass



       def 		_lowerCamelCase			(		self							:     Dict):


           '''simple docstring'''

           pass



       def 		_lowerCamelCase			(		self							:     Dict     ,		__SCREAMING_SNAKE_CASE							:     np.ndarray     ,		__SCREAMING_SNAKE_CASE							:     np.ndarray     ,		__SCREAMING_SNAKE_CASE							:     float):


           '''simple docstring'''

           __a						       =				np.abs((a - b)).max()
           self.assertLessEqual(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		F'Difference between torch and flax is {diff} (>= {tol}).')



       def 		_lowerCamelCase			(		self							:     Union[str, Any]     ,		__SCREAMING_SNAKE_CASE							:     str     ,		__SCREAMING_SNAKE_CASE							:     Tuple     ,		__SCREAMING_SNAKE_CASE							:     Union[str, Any]     ,		__SCREAMING_SNAKE_CASE							:     Union[str, Any]     ,		__SCREAMING_SNAKE_CASE							:     int=None     ,		**__SCREAMING_SNAKE_CASE							:     List[Any]):


           '''simple docstring'''

           __a						       =				VisionTextDualEncoderConfig.from_vision_text_configs(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE)

           __a						       =				FlaxVisionTextDualEncoderModel(__SCREAMING_SNAKE_CASE)

           __a						       =				model(input_ids=__SCREAMING_SNAKE_CASE     ,		pixel_values=__SCREAMING_SNAKE_CASE     ,		attention_mask=__SCREAMING_SNAKE_CASE)

           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							:     Optional[int]     ,		__SCREAMING_SNAKE_CASE							:     List[Any]     ,		__SCREAMING_SNAKE_CASE							:     Union[str, Any]     ,		__SCREAMING_SNAKE_CASE							:     List[str]     ,		__SCREAMING_SNAKE_CASE							:     Any     ,		__SCREAMING_SNAKE_CASE							:     int=None     ,		**__SCREAMING_SNAKE_CASE							:     Dict):


           '''simple docstring'''

           __a       ,   __a						       =				self.get_vision_text_model(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE)
           __a						       =				{'''vision_model''': vision_model, '''text_model''': text_model}
           __a						       =				FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**__SCREAMING_SNAKE_CASE)

           __a						       =				model(input_ids=__SCREAMING_SNAKE_CASE     ,		pixel_values=__SCREAMING_SNAKE_CASE     ,		attention_mask=__SCREAMING_SNAKE_CASE)

           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							:     Optional[Any]     ,		__SCREAMING_SNAKE_CASE							:     Optional[int]     ,		__SCREAMING_SNAKE_CASE							:     str     ,		__SCREAMING_SNAKE_CASE							:     Any     ,		__SCREAMING_SNAKE_CASE							:     Optional[int]     ,		__SCREAMING_SNAKE_CASE							:     int=None     ,		**__SCREAMING_SNAKE_CASE							:     List[str]):


           '''simple docstring'''

           __a       ,   __a						       =				self.get_vision_text_model(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE)
           __a						       =				{'''vision_model''': vision_model, '''text_model''': text_model}
           __a						       =				FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**__SCREAMING_SNAKE_CASE)

           __a						       =				model(input_ids=__SCREAMING_SNAKE_CASE     ,		pixel_values=__SCREAMING_SNAKE_CASE     ,		attention_mask=__SCREAMING_SNAKE_CASE)
           __a						       =				output[0]

           with tempfile.TemporaryDirectory() as tmpdirname:
               model.save_pretrained(__SCREAMING_SNAKE_CASE)
               __a						       =				FlaxVisionTextDualEncoderModel.from_pretrained(__SCREAMING_SNAKE_CASE)

               __a						       =				model(input_ids=__SCREAMING_SNAKE_CASE     ,		pixel_values=__SCREAMING_SNAKE_CASE     ,		attention_mask=__SCREAMING_SNAKE_CASE)
               __a						       =				after_output[0]
               __a						       =				np.amax(np.abs(out_a - out_a))
               self.assertLessEqual(__SCREAMING_SNAKE_CASE     ,		1E-3)



       def 		_lowerCamelCase			(		self							:     int     ,		__SCREAMING_SNAKE_CASE							:     Optional[int]     ,		__SCREAMING_SNAKE_CASE							:     List[str]     ,		__SCREAMING_SNAKE_CASE							:     Optional[Any]     ,		__SCREAMING_SNAKE_CASE							:     Dict     ,		__SCREAMING_SNAKE_CASE							:     List[Any]=None     ,		**__SCREAMING_SNAKE_CASE							:     int):


           '''simple docstring'''

           __a       ,   __a						       =				self.get_vision_text_model(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE)
           __a						       =				{'''vision_model''': vision_model, '''text_model''': text_model}
           __a						       =				FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**__SCREAMING_SNAKE_CASE)

           __a						       =				model(
               input_ids=__SCREAMING_SNAKE_CASE     ,		pixel_values=__SCREAMING_SNAKE_CASE     ,		attention_mask=__SCREAMING_SNAKE_CASE     ,		output_attentions=__SCREAMING_SNAKE_CASE)

           __a						       =				output.vision_model_output.attentions
           self.assertEqual(len(__SCREAMING_SNAKE_CASE)     ,		vision_config.num_hidden_layers)

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

           __a						       =				output.text_model_output.attentions
           self.assertEqual(len(__SCREAMING_SNAKE_CASE)     ,		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							:     Optional[int]     ,		__SCREAMING_SNAKE_CASE							:     Dict     ,		__SCREAMING_SNAKE_CASE							:     int     ,		__SCREAMING_SNAKE_CASE							:     Union[str, Any]):


           '''simple docstring'''

           pt_model.to(__SCREAMING_SNAKE_CASE)
           pt_model.eval()

           # prepare inputs
           __a						       =				inputs_dict
           __a						       =				{k: torch.tensor(v.tolist()) for k, v in flax_inputs.items()}

           with torch.no_grad():
               __a						       =				pt_model(**__SCREAMING_SNAKE_CASE).to_tuple()

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

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

           __a						       =				fx_model_loaded(**__SCREAMING_SNAKE_CASE).to_tuple()
           self.assertEqual(len(__SCREAMING_SNAKE_CASE)     ,		len(__SCREAMING_SNAKE_CASE)     ,		'''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(__SCREAMING_SNAKE_CASE     ,		pt_output.numpy()     ,		4E-2)

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

           pt_model_loaded.to(__SCREAMING_SNAKE_CASE)
           pt_model_loaded.eval()

           with torch.no_grad():
               __a						       =				pt_model_loaded(**__SCREAMING_SNAKE_CASE).to_tuple()

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



       def 		_lowerCamelCase			(		self							:     List[str]     ,		__SCREAMING_SNAKE_CASE							:     Union[str, Any]     ,		__SCREAMING_SNAKE_CASE							:     Any     ,		__SCREAMING_SNAKE_CASE							:     str):


           '''simple docstring'''

           __a						       =				VisionTextDualEncoderConfig.from_vision_text_configs(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE)

           __a						       =				VisionTextDualEncoderModel(__SCREAMING_SNAKE_CASE)
           __a						       =				FlaxVisionTextDualEncoderModel(__SCREAMING_SNAKE_CASE)

           __a						       =				convert_pytorch_state_dict_to_flax(pt_model.state_dict()     ,		__SCREAMING_SNAKE_CASE)
           __a						       =				fx_state

           self.check_pt_flax_equivalence(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE)



       def 		_lowerCamelCase			(		self							:     Optional[int]     ,		__SCREAMING_SNAKE_CASE							:     Dict     ,		__SCREAMING_SNAKE_CASE							:     Union[str, Any]     ,		__SCREAMING_SNAKE_CASE							:     str):


           '''simple docstring'''

           __a						       =				VisionTextDualEncoderConfig.from_vision_text_configs(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE)

           __a						       =				VisionTextDualEncoderModel(__SCREAMING_SNAKE_CASE)
           __a						       =				FlaxVisionTextDualEncoderModel(__SCREAMING_SNAKE_CASE)

           __a						       =				load_flax_weights_in_pytorch_model(__SCREAMING_SNAKE_CASE     ,		fx_model.params)

           self.check_pt_flax_equivalence(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE)



       def 		_lowerCamelCase			(		self							:     Any):


           '''simple docstring'''

           __a						       =				self.prepare_config_and_inputs()
           self.check_model_from_pretrained_configs(**__SCREAMING_SNAKE_CASE)



       def 		_lowerCamelCase			(		self							:     List[str]):


           '''simple docstring'''

           __a						       =				self.prepare_config_and_inputs()
           self.check_vision_text_dual_encoder_from_pretrained(**__SCREAMING_SNAKE_CASE)



       def 		_lowerCamelCase			(		self							:     List[str]):


           '''simple docstring'''

           __a						       =				self.prepare_config_and_inputs()
           self.check_save_load(**__SCREAMING_SNAKE_CASE)



       def 		_lowerCamelCase			(		self							:     Any):


           '''simple docstring'''

           __a						       =				self.prepare_config_and_inputs()
           self.check_vision_text_output_attention(**__SCREAMING_SNAKE_CASE)



       @is_pt_flax_cross_test
       def 		_lowerCamelCase			(		self							:     str):


           '''simple docstring'''

           __a						       =				self.prepare_config_and_inputs()
           __a						       =				config_inputs_dict.pop('''vision_config''')
           __a						       =				config_inputs_dict.pop('''text_config''')

           __a						       =				config_inputs_dict

           self.check_equivalence_pt_to_flax(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE)
           self.check_equivalence_flax_to_pt(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE)





       @slow
       def 		_lowerCamelCase			(		self							:     Dict):


           '''simple docstring'''

           __a       ,   __a						       =				self.get_pretrained_model_and_inputs()

           __a						       =				model_a(**__SCREAMING_SNAKE_CASE)
           __a						       =				outputs[0]

           with tempfile.TemporaryDirectory() as tmp_dirname:
               model_a.save_pretrained(__SCREAMING_SNAKE_CASE)
               __a						       =				FlaxVisionTextDualEncoderModel.from_pretrained(__SCREAMING_SNAKE_CASE)

               __a						       =				model_a(**__SCREAMING_SNAKE_CASE)
               __a						       =				after_outputs[0]
               __a						       =				np.amax(np.abs(out_a - out_a))
               self.assertLessEqual(__SCREAMING_SNAKE_CASE     ,		1E-5)


@require_flax
class     _A       (	__UpperCAmelCase      ,unittest.TestCase    ):



       def 		_lowerCamelCase			(		self							:     Dict):


           '''simple docstring'''

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

           return model, inputs



       def 		_lowerCamelCase			(		self							:     Optional[Any]     ,		__SCREAMING_SNAKE_CASE							:     Tuple     ,		__SCREAMING_SNAKE_CASE							:     Optional[int]):


           '''simple docstring'''

           __a						       =				FlaxViTModel(__SCREAMING_SNAKE_CASE)
           __a						       =				FlaxBertModel(__SCREAMING_SNAKE_CASE)
           return vision_model, text_model





       def 		_lowerCamelCase			(		self							:     Union[str, Any]):


           '''simple docstring'''

           __a						       =				FlaxViTModelTester(self)
           __a						       =				FlaxBertModelTester(self)
           __a						       =				vit_model_tester.prepare_config_and_inputs()
           __a						       =				bert_model_tester.prepare_config_and_inputs()

           __a       ,   __a						       =				vision_config_and_inputs

           __a       ,   __a       ,   __a       ,   __a						       =				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     _A       (	__UpperCAmelCase      ,unittest.TestCase    ):



       def 		_lowerCamelCase			(		self							:     Tuple):


           '''simple docstring'''

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

           return model, inputs



       def 		_lowerCamelCase			(		self							:     List[str]     ,		__SCREAMING_SNAKE_CASE							:     List[Any]     ,		__SCREAMING_SNAKE_CASE							:     int):


           '''simple docstring'''

           __a						       =				FlaxCLIPVisionModel(__SCREAMING_SNAKE_CASE)
           __a						       =				FlaxBertModel(__SCREAMING_SNAKE_CASE)
           return vision_model, text_model





       def 		_lowerCamelCase			(		self							:     List[str]):


           '''simple docstring'''

           __a						       =				FlaxCLIPVisionModelTester(self)
           __a						       =				FlaxBertModelTester(self)
           __a						       =				clip_model_tester.prepare_config_and_inputs()
           __a						       =				bert_model_tester.prepare_config_and_inputs()

           __a       ,   __a						       =				vision_config_and_inputs

           __a       ,   __a       ,   __a       ,   __a						       =				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     _A       (	unittest.TestCase    ):



       @slow
       def 		_lowerCamelCase			(		self							:     List[str]):


           '''simple docstring'''

           __a						       =				FlaxVisionTextDualEncoderModel.from_pretrained('''clip-italian/clip-italian'''     ,		logit_scale_init_value=1.0)
           __a						       =				VisionTextDualEncoderProcessor.from_pretrained('''clip-italian/clip-italian''')

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

           __a						       =				model(**__SCREAMING_SNAKE_CASE)

           # 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						       =				np.array([[1.2_28_47_27, 0.3_10_41_22]])

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







# NOTE: This file is deprecated and will be removed in a future version.
# It only exists so that temporarely `from diffusers.pipelines import DiffusionPipeline` works
from ...utils import deprecate
from ..controlnet.multicontrolnet import MultiControlNetModel  # noqa: F401
from ..controlnet.pipeline_controlnet import StableDiffusionControlNetPipeline  # noqa: F401


deprecate(
    'stable diffusion controlnet',
    '0.22.0',
    'Importing `StableDiffusionControlNetPipeline` or `MultiControlNetModel` from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet is deprecated. Please import `from diffusers import StableDiffusionControlNetPipeline` instead.',
    standard_warn=False,
    stacklevel=3,
)