
"""simple docstring"""





import json
import os
import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers import CLIPTokenizer, CLIPTokenizerFast
from transformers.models.clip.tokenization_clip import VOCAB_FILES_NAMES
from transformers.testing_utils import require_vision
from transformers.utils import IMAGE_PROCESSOR_NAME, is_vision_available


if is_vision_available():
				from PIL import Image

				from transformers import CLIPSegProcessor, ViTImageProcessor






@require_vision
class        _UpperCAmelCase(       unittest.TestCase      ):



		def 		UpperCAmelCase				(   self)			->			Dict:






							'''simple docstring'''
							_UpperCamelCase									=    tempfile.mkdtemp()

							# fmt: off
							_UpperCamelCase									=    ['''l''', '''o''', '''w''', '''e''', '''r''', '''s''', '''t''', '''i''', '''d''', '''n''', '''lo''', '''l</w>''', '''w</w>''', '''r</w>''', '''t</w>''', '''low</w>''', '''er</w>''', '''lowest</w>''', '''newer</w>''', '''wider''', '''<unk>''', '''<|startoftext|>''', '''<|endoftext|>''']
							# fmt: on
							_UpperCamelCase									=    dict(zip(A__							,						range(len(A__))))
							_UpperCamelCase									=    ['''#version: 0.2''', '''l o''', '''lo w</w>''', '''e r</w>''', '''''']
							_UpperCamelCase									=    {'''unk_token''': '''<unk>'''}

							_UpperCamelCase									=    os.path.join(self.tmpdirname							,						VOCAB_FILES_NAMES['''vocab_file'''])
							_UpperCamelCase									=    os.path.join(self.tmpdirname							,						VOCAB_FILES_NAMES['''merges_file'''])
							with open(self.vocab_file							,						'''w'''							,						encoding='''utf-8''') as fp:
												fp.write(json.dumps(A__) + '''\n''')
							with open(self.merges_file							,						'''w'''							,						encoding='''utf-8''') as fp:
												fp.write('''\n'''.join(A__))

							_UpperCamelCase									=    {
							    '''do_resize''': True,
							    '''size''': 20,
							    '''do_center_crop''': True,
							    '''crop_size''': 18,
							    '''do_normalize''': True,
							    '''image_mean''': [0.4814_5466, 0.457_8275, 0.4082_1073],
							    '''image_std''': [0.2686_2954, 0.2613_0258, 0.2757_7711],
							}
							_UpperCamelCase									=    os.path.join(self.tmpdirname							,						A__)
							with open(self.image_processor_file							,						'''w'''							,						encoding='''utf-8''') as fp:
												json.dump(A__							,						A__)



		def 		UpperCAmelCase				(   self							,						**__a)			->			Dict:






							'''simple docstring'''
							return CLIPTokenizer.from_pretrained(self.tmpdirname							,						**A__)



		def 		UpperCAmelCase				(   self							,						**__a)			->			int:






							'''simple docstring'''
							return CLIPTokenizerFast.from_pretrained(self.tmpdirname							,						**A__)



		def 		UpperCAmelCase				(   self							,						**__a)			->			List[str]:






							'''simple docstring'''
							return ViTImageProcessor.from_pretrained(self.tmpdirname							,						**A__)



		def 		UpperCAmelCase				(   self)			->			Any:






							'''simple docstring'''
							shutil.rmtree(self.tmpdirname)



		def 		UpperCAmelCase				(   self)			->			Dict:






							'''simple docstring'''
							_UpperCamelCase									=    [np.random.randint(2_55							,						size=(3, 30, 4_00)							,						dtype=np.uinta)]

							_UpperCamelCase									=    [Image.fromarray(np.moveaxis(A__							,						0							,						-1)) for x in image_inputs]

							return image_inputs



		def 		UpperCAmelCase				(   self)			->			Any:






							'''simple docstring'''
							_UpperCamelCase									=    self.get_tokenizer()
							_UpperCamelCase									=    self.get_rust_tokenizer()
							_UpperCamelCase									=    self.get_image_processor()

							_UpperCamelCase									=    CLIPSegProcessor(tokenizer=A__							,						image_processor=A__)
							processor_slow.save_pretrained(self.tmpdirname)
							_UpperCamelCase									=    CLIPSegProcessor.from_pretrained(self.tmpdirname							,						use_fast=A__)

							_UpperCamelCase									=    CLIPSegProcessor(tokenizer=A__							,						image_processor=A__)
							processor_fast.save_pretrained(self.tmpdirname)
							_UpperCamelCase									=    CLIPSegProcessor.from_pretrained(self.tmpdirname)

							self.assertEqual(processor_slow.tokenizer.get_vocab()							,						tokenizer_slow.get_vocab())
							self.assertEqual(processor_fast.tokenizer.get_vocab()							,						tokenizer_fast.get_vocab())
							self.assertEqual(tokenizer_slow.get_vocab()							,						tokenizer_fast.get_vocab())
							self.assertIsInstance(processor_slow.tokenizer							,						A__)
							self.assertIsInstance(processor_fast.tokenizer							,						A__)

							self.assertEqual(processor_slow.image_processor.to_json_string()							,						image_processor.to_json_string())
							self.assertEqual(processor_fast.image_processor.to_json_string()							,						image_processor.to_json_string())
							self.assertIsInstance(processor_slow.image_processor							,						A__)
							self.assertIsInstance(processor_fast.image_processor							,						A__)



		def 		UpperCAmelCase				(   self)			->			int:






							'''simple docstring'''
							_UpperCamelCase									=    CLIPSegProcessor(tokenizer=self.get_tokenizer()							,						image_processor=self.get_image_processor())
							processor.save_pretrained(self.tmpdirname)

							_UpperCamelCase									=    self.get_tokenizer(bos_token='''(BOS)'''							,						eos_token='''(EOS)''')
							_UpperCamelCase									=    self.get_image_processor(do_normalize=A__							,						padding_value=1.0)

							_UpperCamelCase									=    CLIPSegProcessor.from_pretrained(
							    self.tmpdirname							,						bos_token='''(BOS)'''							,						eos_token='''(EOS)'''							,						do_normalize=A__							,						padding_value=1.0)

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

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



		def 		UpperCAmelCase				(   self)			->			Any:






							'''simple docstring'''
							_UpperCamelCase									=    self.get_image_processor()
							_UpperCamelCase									=    self.get_tokenizer()

							_UpperCamelCase									=    CLIPSegProcessor(tokenizer=A__							,						image_processor=A__)

							_UpperCamelCase									=    self.prepare_image_inputs()

							_UpperCamelCase									=    image_processor(A__							,						return_tensors='''np''')
							_UpperCamelCase									=    processor(images=A__							,						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 		UpperCAmelCase				(   self)			->			Any:






							'''simple docstring'''
							_UpperCamelCase									=    self.get_image_processor()
							_UpperCamelCase									=    self.get_tokenizer()

							_UpperCamelCase									=    CLIPSegProcessor(tokenizer=A__							,						image_processor=A__)

							_UpperCamelCase									=    '''lower newer'''

							_UpperCamelCase									=    processor(text=A__)

							_UpperCamelCase									=    tokenizer(A__)

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



		def 		UpperCAmelCase				(   self)			->			str:






							'''simple docstring'''
							_UpperCamelCase									=    self.get_image_processor()
							_UpperCamelCase									=    self.get_tokenizer()

							_UpperCamelCase									=    CLIPSegProcessor(tokenizer=A__							,						image_processor=A__)

							_UpperCamelCase									=    '''lower newer'''
							_UpperCamelCase									=    self.prepare_image_inputs()

							_UpperCamelCase									=    processor(text=A__							,						images=A__)

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

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



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






							'''simple docstring'''
							_UpperCamelCase									=    self.get_image_processor()
							_UpperCamelCase									=    self.get_tokenizer()

							_UpperCamelCase									=    CLIPSegProcessor(tokenizer=A__							,						image_processor=A__)

							_UpperCamelCase									=    self.prepare_image_inputs()
							_UpperCamelCase									=    self.prepare_image_inputs()

							_UpperCamelCase									=    processor(images=A__							,						visual_prompt=A__)

							self.assertListEqual(list(inputs.keys())							,						['''pixel_values''', '''conditional_pixel_values'''])

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




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






							'''simple docstring'''
							_UpperCamelCase									=    self.get_image_processor()
							_UpperCamelCase									=    self.get_tokenizer()

							_UpperCamelCase									=    CLIPSegProcessor(tokenizer=A__							,						image_processor=A__)

							_UpperCamelCase									=    [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

							_UpperCamelCase									=    processor.batch_decode(A__)
							_UpperCamelCase									=    tokenizer.batch_decode(A__)

							self.assertListEqual(A__							,						A__)














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

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






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







		def  A__      (					self		):
						_A	:				Optional[Any]   =      []
						_A	:				str   =      []

						for i in range(self.num_layers		):
										_A	:				Union[str, Any]   =      self.in_channels if i == 0 else self.out_channels

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

										_A	:				List[str]   =      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(A__		)

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

						if self.add_downsample:
										_A	:				Any   =      FlaxDownsampleaD(self.out_channels       ,dtype=self.dtype		)







		def __call__(					self       ,A__       ,A__       ,A__       ,A__=True		):
						_A	:				List[str]   =      ()

						for resnet, attn in zip(self.resnets       ,self.attentions		):
										_A	:				Optional[Any]   =      resnet(A__       ,A__       ,deterministic=A__		)
										_A	:				Tuple   =      attn(A__       ,A__       ,deterministic=A__		)
										output_states += (hidden_states,)

						if self.add_downsample:
										_A	:				Optional[int]   =      self.downsamplers_a(A__		)
										output_states += (hidden_states,)

						return hidden_states, output_states






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







		def  A__      (					self		):
						_A	:				List[Any]   =      []

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

										_A	:				List[Any]   =      FlaxResnetBlockaD(
										    in_channels=A__       ,out_channels=self.out_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
										resnets.append(A__		)
						_A	:				Tuple   =      resnets

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







		def __call__(					self       ,A__       ,A__       ,A__=True		):
						_A	:				List[Any]   =      ()

						for resnet in self.resnets:
										_A	:				Optional[int]   =      resnet(A__       ,A__       ,deterministic=A__		)
										output_states += (hidden_states,)

						if self.add_downsample:
										_A	:				List[Any]   =      self.downsamplers_a(A__		)
										output_states += (hidden_states,)

						return hidden_states, output_states






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







		def  A__      (					self		):
						_A	:				List[str]   =      []
						_A	:				Optional[Any]   =      []

						for i in range(self.num_layers		):
										_A	:				Dict   =      self.in_channels if (i == self.num_layers - 1) else self.out_channels
										_A	:				List[str]   =      self.prev_output_channel if i == 0 else self.out_channels

										_A	:				Optional[int]   =      FlaxResnetBlockaD(
										    in_channels=resnet_in_channels + res_skip_channels       ,out_channels=self.out_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
										resnets.append(A__		)

										_A	:				Union[str, Any]   =      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(A__		)

						_A	:				Dict   =      resnets
						_A	:				int   =      attentions

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







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

										_A	:				str   =      resnet(A__       ,A__       ,deterministic=A__		)
										_A	:				Optional[int]   =      attn(A__       ,A__       ,deterministic=A__		)

						if self.add_upsample:
										_A	:				Union[str, Any]   =      self.upsamplers_a(A__		)

						return hidden_states






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







		def  A__      (					self		):
						_A	:				Optional[Any]   =      []

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

										_A	:				List[Any]   =      FlaxResnetBlockaD(
										    in_channels=resnet_in_channels + res_skip_channels       ,out_channels=self.out_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
										resnets.append(A__		)

						_A	:				int   =      resnets

						if self.add_upsample:
										_A	:				List[str]   =      FlaxUpsampleaD(self.out_channels       ,dtype=self.dtype		)







		def __call__(					self       ,A__       ,A__       ,A__       ,A__=True		):
						for resnet in self.resnets:
										# pop res hidden states
										_A	:				Tuple   =      res_hidden_states_tuple[-1]
										_A	:				Optional[int]   =      res_hidden_states_tuple[:-1]
										_A	:				Any   =      jnp.concatenate((hidden_states, res_hidden_states)       ,axis=-1		)

										_A	:				Optional[Any]   =      resnet(A__       ,A__       ,deterministic=A__		)

						if self.add_upsample:
										_A	:				Tuple   =      self.upsamplers_a(A__		)

						return hidden_states






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







		def  A__      (					self		):
						# there is always at least one resnet
						_A	:				List[str]   =      [
						    FlaxResnetBlockaD(
						        in_channels=self.in_channels       ,out_channels=self.in_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
						]

						_A	:				Union[str, Any]   =      []

						for _ in range(self.num_layers		):
										_A	:				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(A__		)

										_A	:				Any   =      FlaxResnetBlockaD(
										    in_channels=self.in_channels       ,out_channels=self.in_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
										resnets.append(A__		)

						_A	:				List[str]   =      resnets
						_A	:				str   =      attentions







		def __call__(					self       ,A__       ,A__       ,A__       ,A__=True		):
						_A	:				Optional[int]   =      self.resnets[0](A__       ,A__		)
						for attn, resnet in zip(self.attentions       ,self.resnets[1:]		):
										_A	:				Dict   =      attn(A__       ,A__       ,deterministic=A__		)
										_A	:				Any   =      resnet(A__       ,A__       ,deterministic=A__		)

						return hidden_states



