# This model implementation is heavily inspired by https://github.com/haofanwang/ControlNet-for-Diffusers/

import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    ControlNetModel,
    DDIMScheduler,
    StableDiffusionControlNetImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import MultiControlNetModel
from diffusers.utils import floats_tensor, load_image, load_numpy, randn_tensor, slow, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
)


enable_full_determinism()

class    snake_case       (   UpperCamelCase_ ,    UpperCamelCase_ ,    UpperCamelCase_ ,    unittest.TestCase    ):
  lowercase_								=  StableDiffusionControlNetImgaImgPipeline
  lowercase_								=  TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'height', 'width'}
  lowercase_								=  TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
  lowercase_								=  IMAGE_TO_IMAGE_IMAGE_PARAMS.union({'control_image'}    )
  lowercase_								=  IMAGE_TO_IMAGE_IMAGE_PARAMS




  def 						__lowercase(   self  :						List[Any]					)->							List[Any]:




       """simple docstring"""
       torch.manual_seed(0					)
       SCREAMING_SNAKE_CASE__					:					List[str]      =		UNetaDConditionModel(
           block_out_channels=(32, 64)  ,			layers_per_block=2  ,			sample_size=32  ,			in_channels=4  ,			out_channels=4  ,			down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')  ,			up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')  ,			cross_attention_dim=32  ,			)
       torch.manual_seed(0					)
       SCREAMING_SNAKE_CASE__					:					Any      =		ControlNetModel(
           block_out_channels=(32, 64)  ,			layers_per_block=2  ,			in_channels=4  ,			down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')  ,			cross_attention_dim=32  ,			conditioning_embedding_out_channels=(16, 32)  ,			)
       torch.manual_seed(0					)
       SCREAMING_SNAKE_CASE__					:					Optional[Any]      =		DDIMScheduler(
           beta_start=0.0_0085  ,			beta_end=0.012  ,			beta_schedule='scaled_linear'  ,			clip_sample=a_  ,			set_alpha_to_one=a_  ,			)
       torch.manual_seed(0					)
       SCREAMING_SNAKE_CASE__					:					Optional[int]      =		AutoencoderKL(
           block_out_channels=[32, 64]  ,			in_channels=3  ,			out_channels=3  ,			down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']  ,			up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']  ,			latent_channels=4  ,			)
       torch.manual_seed(0					)
       SCREAMING_SNAKE_CASE__					:					int      =		CLIPTextConfig(
           bos_token_id=0  ,			eos_token_id=2  ,			hidden_size=32  ,			intermediate_size=37  ,			layer_norm_eps=1e-0_5  ,			num_attention_heads=4  ,			num_hidden_layers=5  ,			pad_token_id=1  ,			vocab_size=1000  ,			)
       SCREAMING_SNAKE_CASE__					:					Any      =		CLIPTextModel(a_					)
       SCREAMING_SNAKE_CASE__					:					Dict      =		CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'					)

       SCREAMING_SNAKE_CASE__					:					Tuple      =		{
           'unet': unet,
           'controlnet': controlnet,
           'scheduler': scheduler,
           'vae': vae,
           'text_encoder': text_encoder,
           'tokenizer': tokenizer,
           'safety_checker': None,
           'feature_extractor': None,
       }
       return components




  def 						__lowercase(   self  :						str  ,			a_  :						Optional[int]  ,			a_  :						Optional[int]=0					)->							Any:




       """simple docstring"""
       if str(a_					).startswith('mps'					):
            SCREAMING_SNAKE_CASE__					:					Union[str, Any]      =		torch.manual_seed(a_					)
       else:
            SCREAMING_SNAKE_CASE__					:					List[str]      =		torch.Generator(device=a_					).manual_seed(a_					)

       SCREAMING_SNAKE_CASE__					:					Any      =		2
       SCREAMING_SNAKE_CASE__					:					Tuple      =		randn_tensor(
           (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)  ,			generator=a_  ,			device=torch.device(a_					)  ,			)
       SCREAMING_SNAKE_CASE__					:					List[str]      =		floats_tensor(control_image.shape  ,			rng=random.Random(a_					)					).to(a_					)
       SCREAMING_SNAKE_CASE__					:					List[str]      =		image.cpu().permute(0  ,			2  ,			3  ,			1					)[0]
       SCREAMING_SNAKE_CASE__					:					int      =		Image.fromarray(np.uinta(a_					)					).convert('RGB'					).resize((64, 64)					)
       SCREAMING_SNAKE_CASE__					:					Optional[Any]      =		{
           'prompt': 'A painting of a squirrel eating a burger',
           'generator': generator,
           'num_inference_steps': 2,
           'guidance_scale': 6.0,
           'output_type': 'numpy',
           'image': image,
           'control_image': control_image,
       }

       return inputs




  def 						__lowercase(   self  :						Dict					)->							List[str]:




       """simple docstring"""
       return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3					)




  @unittest.skipIf(
      torch_device != 'cuda' or not is_xformers_available()  ,			reason='XFormers attention is only available with CUDA and `xformers` installed'  ,			)
  def 						__lowercase(   self  :						List[Any]					)->							int:




       """simple docstring"""
       self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3					)




  def 						__lowercase(   self  :						str					)->							int:




       """simple docstring"""
       self._test_inference_batch_single_identical(expected_max_diff=2e-3					)




