





import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyImgaImgPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
from diffusers.utils import floats_tensor, load_image, 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 				A_				(      UpperCAmelCase					,	unittest.TestCase	):







       """simple docstring"""





       SCREAMING_SNAKE_CASE_      : Optional[int]      =  KandinskyImgaImgPipeline
       SCREAMING_SNAKE_CASE_      : str      =  ['''prompt''', '''image_embeds''', '''negative_image_embeds''', '''image''']
       SCREAMING_SNAKE_CASE_      : Dict      =  [
           '''prompt''',
           '''negative_prompt''',
           '''image_embeds''',
           '''negative_image_embeds''',
           '''image''',
       ]
       SCREAMING_SNAKE_CASE_      : Optional[int]      =  [
           '''generator''',
           '''height''',
           '''width''',
           '''strength''',
           '''guidance_scale''',
           '''negative_prompt''',
           '''num_inference_steps''',
           '''return_dict''',
           '''guidance_scale''',
           '''num_images_per_prompt''',
           '''output_type''',
           '''return_dict''',
       ]
       SCREAMING_SNAKE_CASE_      : Dict      =  False







       @property
       def    __UpperCAmelCase (    self	:	List[str]  )						->					str:
         return 32







       @property
       def    __UpperCAmelCase (    self	:	List[str]  )						->					Dict:
         return 32







       @property
       def    __UpperCAmelCase (    self	:	Tuple  )						->					Union[str, Any]:
         return self.time_input_dim







       @property
       def    __UpperCAmelCase (    self	:	Optional[Any]  )						->					Union[str, Any]:
         return self.time_input_dim * 4







       @property
       def    __UpperCAmelCase (    self	:	List[Any]  )						->					Union[str, Any]:
         return 100







       @property
       def    __UpperCAmelCase (    self	:	Optional[int]  )						->					List[Any]:
         _lowercase							= XLMRobertaTokenizerFast.from_pretrained('YiYiXu/tiny-random-mclip-base'  )
         return tokenizer







       @property
       def    __UpperCAmelCase (    self	:	Any  )						->					Any:
         torch.manual_seed(0  )
         _lowercase							= 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=1005						,)

         _lowercase							= MultilingualCLIP(_SCREAMING_SNAKE_CASE  )
         _lowercase							= text_encoder.eval()

         return text_encoder







       @property
       def    __UpperCAmelCase (    self	:	Union[str, Any]  )						->					Union[str, Any]:
         torch.manual_seed(0  )

         _lowercase							= {
             '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,
         }

         _lowercase							= UNetaDConditionModel(**_SCREAMING_SNAKE_CASE  )
         return model







       @property
       def    __UpperCAmelCase (    self	:	Tuple  )						->					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    __UpperCAmelCase (    self	:	List[str]  )						->					Dict:
         torch.manual_seed(0  )
         _lowercase							= VQModel(**self.dummy_movq_kwargs  )
         return model







       def    __UpperCAmelCase (    self	:	List[Any]  )						->					str:
         _lowercase							= self.dummy_text_encoder
         _lowercase							= self.dummy_tokenizer
         _lowercase							= self.dummy_unet
         _lowercase							= self.dummy_movq

         _lowercase							= {
             'num_train_timesteps': 1000,
             'beta_schedule': 'linear',
             'beta_start': 0.00085,
             'beta_end': 0.012,
             'clip_sample': False,
             'set_alpha_to_one': False,
             'steps_offset': 0,
             'prediction_type': 'epsilon',
             'thresholding': False,
         }

         _lowercase							= DDIMScheduler(**_SCREAMING_SNAKE_CASE  )

         _lowercase							= {
             'text_encoder': text_encoder,
             'tokenizer': tokenizer,
             'unet': unet,
             'scheduler': scheduler,
             'movq': movq,
         }

         return components







       def    __UpperCAmelCase (    self	:	List[Any]						,__A	:	Any						,__A	:	int=0  )						->					int:
         _lowercase							= floats_tensor((1, self.cross_attention_dim)						,rng=random.Random(_SCREAMING_SNAKE_CASE  )  ).to(_SCREAMING_SNAKE_CASE  )
         _lowercase							= floats_tensor((1, self.cross_attention_dim)						,rng=random.Random(seed + 1  )  ).to(_SCREAMING_SNAKE_CASE  )
         # create init_image
         _lowercase							= floats_tensor((1, 3, 64, 64)						,rng=random.Random(_SCREAMING_SNAKE_CASE  )  ).to(_SCREAMING_SNAKE_CASE  )
         _lowercase							= image.cpu().permute(0						,2						,3						,1  )[0]
         _lowercase							= Image.fromarray(np.uinta(_SCREAMING_SNAKE_CASE  )  ).convert('RGB'  ).resize((256, 256)  )

         if str(_SCREAMING_SNAKE_CASE  ).startswith('mps'  ):
           _lowercase							= torch.manual_seed(_SCREAMING_SNAKE_CASE  )
         else:
           _lowercase							= torch.Generator(device=_SCREAMING_SNAKE_CASE  ).manual_seed(_SCREAMING_SNAKE_CASE  )
         _lowercase							= {
             'prompt': 'horse',
             'image': init_image,
             'image_embeds': image_embeds,
             'negative_image_embeds': negative_image_embeds,
             'generator': generator,
             'height': 64,
             'width': 64,
             'num_inference_steps': 10,
             'guidance_scale': 7.0,
             'strength': 0.2,
             'output_type': 'np',
         }
         return inputs







       def    __UpperCAmelCase (    self	:	Union[str, Any]  )						->					Optional[Any]:
         _lowercase							= 'cpu'

         _lowercase							= self.get_dummy_components()

         _lowercase							= self.pipeline_class(**_SCREAMING_SNAKE_CASE  )
         _lowercase							= pipe.to(_SCREAMING_SNAKE_CASE  )

         pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE  )

         _lowercase							= pipe(**self.get_dummy_inputs(_SCREAMING_SNAKE_CASE  )  )
         _lowercase							= output.images

         _lowercase							= pipe(
             **self.get_dummy_inputs(_SCREAMING_SNAKE_CASE  )						,return_dict=_SCREAMING_SNAKE_CASE						,)[0]

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

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

         _lowercase							= np.array(
             [0.61474943, 0.6073539, 0.43308544, 0.5928269, 0.47493595, 0.46755973, 0.4613838, 0.45368797, 0.50119233]  )
         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 				A_				(      unittest.TestCase	):







       """simple docstring"""





       def    __UpperCAmelCase (    self	:	str  )						->					int:
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()







       def    __UpperCAmelCase (    self	:	Optional[Any]  )						->					int:
         _lowercase							= load_numpy(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
             '/kandinsky/kandinsky_img2img_frog.npy'  )

         _lowercase							= load_image(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main' '/kandinsky/cat.png'  )
         _lowercase							= 'A red cartoon frog, 4k'

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

         _lowercase							= KandinskyImgaImgPipeline.from_pretrained(
             'kandinsky-community/kandinsky-2-1'						,torch_dtype=torch.floataa  )
         _lowercase							= pipeline.to(_SCREAMING_SNAKE_CASE  )

         pipeline.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE  )

         _lowercase							= torch.Generator(device='cpu'  ).manual_seed(0  )
         _lowercase						,		_lowercase							= pipe_prior(
             _SCREAMING_SNAKE_CASE						,generator=_SCREAMING_SNAKE_CASE						,num_inference_steps=5						,negative_prompt=''						,).to_tuple()

         _lowercase							= pipeline(
             _SCREAMING_SNAKE_CASE						,image=_SCREAMING_SNAKE_CASE						,image_embeds=_SCREAMING_SNAKE_CASE						,negative_image_embeds=_SCREAMING_SNAKE_CASE						,generator=_SCREAMING_SNAKE_CASE						,num_inference_steps=100						,height=768						,width=768						,strength=0.2						,output_type='np'						,)

         _lowercase							= output.images[0]

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

         assert_mean_pixel_difference(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE  )









