


"""simple docstring"""






import inspect
import unittest

from transformers import DataaVecVisionConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
	import torch
	from torch import nn

	from transformers import (
	    MODEL_MAPPING,
	    DataaVecVisionForImageClassification,
	    DataaVecVisionForSemanticSegmentation,
	    DataaVecVisionModel,
	)
	from transformers.models.dataavec.modeling_dataavec_vision import DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
	from PIL import Image

	from transformers import BeitImageProcessor


class       lowerCamelCase_		:





				"""simple docstring"""







				def __init__(    self							,		UpperCAmelCase__							,		UpperCAmelCase__=100							,		UpperCAmelCase__=13							,		UpperCAmelCase__=30							,		UpperCAmelCase__=2							,		UpperCAmelCase__=3							,		UpperCAmelCase__=True							,		UpperCAmelCase__=True							,		UpperCAmelCase__=32							,		UpperCAmelCase__=4							,		UpperCAmelCase__=4							,		UpperCAmelCase__=37							,		UpperCAmelCase__="gelu"							,		UpperCAmelCase__=0.1							,		UpperCAmelCase__=0.1							,		UpperCAmelCase__=10							,		UpperCAmelCase__=0.02							,		UpperCAmelCase__=3							,		UpperCAmelCase__=None							,		UpperCAmelCase__=[0, 1, 2, 3]							,		):
								SCREAMING_SNAKE_CASE__  				=					parent
								SCREAMING_SNAKE_CASE__  				=					100
								SCREAMING_SNAKE_CASE__  				=					batch_size
								SCREAMING_SNAKE_CASE__  				=					image_size
								SCREAMING_SNAKE_CASE__  				=					patch_size
								SCREAMING_SNAKE_CASE__  				=					num_channels
								SCREAMING_SNAKE_CASE__  				=					is_training
								SCREAMING_SNAKE_CASE__  				=					use_labels
								SCREAMING_SNAKE_CASE__  				=					hidden_size
								SCREAMING_SNAKE_CASE__  				=					num_hidden_layers
								SCREAMING_SNAKE_CASE__  				=					num_attention_heads
								SCREAMING_SNAKE_CASE__  				=					intermediate_size
								SCREAMING_SNAKE_CASE__  				=					hidden_act
								SCREAMING_SNAKE_CASE__  				=					hidden_dropout_prob
								SCREAMING_SNAKE_CASE__  				=					attention_probs_dropout_prob
								SCREAMING_SNAKE_CASE__  				=					type_sequence_label_size
								SCREAMING_SNAKE_CASE__  				=					initializer_range
								SCREAMING_SNAKE_CASE__  				=					scope
								SCREAMING_SNAKE_CASE__  				=					out_indices
								SCREAMING_SNAKE_CASE__  				=					num_labels

								# in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
								SCREAMING_SNAKE_CASE__  				=					(image_size // patch_size) ** 2
								SCREAMING_SNAKE_CASE__  				=					num_patches + 1



				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

								SCREAMING_SNAKE_CASE__  				=					None
								SCREAMING_SNAKE_CASE__  				=					None
								if self.use_labels:
												SCREAMING_SNAKE_CASE__  				=					ids_tensor([self.batch_size]							,		self.type_sequence_label_size  )
												SCREAMING_SNAKE_CASE__  				=					ids_tensor([self.batch_size, self.image_size, self.image_size]							,		self.num_labels  )

								SCREAMING_SNAKE_CASE__  				=					self.get_config()

								return config, pixel_values, labels, pixel_labels



				def 	lowerCAmelCase__		(    self  ):
								return DataaVecVisionConfig(
								    vocab_size=self.vocab_size							,		image_size=self.image_size							,		patch_size=self.patch_size							,		num_channels=self.num_channels							,		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							,		is_decoder=UpperCAmelCase__							,		initializer_range=self.initializer_range							,		out_indices=self.out_indices							,		)



				def 	lowerCAmelCase__		(    self							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__  ):
								SCREAMING_SNAKE_CASE__  				=					DataaVecVisionModel(config=UpperCAmelCase__  )
								model.to(UpperCAmelCase__  )
								model.eval()
								SCREAMING_SNAKE_CASE__  				=					model(UpperCAmelCase__  )
								# expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
								SCREAMING_SNAKE_CASE__  				=					(self.image_size // self.patch_size) ** 2
								self.parent.assertEqual(result.last_hidden_state.shape							,		(self.batch_size, num_patches + 1, self.hidden_size)  )



				def 	lowerCAmelCase__		(    self							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__  ):
								SCREAMING_SNAKE_CASE__  				=					self.type_sequence_label_size
								SCREAMING_SNAKE_CASE__  				=					DataaVecVisionForImageClassification(UpperCAmelCase__  )
								model.to(UpperCAmelCase__  )
								model.eval()
								SCREAMING_SNAKE_CASE__  				=					model(UpperCAmelCase__							,		labels=UpperCAmelCase__  )
								self.parent.assertEqual(result.logits.shape							,		(self.batch_size, self.type_sequence_label_size)  )



				def 	lowerCAmelCase__		(    self							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__  ):
								SCREAMING_SNAKE_CASE__  				=					self.num_labels
								SCREAMING_SNAKE_CASE__  				=					DataaVecVisionForSemanticSegmentation(UpperCAmelCase__  )
								model.to(UpperCAmelCase__  )
								model.eval()
								SCREAMING_SNAKE_CASE__  				=					model(UpperCAmelCase__  )
								self.parent.assertEqual(
								    result.logits.shape							,		(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)  )
								SCREAMING_SNAKE_CASE__  				=					model(UpperCAmelCase__							,		labels=UpperCAmelCase__  )
								self.parent.assertEqual(
								    result.logits.shape							,		(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)  )






				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					self.prepare_config_and_inputs()
								SCREAMING_SNAKE_CASE__				,   SCREAMING_SNAKE_CASE__				,   SCREAMING_SNAKE_CASE__				,   SCREAMING_SNAKE_CASE__  				=					config_and_inputs
								SCREAMING_SNAKE_CASE__  				=					{"pixel_values": pixel_values}
								return config, inputs_dict


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





				"""simple docstring"""







				_lowerCAmelCase   :   Union[str, Any]   				=       (
				    (DataaVecVisionModel, DataaVecVisionForImageClassification, DataaVecVisionForSemanticSegmentation)
				    if is_torch_available()
				    else ()
				)
				_lowerCAmelCase   :   Any   				=       (
				    {
				        "feature-extraction": DataaVecVisionModel,
				        "image-classification": DataaVecVisionForImageClassification,
				        "image-segmentation": DataaVecVisionForSemanticSegmentation,
				    }
				    if is_torch_available()
				    else {}
				)

				_lowerCAmelCase   :   Dict   				=       False
				_lowerCAmelCase   :   Optional[int]   				=       False
				_lowerCAmelCase   :   Any   				=       False



				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					DataaVecVisionModelTester(self  )
								SCREAMING_SNAKE_CASE__  				=					ConfigTester(
								    self							,		config_class=UpperCAmelCase__							,		has_text_modality=UpperCAmelCase__							,		hidden_size=37  )



				def 	lowerCAmelCase__		(    self  ):
								self.config_tester.run_common_tests()



				def 	lowerCAmelCase__		(    self  ):
								# Data2VecVision does not use inputs_embeds
								pass



				@require_torch_multi_gpu
				@unittest.skip(
				    reason="Data2VecVision has some layers using `add_module` which doesn't work well with `nn.DataParallel`"  )
				def 	lowerCAmelCase__		(    self  ):
								pass



				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__				,   SCREAMING_SNAKE_CASE__  				=					self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
												SCREAMING_SNAKE_CASE__  				=					model_class(UpperCAmelCase__  )
												self.assertIsInstance(model.get_input_embeddings()							,		(nn.Module)  )
												SCREAMING_SNAKE_CASE__  				=					model.get_output_embeddings()
												self.assertTrue(x is None or isinstance(UpperCAmelCase__							,		nn.Linear  )  )



				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__				,   SCREAMING_SNAKE_CASE__  				=					self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
												SCREAMING_SNAKE_CASE__  				=					model_class(UpperCAmelCase__  )
												SCREAMING_SNAKE_CASE__  				=					inspect.signature(model.forward  )
												# signature.parameters is an OrderedDict => so arg_names order is deterministic
												SCREAMING_SNAKE_CASE__  				=					[*signature.parameters.keys()]

												SCREAMING_SNAKE_CASE__  				=					["pixel_values"]
												self.assertListEqual(arg_names[:1]							,		UpperCAmelCase__  )



				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*UpperCAmelCase__  )



				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_segmentation(*UpperCAmelCase__  )



				def 	lowerCAmelCase__		(    self  ):
								if not self.model_tester.is_training:
												return

								SCREAMING_SNAKE_CASE__				,   SCREAMING_SNAKE_CASE__  				=					self.model_tester.prepare_config_and_inputs_for_common()
								SCREAMING_SNAKE_CASE__  				=					True

								for model_class in self.all_model_classes:
												if model_class in [*get_values(UpperCAmelCase__  )]:
																continue

												SCREAMING_SNAKE_CASE__  				=					model_class(UpperCAmelCase__  )
												model.to(UpperCAmelCase__  )
												model.train()
												SCREAMING_SNAKE_CASE__  				=					self._prepare_for_class(UpperCAmelCase__							,		UpperCAmelCase__							,		return_labels=UpperCAmelCase__  )
												SCREAMING_SNAKE_CASE__  				=					model(**UpperCAmelCase__  ).loss
												loss.backward()



				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__				,   SCREAMING_SNAKE_CASE__  				=					self.model_tester.prepare_config_and_inputs_for_common()
								if not self.model_tester.is_training:
												return

								SCREAMING_SNAKE_CASE__  				=					False
								SCREAMING_SNAKE_CASE__  				=					True

								for model_class in self.all_model_classes:
												if model_class in [*get_values(UpperCAmelCase__  )] or not model_class.supports_gradient_checkpointing:
																continue
												# TODO: remove the following 3 lines once we have a MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING
												# this can then be incorporated into _prepare_for_class in test_modeling_common.py
												elif model_class.__name__ == "Data2VecVisionForSemanticSegmentation":
																SCREAMING_SNAKE_CASE__				,   SCREAMING_SNAKE_CASE__				,   SCREAMING_SNAKE_CASE__				,   SCREAMING_SNAKE_CASE__  				=					inputs_dict["pixel_values"].shape
																SCREAMING_SNAKE_CASE__  				=					torch.zeros(
																    [self.model_tester.batch_size, height, width]							,		device=UpperCAmelCase__  ).long()
												SCREAMING_SNAKE_CASE__  				=					model_class(UpperCAmelCase__  )
												model.gradient_checkpointing_enable()
												model.to(UpperCAmelCase__  )
												model.train()
												SCREAMING_SNAKE_CASE__  				=					self._prepare_for_class(UpperCAmelCase__							,		UpperCAmelCase__							,		return_labels=UpperCAmelCase__  )
												SCREAMING_SNAKE_CASE__  				=					model(**UpperCAmelCase__  ).loss
												loss.backward()



				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__				,   SCREAMING_SNAKE_CASE__  				=					self.model_tester.prepare_config_and_inputs_for_common()

								SCREAMING_SNAKE_CASE__  				=					_config_zero_init(UpperCAmelCase__  )
								for model_class in self.all_model_classes:
												SCREAMING_SNAKE_CASE__  				=					model_class(config=UpperCAmelCase__  )
												for name, param in model.named_parameters():
																# we skip lambda parameters as these require special initial values
																# determined by config.layer_scale_init_value
																if "lambda" in name:
																				continue
																if param.requires_grad:
																				self.assertIn(
																				    ((param.data.mean() * 1e9).round() / 1e9).item()							,		[0.0, 1.0]							,		msg=f'''Parameter {name} of model {model_class} seems not properly initialized'''							,		)



				def 	lowerCAmelCase__		(    self							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__=2e-4							,		UpperCAmelCase__="outputs"							,		UpperCAmelCase__=None  ):
								# We override with a slightly higher tol value, as semseg models tend to diverge a bit more
								super().check_pt_tf_outputs(UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__  )



				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase__  )






				@slow
				def 	lowerCAmelCase__		(    self  ):
								for model_name in DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												SCREAMING_SNAKE_CASE__  				=					DataaVecVisionModel.from_pretrained(UpperCAmelCase__  )
												self.assertIsNotNone(UpperCAmelCase__  )