class    snake_case       (   UpperCamelCase_ ,    UpperCamelCase_ ,    unittest.TestCase    ):
  lowercase_								=  StableDiffusionControlNetImgaImgPipeline
  lowercase_								=  TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'height', 'width'}
  lowercase_								=  TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
  lowercase_								=  frozenset([]    )  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess




  def 						__lowercase(   self  :						str					)->							Any:




       """simple docstring"""
       torch.manual_seed(0					)
       SCREAMING_SNAKE_CASE__					:					Union[str, Any]      =		UNetaDConditionModel(
           block_out_channels=(32, 64)  ,			layers_per_block=2  ,			sample_size=32  ,			in_channels=4  ,			out_channels=4  ,			down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')  ,			up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')  ,			cross_attention_dim=32  ,			)
       torch.manual_seed(0					)

       def init_weights(a_  :						Tuple					):
            if isinstance(a_  ,			torch.nn.Convad					):
                 torch.nn.init.normal(m.weight					)
                 m.bias.data.fill_(1.0					)

       SCREAMING_SNAKE_CASE__					:					Optional[int]      =		ControlNetModel(
           block_out_channels=(32, 64)  ,			layers_per_block=2  ,			in_channels=4  ,			down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')  ,			cross_attention_dim=32  ,			conditioning_embedding_out_channels=(16, 32)  ,			)
       controlneta.controlnet_down_blocks.apply(a_					)

       torch.manual_seed(0					)
       SCREAMING_SNAKE_CASE__					:					str      =		ControlNetModel(
           block_out_channels=(32, 64)  ,			layers_per_block=2  ,			in_channels=4  ,			down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')  ,			cross_attention_dim=32  ,			conditioning_embedding_out_channels=(16, 32)  ,			)
       controlneta.controlnet_down_blocks.apply(a_					)

       torch.manual_seed(0					)
       SCREAMING_SNAKE_CASE__					:					Tuple      =		DDIMScheduler(
           beta_start=0.0_0085  ,			beta_end=0.012  ,			beta_schedule='scaled_linear'  ,			clip_sample=a_  ,			set_alpha_to_one=a_  ,			)
       torch.manual_seed(0					)
       SCREAMING_SNAKE_CASE__					:					List[Any]      =		AutoencoderKL(
           block_out_channels=[32, 64]  ,			in_channels=3  ,			out_channels=3  ,			down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']  ,			up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']  ,			latent_channels=4  ,			)
       torch.manual_seed(0					)
       SCREAMING_SNAKE_CASE__					:					List[str]      =		CLIPTextConfig(
           bos_token_id=0  ,			eos_token_id=2  ,			hidden_size=32  ,			intermediate_size=37  ,			layer_norm_eps=1e-0_5  ,			num_attention_heads=4  ,			num_hidden_layers=5  ,			pad_token_id=1  ,			vocab_size=1000  ,			)
       SCREAMING_SNAKE_CASE__					:					Optional[int]      =		CLIPTextModel(a_					)
       SCREAMING_SNAKE_CASE__					:					Dict      =		CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'					)

       SCREAMING_SNAKE_CASE__					:					List[Any]      =		MultiControlNetModel([controlneta, controlneta]					)

       SCREAMING_SNAKE_CASE__					:					Tuple      =		{
           'unet': unet,
           'controlnet': controlnet,
           'scheduler': scheduler,
           'vae': vae,
           'text_encoder': text_encoder,
           'tokenizer': tokenizer,
           'safety_checker': None,
           'feature_extractor': None,
       }
       return components




  def 						__lowercase(   self  :						Tuple  ,			a_  :						Any  ,			a_  :						Union[str, Any]=0					)->							Optional[Any]:




       """simple docstring"""
       if str(a_					).startswith('mps'					):
            SCREAMING_SNAKE_CASE__					:					Optional[int]      =		torch.manual_seed(a_					)
       else:
            SCREAMING_SNAKE_CASE__					:					Any      =		torch.Generator(device=a_					).manual_seed(a_					)

       SCREAMING_SNAKE_CASE__					:					List[str]      =		2

       SCREAMING_SNAKE_CASE__					:					Optional[int]      =		[
           randn_tensor(
               (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)  ,			generator=a_  ,			device=torch.device(a_					)  ,			),
           randn_tensor(
               (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)  ,			generator=a_  ,			device=torch.device(a_					)  ,			),
       ]

       SCREAMING_SNAKE_CASE__					:					List[Any]      =		floats_tensor(control_image[0].shape  ,			rng=random.Random(a_					)					).to(a_					)
       SCREAMING_SNAKE_CASE__					:					List[Any]      =		image.cpu().permute(0  ,			2  ,			3  ,			1					)[0]
       SCREAMING_SNAKE_CASE__					:					Optional[Any]      =		Image.fromarray(np.uinta(a_					)					).convert('RGB'					).resize((64, 64)					)
       SCREAMING_SNAKE_CASE__					:					Tuple      =		{
           'prompt': 'A painting of a squirrel eating a burger',
           'generator': generator,
           'num_inference_steps': 2,
           'guidance_scale': 6.0,
           'output_type': 'numpy',
           'image': image,
           'control_image': control_image,
       }

       return inputs




  def 						__lowercase(   self  :						Optional[int]					)->							List[Any]:




       """simple docstring"""
       SCREAMING_SNAKE_CASE__					:					str      =		self.get_dummy_components()
       SCREAMING_SNAKE_CASE__					:					str      =		self.pipeline_class(**a_					)
       pipe.to(a_					)

       SCREAMING_SNAKE_CASE__					:					List[Any]      =		10.0
       SCREAMING_SNAKE_CASE__					:					Any      =		4

       SCREAMING_SNAKE_CASE__					:					Optional[int]      =		self.get_dummy_inputs(a_					)
       SCREAMING_SNAKE_CASE__					:					Any      =		steps
       SCREAMING_SNAKE_CASE__					:					Optional[int]      =		scale
       SCREAMING_SNAKE_CASE__					:					Optional[int]      =		pipe(**a_					)[0]

       SCREAMING_SNAKE_CASE__					:					Tuple      =		self.get_dummy_inputs(a_					)
       SCREAMING_SNAKE_CASE__					:					str      =		steps
       SCREAMING_SNAKE_CASE__					:					int      =		scale
       SCREAMING_SNAKE_CASE__					:					Optional[int]      =		pipe(**a_  ,			control_guidance_start=0.1  ,			control_guidance_end=0.2					)[0]

       SCREAMING_SNAKE_CASE__					:					Any      =		self.get_dummy_inputs(a_					)
       SCREAMING_SNAKE_CASE__					:					Optional[int]      =		steps
       SCREAMING_SNAKE_CASE__					:					Optional[Any]      =		scale
       SCREAMING_SNAKE_CASE__					:					List[Any]      =		pipe(**a_  ,			control_guidance_start=[0.1, 0.3]  ,			control_guidance_end=[0.2, 0.7]					)[0]

       SCREAMING_SNAKE_CASE__					:					Optional[int]      =		self.get_dummy_inputs(a_					)
       SCREAMING_SNAKE_CASE__					:					Dict      =		steps
       SCREAMING_SNAKE_CASE__					:					List[Any]      =		scale
       SCREAMING_SNAKE_CASE__					:					str      =		pipe(**a_  ,			control_guidance_start=0.4  ,			control_guidance_end=[0.5, 0.8]					)[0]

       # make sure that all outputs are different
       assert np.sum(np.abs(output_a - output_a					)					) > 1e-3
       assert np.sum(np.abs(output_a - output_a					)					) > 1e-3
       assert np.sum(np.abs(output_a - output_a					)					) > 1e-3




  def 						__lowercase(   self  :						Any					)->							int:




       """simple docstring"""
       return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3					)




  @unittest.skipIf(
      torch_device != 'cuda' or not is_xformers_available()  ,			reason='XFormers attention is only available with CUDA and `xformers` installed'  ,			)
  def 						__lowercase(   self  :						List[str]					)->							int:




       """simple docstring"""
       self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3					)




  def 						__lowercase(   self  :						int					)->							List[Any]:




       """simple docstring"""
       self._test_inference_batch_single_identical(expected_max_diff=2e-3					)




  def 						__lowercase(   self  :						Optional[int]					)->							Union[str, Any]:




       """simple docstring"""
       SCREAMING_SNAKE_CASE__					:					str      =		self.get_dummy_components()
       SCREAMING_SNAKE_CASE__					:					str      =		self.pipeline_class(**a_					)
       pipe.to(a_					)
       pipe.set_progress_bar_config(disable=a_					)
       with tempfile.TemporaryDirectory() as tmpdir:
            try:
                 # save_pretrained is not implemented for Multi-ControlNet
                 pipe.save_pretrained(a_					)
            except NotImplementedError:
                 pass




