



"""simple docstring"""







import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import MBartConfig, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, slow
from transformers.utils import cached_property

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
		import os

		# The slow tests are often failing with OOM error on GPU
		# This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
		# but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
		UpperCAmelCase:  int			     =				"""platform"""

		import jax
		import jax.numpy as jnp

		from transformers import AutoTokenizer
		from transformers.models.mbart.modeling_flax_mbart import (
		    FlaxMBartForConditionalGeneration,
		    FlaxMBartForQuestionAnswering,
		    FlaxMBartForSequenceClassification,
		    FlaxMBartModel,
		    shift_tokens_right,
		)


def  __SCREAMING_SNAKE_CASE (  __UpperCAmelCase   ,				__UpperCAmelCase   ,				__UpperCAmelCase=None   ,				__UpperCAmelCase=None   ,				__UpperCAmelCase=None   ,				__UpperCAmelCase=None   ,				__UpperCAmelCase=None   ,				__UpperCAmelCase=None   ,				):
		if attention_mask is None:
				_lowercase				: int     =   np.where(input_ids != config.pad_token_id   ,				1   ,				0      )
		if decoder_attention_mask is None:
				_lowercase				: Any     =   np.where(decoder_input_ids != config.pad_token_id   ,				1   ,				0      )
		if head_mask is None:
				_lowercase				: Tuple     =   np.ones((config.encoder_layers, config.encoder_attention_heads)      )
		if decoder_head_mask is None:
				_lowercase				: Tuple     =   np.ones((config.decoder_layers, config.decoder_attention_heads)      )
		if cross_attn_head_mask is None:
				_lowercase				: str     =   np.ones((config.decoder_layers, config.decoder_attention_heads)      )
		return {
		    "input_ids": input_ids,
		    "decoder_input_ids": decoder_input_ids,
		    "attention_mask": attention_mask,
		    "decoder_attention_mask": decoder_attention_mask,
		}