def __lowercase  (   ):
				SCREAMING_SNAKE_CASE__  				=					Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"	)
				return image


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





				"""simple docstring"""







				@cached_property
				def 	lowerCAmelCase__		(    self  ):
								return (
								    BeitImageProcessor.from_pretrained("facebook/data2vec-vision-base-ft1k"  ) if is_vision_available() else None
								)






				@slow
				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					DataaVecVisionForImageClassification.from_pretrained("facebook/data2vec-vision-base-ft1k"  ).to(
								    UpperCAmelCase__  )

								SCREAMING_SNAKE_CASE__  				=					self.default_image_processor
								SCREAMING_SNAKE_CASE__  				=					prepare_img()
								SCREAMING_SNAKE_CASE__  				=					image_processor(images=UpperCAmelCase__							,		return_tensors="pt"  ).to(UpperCAmelCase__  )

								# forward pass
								with torch.no_grad():
												SCREAMING_SNAKE_CASE__  				=					model(**UpperCAmelCase__  )
								SCREAMING_SNAKE_CASE__  				=					outputs.logits

								# verify the logits
								SCREAMING_SNAKE_CASE__  				=					torch.Size((1, 1000)  )
								self.assertEqual(logits.shape							,		UpperCAmelCase__  )

								SCREAMING_SNAKE_CASE__  				=					torch.tensor([0.3_277, -0.1_395, 0.0_911]  ).to(UpperCAmelCase__  )

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

								SCREAMING_SNAKE_CASE__  				=					[model.config.labelaid[i] for i in ["remote control, remote", "tabby, tabby cat"]]
								self.assertEqual(logits[0].topk(2  ).indices.cpu().tolist()							,		UpperCAmelCase__  )




