





'''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 				_lowerCAmelCase  (			lowerCamelCase_     :					List[str]			):






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






@require_flax
class 							__lowercase    :







 '''simple docstring'''

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







    '''simple docstring'''





    pass


 def   _UpperCAmelCase       (self	)							->							Optional[int]:







    '''simple docstring'''





    pass


 def   _UpperCAmelCase       (self	)							->							Dict:







    '''simple docstring'''





    pass


 def   _UpperCAmelCase       (self       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase	)							->							Optional[Any]:







    '''simple docstring'''





    __lowercase			  =							np.abs((a - b)	).max()
    self.assertLessEqual(_lowercase       ,_lowercase       ,f"Difference between torch and flax is {diff} (>= {tol})."	)


 def   _UpperCAmelCase       (self       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase=None       ,**_lowerCamelCase	)							->							List[Any]:







    '''simple docstring'''





    __lowercase			  =							VisionTextDualEncoderConfig.from_vision_text_configs(_lowercase       ,_lowercase	)

    __lowercase			  =							FlaxVisionTextDualEncoderModel(_lowercase	)

    __lowercase			  =							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   _UpperCAmelCase       (self       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase=None       ,**_lowerCamelCase	)							->							List[Any]:







    '''simple docstring'''





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

    __lowercase			  =							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   _UpperCAmelCase       (self       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase=None       ,**_lowerCamelCase	)							->							int:







    '''simple docstring'''





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

    __lowercase			  =							model(input_ids=_lowercase       ,pixel_values=_lowercase       ,attention_mask=_lowercase	)
    __lowercase			  =							output[0]

    with tempfile.TemporaryDirectory() as tmpdirname:
       model.save_pretrained(_lowercase	)
       __lowercase			  =							FlaxVisionTextDualEncoderModel.from_pretrained(_lowercase	)

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


 def   _UpperCAmelCase       (self       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase=None       ,**_lowerCamelCase	)							->							Optional[Any]:







    '''simple docstring'''





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

    __lowercase			  =							model(
        input_ids=_lowercase       ,pixel_values=_lowercase       ,attention_mask=_lowercase       ,output_attentions=_lowercase	)

    __lowercase			  =							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)
    __lowercase			  =							to_atuple(vision_model.config.image_size	)
    __lowercase			  =							to_atuple(vision_model.config.patch_size	)
    __lowercase			  =							(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
    __lowercase			  =							num_patches + 1
    self.assertEqual(vision_attentions[0].shape[-3:]       ,(vision_config.num_attention_heads, seq_len, seq_len)	)

    __lowercase			  =							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   _UpperCAmelCase       (self       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase	)							->							Optional[Any]:







    '''simple docstring'''





    pt_model.to(_lowercase	)
    pt_model.eval()

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

    with torch.no_grad():
       __lowercase			  =							pt_model(**_lowercase	).to_tuple()

    __lowercase			  =							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	)
       __lowercase			  =							FlaxVisionTextDualEncoderModel.from_pretrained(_lowercase       ,from_pt=_lowercase	)

    __lowercase			  =							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	)
       __lowercase			  =							VisionTextDualEncoderModel.from_pretrained(_lowercase       ,from_flax=_lowercase	)

    pt_model_loaded.to(_lowercase	)
    pt_model_loaded.eval()

    with torch.no_grad():
       __lowercase			  =							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   _UpperCAmelCase       (self       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase	)							->							Union[str, Any]:







    '''simple docstring'''





    __lowercase			  =							VisionTextDualEncoderConfig.from_vision_text_configs(_lowercase       ,_lowercase	)

    __lowercase			  =							VisionTextDualEncoderModel(_lowercase	)
    __lowercase			  =							FlaxVisionTextDualEncoderModel(_lowercase	)

    __lowercase			  =							convert_pytorch_state_dict_to_flax(pt_model.state_dict()       ,_lowercase	)
    __lowercase			  =							fx_state

    self.check_pt_flax_equivalence(_lowercase       ,_lowercase       ,_lowercase	)


 def   _UpperCAmelCase       (self       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase	)							->							List[Any]:







    '''simple docstring'''





    __lowercase			  =							VisionTextDualEncoderConfig.from_vision_text_configs(_lowercase       ,_lowercase	)

    __lowercase			  =							VisionTextDualEncoderModel(_lowercase	)
    __lowercase			  =							FlaxVisionTextDualEncoderModel(_lowercase	)

    __lowercase			  =							load_flax_weights_in_pytorch_model(_lowercase       ,fx_model.params	)

    self.check_pt_flax_equivalence(_lowercase       ,_lowercase       ,_lowercase	)


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







    '''simple docstring'''





    __lowercase			  =							self.prepare_config_and_inputs()
    self.check_model_from_pretrained_configs(**_lowercase	)


 def   _UpperCAmelCase       (self	)							->							Dict:







    '''simple docstring'''





    __lowercase			  =							self.prepare_config_and_inputs()
    self.check_vision_text_dual_encoder_from_pretrained(**_lowercase	)


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







    '''simple docstring'''





    __lowercase			  =							self.prepare_config_and_inputs()
    self.check_save_load(**_lowercase	)


 def   _UpperCAmelCase       (self	)							->							Tuple:







    '''simple docstring'''





    __lowercase			  =							self.prepare_config_and_inputs()
    self.check_vision_text_output_attention(**_lowercase	)


 @is_pt_flax_cross_test
 def   _UpperCAmelCase       (self	)							->							List[Any]:







    '''simple docstring'''





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

    __lowercase			  =							config_inputs_dict

    self.check_equivalence_pt_to_flax(_lowercase       ,_lowercase       ,_lowercase	)
    self.check_equivalence_flax_to_pt(_lowercase       ,_lowercase       ,_lowercase	)


 @slow
 def   _UpperCAmelCase       (self	)							->							Tuple:







    '''simple docstring'''





    __lowercase  ,						__lowercase			  =							self.get_pretrained_model_and_inputs()

    __lowercase			  =							model_a(**_lowercase	)
    __lowercase			  =							outputs[0]

    with tempfile.TemporaryDirectory() as tmp_dirname:
       model_a.save_pretrained(_lowercase	)
       __lowercase			  =							FlaxVisionTextDualEncoderModel.from_pretrained(_lowercase	)

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






@require_flax
class 							__lowercase    ( lowerCAmelCase__	, unittest.TestCase  ):







 '''simple docstring'''

 def   _UpperCAmelCase       (self	)							->							Any:







    '''simple docstring'''





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

    return model, inputs


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







    '''simple docstring'''





    __lowercase			  =							FlaxViTModel(_lowercase	)
    __lowercase			  =							FlaxBertModel(_lowercase	)
    return vision_model, text_model


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







    '''simple docstring'''





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

    __lowercase  ,						__lowercase			  =							vision_config_and_inputs

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







 '''simple docstring'''

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







    '''simple docstring'''





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

    return model, inputs


 def   _UpperCAmelCase       (self       ,_lowerCamelCase       ,_lowerCamelCase	)							->							str:







    '''simple docstring'''





    __lowercase			  =							FlaxCLIPVisionModel(_lowercase	)
    __lowercase			  =							FlaxBertModel(_lowercase	)
    return vision_model, text_model


 def   _UpperCAmelCase       (self	)							->							Any:







    '''simple docstring'''





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

    __lowercase  ,						__lowercase			  =							vision_config_and_inputs

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







 '''simple docstring'''

 @slow
 def   _UpperCAmelCase       (self	)							->							str:







    '''simple docstring'''





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

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

    __lowercase			  =							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])       ,)

    __lowercase			  =							np.array([[1.2_2_8_4_7_2_7, 0.3_1_0_4_1_2_2]]	)

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












'''simple docstring'''



from dataclasses import dataclass
from enum import Enum
from typing import List, Optional, Union

import numpy as np
import PIL
from PIL import Image

from ...utils import BaseOutput, is_torch_available, is_transformers_available






@dataclass
class 							__lowercase    ( lowerCAmelCase__  ):







	'''simple docstring'''

	a      :			Union[List[PIL.Image.Image], np.ndarray]
	a      :			Optional[List[bool]]


if is_transformers_available() and is_torch_available():
	from .pipeline_semantic_stable_diffusion import SemanticStableDiffusionPipeline