class   UpperCamelCase	(     unittest.TestCase       ):




	"""simple docstring"""






	def __init__(     self				,UpperCAmelCase_				,UpperCAmelCase_=13				,UpperCAmelCase_=7				,UpperCAmelCase_=True				,UpperCAmelCase_=False				,UpperCAmelCase_=99				,UpperCAmelCase_=16				,UpperCAmelCase_=2				,UpperCAmelCase_=4				,UpperCAmelCase_=4				,UpperCAmelCase_="gelu"				,UpperCAmelCase_=0.1				,UpperCAmelCase_=0.1				,UpperCAmelCase_=32				,UpperCAmelCase_=2				,UpperCAmelCase_=1				,UpperCAmelCase_=0				,UpperCAmelCase_=2				,UpperCAmelCase_=0.02				,):
			_lowercase				: Optional[int]     =   parent
			_lowercase				: Optional[Any]     =   batch_size
			_lowercase				: Optional[int]     =   seq_length
			_lowercase				: Tuple     =   is_training
			_lowercase				: Any     =   use_labels
			_lowercase				: str     =   vocab_size
			_lowercase				: Any     =   hidden_size
			_lowercase				: Dict     =   num_hidden_layers
			_lowercase				: Any     =   num_attention_heads
			_lowercase				: Optional[Any]     =   intermediate_size
			_lowercase				: Optional[int]     =   hidden_act
			_lowercase				: Optional[int]     =   hidden_dropout_prob
			_lowercase				: Optional[Any]     =   attention_probs_dropout_prob
			_lowercase				: Tuple     =   max_position_embeddings
			_lowercase				: int     =   eos_token_id
			_lowercase				: int     =   pad_token_id
			_lowercase				: Any     =   bos_token_id
			_lowercase				: List[str]     =   decoder_start_token_id
			_lowercase				: Tuple     =   initializer_range






	def    lowerCamelCase__					(     self   ):
			_lowercase				: str     =   np.clip(ids_tensor([self.batch_size, self.seq_length - 1]				,self.vocab_size   )				,3				,self.vocab_size   )
			_lowercase				: int     =   np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)				,dtype=np.intaa   ))				,-1   )

			_lowercase				: List[str]     =   shift_tokens_right(UpperCAmelCase_				,1   )

			_lowercase				: List[str]     =   MBartConfig(
			    vocab_size=self.vocab_size				,d_model=self.hidden_size				,encoder_layers=self.num_hidden_layers				,decoder_layers=self.num_hidden_layers				,encoder_attention_heads=self.num_attention_heads				,decoder_attention_heads=self.num_attention_heads				,encoder_ffn_dim=self.intermediate_size				,decoder_ffn_dim=self.intermediate_size				,dropout=self.hidden_dropout_prob				,attention_dropout=self.attention_probs_dropout_prob				,max_position_embeddings=self.max_position_embeddings				,eos_token_id=self.eos_token_id				,bos_token_id=self.bos_token_id				,pad_token_id=self.pad_token_id				,decoder_start_token_id=self.decoder_start_token_id				,initializer_range=self.initializer_range				,use_cache=UpperCAmelCase_				,)
			_lowercase				: Tuple     =   prepare_mbart_inputs_dict(UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_   )
			return config, inputs_dict






	def    lowerCamelCase__					(     self   ):
			_lowercase     ,  _lowercase				: Optional[Any]     =   self.prepare_config_and_inputs()
			return config, inputs_dict






	def    lowerCamelCase__					(     self				,UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_   ):
			_lowercase				: Optional[int]     =   20
			_lowercase				: Dict     =   model_class_name(UpperCAmelCase_   )

			_lowercase				: str     =   model.encode(inputs_dict["""input_ids"""]   )

			_lowercase     ,  _lowercase				: List[str]     =   (
			    inputs_dict["""decoder_input_ids"""],
			    inputs_dict["""decoder_attention_mask"""],
			)

			_lowercase				: Union[str, Any]     =   model.init_cache(decoder_input_ids.shape[0]				,UpperCAmelCase_				,UpperCAmelCase_   )
			_lowercase				: int     =   jnp.ones((decoder_input_ids.shape[0], max_decoder_length)				,dtype="""i4"""   )

			_lowercase				: Dict     =   jnp.broadcast_to(
			    jnp.arange(decoder_input_ids.shape[-1] - 1   )[None, :]				,(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)				,)
			_lowercase				: int     =   model.decode(
			    decoder_input_ids[:, :-1]				,UpperCAmelCase_				,decoder_attention_mask=UpperCAmelCase_				,past_key_values=UpperCAmelCase_				,decoder_position_ids=UpperCAmelCase_				,)

			_lowercase				: List[Any]     =   jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]				,dtype="""i4"""   )
			_lowercase				: Optional[Any]     =   model.decode(
			    decoder_input_ids[:, -1:]				,UpperCAmelCase_				,decoder_attention_mask=UpperCAmelCase_				,past_key_values=outputs_cache.past_key_values				,decoder_position_ids=UpperCAmelCase_				,)

			_lowercase				: str     =   model.decode(UpperCAmelCase_				,UpperCAmelCase_   )

			_lowercase				: List[str]     =   np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])   )   )
			self.parent.assertTrue(diff < 1E-3				,msg=f"""Max diff is {diff}"""   )






	def    lowerCamelCase__					(     self				,UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_   ):
			_lowercase				: Tuple     =   20
			_lowercase				: List[Any]     =   model_class_name(UpperCAmelCase_   )

			_lowercase				: Union[str, Any]     =   model.encode(inputs_dict["""input_ids"""]   )

			_lowercase     ,  _lowercase				: Any     =   (
			    inputs_dict["""decoder_input_ids"""],
			    inputs_dict["""decoder_attention_mask"""],
			)

			_lowercase				: List[str]     =   jnp.concatenate(
			    [
			        decoder_attention_mask,
			        jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])   ),
			    ]				,axis=-1				,)

			_lowercase				: Union[str, Any]     =   model.init_cache(decoder_input_ids.shape[0]				,UpperCAmelCase_				,UpperCAmelCase_   )
			_lowercase				: List[str]     =   jnp.broadcast_to(
			    jnp.arange(decoder_input_ids.shape[-1] - 1   )[None, :]				,(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)				,)

			_lowercase				: List[Any]     =   model.decode(
			    decoder_input_ids[:, :-1]				,UpperCAmelCase_				,decoder_attention_mask=UpperCAmelCase_				,past_key_values=UpperCAmelCase_				,decoder_position_ids=UpperCAmelCase_				,)
			_lowercase				: List[Any]     =   jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]				,dtype="""i4"""   )
			_lowercase				: Optional[Any]     =   model.decode(
			    decoder_input_ids[:, -1:]				,UpperCAmelCase_				,past_key_values=outputs_cache.past_key_values				,decoder_attention_mask=UpperCAmelCase_				,decoder_position_ids=UpperCAmelCase_				,)

			_lowercase				: Union[str, Any]     =   model.decode(UpperCAmelCase_				,UpperCAmelCase_				,decoder_attention_mask=UpperCAmelCase_   )

			_lowercase				: List[str]     =   np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])   )   )
			self.parent.assertTrue(diff < 1E-3				,msg=f"""Max diff is {diff}"""   )