import unittest

import numpy as np

from transformers import RobertaPreLayerNormConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
      import jax.numpy as jnp

      from transformers.models.roberta_prelayernorm.modeling_flax_roberta_prelayernorm import (
          FlaxRobertaPreLayerNormForCausalLM,
          FlaxRobertaPreLayerNormForMaskedLM,
          FlaxRobertaPreLayerNormForMultipleChoice,
          FlaxRobertaPreLayerNormForQuestionAnswering,
          FlaxRobertaPreLayerNormForSequenceClassification,
          FlaxRobertaPreLayerNormForTokenClassification,
          FlaxRobertaPreLayerNormModel,
      )





class 				A_				(      unittest.TestCase	):







       """simple docstring"""





       def __init__(    self	:	List[Any]						,__A	:	int						,__A	:	str=13						,__A	:	str=7						,__A	:	Optional[Any]=True						,__A	:	Optional[int]=True						,__A	:	int=True						,__A	:	List[str]=True						,__A	:	Dict=99						,__A	:	Dict=32						,__A	:	Dict=5						,__A	:	Dict=4						,__A	:	int=37						,__A	:	Any="gelu"						,__A	:	int=0.1						,__A	:	int=0.1						,__A	:	Dict=512						,__A	:	List[str]=16						,__A	:	Tuple=2						,__A	:	Dict=0.02						,__A	:	int=4						,)						->					Optional[int]:
         _lowercase							= parent
         _lowercase							= batch_size
         _lowercase							= seq_length
         _lowercase							= is_training
         _lowercase							= use_attention_mask
         _lowercase							= use_token_type_ids
         _lowercase							= use_labels
         _lowercase							= vocab_size
         _lowercase							= hidden_size
         _lowercase							= num_hidden_layers
         _lowercase							= num_attention_heads
         _lowercase							= intermediate_size
         _lowercase							= hidden_act
         _lowercase							= hidden_dropout_prob
         _lowercase							= attention_probs_dropout_prob
         _lowercase							= max_position_embeddings
         _lowercase							= type_vocab_size
         _lowercase							= type_sequence_label_size
         _lowercase							= initializer_range
         _lowercase							= num_choices







       def    __UpperCAmelCase (    self	:	int  )						->					str:
         _lowercase							= ids_tensor([self.batch_size, self.seq_length]						,self.vocab_size  )

         _lowercase							= None
         if self.use_attention_mask:
           _lowercase							= random_attention_mask([self.batch_size, self.seq_length]  )

         _lowercase							= None
         if self.use_token_type_ids:
           _lowercase							= ids_tensor([self.batch_size, self.seq_length]						,self.type_vocab_size  )

         _lowercase							= RobertaPreLayerNormConfig(
             vocab_size=self.vocab_size						,hidden_size=self.hidden_size						,num_hidden_layers=self.num_hidden_layers						,num_attention_heads=self.num_attention_heads						,intermediate_size=self.intermediate_size						,hidden_act=self.hidden_act						,hidden_dropout_prob=self.hidden_dropout_prob						,attention_probs_dropout_prob=self.attention_probs_dropout_prob						,max_position_embeddings=self.max_position_embeddings						,type_vocab_size=self.type_vocab_size						,is_decoder=__A						,initializer_range=self.initializer_range						,)

         return config, input_ids, token_type_ids, attention_mask







       def    __UpperCAmelCase (    self	:	Optional[int]  )						->					Tuple:
         _lowercase							= self.prepare_config_and_inputs()
         _lowercase						,		_lowercase						,		_lowercase						,		_lowercase							= config_and_inputs
         _lowercase							= {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': attention_mask}
         return config, inputs_dict







       def    __UpperCAmelCase (    self	:	Tuple  )						->					List[Any]:
         _lowercase							= self.prepare_config_and_inputs()
         _lowercase						,		_lowercase						,		_lowercase						,		_lowercase							= config_and_inputs

         _lowercase							= True
         _lowercase							= floats_tensor([self.batch_size, self.seq_length, self.hidden_size]  )
         _lowercase							= ids_tensor([self.batch_size, self.seq_length]						,vocab_size=2  )

         return (
             config,
             input_ids,
             token_type_ids,
             encoder_hidden_states,
             encoder_attention_mask,
         )





@require_flax
# Copied from tests.models.roberta.test_modelling_flax_roberta.FlaxRobertaPreLayerNormModelTest with ROBERTA->ROBERTA_PRELAYERNORM,Roberta->RobertaPreLayerNorm,roberta-base->andreasmadsen/efficient_mlm_m0.40
class 				A_				(      UpperCAmelCase					,	unittest.TestCase	):







       """simple docstring"""





       SCREAMING_SNAKE_CASE_      : Optional[int]      =  True

       SCREAMING_SNAKE_CASE_      : Dict      =  (
           (
               FlaxRobertaPreLayerNormModel,
               FlaxRobertaPreLayerNormForCausalLM,
               FlaxRobertaPreLayerNormForMaskedLM,
               FlaxRobertaPreLayerNormForSequenceClassification,
               FlaxRobertaPreLayerNormForTokenClassification,
               FlaxRobertaPreLayerNormForMultipleChoice,
               FlaxRobertaPreLayerNormForQuestionAnswering,
           )
           if is_flax_available()
           else ()
       )







       def    __UpperCAmelCase (    self	:	List[Any]  )						->					Any:
         _lowercase							= FlaxRobertaPreLayerNormModelTester(self  )







       @slow
       def    __UpperCAmelCase (    self	:	Dict  )						->					Dict:
         for model_class_name in self.all_model_classes:
           _lowercase							= model_class_name.from_pretrained('andreasmadsen/efficient_mlm_m0.40'						,from_pt=__A  )
           _lowercase							= model(np.ones((1, 1)  )  )
           self.assertIsNotNone(__A  )





@require_flax
class 				A_				(      unittest.TestCase	):







       """simple docstring"""





       @slow
       def    __UpperCAmelCase (    self	:	int  )						->					Optional[Any]:
         _lowercase							= FlaxRobertaPreLayerNormForMaskedLM.from_pretrained('andreasmadsen/efficient_mlm_m0.40'						,from_pt=__A  )

         _lowercase							= np.array([[0, 3_1414, 232, 328, 740, 1140, 1_2695, 69, 4_6078, 1588, 2]]						,dtype=jnp.intaa  )
         _lowercase							= model(__A  )[0]
         _lowercase							= [1, 11, 5_0265]
         self.assertEqual(list(output.shape  )						,__A  )
         # compare the actual values for a slice.
         _lowercase							= np.array(
             [[[40.4880, 18.0199, -5.2367], [-1.8877, -4.0885, 10.7085], [-2.2613, -5.6110, 7.2665]]]						,dtype=np.floataa  )
         self.assertTrue(np.allclose(output[:, :3, :3]						,__A						,atol=1e-4  )  )







       @slow
       def    __UpperCAmelCase (    self	:	Union[str, Any]  )						->					List[Any]:
         _lowercase							= FlaxRobertaPreLayerNormModel.from_pretrained('andreasmadsen/efficient_mlm_m0.40'						,from_pt=__A  )

         _lowercase							= np.array([[0, 3_1414, 232, 328, 740, 1140, 1_2695, 69, 4_6078, 1588, 2]]						,dtype=jnp.intaa  )
         _lowercase							= model(__A  )[0]
         # compare the actual values for a slice.
         _lowercase							= np.array(
             [[[0.0208, -0.0356, 0.0237], [-0.1569, -0.0411, -0.2626], [0.1879, 0.0125, -0.0089]]]						,dtype=np.floataa  )
         self.assertTrue(np.allclose(output[:, :3, :3]						,__A						,atol=1e-4  )  )