import copy
import inspect
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import VideoMAEConfig
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
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch
		from torch import nn

		from transformers import (
		    MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING,
		    VideoMAEForPreTraining,
		    VideoMAEForVideoClassification,
		    VideoMAEModel,
		)
		from transformers.models.videomae.modeling_videomae import VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
		from transformers import VideoMAEImageProcessor




class   _UpperCamelCase :


			"""simple docstring"""






			def __init__(  self ,  lowerCAmelCase__ ,  lowerCAmelCase__=13 ,  lowerCAmelCase__=10 ,  lowerCAmelCase__=3 ,  lowerCAmelCase__=2 ,  lowerCAmelCase__=2 ,  lowerCAmelCase__=2 ,  lowerCAmelCase__=True ,  lowerCAmelCase__=True ,  lowerCAmelCase__=32 ,  lowerCAmelCase__=5 ,  lowerCAmelCase__=4 ,  lowerCAmelCase__=37 ,  lowerCAmelCase__="gelu" ,  lowerCAmelCase__=0.1 ,  lowerCAmelCase__=0.1 ,  lowerCAmelCase__=10 ,  lowerCAmelCase__=0.02 ,  lowerCAmelCase__=0.9 ,  lowerCAmelCase__=None ,  )	->   Union[str, Any]:

							'''simple docstring'''


							__lowercase					     =				parent
							__lowercase					     =				batch_size
							__lowercase					     =				image_size
							__lowercase					     =				num_channels
							__lowercase					     =				patch_size
							__lowercase					     =				tubelet_size
							__lowercase					     =				num_frames
							__lowercase					     =				is_training
							__lowercase					     =				use_labels
							__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					     =				type_sequence_label_size
							__lowercase					     =				initializer_range
							__lowercase					     =				mask_ratio
							__lowercase					     =				scope

							# in VideoMAE, the number of tokens equals num_frames/tubelet_size * num_patches per frame
							__lowercase					     =				(image_size // patch_size) ** 2
							__lowercase					     =				(num_frames // tubelet_size) * self.num_patches_per_frame

							# use this variable to define bool_masked_pos
							__lowercase					     =				int(mask_ratio * self.seq_length		)





			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   Optional[int]:

							'''simple docstring'''


							__lowercase					     =				floats_tensor(
							    [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]		)

							__lowercase					     =				None
							if self.use_labels:
											__lowercase					     =				ids_tensor([self.batch_size] ,  self.type_sequence_label_size		)

							__lowercase					     =				self.get_config()

							return config, pixel_values, labels





			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   Union[str, Any]:

							'''simple docstring'''


							return VideoMAEConfig(
							    image_size=self.image_size ,  patch_size=self.patch_size ,  num_channels=self.num_channels ,  num_frames=self.num_frames ,  tubelet_size=self.tubelet_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 ,  is_decoder=lowerCAmelCase__ ,  initializer_range=self.initializer_range ,  )





			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__		)	->   List[str]:

							'''simple docstring'''


							__lowercase					     =				VideoMAEModel(config=lowerCAmelCase__		)
							model.to(lowerCAmelCase__		)
							model.eval()
							__lowercase					     =				model(lowerCAmelCase__		)
							self.parent.assertEqual(result.last_hidden_state.shape ,  (self.batch_size, self.seq_length, self.hidden_size)		)





			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__		)	->   str:

							'''simple docstring'''


							__lowercase					     =				VideoMAEForPreTraining(lowerCAmelCase__		)
							model.to(lowerCAmelCase__		)
							model.eval()
							# important: each video needs to have the same number of masked patches
							# hence we define a single mask, which we then repeat for each example in the batch
							__lowercase					     =				torch.ones((self.num_masks,)		)
							__lowercase					     =				torch.cat([mask, torch.zeros(self.seq_length - mask.size(0		)		)]		)
							__lowercase					     =				mask.expand(self.batch_size ,  -1		).bool()

							__lowercase					     =				model(lowerCAmelCase__ ,  lowerCAmelCase__		)
							# model only returns predictions for masked patches
							__lowercase					     =				mask.sum().item()
							__lowercase					     =				3 * self.tubelet_size * self.patch_size**2
							self.parent.assertEqual(result.logits.shape ,  (self.batch_size, num_masked_patches, decoder_num_labels)		)





			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   Union[str, Any]:

							'''simple docstring'''


							__lowercase					     =				self.prepare_config_and_inputs()
							__lowercase     ,	__lowercase     ,	__lowercase					     =				config_and_inputs
							__lowercase					     =				{'''pixel_values''': pixel_values}
							return config, inputs_dict