@require_flax
class   UpperCamelCase	(     unittest.TestCase       ):




	"""simple docstring"""
	SCREAMING_SNAKE_CASE_							:       List[Any]						       =  9_9






	def    lowerCamelCase__					(     self   ):
			_lowercase				: List[str]     =   np.array(
			    [
			        [71, 82, 18, 33, 46, 91, 2],
			        [68, 34, 26, 58, 30, 82, 2],
			        [5, 97, 17, 39, 94, 40, 2],
			        [76, 83, 94, 25, 70, 78, 2],
			        [87, 59, 41, 35, 48, 66, 2],
			        [55, 13, 16, 58, 5, 2, 1],  # note padding
			        [64, 27, 31, 51, 12, 75, 2],
			        [52, 64, 86, 17, 83, 39, 2],
			        [48, 61, 9, 24, 71, 82, 2],
			        [26, 1, 60, 48, 22, 13, 2],
			        [21, 5, 62, 28, 14, 76, 2],
			        [45, 98, 37, 86, 59, 48, 2],
			        [70, 70, 50, 9, 28, 0, 2],
			    ]				,dtype=np.intaa				,)

			_lowercase				: List[str]     =   input_ids.shape[0]
			_lowercase				: List[Any]     =   MBartConfig(
			    vocab_size=self.vocab_size				,d_model=24				,encoder_layers=2				,decoder_layers=2				,encoder_attention_heads=2				,decoder_attention_heads=2				,encoder_ffn_dim=32				,decoder_ffn_dim=32				,max_position_embeddings=48				,eos_token_id=2				,pad_token_id=1				,bos_token_id=0				,)
			return config, input_ids, batch_size






	def    lowerCamelCase__					(     self   ):
			_lowercase     ,  _lowercase     ,  _lowercase				: int     =   self._get_config_and_data()
			_lowercase				: Optional[Any]     =   FlaxMBartForSequenceClassification(UpperCAmelCase_   )
			_lowercase				: Any     =   model(input_ids=UpperCAmelCase_				,decoder_input_ids=UpperCAmelCase_   )
			_lowercase				: Any     =   (batch_size, config.num_labels)
			self.assertEqual(outputs["""logits"""].shape				,UpperCAmelCase_   )






	def    lowerCamelCase__					(     self   ):
			_lowercase     ,  _lowercase     ,  _lowercase				: Optional[int]     =   self._get_config_and_data()
			_lowercase				: Tuple     =   FlaxMBartForQuestionAnswering(UpperCAmelCase_   )
			_lowercase				: Optional[int]     =   model(input_ids=UpperCAmelCase_   )

			self.assertEqual(outputs["""start_logits"""].shape				,input_ids.shape   )
			self.assertEqual(outputs["""end_logits"""].shape				,input_ids.shape   )






	def    lowerCamelCase__					(     self   ):
			_lowercase     ,  _lowercase     ,  _lowercase				: Any     =   self._get_config_and_data()
			_lowercase				: List[str]     =   FlaxMBartForConditionalGeneration(UpperCAmelCase_   )
			_lowercase				: Dict     =   lm_model(input_ids=UpperCAmelCase_   )
			_lowercase				: Any     =   (batch_size, input_ids.shape[1], config.vocab_size)
			self.assertEqual(outputs["""logits"""].shape				,UpperCAmelCase_   )






	def    lowerCamelCase__					(     self   ):
			_lowercase				: List[str]     =   MBartConfig(
			    vocab_size=self.vocab_size				,d_model=14				,encoder_layers=2				,decoder_layers=2				,encoder_attention_heads=2				,decoder_attention_heads=2				,encoder_ffn_dim=8				,decoder_ffn_dim=8				,max_position_embeddings=48				,)
			_lowercase				: Optional[int]     =   FlaxMBartForConditionalGeneration(UpperCAmelCase_   )
			_lowercase				: str     =   np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]				,dtype=np.intaa   )
			_lowercase				: Dict     =   np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]				,dtype=np.intaa   )
			_lowercase				: str     =   lm_model(input_ids=UpperCAmelCase_				,decoder_input_ids=UpperCAmelCase_   )
			_lowercase				: Optional[int]     =   (*summary.shape, config.vocab_size)
			self.assertEqual(outputs["""logits"""].shape				,UpperCAmelCase_   )






	def    lowerCamelCase__					(     self   ):
			_lowercase				: Any     =   np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]				,dtype=np.intaa   )
			_lowercase				: Any     =   shift_tokens_right(UpperCAmelCase_				,1   )
			_lowercase				: List[str]     =   np.equal(UpperCAmelCase_				,1   ).astype(np.floataa   ).sum()
			_lowercase				: Dict     =   np.equal(UpperCAmelCase_				,1   ).astype(np.floataa   ).sum()
			self.assertEqual(shifted.shape				,input_ids.shape   )
			self.assertEqual(UpperCAmelCase_				,n_pad_before - 1   )
			self.assertTrue(np.equal(shifted[:, 0]				,2   ).all()   )







