

import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaInpaintPipeline,
    KandinskyVaaPriorPipeline,
    UNetaDConditionModel,
    VQModel,
)
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 	lowercase    (				_lowercase    ,							unittest.TestCase							):


		"""simple docstring"""




		a__					       =  KandinskyVaaInpaintPipeline
		a__					       =  ["image_embeds", "negative_image_embeds", "image", "mask_image"]
		a__					       =  [
		    "image_embeds",
		    "negative_image_embeds",
		    "image",
		    "mask_image",
		]
		a__					       =  [
		    "generator",
		    "height",
		    "width",
		    "latents",
		    "guidance_scale",
		    "num_inference_steps",
		    "return_dict",
		    "guidance_scale",
		    "num_images_per_prompt",
		    "output_type",
		    "return_dict",
		]
		a__					       =  False


		@property
		def 		A__					(				self):
			return 32


		@property
		def 		A__					(				self):
			return 32


		@property
		def 		A__					(				self):
			return self.time_input_dim


		@property
		def 		A__					(				self):
			return self.time_input_dim * 4


		@property
		def 		A__					(				self):
			return 1_00


		@property
		def 		A__					(				self):
			torch.manual_seed(0)

			_UpperCamelCase   :		Optional[Any]			=						{
			    'in_channels': 9,
			    # Out channels is double in channels because predicts mean and variance
			    'out_channels': 8,
			    'addition_embed_type': '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': 'image_proj',
			    'cross_attention_dim': self.cross_attention_dim,
			    'attention_head_dim': 4,
			    'resnet_time_scale_shift': 'scale_shift',
			    'class_embed_type': None,
			}

			_UpperCamelCase   :		Tuple			=						UNetaDConditionModel(**__snake_case)
			return model


		@property
		def 		A__					(				self):
			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):
			torch.manual_seed(0)
			_UpperCamelCase   :		int			=						VQModel(**self.dummy_movq_kwargs)
			return model


		def 		A__					(				self):
			_UpperCamelCase   :		Optional[Any]			=						self.dummy_unet
			_UpperCamelCase   :		List[str]			=						self.dummy_movq

			_UpperCamelCase   :		List[Any]			=						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=__snake_case	,       set_alpha_to_one=__snake_case	,       steps_offset=1	,       prediction_type='epsilon'	,       thresholding=__snake_case	,       )

			_UpperCamelCase   :		str			=						{
			    'unet': unet,
			    'scheduler': scheduler,
			    'movq': movq,
			}

			return components


		def 		A__					(				self	,       __snake_case	,       __snake_case=0):
			_UpperCamelCase   :		Any			=						floats_tensor((1, self.text_embedder_hidden_size)	,       rng=random.Random(__snake_case)).to(__snake_case)
			_UpperCamelCase   :		Dict			=						floats_tensor((1, self.text_embedder_hidden_size)	,       rng=random.Random(seed + 1)).to(
			    __snake_case)
			# create init_image
			_UpperCamelCase   :		Dict			=						floats_tensor((1, 3, 64, 64)	,       rng=random.Random(__snake_case)).to(__snake_case)
			_UpperCamelCase   :		Tuple			=						image.cpu().permute(0	,       2	,       3	,       1)[0]
			_UpperCamelCase   :		Dict			=						Image.fromarray(np.uinta(__snake_case)).convert('RGB').resize((2_56, 2_56))
			# create mask
			_UpperCamelCase   :		Tuple			=						np.ones((64, 64)	,       dtype=np.floataa)
			_UpperCamelCase   :		int			=						0

			if str(__snake_case).startswith('mps'):
				_UpperCamelCase   :		List[str]			=						torch.manual_seed(__snake_case)
			else:
				_UpperCamelCase   :		Tuple			=						torch.Generator(device=__snake_case).manual_seed(__snake_case)
			_UpperCamelCase   :		Optional[Any]			=						{
			    '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):
			_UpperCamelCase   :		List[str]			=						'cpu'

			_UpperCamelCase   :		Union[str, Any]			=						self.get_dummy_components()

			_UpperCamelCase   :		str			=						self.pipeline_class(**__snake_case)
			_UpperCamelCase   :		int			=						pipe.to(__snake_case)

			pipe.set_progress_bar_config(disable=__snake_case)

			_UpperCamelCase   :		Optional[Any]			=						pipe(**self.get_dummy_inputs(__snake_case))
			_UpperCamelCase   :		List[str]			=						output.images

			_UpperCamelCase   :		Union[str, Any]			=						pipe(
			    **self.get_dummy_inputs(__snake_case)	,       return_dict=__snake_case	,       )[0]

			_UpperCamelCase   :		Union[str, Any]			=						image[0, -3:, -3:, -1]
			_UpperCamelCase   :		List[Any]			=						image_from_tuple[0, -3:, -3:, -1]

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

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

			_UpperCamelCase   :		List[Any]			=						np.array(
			    [0.5_0_7_7_5_9_0_3, 0.4_9_5_2_7_1_9_5, 0.4_8_8_2_4_5_4_3, 0.5_0_1_9_2_2_3_7, 0.4_8_6_4_4_9_0_6, 0.4_9_3_7_3_8_1_4, 0.4_7_8_0_5_9_8, 0.4_7_2_3_4_8_2_7, 0.4_8_3_2_7_8_4_8])

			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):
			super().test_inference_batch_single_identical(expected_max_diff=3e-3)







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


		"""simple docstring"""




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






		def 		A__					(				self):
			_UpperCamelCase   :		str			=						load_numpy(
			    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
			    '/kandinskyv22/kandinskyv22_inpaint_cat_with_hat_fp16.npy')

			_UpperCamelCase   :		List[str]			=						load_image(
			    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main' '/kandinsky/cat.png')
			_UpperCamelCase   :		List[Any]			=						np.ones((7_68, 7_68)	,       dtype=np.floataa)
			_UpperCamelCase   :		Optional[int]			=						0

			_UpperCamelCase   :		List[str]			=						'a hat'

			_UpperCamelCase   :		List[Any]			=						KandinskyVaaPriorPipeline.from_pretrained(
			    'kandinsky-community/kandinsky-2-2-prior'	,       torch_dtype=torch.floataa)
			pipe_prior.to(__snake_case)

			_UpperCamelCase   :		List[str]			=						KandinskyVaaInpaintPipeline.from_pretrained(
			    'kandinsky-community/kandinsky-2-2-decoder-inpaint'	,       torch_dtype=torch.floataa)
			_UpperCamelCase   :		Union[str, Any]			=						pipeline.to(__snake_case)
			pipeline.set_progress_bar_config(disable=__snake_case)

			_UpperCamelCase   :		int			=						torch.Generator(device='cpu').manual_seed(0)
			_UpperCamelCase   :		Union[str, Any]			=						pipe_prior(
			    __snake_case	,       generator=__snake_case	,       num_inference_steps=5	,       negative_prompt=''	,       ).to_tuple()

			_UpperCamelCase   :		List[Any]			=						pipeline(
			    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'	,       )

			_UpperCamelCase   :		Union[str, Any]			=						output.images[0]

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

			assert_mean_pixel_difference(__snake_case	,       __snake_case)







import inspect
import unittest

from transformers import BitConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch
							from torch import nn

							from transformers import BitBackbone, BitForImageClassification, BitImageProcessor, BitModel
							from transformers.models.bit.modeling_bit import BIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
							from PIL import Image







class 	lowercase    :


		"""simple docstring"""




		def __init__(				self	,       __snake_case	,       __snake_case=3	,       __snake_case=32	,       __snake_case=3	,       __snake_case=10	,       __snake_case=[8, 16, 32, 64]	,       __snake_case=[1, 1, 2, 1]	,       __snake_case=True	,       __snake_case=True	,       __snake_case="relu"	,       __snake_case=3	,       __snake_case=None	,       __snake_case=["stage2", "stage3", "stage4"]	,       __snake_case=[2, 3, 4]	,       __snake_case=1	,       ):
			_UpperCamelCase   :		List[Any]			=						parent
			_UpperCamelCase   :		Dict			=						batch_size
			_UpperCamelCase   :		Optional[int]			=						image_size
			_UpperCamelCase   :		str			=						num_channels
			_UpperCamelCase   :		Optional[Any]			=						embeddings_size
			_UpperCamelCase   :		Tuple			=						hidden_sizes
			_UpperCamelCase   :		Dict			=						depths
			_UpperCamelCase   :		str			=						is_training
			_UpperCamelCase   :		Optional[int]			=						use_labels
			_UpperCamelCase   :		str			=						hidden_act
			_UpperCamelCase   :		Optional[int]			=						num_labels
			_UpperCamelCase   :		Optional[int]			=						scope
			_UpperCamelCase   :		Tuple			=						len(__snake_case)
			_UpperCamelCase   :		Dict			=						out_features
			_UpperCamelCase   :		Union[str, Any]			=						out_indices
			_UpperCamelCase   :		int			=						num_groups


		def 		A__					(				self):
			_UpperCamelCase   :		Tuple			=						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

			_UpperCamelCase   :		str			=						None
			if self.use_labels:
				_UpperCamelCase   :		Optional[int]			=						ids_tensor([self.batch_size]	,       self.num_labels)

			_UpperCamelCase   :		str			=						self.get_config()

			return config, pixel_values, labels


		def 		A__					(				self):
			return BitConfig(
			    num_channels=self.num_channels	,       embeddings_size=self.embeddings_size	,       hidden_sizes=self.hidden_sizes	,       depths=self.depths	,       hidden_act=self.hidden_act	,       num_labels=self.num_labels	,       out_features=self.out_features	,       out_indices=self.out_indices	,       num_groups=self.num_groups	,       )


		def 		A__					(				self	,       __snake_case	,       __snake_case	,       __snake_case):
			_UpperCamelCase   :		str			=						BitModel(config=__snake_case)
			model.to(__snake_case)
			model.eval()
			_UpperCamelCase   :		Optional[Any]			=						model(__snake_case)
			self.parent.assertEqual(
			    result.last_hidden_state.shape	,       (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)	,       )


		def 		A__					(				self	,       __snake_case	,       __snake_case	,       __snake_case):
			_UpperCamelCase   :		Dict			=						self.num_labels
			_UpperCamelCase   :		Dict			=						BitForImageClassification(__snake_case)
			model.to(__snake_case)
			model.eval()
			_UpperCamelCase   :		Dict			=						model(__snake_case	,       labels=__snake_case)
			self.parent.assertEqual(result.logits.shape	,       (self.batch_size, self.num_labels))


		def 		A__					(				self	,       __snake_case	,       __snake_case	,       __snake_case):
			_UpperCamelCase   :		Optional[Any]			=						BitBackbone(config=__snake_case)
			model.to(__snake_case)
			model.eval()
			_UpperCamelCase   :		List[Any]			=						model(__snake_case)

			# verify feature maps
			self.parent.assertEqual(len(result.feature_maps)	,       len(config.out_features))
			self.parent.assertListEqual(list(result.feature_maps[0].shape)	,       [self.batch_size, self.hidden_sizes[1], 4, 4])

			# verify channels
			self.parent.assertEqual(len(model.channels)	,       len(config.out_features))
			self.parent.assertListEqual(model.channels	,       config.hidden_sizes[1:])

			# verify backbone works with out_features=None
			_UpperCamelCase   :		Any			=						None
			_UpperCamelCase   :		str			=						BitBackbone(config=__snake_case)
			model.to(__snake_case)
			model.eval()
			_UpperCamelCase   :		Any			=						model(__snake_case)

			# verify feature maps
			self.parent.assertEqual(len(result.feature_maps)	,       1)
			self.parent.assertListEqual(list(result.feature_maps[0].shape)	,       [self.batch_size, self.hidden_sizes[-1], 1, 1])

			# verify channels
			self.parent.assertEqual(len(model.channels)	,       1)
			self.parent.assertListEqual(model.channels	,       [config.hidden_sizes[-1]])






		def 		A__					(				self):
			_UpperCamelCase   :		Optional[int]			=						self.prepare_config_and_inputs()
			_UpperCamelCase	,		_UpperCamelCase	,		_UpperCamelCase   :		int			=						config_and_inputs
			_UpperCamelCase   :		int			=						{'pixel_values': pixel_values}
			return config, inputs_dict







@require_torch
class 	lowercase    (				_lowercase    ,							_lowercase    ,							unittest.TestCase							):


		"""simple docstring"""




		a__					       =  (BitModel, BitForImageClassification, BitBackbone) if is_torch_available() else ()
		a__					       =  (
		    {"feature-extraction": BitModel, "image-classification": BitForImageClassification}
		    if is_torch_available()
		    else {}
		)

		a__					       =  False
		a__					       =  False
		a__					       =  False
		a__					       =  False
		a__					       =  False


		def 		A__					(				self):
			_UpperCamelCase   :		Dict			=						BitModelTester(self)
			_UpperCamelCase   :		Optional[Any]			=						ConfigTester(self	,       config_class=__snake_case	,       has_text_modality=__snake_case)


		def 		A__					(				self):
			self.create_and_test_config_common_properties()
			self.config_tester.create_and_test_config_to_json_string()
			self.config_tester.create_and_test_config_to_json_file()
			self.config_tester.create_and_test_config_from_and_save_pretrained()
			self.config_tester.create_and_test_config_with_num_labels()
			self.config_tester.check_config_can_be_init_without_params()
			self.config_tester.check_config_arguments_init()


		def 		A__					(				self):
			return


		@unittest.skip(reason='Bit does not output attentions')
		def 		A__					(				self):
			pass


		@unittest.skip(reason='Bit does not use inputs_embeds')
		def 		A__					(				self):
			pass


		@unittest.skip(reason='Bit does not support input and output embeddings')
		def 		A__					(				self):
			pass


		def 		A__					(				self):
			_UpperCamelCase	,		_UpperCamelCase   :		List[Any]			=						self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
				_UpperCamelCase   :		int			=						model_class(__snake_case)
				_UpperCamelCase   :		List[Any]			=						inspect.signature(model.forward)
				# signature.parameters is an OrderedDict => so arg_names order is deterministic
				_UpperCamelCase   :		Optional[int]			=						[*signature.parameters.keys()]

				_UpperCamelCase   :		List[Any]			=						['pixel_values']
				self.assertListEqual(arg_names[:1]	,       __snake_case)


		def 		A__					(				self):
			_UpperCamelCase   :		Union[str, Any]			=						self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_model(*__snake_case)


		def 		A__					(				self):
			_UpperCamelCase   :		Tuple			=						self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_backbone(*__snake_case)


		def 		A__					(				self):
			_UpperCamelCase	,		_UpperCamelCase   :		Optional[int]			=						self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
				_UpperCamelCase   :		Union[str, Any]			=						model_class(config=__snake_case)
				for name, module in model.named_modules():
					if isinstance(__snake_case	,       (nn.BatchNormad, nn.GroupNorm)):
						self.assertTrue(
						    torch.all(module.weight == 1)	,       msg=f'''Parameter {name} of model {model_class} seems not properly initialized'''	,       )
						self.assertTrue(
						    torch.all(module.bias == 0)	,       msg=f'''Parameter {name} of model {model_class} seems not properly initialized'''	,       )


		def 		A__					(				self):
			def check_hidden_states_output(__snake_case	,       __snake_case	,       __snake_case):
				_UpperCamelCase   :		str			=						model_class(__snake_case)
				model.to(__snake_case)
				model.eval()

				with torch.no_grad():
					_UpperCamelCase   :		Union[str, Any]			=						model(**self._prepare_for_class(__snake_case	,       __snake_case))

				_UpperCamelCase   :		List[Any]			=						outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

				_UpperCamelCase   :		str			=						self.model_tester.num_stages
				self.assertEqual(len(__snake_case)	,       expected_num_stages + 1)

				# Bit's feature maps are of shape (batch_size, num_channels, height, width)
				self.assertListEqual(
				    list(hidden_states[0].shape[-2:])	,       [self.model_tester.image_size // 4, self.model_tester.image_size // 4]	,       )

			_UpperCamelCase	,		_UpperCamelCase   :		int			=						self.model_tester.prepare_config_and_inputs_for_common()
			_UpperCamelCase   :		List[str]			=						['preactivation', 'bottleneck']
			for model_class in self.all_model_classes:
				for layer_type in layers_type:
					_UpperCamelCase   :		Any			=						layer_type
					_UpperCamelCase   :		Tuple			=						True
					check_hidden_states_output(__snake_case	,       __snake_case	,       __snake_case)

					# check that output_hidden_states also work using config
					del inputs_dict["output_hidden_states"]
					_UpperCamelCase   :		List[str]			=						True

					check_hidden_states_output(__snake_case	,       __snake_case	,       __snake_case)


		@unittest.skip(reason='Bit does not use feedforward chunking')
		def 		A__					(				self):
			pass


		def 		A__					(				self):
			_UpperCamelCase   :		Union[str, Any]			=						self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_image_classification(*__snake_case)






		@slow
		def 		A__					(				self):
			for model_name in BIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
				_UpperCamelCase   :		Optional[Any]			=						BitModel.from_pretrained(__snake_case)
				self.assertIsNotNone(__snake_case)






def 					lowerCamelCase_     (							)						->				Optional[int]:




	'''simple docstring'''




	_UpperCamelCase   :		List[Any]			=						Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'	)
	return image







@require_torch
@require_vision
class 	lowercase    (				unittest.TestCase							):


		"""simple docstring"""




		@cached_property
		def 		A__					(				self):
			return (
			    BitImageProcessor.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]) if is_vision_available() else None
			)






		@slow
		def 		A__					(				self):
			_UpperCamelCase   :		Optional[Any]			=						BitForImageClassification.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]).to(__snake_case)

			_UpperCamelCase   :		str			=						self.default_image_processor
			_UpperCamelCase   :		List[str]			=						prepare_img()
			_UpperCamelCase   :		int			=						image_processor(images=__snake_case	,       return_tensors='pt').to(__snake_case)

			# forward pass
			with torch.no_grad():
				_UpperCamelCase   :		Any			=						model(**__snake_case)

			# verify the logits
			_UpperCamelCase   :		Dict			=						torch.Size((1, 10_00))
			self.assertEqual(outputs.logits.shape	,       __snake_case)

			_UpperCamelCase   :		Optional[int]			=						torch.tensor([[-0.6_5_2_6, -0.5_2_6_3, -1.4_3_9_8]]).to(__snake_case)

			self.assertTrue(torch.allclose(outputs.logits[0, :3]	,       __snake_case	,       atol=1e-4))







@require_torch
class 	lowercase    (				_lowercase    ,							unittest.TestCase							):


		"""simple docstring"""




		a__					       =  (BitBackbone,) if is_torch_available() else ()
		a__					       =  BitConfig

		a__					       =  False


		def 		A__					(				self):
			_UpperCamelCase   :		List[str]			=						BitModelTester(self)




