


"""simple docstring"""






import copy
import inspect
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import VivitConfig
from transformers.models.auto import get_values
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_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor


if is_torch_available():
	import torch
	from torch import nn

	from transformers import MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING, VivitForVideoClassification, VivitModel
	from transformers.models.vivit.modeling_vivit import VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
	from transformers import VivitImageProcessor


class       lowerCamelCase_		:





				"""simple docstring"""







				def __init__(    self							,		UpperCAmelCase__							,		UpperCAmelCase__=2							,		UpperCAmelCase__=True							,		UpperCAmelCase__=True							,		UpperCAmelCase__=10							,		UpperCAmelCase__=10							,		UpperCAmelCase__=8							,		UpperCAmelCase__=[2, 4, 4]							,		UpperCAmelCase__=3							,		UpperCAmelCase__=768							,		UpperCAmelCase__=5							,		UpperCAmelCase__=4							,		UpperCAmelCase__=37							,		UpperCAmelCase__="gelu_fast"							,		UpperCAmelCase__=0.0							,		UpperCAmelCase__=0.0							,		UpperCAmelCase__=0.02							,		UpperCAmelCase__=1e-06							,		UpperCAmelCase__=True							,		UpperCAmelCase__=None							,		):
								SCREAMING_SNAKE_CASE__  				=					parent
								SCREAMING_SNAKE_CASE__  				=					batch_size
								SCREAMING_SNAKE_CASE__  				=					is_training
								SCREAMING_SNAKE_CASE__  				=					use_labels
								SCREAMING_SNAKE_CASE__  				=					num_labels
								SCREAMING_SNAKE_CASE__  				=					image_size
								SCREAMING_SNAKE_CASE__  				=					num_frames
								SCREAMING_SNAKE_CASE__  				=					tubelet_size
								SCREAMING_SNAKE_CASE__  				=					num_channels
								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__  				=					initializer_range
								SCREAMING_SNAKE_CASE__  				=					layer_norm_eps
								SCREAMING_SNAKE_CASE__  				=					qkv_bias
								SCREAMING_SNAKE_CASE__  				=					scope

								SCREAMING_SNAKE_CASE__  				=					(
								    (self.image_size // self.tubelet_size[2])
								    * (self.image_size // self.tubelet_size[1])
								    * (self.num_frames // self.tubelet_size[0])
								) + 1  # CLS token



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

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

								SCREAMING_SNAKE_CASE__  				=					self.get_config()

								return config, pixel_values, labels



				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					VivitConfig(
								    num_frames=self.num_frames							,		image_size=self.image_size							,		tubelet_size=self.tubelet_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							,		initializer_range=self.initializer_range							,		layer_norm_eps=self.layer_norm_eps							,		qkv_bias=self.qkv_bias							,		)
								SCREAMING_SNAKE_CASE__  				=					self.num_labels
								return config



				def 	lowerCAmelCase__		(    self							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__  ):
								SCREAMING_SNAKE_CASE__  				=					VivitModel(config=UpperCAmelCase__  )
								model.to(UpperCAmelCase__  )
								model.eval()
								SCREAMING_SNAKE_CASE__  				=					model(UpperCAmelCase__  )
								self.parent.assertEqual(result.last_hidden_state.shape							,		(self.batch_size, self.seq_length, self.hidden_size)  )



				def 	lowerCAmelCase__		(    self							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__  ):
								SCREAMING_SNAKE_CASE__  				=					VivitForVideoClassification(UpperCAmelCase__  )
								model.to(UpperCAmelCase__  )
								model.eval()

								SCREAMING_SNAKE_CASE__  				=					model(UpperCAmelCase__  )

								# verify the logits shape
								SCREAMING_SNAKE_CASE__  				=					torch.Size((self.batch_size, self.num_labels)  )
								self.parent.assertEqual(result.logits.shape							,		UpperCAmelCase__  )






				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					self.prepare_config_and_inputs()
								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							,	unittest.TestCase ):





				"""simple docstring"""







				_lowerCAmelCase   :   Tuple   				=       (VivitModel, VivitForVideoClassification) if is_torch_available() else ()

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



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



				def 	lowerCAmelCase__		(    self							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__=False  ):
								SCREAMING_SNAKE_CASE__  				=					copy.deepcopy(UpperCAmelCase__  )

								if return_labels:
												if model_class in get_values(UpperCAmelCase__  ):
																SCREAMING_SNAKE_CASE__  				=					torch.zeros(
																    self.model_tester.batch_size							,		dtype=torch.long							,		device=UpperCAmelCase__  )

								return inputs_dict



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



				@unittest.skip(reason="Vivit does not use inputs_embeds"  )
				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", "head_mask"]
												self.assertListEqual(arg_names[:2]							,		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_video_classification(*UpperCAmelCase__  )



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



				def 	lowerCAmelCase__		(    self  ):
								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:
												SCREAMING_SNAKE_CASE__  				=					self.model_tester.seq_length

												SCREAMING_SNAKE_CASE__  				=					True
												SCREAMING_SNAKE_CASE__  				=					False
												SCREAMING_SNAKE_CASE__  				=					True
												SCREAMING_SNAKE_CASE__  				=					model_class(UpperCAmelCase__  )
												model.to(UpperCAmelCase__  )
												model.eval()
												with torch.no_grad():
																SCREAMING_SNAKE_CASE__  				=					model(**self._prepare_for_class(UpperCAmelCase__							,		UpperCAmelCase__  )  )
												SCREAMING_SNAKE_CASE__  				=					outputs.attentions
												self.assertEqual(len(UpperCAmelCase__  )							,		self.model_tester.num_hidden_layers  )

												# check that output_attentions also work using config
												del inputs_dict["output_attentions"]
												SCREAMING_SNAKE_CASE__  				=					True
												SCREAMING_SNAKE_CASE__  				=					model_class(UpperCAmelCase__  )
												model.to(UpperCAmelCase__  )
												model.eval()
												with torch.no_grad():
																SCREAMING_SNAKE_CASE__  				=					model(**self._prepare_for_class(UpperCAmelCase__							,		UpperCAmelCase__  )  )
												SCREAMING_SNAKE_CASE__  				=					outputs.attentions
												self.assertEqual(len(UpperCAmelCase__  )							,		self.model_tester.num_hidden_layers  )

												self.assertListEqual(
												    list(attentions[0].shape[-3:]  )							,		[self.model_tester.num_attention_heads, seq_len, seq_len]							,		)
												SCREAMING_SNAKE_CASE__  				=					len(UpperCAmelCase__  )

												# Check attention is always last and order is fine
												SCREAMING_SNAKE_CASE__  				=					True
												SCREAMING_SNAKE_CASE__  				=					True
												SCREAMING_SNAKE_CASE__  				=					model_class(UpperCAmelCase__  )
												model.to(UpperCAmelCase__  )
												model.eval()
												with torch.no_grad():
																SCREAMING_SNAKE_CASE__  				=					model(**self._prepare_for_class(UpperCAmelCase__							,		UpperCAmelCase__  )  )

												self.assertEqual(out_len + 1							,		len(UpperCAmelCase__  )  )

												SCREAMING_SNAKE_CASE__  				=					outputs.attentions

												self.assertEqual(len(UpperCAmelCase__  )							,		self.model_tester.num_hidden_layers  )

												self.assertListEqual(
												    list(self_attentions[0].shape[-3:]  )							,		[self.model_tester.num_attention_heads, seq_len, seq_len]							,		)



				def 	lowerCAmelCase__		(    self  ):
								def check_hidden_states_output(UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__  ):
												SCREAMING_SNAKE_CASE__  				=					model_class(UpperCAmelCase__  )
												model.to(UpperCAmelCase__  )
												model.eval()

												with torch.no_grad():
																SCREAMING_SNAKE_CASE__  				=					model(**self._prepare_for_class(UpperCAmelCase__							,		UpperCAmelCase__  )  )

												SCREAMING_SNAKE_CASE__  				=					outputs.hidden_states
												SCREAMING_SNAKE_CASE__  				=					self.model_tester.num_hidden_layers + 1
												self.assertEqual(len(UpperCAmelCase__  )							,		UpperCAmelCase__  )

												SCREAMING_SNAKE_CASE__  				=					self.model_tester.seq_length

												self.assertListEqual(
												    list(hidden_states[0].shape[-2:]  )							,		[seq_length, self.model_tester.hidden_size]							,		)

								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__  				=					True
												check_hidden_states_output(UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__  )

												# check that output_hidden_states also work using config
												del inputs_dict["output_hidden_states"]
												SCREAMING_SNAKE_CASE__  				=					True

												check_hidden_states_output(UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__  )






				@unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."  )
				def 	lowerCAmelCase__		(    self  ):
								pass


def __lowercase  (   ):
				SCREAMING_SNAKE_CASE__  				=					hf_hub_download(
				    repo_id="hf-internal-testing/spaghetti-video"    ,       filename="eating_spaghetti_32_frames.npy"    ,       repo_type="dataset"	)
				SCREAMING_SNAKE_CASE__  				=					np.load(lowerCamelCase_	)
				return list(lowerCamelCase_	)


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





				"""simple docstring"""







				@cached_property
				def 	lowerCAmelCase__		(    self  ):
								return VivitImageProcessor() if is_vision_available() else None






				@slow
				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					VivitForVideoClassification.from_pretrained("google/vivit-b-16x2-kinetics400"  ).to(UpperCAmelCase__  )

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

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

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

								# taken from original model
								SCREAMING_SNAKE_CASE__  				=					torch.tensor([-1.0_543, 2.0_764, -0.2_104, 0.4_439, -0.9_658]  ).to(UpperCAmelCase__  )

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




"""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