@require_flax
class   UpperCamelCase	(     snake_case					,   unittest.TestCase					,   snake_case       ):




	"""simple docstring"""
	SCREAMING_SNAKE_CASE_							:       Dict						       =  True
	SCREAMING_SNAKE_CASE_							:       int						       =  (
	    (
	        FlaxMBartModel,
	        FlaxMBartForConditionalGeneration,
	        FlaxMBartForSequenceClassification,
	        FlaxMBartForQuestionAnswering,
	    )
	    if is_flax_available()
	    else ()
	)
	SCREAMING_SNAKE_CASE_							:       Dict						       =  (FlaxMBartForConditionalGeneration,) if is_flax_available() else ()






	def    lowerCamelCase__					(     self   ):
			_lowercase				: Any     =   FlaxMBartModelTester(self   )






	def    lowerCamelCase__					(     self   ):
			_lowercase     ,  _lowercase				: List[str]     =   self.model_tester.prepare_config_and_inputs()
			for model_class in self.all_model_classes:
					self.model_tester.check_use_cache_forward(UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_   )






	def    lowerCamelCase__					(     self   ):
			_lowercase     ,  _lowercase				: Optional[int]     =   self.model_tester.prepare_config_and_inputs()
			for model_class in self.all_model_classes:
					self.model_tester.check_use_cache_forward_with_attn_mask(UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_   )






	def    lowerCamelCase__					(     self   ):
			_lowercase     ,  _lowercase				: Dict     =   self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
					with self.subTest(model_class.__name__   ):
							_lowercase				: Any     =   self._prepare_for_class(UpperCAmelCase_				,UpperCAmelCase_   )
							_lowercase				: Union[str, Any]     =   model_class(UpperCAmelCase_   )

							@jax.jit
							def encode_jitted(UpperCAmelCase_				,UpperCAmelCase_=None				,**UpperCAmelCase_   ):
									return model.encode(input_ids=UpperCAmelCase_				,attention_mask=UpperCAmelCase_   )

							with self.subTest("""JIT Enabled"""   ):
									_lowercase				: str     =   encode_jitted(**UpperCAmelCase_   ).to_tuple()

							with self.subTest("""JIT Disabled"""   ):
									with jax.disable_jit():
											_lowercase				: Optional[Any]     =   encode_jitted(**UpperCAmelCase_   ).to_tuple()

							self.assertEqual(len(UpperCAmelCase_   )				,len(UpperCAmelCase_   )   )
							for jitted_output, output in zip(UpperCAmelCase_				,UpperCAmelCase_   ):
									self.assertEqual(jitted_output.shape				,output.shape   )






	def    lowerCamelCase__					(     self   ):
			_lowercase     ,  _lowercase				: Any     =   self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
					with self.subTest(model_class.__name__   ):
							_lowercase				: List[Any]     =   model_class(UpperCAmelCase_   )
							_lowercase				: List[Any]     =   model.encode(inputs_dict["""input_ids"""]				,inputs_dict["""attention_mask"""]   )

							_lowercase				: Optional[int]     =   {
							    """decoder_input_ids""": inputs_dict["""decoder_input_ids"""],
							    """decoder_attention_mask""": inputs_dict["""decoder_attention_mask"""],
							    """encoder_outputs""": encoder_outputs,
							}

							@jax.jit
							def decode_jitted(UpperCAmelCase_				,UpperCAmelCase_				,UpperCAmelCase_   ):
									return model.decode(
									    decoder_input_ids=UpperCAmelCase_				,decoder_attention_mask=UpperCAmelCase_				,encoder_outputs=UpperCAmelCase_				,)

							with self.subTest("""JIT Enabled"""   ):
									_lowercase				: List[Any]     =   decode_jitted(**UpperCAmelCase_   ).to_tuple()

							with self.subTest("""JIT Disabled"""   ):
									with jax.disable_jit():
											_lowercase				: int     =   decode_jitted(**UpperCAmelCase_   ).to_tuple()

							self.assertEqual(len(UpperCAmelCase_   )				,len(UpperCAmelCase_   )   )
							for jitted_output, output in zip(UpperCAmelCase_				,UpperCAmelCase_   ):
									self.assertEqual(jitted_output.shape				,output.shape   )






	@slow
	def    lowerCamelCase__					(     self   ):
			for model_class_name in self.all_model_classes:
					_lowercase				: Tuple     =   model_class_name.from_pretrained("""facebook/mbart-large-cc25"""				,from_pt=UpperCAmelCase_   )
					# FlaxMBartForSequenceClassification expects eos token in input_ids
					_lowercase				: Any     =   np.ones((1, 1)   ) * model.config.eos_token_id
					_lowercase				: Union[str, Any]     =   model(UpperCAmelCase_   )
					self.assertIsNotNone(UpperCAmelCase_   )







