



import gc
import random
import unittest

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

from diffusers import DDIMScheduler, KandinskyInpaintPipeline, 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   __lowerCAmelCase  (		UpperCAmelCase_				,  unittest.TestCase   ):
					"""simple docstring"""






					A__					:		Optional[Any]      							=     KandinskyInpaintPipeline
					A__					:		Dict      							=     ["prompt", "image_embeds", "negative_image_embeds", "image", "mask_image"]
					A__					:		Any      							=     [
					    "prompt",
					    "negative_prompt",
					    "image_embeds",
					    "negative_image_embeds",
					    "image",
					    "mask_image",
					]
					A__					:		Tuple      							=     [
					    "generator",
					    "height",
					    "width",
					    "latents",
					    "guidance_scale",
					    "negative_prompt",
					    "num_inference_steps",
					    "return_dict",
					    "guidance_scale",
					    "num_images_per_prompt",
					    "output_type",
					    "return_dict",
					]
					A__					:		Dict      							=     False



					@property
					def      _a				(  self       :							Optional[Any]       ):






								"""simple docstring"""


								return 32



					@property
					def      _a				(  self       :							int       ):






								"""simple docstring"""


								return 32



					@property
					def      _a				(  self       :							List[Any]       ):






								"""simple docstring"""


								return self.time_input_dim



					@property
					def      _a				(  self       :							str       ):






								"""simple docstring"""


								return self.time_input_dim * 4



					@property
					def      _a				(  self       :							List[str]       ):






								"""simple docstring"""


								return 1_00



					@property
					def      _a				(  self       :							Any       ):






								"""simple docstring"""


								A__	  =     XLMRobertaTokenizerFast.from_pretrained('YiYiXu/tiny-random-mclip-base'       )
								return tokenizer



					@property
					def      _a				(  self       :							str       ):






								"""simple docstring"""


								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(_snake_case       )
								A__	  =     text_encoder.eval()

								return text_encoder



					@property
					def      _a				(  self       :							Dict       ):






								"""simple docstring"""


								torch.manual_seed(0       )

								A__	  =     {
								    'in_channels': 9,
								    # 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(**_snake_case       )
								return model



					@property
					def      _a				(  self       :							int       ):






								"""simple docstring"""


								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      _a				(  self       :							Optional[Any]       ):






								"""simple docstring"""


								torch.manual_seed(0       )
								A__	  =     VQModel(**self.dummy_movq_kwargs       )
								return model



					def      _a				(  self       :							List[Any]       ):






								"""simple docstring"""


								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_0085					,       beta_end=0.012					,       clip_sample=_snake_case					,       set_alpha_to_one=_snake_case					,       steps_offset=1					,       prediction_type='epsilon'					,       thresholding=_snake_case					,       )

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

								return components



					def      _a				(  self       :							str					,       _snake_case       :							int					,       _snake_case       :							List[str]=0       ):






								"""simple docstring"""


								A__	  =     floats_tensor((1, self.cross_attention_dim)					,       rng=random.Random(_snake_case       )       ).to(_snake_case       )
								A__	  =     floats_tensor((1, self.cross_attention_dim)					,       rng=random.Random(seed + 1       )       ).to(_snake_case       )
								# create init_image
								A__	  =     floats_tensor((1, 3, 64, 64)					,       rng=random.Random(_snake_case       )       ).to(_snake_case       )
								A__	  =     image.cpu().permute(0					,       2					,       3					,       1       )[0]
								A__	  =     Image.fromarray(np.uinta(_snake_case       )       ).convert('RGB'       ).resize((2_56, 2_56)       )
								# create mask
								A__	  =     np.ones((64, 64)					,       dtype=np.floataa       )
								A__	  =     0

								if str(_snake_case       ).startswith('mps'       ):
											A__	  =     torch.manual_seed(_snake_case       )
								else:
											A__	  =     torch.Generator(device=_snake_case       ).manual_seed(_snake_case       )
								A__	  =     {
								    'prompt': 'horse',
								    'image': init_image,
								    'mask_image': mask,
								    'image_embeds': image_embeds,
								    'negative_image_embeds': negative_image_embeds,
								    'generator': generator,
								    'height': 64,
								    'width': 64,
								    'num_inference_steps': 2,
								    'guidance_scale': 4.0,
								    'output_type': 'np',
								}
								return inputs



					def      _a				(  self       :							List[str]       ):






								"""simple docstring"""


								A__	  =     'cpu'

								A__	  =     self.get_dummy_components()

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

								pipe.set_progress_bar_config(disable=_snake_case       )

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

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

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

								print(F'''image.shape {image.shape}'''       )

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

								A__	  =     np.array(
								    [0.832_6919, 0.7379_0467, 0.2091_8581, 0.930_9612, 0.551_1791, 0.4371_3328, 0.551_3321, 0.4992_2934, 0.5949_7786]       )

								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()}'''



					def      _a				(  self       :							Dict       ):






								"""simple docstring"""


								super().test_inference_batch_single_identical(expected_max_diff=3E-3       )







@slow
@require_torch_gpu
class   __lowerCAmelCase  (		unittest.TestCase   ):
					"""simple docstring"""






					def      _a				(  self       :							Optional[int]       ):






								"""simple docstring"""


								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()



					def      _a				(  self       :							Union[str, Any]       ):






								"""simple docstring"""


								A__	  =     load_numpy(
								    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
								    '/kandinsky/kandinsky_inpaint_cat_with_hat_fp16.npy'       )

								A__	  =     load_image(
								    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main' '/kandinsky/cat.png'       )
								A__	  =     np.ones((7_68, 7_68)					,       dtype=np.floataa       )
								A__	  =     0

								A__	  =     'a hat'

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

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

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

								A__	  =     pipeline(
								    _snake_case					,       image=_snake_case					,       mask_image=_snake_case					,       image_embeds=_snake_case					,       negative_image_embeds=_snake_case					,       generator=_snake_case					,       num_inference_steps=1_00					,       height=7_68					,       width=7_68					,       output_type='np'					,       )

								A__	  =     output.images[0]

								assert image.shape == (7_68, 7_68, 3)

								assert_mean_pixel_difference(_snake_case					,       _snake_case       )









from __future__ import annotations

import unittest

import numpy as np

from transformers import BlipTextConfig
from transformers.testing_utils import require_tf, slow
from transformers.utils import is_tf_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask


if is_tf_available():
							import tensorflow as tf

							from transformers import TFBlipTextModel
							from transformers.models.blip.modeling_tf_blip import TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST


class   __lowerCAmelCase  :
					"""simple docstring"""






					def __init__(  self       :							Dict					,       _snake_case       :							Union[str, Any]					,       _snake_case       :							Optional[Any]=12					,       _snake_case       :							Any=7					,       _snake_case       :							List[str]=True					,       _snake_case       :							int=True					,       _snake_case       :							int=True					,       _snake_case       :							Tuple=99					,       _snake_case       :							List[Any]=32					,       _snake_case       :							Optional[int]=32					,       _snake_case       :							List[str]=2					,       _snake_case       :							List[str]=4					,       _snake_case       :							List[Any]=37					,       _snake_case       :							Union[str, Any]=0.1					,       _snake_case       :							Tuple=0.1					,       _snake_case       :							Dict=5_12					,       _snake_case       :							Union[str, Any]=0.02					,       _snake_case       :							Any=0					,       _snake_case       :							Optional[Any]=None					,       ):






								"""simple docstring"""


								A__	  =     parent
								A__	  =     batch_size
								A__	  =     seq_length
								A__	  =     is_training
								A__	  =     use_input_mask
								A__	  =     use_labels
								A__	  =     vocab_size
								A__	  =     hidden_size
								A__	  =     projection_dim
								A__	  =     num_hidden_layers
								A__	  =     num_attention_heads
								A__	  =     intermediate_size
								A__	  =     dropout
								A__	  =     attention_dropout
								A__	  =     max_position_embeddings
								A__	  =     initializer_range
								A__	  =     scope
								A__	  =     bos_token_id



					def      _a				(  self       :							Optional[Any]       ):






								"""simple docstring"""


								A__	  =     ids_tensor([self.batch_size, self.seq_length]					,       self.vocab_size       )

								A__	  =     None
								if self.use_input_mask:
											A__	  =     random_attention_mask([self.batch_size, self.seq_length]       )

								if input_mask is not None:
											A__	  =     input_mask.numpy()
											A__  ,			A__	  =     input_mask.shape
											A__	  =     np.random.randint(1					,       seq_length - 1					,       size=(batch_size,)       )
											for batch_idx, start_index in enumerate(_snake_case       ):
														A__	  =     1
														A__	  =     0

								A__	  =     self.get_config()

								return config, input_ids, tf.convert_to_tensor(_snake_case       )



					def      _a				(  self       :							Tuple       ):






								"""simple docstring"""


								return BlipTextConfig(
								    vocab_size=self.vocab_size					,       hidden_size=self.hidden_size					,       projection_dim=self.projection_dim					,       num_hidden_layers=self.num_hidden_layers					,       num_attention_heads=self.num_attention_heads					,       intermediate_size=self.intermediate_size					,       dropout=self.dropout					,       attention_dropout=self.attention_dropout					,       max_position_embeddings=self.max_position_embeddings					,       initializer_range=self.initializer_range					,       bos_token_id=self.bos_token_id					,       )



					def      _a				(  self       :							int					,       _snake_case       :							Union[str, Any]					,       _snake_case       :							Any					,       _snake_case       :							List[str]       ):






								"""simple docstring"""


								A__	  =     TFBlipTextModel(config=_snake_case       )
								A__	  =     model(_snake_case					,       attention_mask=_snake_case					,       training=_snake_case       )
								A__	  =     model(_snake_case					,       training=_snake_case       )
								self.parent.assertEqual(result.last_hidden_state.shape					,       (self.batch_size, self.seq_length, self.hidden_size)       )
								self.parent.assertEqual(result.pooler_output.shape					,       (self.batch_size, self.hidden_size)       )



					def      _a				(  self       :							str       ):






								"""simple docstring"""


								A__	  =     self.prepare_config_and_inputs()
								A__  ,			A__  ,			A__	  =     config_and_inputs
								A__	  =     {'input_ids': input_ids, 'attention_mask': input_mask}
								return config, inputs_dict







@require_tf
class   __lowerCAmelCase  (		UpperCAmelCase_				,  unittest.TestCase   ):
					"""simple docstring"""






					A__					:		Tuple      							=     (TFBlipTextModel,) if is_tf_available() else ()
					A__					:		Optional[int]      							=     False
					A__					:		Union[str, Any]      							=     False
					A__					:		Union[str, Any]      							=     False



					def      _a				(  self       :							Any       ):






								"""simple docstring"""


								A__	  =     BlipTextModelTester(self       )
								A__	  =     ConfigTester(self					,       config_class=_snake_case					,       hidden_size=37       )



					def      _a				(  self       :							List[str]       ):






								"""simple docstring"""


								self.config_tester.run_common_tests()



					def      _a				(  self       :							Union[str, Any]       ):






								"""simple docstring"""


								A__	  =     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_snake_case       )



					def      _a				(  self       :							Tuple       ):






								"""simple docstring"""


								pass



					def      _a				(  self       :							int       ):






								"""simple docstring"""


								pass



					@unittest.skip(reason='Blip does not use inputs_embeds'       )
					def      _a				(  self       :							Any       ):






								"""simple docstring"""


								pass



					@unittest.skip(reason='BlipTextModel has no base class and is not available in MODEL_MAPPING'       )
					def      _a				(  self       :							str       ):






								"""simple docstring"""


								pass



					@unittest.skip(reason='BlipTextModel has no base class and is not available in MODEL_MAPPING'       )
					def      _a				(  self       :							Optional[Any]       ):






								"""simple docstring"""


								pass



					@slow
					def      _a				(  self       :							Union[str, Any]       ):






								"""simple docstring"""


								for model_name in TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											A__	  =     TFBlipTextModel.from_pretrained(_snake_case       )
											self.assertIsNotNone(_snake_case       )



					def      _a				(  self       :							int					,       _snake_case       :							int=True       ):






								"""simple docstring"""


								super().test_pt_tf_model_equivalence(allow_missing_keys=_snake_case       )