"""simple docstring"""






import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionDualGuidedPipeline
from diffusers.utils.testing_utils import load_image, nightly, require_torch_gpu, torch_device


_lowerCamelCase  =	False


@nightly
@require_torch_gpu
class       lowerCamelCase_		(      unittest.TestCase ):





				"""simple docstring"""







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



				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					VersatileDiffusionDualGuidedPipeline.from_pretrained("shi-labs/versatile-diffusion"  )
								# remove text_unet
								pipe.remove_unused_weights()
								pipe.to(UpperCAmelCase__  )
								pipe.set_progress_bar_config(disable=UpperCAmelCase__  )

								SCREAMING_SNAKE_CASE__  				=					load_image(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg"  )

								SCREAMING_SNAKE_CASE__  				=					torch.manual_seed(0  )
								SCREAMING_SNAKE_CASE__  				=					pipe(
								    prompt="first prompt"							,		image=UpperCAmelCase__							,		text_to_image_strength=0.75							,		generator=UpperCAmelCase__							,		guidance_scale=7.5							,		num_inference_steps=2							,		output_type="numpy"							,		).images

								with tempfile.TemporaryDirectory() as tmpdirname:
												pipe.save_pretrained(UpperCAmelCase__  )
												SCREAMING_SNAKE_CASE__  				=					VersatileDiffusionDualGuidedPipeline.from_pretrained(UpperCAmelCase__  )

								pipe.to(UpperCAmelCase__  )
								pipe.set_progress_bar_config(disable=UpperCAmelCase__  )

								SCREAMING_SNAKE_CASE__  				=					generator.manual_seed(0  )
								SCREAMING_SNAKE_CASE__  				=					pipe(
								    prompt="first prompt"							,		image=UpperCAmelCase__							,		text_to_image_strength=0.75							,		generator=UpperCAmelCase__							,		guidance_scale=7.5							,		num_inference_steps=2							,		output_type="numpy"							,		).images

								assert np.abs(image - new_image  ).sum() < 1e-5, "Models don't have the same forward pass"






				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					VersatileDiffusionDualGuidedPipeline.from_pretrained("shi-labs/versatile-diffusion"  )
								pipe.remove_unused_weights()
								pipe.to(UpperCAmelCase__  )
								pipe.set_progress_bar_config(disable=UpperCAmelCase__  )

								SCREAMING_SNAKE_CASE__  				=					"cyberpunk 2077"
								SCREAMING_SNAKE_CASE__  				=					load_image(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg"  )
								SCREAMING_SNAKE_CASE__  				=					torch.manual_seed(0  )
								SCREAMING_SNAKE_CASE__  				=					pipe(
								    prompt=UpperCAmelCase__							,		image=UpperCAmelCase__							,		text_to_image_strength=0.75							,		generator=UpperCAmelCase__							,		guidance_scale=7.5							,		num_inference_steps=50							,		output_type="numpy"							,		).images

								SCREAMING_SNAKE_CASE__  				=					image[0, 253:256, 253:256, -1]

								assert image.shape == (1, 512, 512, 3)
								SCREAMING_SNAKE_CASE__  				=					np.array([0.0_787, 0.0_849, 0.0_826, 0.0_812, 0.0_807, 0.0_795, 0.0_818, 0.0_798, 0.0_779]  )

								assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-2
