


"""simple docstring"""



import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers.testing_utils import require_vision
from transformers.utils import is_vision_available


if is_vision_available():
  from PIL import Image

  from transformers import AutoProcessor, BertTokenizer, BlipImageProcessor, BlipProcessor, PreTrainedTokenizerFast







@require_vision
class _lowerCAmelCase       (   unittest.TestCase			):







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

       '''simple docstring'''




       snake_case					:   int         =					tempfile.mkdtemp()

       snake_case					:   Union[str, Any]         =					BlipImageProcessor()
       snake_case					:   int         =					BertTokenizer.from_pretrained("hf-internal-testing/tiny-random-BertModel"		)

       snake_case					:   int         =					BlipProcessor(UpperCamelCase__		,			UpperCamelCase__		)

       processor.save_pretrained(self.tmpdirname		)







   def 				lowerCamelCase	(			self		,			**UpperCamelCase__		)			->	Tuple:

       '''simple docstring'''




       return AutoProcessor.from_pretrained(self.tmpdirname		,			**UpperCamelCase__		).tokenizer







   def 				lowerCamelCase	(			self		,			**UpperCamelCase__		)			->	str:

       '''simple docstring'''




       return AutoProcessor.from_pretrained(self.tmpdirname		,			**UpperCamelCase__		).image_processor







   def 				lowerCamelCase	(			self		)			->	Any:

       '''simple docstring'''




       shutil.rmtree(self.tmpdirname		)







   def 				lowerCamelCase	(			self		)			->	str:

       '''simple docstring'''




       snake_case					:   Union[str, Any]         =					[np.random.randint(255		,			size=(3, 30, 400)		,			dtype=np.uinta		)]

       snake_case					:   List[Any]         =					[Image.fromarray(np.moveaxis(UpperCamelCase__		,			0		,			-1		)		) for x in image_inputs]

       return image_inputs







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

       '''simple docstring'''




       snake_case					:   List[Any]         =					BlipProcessor(tokenizer=self.get_tokenizer()		,			image_processor=self.get_image_processor()		)
       processor.save_pretrained(self.tmpdirname		)

       snake_case					:   str         =					self.get_tokenizer(bos_token="(BOS)"		,			eos_token="(EOS)"		)
       snake_case					:   Tuple         =					self.get_image_processor(do_normalize=UpperCamelCase__		,			padding_value=1.0		)

       snake_case					:   Optional[int]         =					BlipProcessor.from_pretrained(
           self.tmpdirname		,			bos_token="(BOS)"		,			eos_token="(EOS)"		,			do_normalize=UpperCamelCase__		,			padding_value=1.0		)

       self.assertEqual(processor.tokenizer.get_vocab()		,			tokenizer_add_kwargs.get_vocab()		)
       self.assertIsInstance(processor.tokenizer		,			UpperCamelCase__		)

       self.assertEqual(processor.image_processor.to_json_string()		,			image_processor_add_kwargs.to_json_string()		)
       self.assertIsInstance(processor.image_processor		,			UpperCamelCase__		)







   def 				lowerCamelCase	(			self		)			->	Any:

       '''simple docstring'''




       snake_case					:   Optional[Any]         =					self.get_image_processor()
       snake_case					:   List[str]         =					self.get_tokenizer()

       snake_case					:   List[Any]         =					BlipProcessor(tokenizer=UpperCamelCase__		,			image_processor=UpperCamelCase__		)

       snake_case					:   Union[str, Any]         =					self.prepare_image_inputs()

       snake_case					:   int         =					image_processor(UpperCamelCase__		,			return_tensors="np"		)
       snake_case					:   Tuple         =					processor(images=UpperCamelCase__		,			return_tensors="np"		)

       for key in input_feat_extract.keys():
           self.assertAlmostEqual(input_feat_extract[key].sum()		,			input_processor[key].sum()		,			delta=1e-2		)







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

       '''simple docstring'''




       snake_case					:   Union[str, Any]         =					self.get_image_processor()
       snake_case					:   str         =					self.get_tokenizer()

       snake_case					:   List[str]         =					BlipProcessor(tokenizer=UpperCamelCase__		,			image_processor=UpperCamelCase__		)

       snake_case					:   Tuple         =					"lower newer"

       snake_case					:   Union[str, Any]         =					processor(text=UpperCamelCase__		)

       snake_case					:   str         =					tokenizer(UpperCamelCase__		,			return_token_type_ids=UpperCamelCase__		)

       for key in encoded_tok.keys():
           self.assertListEqual(encoded_tok[key]		,			encoded_processor[key]		)







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

       '''simple docstring'''




       snake_case					:   Dict         =					self.get_image_processor()
       snake_case					:   Any         =					self.get_tokenizer()

       snake_case					:   str         =					BlipProcessor(tokenizer=UpperCamelCase__		,			image_processor=UpperCamelCase__		)

       snake_case					:   Any         =					"lower newer"
       snake_case					:   Optional[int]         =					self.prepare_image_inputs()

       snake_case					:   Dict         =					processor(text=UpperCamelCase__		,			images=UpperCamelCase__		)

       self.assertListEqual(list(inputs.keys()		)		,			["pixel_values", "input_ids", "attention_mask"]		)

       # test if it raises when no input is passed
       with pytest.raises(UpperCamelCase__		):
           processor()







   def 				lowerCamelCase	(			self		)			->	Any:

       '''simple docstring'''




       snake_case					:   Optional[Any]         =					self.get_image_processor()
       snake_case					:   Optional[int]         =					self.get_tokenizer()

       snake_case					:   Any         =					BlipProcessor(tokenizer=UpperCamelCase__		,			image_processor=UpperCamelCase__		)

       snake_case					:   List[Any]         =					[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

       snake_case					:   List[str]         =					processor.batch_decode(UpperCamelCase__		)
       snake_case					:   List[str]         =					tokenizer.batch_decode(UpperCamelCase__		)

       self.assertListEqual(UpperCamelCase__		,			UpperCamelCase__		)







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

       '''simple docstring'''




       snake_case					:   List[Any]         =					self.get_image_processor()
       snake_case					:   List[Any]         =					self.get_tokenizer()

       snake_case					:   str         =					BlipProcessor(tokenizer=UpperCamelCase__		,			image_processor=UpperCamelCase__		)

       snake_case					:   str         =					"lower newer"
       snake_case					:   Optional[int]         =					self.prepare_image_inputs()

       snake_case					:   Optional[int]         =					processor(text=UpperCamelCase__		,			images=UpperCamelCase__		)

       # For now the processor supports only ['pixel_values', 'input_ids', 'attention_mask']
       self.assertListEqual(list(inputs.keys()		)		,			["pixel_values", "input_ids", "attention_mask"]		)









"""simple docstring"""



import flax.linen as nn
import jax.numpy as jnp

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD







class _lowerCAmelCase       (   nn.Module			):
   __UpperCAmelCase				: int
   __UpperCAmelCase				: int
   __UpperCAmelCase				: float													=						0.0
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: bool													=						True
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: jnp.dtype													=						jnp.floataa







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

       '''simple docstring'''




       snake_case					:   Optional[int]         =					[]
       snake_case					:   Tuple         =					[]

       for i in range(self.num_layers		):
           snake_case					:   Any         =					self.in_channels if i == 0 else self.out_channels

           snake_case					:   Optional[int]         =					FlaxResnetBlockaD(
               in_channels=UpperCamelCase__		,			out_channels=self.out_channels		,			dropout_prob=self.dropout		,			dtype=self.dtype		,			)
           resnets.append(UpperCamelCase__		)

           snake_case					:   Dict         =					FlaxTransformeraDModel(
               in_channels=self.out_channels		,			n_heads=self.num_attention_heads		,			d_head=self.out_channels // self.num_attention_heads		,			depth=1		,			use_linear_projection=self.use_linear_projection		,			only_cross_attention=self.only_cross_attention		,			use_memory_efficient_attention=self.use_memory_efficient_attention		,			dtype=self.dtype		,			)
           attentions.append(UpperCamelCase__		)

       snake_case					:   List[Any]         =					resnets
       snake_case					:   List[Any]         =					attentions

       if self.add_downsample:
           snake_case					:   Optional[Any]         =					FlaxDownsampleaD(self.out_channels		,			dtype=self.dtype		)







   def __call__(			self		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__=True		)			->	List[str]:

       '''simple docstring'''




       snake_case					:   Tuple         =					()

       for resnet, attn in zip(self.resnets		,			self.attentions		):
           snake_case					:   str         =					resnet(UpperCamelCase__		,			UpperCamelCase__		,			deterministic=UpperCamelCase__		)
           snake_case					:   int         =					attn(UpperCamelCase__		,			UpperCamelCase__		,			deterministic=UpperCamelCase__		)
           output_states += (hidden_states,)

       if self.add_downsample:
           snake_case					:   Any         =					self.downsamplers_a(UpperCamelCase__		)
           output_states += (hidden_states,)

       return hidden_states, output_states







class _lowerCAmelCase       (   nn.Module			):
   __UpperCAmelCase				: int
   __UpperCAmelCase				: int
   __UpperCAmelCase				: float													=						0.0
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: bool													=						True
   __UpperCAmelCase				: jnp.dtype													=						jnp.floataa







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

       '''simple docstring'''




       snake_case					:   Dict         =					[]

       for i in range(self.num_layers		):
           snake_case					:   List[str]         =					self.in_channels if i == 0 else self.out_channels

           snake_case					:   str         =					FlaxResnetBlockaD(
               in_channels=UpperCamelCase__		,			out_channels=self.out_channels		,			dropout_prob=self.dropout		,			dtype=self.dtype		,			)
           resnets.append(UpperCamelCase__		)
       snake_case					:   str         =					resnets

       if self.add_downsample:
           snake_case					:   str         =					FlaxDownsampleaD(self.out_channels		,			dtype=self.dtype		)







   def __call__(			self		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__=True		)			->	List[Any]:

       '''simple docstring'''




       snake_case					:   Tuple         =					()

       for resnet in self.resnets:
           snake_case					:   Optional[Any]         =					resnet(UpperCamelCase__		,			UpperCamelCase__		,			deterministic=UpperCamelCase__		)
           output_states += (hidden_states,)

       if self.add_downsample:
           snake_case					:   Tuple         =					self.downsamplers_a(UpperCamelCase__		)
           output_states += (hidden_states,)

       return hidden_states, output_states







class _lowerCAmelCase       (   nn.Module			):
   __UpperCAmelCase				: int
   __UpperCAmelCase				: int
   __UpperCAmelCase				: int
   __UpperCAmelCase				: float													=						0.0
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: bool													=						True
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: jnp.dtype													=						jnp.floataa







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

       '''simple docstring'''




       snake_case					:   Dict         =					[]
       snake_case					:   Union[str, Any]         =					[]

       for i in range(self.num_layers		):
           snake_case					:   Optional[Any]         =					self.in_channels if (i == self.num_layers - 1) else self.out_channels
           snake_case					:   Union[str, Any]         =					self.prev_output_channel if i == 0 else self.out_channels

           snake_case					:   Optional[Any]         =					FlaxResnetBlockaD(
               in_channels=resnet_in_channels + res_skip_channels		,			out_channels=self.out_channels		,			dropout_prob=self.dropout		,			dtype=self.dtype		,			)
           resnets.append(UpperCamelCase__		)

           snake_case					:   int         =					FlaxTransformeraDModel(
               in_channels=self.out_channels		,			n_heads=self.num_attention_heads		,			d_head=self.out_channels // self.num_attention_heads		,			depth=1		,			use_linear_projection=self.use_linear_projection		,			only_cross_attention=self.only_cross_attention		,			use_memory_efficient_attention=self.use_memory_efficient_attention		,			dtype=self.dtype		,			)
           attentions.append(UpperCamelCase__		)

       snake_case					:   str         =					resnets
       snake_case					:   List[Any]         =					attentions

       if self.add_upsample:
           snake_case					:   str         =					FlaxUpsampleaD(self.out_channels		,			dtype=self.dtype		)







   def __call__(			self		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__=True		)			->	Optional[int]:

       '''simple docstring'''




       for resnet, attn in zip(self.resnets		,			self.attentions		):
           # pop res hidden states
           snake_case					:   Dict         =					res_hidden_states_tuple[-1]
           snake_case					:   Optional[int]         =					res_hidden_states_tuple[:-1]
           snake_case					:   List[Any]         =					jnp.concatenate((hidden_states, res_hidden_states)		,			axis=-1		)

           snake_case					:   List[Any]         =					resnet(UpperCamelCase__		,			UpperCamelCase__		,			deterministic=UpperCamelCase__		)
           snake_case					:   Any         =					attn(UpperCamelCase__		,			UpperCamelCase__		,			deterministic=UpperCamelCase__		)

       if self.add_upsample:
           snake_case					:   int         =					self.upsamplers_a(UpperCamelCase__		)

       return hidden_states







class _lowerCAmelCase       (   nn.Module			):
   __UpperCAmelCase				: int
   __UpperCAmelCase				: int
   __UpperCAmelCase				: int
   __UpperCAmelCase				: float													=						0.0
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: bool													=						True
   __UpperCAmelCase				: jnp.dtype													=						jnp.floataa







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

       '''simple docstring'''




       snake_case					:   str         =					[]

       for i in range(self.num_layers		):
           snake_case					:   Union[str, Any]         =					self.in_channels if (i == self.num_layers - 1) else self.out_channels
           snake_case					:   Optional[Any]         =					self.prev_output_channel if i == 0 else self.out_channels

           snake_case					:   Dict         =					FlaxResnetBlockaD(
               in_channels=resnet_in_channels + res_skip_channels		,			out_channels=self.out_channels		,			dropout_prob=self.dropout		,			dtype=self.dtype		,			)
           resnets.append(UpperCamelCase__		)

       snake_case					:   Optional[int]         =					resnets

       if self.add_upsample:
           snake_case					:   Dict         =					FlaxUpsampleaD(self.out_channels		,			dtype=self.dtype		)







   def __call__(			self		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__=True		)			->	int:

       '''simple docstring'''




       for resnet in self.resnets:
           # pop res hidden states
           snake_case					:   int         =					res_hidden_states_tuple[-1]
           snake_case					:   List[Any]         =					res_hidden_states_tuple[:-1]
           snake_case					:   Optional[int]         =					jnp.concatenate((hidden_states, res_hidden_states)		,			axis=-1		)

           snake_case					:   Dict         =					resnet(UpperCamelCase__		,			UpperCamelCase__		,			deterministic=UpperCamelCase__		)

       if self.add_upsample:
           snake_case					:   Optional[int]         =					self.upsamplers_a(UpperCamelCase__		)

       return hidden_states







class _lowerCAmelCase       (   nn.Module			):
   __UpperCAmelCase				: int
   __UpperCAmelCase				: float													=						0.0
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: jnp.dtype													=						jnp.floataa







   def 				lowerCamelCase	(			self		)			->	Tuple:

       '''simple docstring'''




       snake_case					:   Dict         =					[
           FlaxResnetBlockaD(
               in_channels=self.in_channels		,			out_channels=self.in_channels		,			dropout_prob=self.dropout		,			dtype=self.dtype		,			)
       ]

       snake_case					:   Any         =					[]

       for _ in range(self.num_layers		):
           snake_case					:   Union[str, Any]         =					FlaxTransformeraDModel(
               in_channels=self.in_channels		,			n_heads=self.num_attention_heads		,			d_head=self.in_channels // self.num_attention_heads		,			depth=1		,			use_linear_projection=self.use_linear_projection		,			use_memory_efficient_attention=self.use_memory_efficient_attention		,			dtype=self.dtype		,			)
           attentions.append(UpperCamelCase__		)

           snake_case					:   List[Any]         =					FlaxResnetBlockaD(
               in_channels=self.in_channels		,			out_channels=self.in_channels		,			dropout_prob=self.dropout		,			dtype=self.dtype		,			)
           resnets.append(UpperCamelCase__		)

       snake_case					:   Union[str, Any]         =					resnets
       snake_case					:   str         =					attentions







   def __call__(			self		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__=True		)			->	Any:

       '''simple docstring'''




       snake_case					:   Optional[Any]         =					self.resnets[0](UpperCamelCase__		,			UpperCamelCase__		)
       for attn, resnet in zip(self.attentions		,			self.resnets[1:]		):
           snake_case					:   Union[str, Any]         =					attn(UpperCamelCase__		,			UpperCamelCase__		,			deterministic=UpperCamelCase__		)
           snake_case					:   Optional[Any]         =					resnet(UpperCamelCase__		,			UpperCamelCase__		,			deterministic=UpperCamelCase__		)

       return hidden_states