@require_torch
class   _UpperCamelCase (		_UpperCAmelCase		,_UpperCAmelCase		,unittest.TestCase					):


			"""simple docstring"""






			__a    :					str     					=		(
			    (VideoMAEModel, VideoMAEForPreTraining, VideoMAEForVideoClassification) if is_torch_available() else ()
			)
			__a    :					Optional[Any]     					=		(
			    {'''feature-extraction''': VideoMAEModel, '''video-classification''': VideoMAEForVideoClassification}
			    if is_torch_available()
			    else {}
			)

			__a    :					Union[str, Any]     					=		False
			__a    :					List[str]     					=		False
			__a    :					Any     					=		False
			__a    :					Union[str, Any]     					=		False





			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   List[Any]:

							'''simple docstring'''


							__lowercase					     =				VideoMAEModelTester(self		)
							__lowercase					     =				ConfigTester(self ,  config_class=lowerCAmelCase__ ,  has_text_modality=lowerCAmelCase__ ,  hidden_size=37		)





			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__=False		)	->   Optional[Any]:

							'''simple docstring'''


							__lowercase					     =				copy.deepcopy(lowerCAmelCase__		)

							if model_class == VideoMAEForPreTraining:
											# important: each video needs to have the same number of masked patches
											# hence we define a single mask, which we then repeat for each example in the batch
											__lowercase					     =				torch.ones((self.model_tester.num_masks,)		)
											__lowercase					     =				torch.cat([mask, torch.zeros(self.model_tester.seq_length - mask.size(0		)		)]		)
											__lowercase					     =				mask.expand(self.model_tester.batch_size ,  -1		).bool()
											__lowercase					     =				bool_masked_pos.to(lowerCAmelCase__		)

							if return_labels:
											if model_class in [
											    *get_values(lowerCAmelCase__		),
											]:
															__lowercase					     =				torch.zeros(
															    self.model_tester.batch_size ,  dtype=torch.long ,  device=lowerCAmelCase__		)

							return inputs_dict





			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   str:

							'''simple docstring'''


							self.config_tester.run_common_tests()





			@unittest.skip(reason='''VideoMAE does not use inputs_embeds'''		)
			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   Optional[int]:

							'''simple docstring'''


							pass





			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   Tuple:

							'''simple docstring'''


							__lowercase     ,	__lowercase					     =				self.model_tester.prepare_config_and_inputs_for_common()

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





			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   Dict:

							'''simple docstring'''


							__lowercase     ,	__lowercase					     =				self.model_tester.prepare_config_and_inputs_for_common()

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

											__lowercase					     =				['''pixel_values''']
											self.assertListEqual(arg_names[:1] ,  lowerCAmelCase__		)





			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   Dict:

							'''simple docstring'''


							__lowercase					     =				self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*lowerCAmelCase__		)





			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   Dict:

							'''simple docstring'''


							__lowercase					     =				self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_pretraining(*lowerCAmelCase__		)





			@slow
			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   Dict:

							'''simple docstring'''


							for model_name in VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											__lowercase					     =				VideoMAEModel.from_pretrained(lowerCAmelCase__		)
											self.assertIsNotNone(lowerCAmelCase__		)





			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   Tuple:

							'''simple docstring'''


							if not self.has_attentions:
											pass

							else:
											__lowercase     ,	__lowercase					     =				self.model_tester.prepare_config_and_inputs_for_common()
											__lowercase					     =				True

											for model_class in self.all_model_classes:
															__lowercase					     =				self.model_tester.seq_length - self.model_tester.num_masks
															__lowercase					     =				(
															    num_visible_patches if model_class == VideoMAEForPreTraining else self.model_tester.seq_length
															)

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

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

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

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

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

															__lowercase					     =				outputs.attentions

															self.assertEqual(len(lowerCAmelCase__		) ,  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 			_SCREAMING_SNAKE_CASE     (  self		)	->   Dict:

							'''simple docstring'''


							def check_hidden_states_output(lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__		):
											__lowercase					     =				model_class(lowerCAmelCase__		)
											model.to(lowerCAmelCase__		)
											model.eval()

											with torch.no_grad():
															__lowercase					     =				model(**self._prepare_for_class(lowerCAmelCase__ ,  lowerCAmelCase__		)		)

											__lowercase					     =				outputs.hidden_states
											__lowercase					     =				self.model_tester.num_hidden_layers + 1
											self.assertEqual(len(lowerCAmelCase__		) ,  lowerCAmelCase__		)

											__lowercase					     =				self.model_tester.seq_length - self.model_tester.num_masks
											__lowercase					     =				num_visible_patches if model_class == VideoMAEForPreTraining else self.model_tester.seq_length

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

							__lowercase     ,	__lowercase					     =				self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
											__lowercase					     =				True
											check_hidden_states_output(lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__		)

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

											check_hidden_states_output(lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__		)





			@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''		)
			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   str:

							'''simple docstring'''


							pass