@require_flax
@require_sentencepiece
@require_tokenizers
class   UpperCamelCase	(     unittest.TestCase       ):




	"""simple docstring"""
	SCREAMING_SNAKE_CASE_							:       Union[str, Any]						       =  [
	    " UN Chief Says There Is No Military Solution in Syria",
	]
	SCREAMING_SNAKE_CASE_							:       List[Any]						       =  [
	    "Şeful ONU declară că nu există o soluţie militară în Siria",
	]
	SCREAMING_SNAKE_CASE_							:       List[Any]						       =  "facebook/mbart-large-en-ro"






	@cached_property
	def    lowerCamelCase__					(     self   ):
			return AutoTokenizer.from_pretrained(self.model_name   )






	@cached_property
	def    lowerCamelCase__					(     self   ):
			_lowercase				: Any     =   FlaxMBartForConditionalGeneration.from_pretrained(self.model_name				,from_pt=UpperCAmelCase_   )
			return model






	def    lowerCamelCase__					(     self				,**UpperCAmelCase_   ):
			_lowercase				: Optional[int]     =   self.translate_src_text(**UpperCAmelCase_   )
			self.assertListEqual(self.expected_text				,UpperCAmelCase_   )






	def    lowerCamelCase__					(     self				,**UpperCAmelCase_   ):
			_lowercase				: Tuple     =   self.tokenizer(self.src_text				,**UpperCAmelCase_				,return_tensors="""np"""   )
			_lowercase				: Optional[Any]     =   self.model.generate(
			    model_inputs.input_ids				,attention_mask=model_inputs.attention_mask				,decoder_start_token_id=self.tokenizer.lang_code_to_id["""ro_RO"""]				,early_stopping=UpperCAmelCase_				,num_beams=2				,).sequences
			_lowercase				: Optional[Any]     =   self.tokenizer.batch_decode(UpperCAmelCase_				,skip_special_tokens=UpperCAmelCase_   )
			return generated_words






	@slow
	def    lowerCamelCase__					(     self   ):
			self._assert_generated_batch_equal_expected()




"""simple docstring"""