@slow
@require_torch_gpu
class    snake_case       (   unittest.TestCase    ):




  def 						__lowercase(   self  :						Dict					)->							Dict:




       """simple docstring"""
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()




  def 						__lowercase(   self  :						Optional[Any]					)->							int:




       """simple docstring"""
       SCREAMING_SNAKE_CASE__					:					Union[str, Any]      =		ControlNetModel.from_pretrained('lllyasviel/sd-controlnet-canny'					)

       SCREAMING_SNAKE_CASE__					:					Optional[Any]      =		StableDiffusionControlNetImgaImgPipeline.from_pretrained(
           'runwayml/stable-diffusion-v1-5'  ,			safety_checker=a_  ,			controlnet=a_					)
       pipe.enable_model_cpu_offload()
       pipe.set_progress_bar_config(disable=a_					)

       SCREAMING_SNAKE_CASE__					:					List[Any]      =		torch.Generator(device='cpu'					).manual_seed(0					)
       SCREAMING_SNAKE_CASE__					:					Optional[int]      =		'evil space-punk bird'
       SCREAMING_SNAKE_CASE__					:					int      =		load_image(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png'					).resize((512, 512)					)
       SCREAMING_SNAKE_CASE__					:					Optional[int]      =		load_image(
           'https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird.png'					).resize((512, 512)					)

       SCREAMING_SNAKE_CASE__					:					Optional[Any]      =		pipe(
           a_  ,			a_  ,			control_image=a_  ,			generator=a_  ,			output_type='np'  ,			num_inference_steps=50  ,			strength=0.6  ,			)

       SCREAMING_SNAKE_CASE__					:					Optional[int]      =		output.images[0]

       assert image.shape == (512, 512, 3)

       SCREAMING_SNAKE_CASE__					:					Optional[int]      =		load_numpy(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/img2img.npy'					)

       assert np.abs(expected_image - image					).max() < 9e-2







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







@require_flax
class 	__lowercase    :





							"""simple docstring"""







							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__       ):





										"""simple docstring"""




										pass






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										pass






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										pass






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					np.abs((a - b)       ).max()
										self.assertLessEqual(lowerCAmelCase__			,						lowerCAmelCase__			,						F'''Difference between torch and flax is {diff} (>= {tol}).'''       )






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__=None			,						**lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Tuple					=					VisionTextDualEncoderConfig.from_vision_text_configs(lowerCAmelCase__			,						lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     List[str]					=					FlaxVisionTextDualEncoderModel(lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Tuple					=					model(input_ids=lowerCAmelCase__			,						pixel_values=lowerCAmelCase__			,						attention_mask=lowerCAmelCase__       )

										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__       ):





										"""simple docstring"""




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

										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					model(input_ids=lowerCAmelCase__			,						pixel_values=lowerCAmelCase__			,						attention_mask=lowerCAmelCase__       )

										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__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     int					=					self.get_vision_text_model(lowerCAmelCase__			,						lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Dict					=					{'vision_model': vision_model, 'text_model': text_model}
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					model(input_ids=lowerCAmelCase__			,						pixel_values=lowerCAmelCase__			,						attention_mask=lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					output[0]

										with tempfile.TemporaryDirectory() as tmpdirname:
													model.save_pretrained(lowerCAmelCase__       )
													SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					FlaxVisionTextDualEncoderModel.from_pretrained(lowerCAmelCase__       )

													SCREAMING_SNAKE_CASE_		:     Tuple					=					model(input_ids=lowerCAmelCase__			,						pixel_values=lowerCAmelCase__			,						attention_mask=lowerCAmelCase__       )
													SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					after_output[0]
													SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					np.amax(np.abs(out_a - out_a       )       )
													self.assertLessEqual(lowerCAmelCase__			,						1E-3       )






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__=None			,						**lowerCAmelCase__       ):





										"""simple docstring"""




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

										SCREAMING_SNAKE_CASE_		:     List[str]					=					model(
										    input_ids=lowerCAmelCase__			,						pixel_values=lowerCAmelCase__			,						attention_mask=lowerCAmelCase__			,						output_attentions=lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     List[str]					=					output.vision_model_output.attentions
										self.assertEqual(len(lowerCAmelCase__       )			,						vision_config.num_hidden_layers       )

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

										SCREAMING_SNAKE_CASE_		:     Tuple					=					output.text_model_output.attentions
										self.assertEqual(len(lowerCAmelCase__       )			,						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__       ):





										"""simple docstring"""




										pt_model.to(lowerCAmelCase__       )
										pt_model.eval()

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

										with torch.no_grad():
													SCREAMING_SNAKE_CASE_		:     List[Any]					=					pt_model(**lowerCAmelCase__       ).to_tuple()

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

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

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

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

										pt_model_loaded.to(lowerCAmelCase__       )
										pt_model_loaded.eval()

										with torch.no_grad():
													SCREAMING_SNAKE_CASE_		:     Optional[int]					=					pt_model_loaded(**lowerCAmelCase__       ).to_tuple()

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






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     List[Any]					=					VisionTextDualEncoderConfig.from_vision_text_configs(lowerCAmelCase__			,						lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					VisionTextDualEncoderModel(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					FlaxVisionTextDualEncoderModel(lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     str					=					convert_pytorch_state_dict_to_flax(pt_model.state_dict()			,						lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					fx_state

										self.check_pt_flax_equivalence(lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					VisionTextDualEncoderConfig.from_vision_text_configs(lowerCAmelCase__			,						lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					VisionTextDualEncoderModel(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     str					=					FlaxVisionTextDualEncoderModel(lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     int					=					load_flax_weights_in_pytorch_model(lowerCAmelCase__			,						fx_model.params       )

										self.check_pt_flax_equivalence(lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     str					=					self.prepare_config_and_inputs()
										self.check_model_from_pretrained_configs(**lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Dict					=					self.prepare_config_and_inputs()
										self.check_vision_text_dual_encoder_from_pretrained(**lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Tuple					=					self.prepare_config_and_inputs()
										self.check_save_load(**lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					self.prepare_config_and_inputs()
										self.check_vision_text_output_attention(**lowerCAmelCase__       )






							@is_pt_flax_cross_test
							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					self.prepare_config_and_inputs()
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					config_inputs_dict.pop('vision_config'       )
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					config_inputs_dict.pop('text_config'       )

										SCREAMING_SNAKE_CASE_		:     Any					=					config_inputs_dict

										self.check_equivalence_pt_to_flax(lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       )
										self.check_equivalence_flax_to_pt(lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       )






							@slow
							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     Dict					=					self.get_pretrained_model_and_inputs()

										SCREAMING_SNAKE_CASE_		:     int					=					model_a(**lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					outputs[0]

										with tempfile.TemporaryDirectory() as tmp_dirname:
													model_a.save_pretrained(lowerCAmelCase__       )
													SCREAMING_SNAKE_CASE_		:     int					=					FlaxVisionTextDualEncoderModel.from_pretrained(lowerCAmelCase__       )

													SCREAMING_SNAKE_CASE_		:     str					=					model_a(**lowerCAmelCase__       )
													SCREAMING_SNAKE_CASE_		:     Any					=					after_outputs[0]
													SCREAMING_SNAKE_CASE_		:     str					=					np.amax(np.abs(out_a - out_a       )       )
													self.assertLessEqual(lowerCAmelCase__			,						1E-5       )







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





							"""simple docstring"""







							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




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

										return model, inputs






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     int					=					FlaxViTModel(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					FlaxBertModel(lowerCAmelCase__       )
										return vision_model, text_model






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     int					=					FlaxViTModelTester(self       )
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					FlaxBertModelTester(self       )
										SCREAMING_SNAKE_CASE_		:     str					=					vit_model_tester.prepare_config_and_inputs()
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					bert_model_tester.prepare_config_and_inputs()

										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     Any					=					vision_config_and_inputs

										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     List[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 	__lowercase    (__SCREAMING_SNAKE_CASE	,       unittest.TestCase  ):





							"""simple docstring"""







							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




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

										return model, inputs






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Dict					=					FlaxCLIPVisionModel(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Any					=					FlaxBertModel(lowerCAmelCase__       )
										return vision_model, text_model






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					FlaxCLIPVisionModelTester(self       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					FlaxBertModelTester(self       )
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					clip_model_tester.prepare_config_and_inputs()
										SCREAMING_SNAKE_CASE_		:     int					=					bert_model_tester.prepare_config_and_inputs()

										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     int					=					vision_config_and_inputs

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





							"""simple docstring"""







							@slow
							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     List[str]					=					FlaxVisionTextDualEncoderModel.from_pretrained('clip-italian/clip-italian'			,						logit_scale_init_value=1.0       )
										SCREAMING_SNAKE_CASE_		:     str					=					VisionTextDualEncoderProcessor.from_pretrained('clip-italian/clip-italian'       )

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

										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					model(**lowerCAmelCase__       )

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

										SCREAMING_SNAKE_CASE_		:     Tuple					=					np.array([[1.2_284_727, 0.3_104_122]]       )

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