def 							UpperCAmelCase		(       ):







				"""simple docstring"""



				__lowercase					     =				hf_hub_download(
				    repo_id='''hf-internal-testing/spaghetti-video'''				,     filename='''eating_spaghetti.npy'''				,     repo_type='''dataset'''				)
				__lowercase					     =				np.load(lowercase				)
				return list(lowercase				)




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


			"""simple docstring"""






			@cached_property
			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   Union[str, Any]:

							'''simple docstring'''


							return (
							    VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5] ,  image_std=[0.5, 0.5, 0.5]		)
							    if is_vision_available()
							    else None
							)





			@slow
			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   Union[str, Any]:

							'''simple docstring'''


							__lowercase					     =				VideoMAEForVideoClassification.from_pretrained('''MCG-NJU/videomae-base-finetuned-kinetics'''		).to(
							    lowerCAmelCase__		)

							__lowercase					     =				self.default_image_processor
							__lowercase					     =				prepare_video()
							__lowercase					     =				image_processor(lowerCAmelCase__ ,  return_tensors='''pt'''		).to(lowerCAmelCase__		)

							# forward pass
							with torch.no_grad():
											__lowercase					     =				model(**lowerCAmelCase__		)

							# verify the logits
							__lowercase					     =				torch.Size((1, 4_00)		)
							self.assertEqual(outputs.logits.shape ,  lowerCAmelCase__		)

							__lowercase					     =				torch.tensor([0.3669, -0.0688, -0.2421]		).to(lowerCAmelCase__		)

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





			@slow
			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   List[str]:

							'''simple docstring'''


							__lowercase					     =				VideoMAEForPreTraining.from_pretrained('''MCG-NJU/videomae-base-short'''		).to(lowerCAmelCase__		)

							__lowercase					     =				self.default_image_processor
							__lowercase					     =				prepare_video()
							__lowercase					     =				image_processor(lowerCAmelCase__ ,  return_tensors='''pt'''		).to(lowerCAmelCase__		)

							# add boolean mask, indicating which patches to mask
							__lowercase					     =				hf_hub_download(repo_id='''hf-internal-testing/bool-masked-pos''' ,  filename='''bool_masked_pos.pt'''		)
							__lowercase					     =				torch.load(lowerCAmelCase__		)

							# forward pass
							with torch.no_grad():
											__lowercase					     =				model(**lowerCAmelCase__		)

							# verify the logits
							__lowercase					     =				torch.Size([1, 14_08, 15_36]		)
							__lowercase					     =				torch.tensor(
							    [[0.7994, 0.9612, 0.8508], [0.7401, 0.8958, 0.8302], [0.5862, 0.7468, 0.7325]] ,  device=lowerCAmelCase__		)
							self.assertEqual(outputs.logits.shape ,  lowerCAmelCase__		)
							self.assertTrue(torch.allclose(outputs.logits[0, :3, :3] ,  lowerCAmelCase__ ,  atol=1E-4		)		)

							# verify the loss (`config.norm_pix_loss` = `True`)
							__lowercase					     =				torch.tensor([0.5142] ,  device=lowerCAmelCase__		)
							self.assertTrue(torch.allclose(outputs.loss ,  lowerCAmelCase__ ,  atol=1E-4		)		)

							# verify the loss (`config.norm_pix_loss` = `False`)
							__lowercase					     =				VideoMAEForPreTraining.from_pretrained('''MCG-NJU/videomae-base-short''' ,  norm_pix_loss=lowerCAmelCase__		).to(
							    lowerCAmelCase__		)

							with torch.no_grad():
											__lowercase					     =				model(**lowerCAmelCase__		)

							__lowercase					     =				torch.tensor(torch.tensor([0.6469]		) ,  device=lowerCAmelCase__		)
							self.assertTrue(torch.allclose(outputs.loss ,  lowerCAmelCase__ ,  atol=1E-4		)		)
from argparse import ArgumentParser

from .env import EnvironmentCommand





def 							UpperCAmelCase		(       ):







				"""simple docstring"""



				__lowercase					     =				ArgumentParser('''Diffusers CLI tool'''				,     usage='''diffusers-cli <command> [<args>]'''				)
				__lowercase					     =				parser.add_subparsers(help='''diffusers-cli command helpers'''				)

				# Register commands
				EnvironmentCommand.register_subcommand(lowercase				)

				# Let's go
				__lowercase					     =				parser.parse_args()

				if not hasattr(lowercase				,     '''func'''				):
								parser.print_help()
								exit(1				)

				# Run
				__lowercase					     =				args.func(lowercase				)
				service.run()


if __name__ == "__main__":
		main()