import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()
class      UpperCamelCase__(      __A		,    unittest.TestCase			):
						lowerCAmelCase__	:    Tuple    				=		KandinskyPipeline
						lowerCAmelCase__	:    int    				=		[
						    'prompt',
						    'image_embeds',
						    'negative_image_embeds',
						]
						lowerCAmelCase__	:    str    				=		['prompt', 'negative_prompt', 'image_embeds', 'negative_image_embeds']
						lowerCAmelCase__	:    Any    				=		[
						    'generator',
						    'height',
						    'width',
						    'latents',
						    'guidance_scale',
						    'negative_prompt',
						    'num_inference_steps',
						    'return_dict',
						    'guidance_scale',
						    'num_images_per_prompt',
						    'output_type',
						    'return_dict',
						]
						lowerCAmelCase__	:    Dict    				=		False





						@property
						def 					snake_case__						(							self  )       ->					List[Any]:
													return 32





						@property
						def 					snake_case__						(							self  )       ->					Any:
													return 32





						@property
						def 					snake_case__						(							self  )       ->					Dict:
													return self.time_input_dim





						@property
						def 					snake_case__						(							self  )       ->					Optional[Any]:
													return self.time_input_dim * 4





						@property
						def 					snake_case__						(							self  )       ->					List[str]:
													return 1_00





						@property
						def 					snake_case__						(							self  )       ->					Any:
													A__						  =				XLMRobertaTokenizerFast.from_pretrained('YiYiXu/tiny-random-mclip-base'  )
													return tokenizer





						@property
						def 					snake_case__						(							self  )       ->					Tuple:
													torch.manual_seed(0  )
													A__						  =				MCLIPConfig(
													    numDims=self.cross_attention_dim		,transformerDimensions=self.text_embedder_hidden_size		,hidden_size=self.text_embedder_hidden_size		,intermediate_size=37		,num_attention_heads=4		,num_hidden_layers=5		,vocab_size=10_05		,)

													A__						  =				MultilingualCLIP(__UpperCAmelCase  )
													A__						  =				text_encoder.eval()

													return text_encoder





						@property
						def 					snake_case__						(							self  )       ->					Optional[Any]:
													torch.manual_seed(0  )

													A__						  =				{
													    'in_channels': 4,
													    # Out channels is double in channels because predicts mean and variance
													    'out_channels': 8,
													    'addition_embed_type': 'text_image',
													    'down_block_types': ('ResnetDownsampleBlock2D', 'SimpleCrossAttnDownBlock2D'),
													    'up_block_types': ('SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D'),
													    'mid_block_type': 'UNetMidBlock2DSimpleCrossAttn',
													    'block_out_channels': (self.block_out_channels_a, self.block_out_channels_a * 2),
													    'layers_per_block': 1,
													    'encoder_hid_dim': self.text_embedder_hidden_size,
													    'encoder_hid_dim_type': 'text_image_proj',
													    'cross_attention_dim': self.cross_attention_dim,
													    'attention_head_dim': 4,
													    'resnet_time_scale_shift': 'scale_shift',
													    'class_embed_type': None,
													}

													A__						  =				UNetaDConditionModel(**__UpperCAmelCase  )
													return model





						@property
						def 					snake_case__						(							self  )       ->					Tuple:
													return {
													    "block_out_channels": [32, 64],
													    "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
													    "in_channels": 3,
													    "latent_channels": 4,
													    "layers_per_block": 1,
													    "norm_num_groups": 8,
													    "norm_type": "spatial",
													    "num_vq_embeddings": 12,
													    "out_channels": 3,
													    "up_block_types": [
													        "AttnUpDecoderBlock2D",
													        "UpDecoderBlock2D",
													    ],
													    "vq_embed_dim": 4,
													}





						@property
						def 					snake_case__						(							self  )       ->					Any:
													torch.manual_seed(0  )
													A__						  =				VQModel(**self.dummy_movq_kwargs  )
													return model





						def 					snake_case__						(							self  )       ->					Optional[int]:
													A__						  =				self.dummy_text_encoder
													A__						  =				self.dummy_tokenizer
													A__						  =				self.dummy_unet
													A__						  =				self.dummy_movq

													A__						  =				DDIMScheduler(
													    num_train_timesteps=10_00		,beta_schedule='linear'		,beta_start=0.0_0_0_8_5		,beta_end=0.0_1_2		,clip_sample=__UpperCAmelCase		,set_alpha_to_one=__UpperCAmelCase		,steps_offset=1		,prediction_type='epsilon'		,thresholding=__UpperCAmelCase		,)

													A__						  =				{
													    'text_encoder': text_encoder,
													    'tokenizer': tokenizer,
													    'unet': unet,
													    'scheduler': scheduler,
													    'movq': movq,
													}
													return components





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase=0  )       ->					Optional[Any]:
													A__						  =				floats_tensor((1, self.cross_attention_dim)		,rng=random.Random(__UpperCAmelCase  )  ).to(__UpperCAmelCase  )
													A__						  =				floats_tensor((1, self.cross_attention_dim)		,rng=random.Random(seed + 1  )  ).to(__UpperCAmelCase  )
													if str(__UpperCAmelCase  ).startswith('mps'  ):
																				A__						  =				torch.manual_seed(__UpperCAmelCase  )
													else:
																				A__						  =				torch.Generator(device=__UpperCAmelCase  ).manual_seed(__UpperCAmelCase  )
													A__						  =				{
													    'prompt': 'horse',
													    'image_embeds': image_embeds,
													    'negative_image_embeds': negative_image_embeds,
													    'generator': generator,
													    'height': 64,
													    'width': 64,
													    'guidance_scale': 4.0,
													    'num_inference_steps': 2,
													    'output_type': 'np',
													}
													return inputs





						def 					snake_case__						(							self  )       ->					Optional[Any]:
													A__						  =				'cpu'

													A__						  =				self.get_dummy_components()

													A__						  =				self.pipeline_class(**__UpperCAmelCase  )
													A__						  =				pipe.to(__UpperCAmelCase  )

													pipe.set_progress_bar_config(disable=__UpperCAmelCase  )

													A__						  =				pipe(**self.get_dummy_inputs(__UpperCAmelCase  )  )
													A__						  =				output.images

													A__						  =				pipe(
													    **self.get_dummy_inputs(__UpperCAmelCase  )		,return_dict=__UpperCAmelCase		,)[0]

													A__						  =				image[0, -3:, -3:, -1]
													A__						  =				image_from_tuple[0, -3:, -3:, -1]

													assert image.shape == (1, 64, 64, 3)

													A__						  =				np.array(
													    [0.3_2_8_6_6_3, 1.0, 0.2_3_2_1_6_8_7_3, 1.0, 0.9_2_7_1_7_5_6_4, 0.4_6_3_9_0_4_6, 0.9_6_8_9_4_7_7_7, 0.3_1_7_1_3_3_7_8, 0.6_2_9_3_9_5_3]  )

													assert (
													    np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-2
													), f''' expected_slice {expected_slice}, but got {image_slice.flatten()}'''
													assert (
													    np.abs(image_from_tuple_slice.flatten() - expected_slice  ).max() < 1e-2
													), f''' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'''






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





						def 					snake_case__						(							self  )       ->					Tuple:
													# clean up the VRAM after each test
													super().tearDown()
													gc.collect()
													torch.cuda.empty_cache()





						def 					snake_case__						(							self  )       ->					List[Any]:
													A__						  =				load_numpy(
													    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
													    '/kandinsky/kandinsky_text2img_cat_fp16.npy'  )

													A__						  =				KandinskyPriorPipeline.from_pretrained(
													    'kandinsky-community/kandinsky-2-1-prior'		,torch_dtype=torch.floataa  )
													pipe_prior.to(__UpperCAmelCase  )

													A__						  =				KandinskyPipeline.from_pretrained('kandinsky-community/kandinsky-2-1'		,torch_dtype=torch.floataa  )
													A__						  =				pipeline.to(__UpperCAmelCase  )
													pipeline.set_progress_bar_config(disable=__UpperCAmelCase  )

													A__						  =				'red cat, 4k photo'

													A__						  =				torch.Generator(device='cuda'  ).manual_seed(0  )
													A__  ,	A__						  =				pipe_prior(
													    __UpperCAmelCase		,generator=__UpperCAmelCase		,num_inference_steps=5		,negative_prompt=''		,).to_tuple()

													A__						  =				torch.Generator(device='cuda'  ).manual_seed(0  )
													A__						  =				pipeline(
													    __UpperCAmelCase		,image_embeds=__UpperCAmelCase		,negative_image_embeds=__UpperCAmelCase		,generator=__UpperCAmelCase		,num_inference_steps=1_00		,output_type='np'		,)

													A__						  =				output.images[0]

													assert image.shape == (5_12, 5_12, 3)

													assert_mean_pixel_difference(__UpperCAmelCase		,__UpperCAmelCase  )





