






"""simple docstring"""



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 							__SCREAMING_SNAKE_CASE						:


							def __init__(			self      :						Optional[Any]			,    snake_case      :						Optional[Any]			,    snake_case      :						Dict=13			,    snake_case      :						Optional[int]=10			,    snake_case      :						str=3			,    snake_case      :						str=2			,    snake_case      :						List[str]=2			,    snake_case      :						int=2			,    snake_case      :						List[Any]=True			,    snake_case      :						Optional[int]=True			,    snake_case      :						Union[str, Any]=32			,    snake_case      :						Optional[int]=5			,    snake_case      :						Optional[Any]=4			,    snake_case      :						Tuple=37			,    snake_case      :						List[str]="gelu"			,    snake_case      :						Any=0.1			,    snake_case      :						Dict=0.1			,    snake_case      :						Dict=10			,    snake_case      :						Union[str, Any]=0.02			,    snake_case      :						Union[str, Any]=0.9			,    snake_case      :						Optional[int]=None			,    ):
													'''simple docstring'''





													A__						:  List[Any]		  =  parent
													A__						:  Dict		  =  batch_size
													A__						:  Union[str, Any]		  =  image_size
													A__						:  Any		  =  num_channels
													A__						:  Optional[int]		  =  patch_size
													A__						:  Union[str, Any]		  =  tubelet_size
													A__						:  str		  =  num_frames
													A__						:  Optional[Any]		  =  is_training
													A__						:  List[Any]		  =  use_labels
													A__						:  List[str]		  =  hidden_size
													A__						:  List[str]		  =  num_hidden_layers
													A__						:  List[Any]		  =  num_attention_heads
													A__						:  List[Any]		  =  intermediate_size
													A__						:  List[str]		  =  hidden_act
													A__						:  List[Any]		  =  hidden_dropout_prob
													A__						:  Optional[int]		  =  attention_probs_dropout_prob
													A__						:  List[Any]		  =  type_sequence_label_size
													A__						:  Dict		  =  initializer_range
													A__						:  List[str]		  =  mask_ratio
													A__						:  str		  =  scope

													# in VideoMAE, the number of tokens equals num_frames/tubelet_size * num_patches per frame
													A__						:  List[str]		  =  (image_size // patch_size) ** 2
													A__						:  List[str]		  =  (num_frames // tubelet_size) * self.num_patches_per_frame

													# use this variable to define bool_masked_pos
													A__						:  List[str]		  =  int(mask_ratio * self.seq_length							)


							def 						_UpperCamelCase       (			self      :						str							):
													'''simple docstring'''





													A__						:  Dict		  =  floats_tensor(
													    [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]							)

													A__						:  Union[str, Any]		  =  None
													if self.use_labels:
																			A__						:  Optional[Any]		  =  ids_tensor([self.batch_size]			,    self.type_sequence_label_size							)

													A__						:  str		  =  self.get_config()

													return config, pixel_values, labels


							def 						_UpperCamelCase       (			self      :						str							):
													'''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=snake_case			,    initializer_range=self.initializer_range			,    )


							def 						_UpperCamelCase       (			self      :						int			,    snake_case      :						str			,    snake_case      :						List[str]			,    snake_case      :						int							):
													'''simple docstring'''





													A__						:  Optional[int]		  =  VideoMAEModel(config=snake_case							)
													model.to(snake_case							)
													model.eval()
													A__						:  Optional[Any]		  =  model(snake_case							)
													self.parent.assertEqual(result.last_hidden_state.shape			,    (self.batch_size, self.seq_length, self.hidden_size)							)


							def 						_UpperCamelCase       (			self      :						Dict			,    snake_case      :						Dict			,    snake_case      :						Optional[Any]			,    snake_case      :						Tuple							):
													'''simple docstring'''





													A__						:  str		  =  VideoMAEForPreTraining(snake_case							)
													model.to(snake_case							)
													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
													A__						:  Any		  =  torch.ones((self.num_masks,)							)
													A__						:  List[Any]		  =  torch.cat([mask, torch.zeros(self.seq_length - mask.size(0							)							)]							)
													A__						:  Any		  =  mask.expand(self.batch_size			,    -1							).bool()

													A__						:  Dict		  =  model(snake_case			,    snake_case							)
													# model only returns predictions for masked patches
													A__						:  Tuple		  =  mask.sum().item()
													A__						:  List[str]		  =  3 * self.tubelet_size * self.patch_size**2
													self.parent.assertEqual(result.logits.shape			,    (self.batch_size, num_masked_patches, decoder_num_labels)							)


							def 						_UpperCamelCase       (			self      :						Any							):
													'''simple docstring'''





													A__						:  Any		  =  self.prepare_config_and_inputs()
													A__  ,  A__  ,  A__						:  Dict		  =  config_and_inputs
													A__						:  Dict		  =  {"""pixel_values""": pixel_values}
													return config, inputs_dict






@require_torch
class 							__SCREAMING_SNAKE_CASE						(				UpperCamelCase       ,		UpperCamelCase       ,		unittest.TestCase		):
							snake_case_     =							(
							    (VideoMAEModel, VideoMAEForPreTraining, VideoMAEForVideoClassification) if is_torch_available() else ()
							)
							snake_case_     =							(
							    {'feature-extraction': VideoMAEModel, 'video-classification': VideoMAEForVideoClassification}
							    if is_torch_available()
							    else {}
							)

							snake_case_     =							False
							snake_case_     =							False
							snake_case_     =							False
							snake_case_     =							False


							def 						_UpperCamelCase       (			self      :						List[str]							):
													'''simple docstring'''





													A__						:  Dict		  =  VideoMAEModelTester(self							)
													A__						:  Any		  =  ConfigTester(self			,    config_class=snake_case			,    has_text_modality=snake_case			,    hidden_size=37							)


							def 						_UpperCamelCase       (			self      :						Optional[int]			,    snake_case      :						Optional[Any]			,    snake_case      :						Union[str, Any]			,    snake_case      :						Optional[Any]=False							):
													'''simple docstring'''





													A__						:  Dict		  =  copy.deepcopy(snake_case							)

													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
																			A__						:  List[Any]		  =  torch.ones((self.model_tester.num_masks,)							)
																			A__						:  Optional[Any]		  =  torch.cat([mask, torch.zeros(self.model_tester.seq_length - mask.size(0							)							)]							)
																			A__						:  Optional[int]		  =  mask.expand(self.model_tester.batch_size			,    -1							).bool()
																			A__						:  Union[str, Any]		  =  bool_masked_pos.to(snake_case							)

													if return_labels:
																			if model_class in [
																			    *get_values(snake_case							),
																			]:
																									A__						:  Optional[Any]		  =  torch.zeros(
																									    self.model_tester.batch_size			,    dtype=torch.long			,    device=snake_case							)

													return inputs_dict


							def 						_UpperCamelCase       (			self      :						Any							):
													'''simple docstring'''





													self.config_tester.run_common_tests()


							@unittest.skip(reason="""VideoMAE does not use inputs_embeds"""							)
							def 						_UpperCamelCase       (			self      :						Any							):
													'''simple docstring'''





													pass


							def 						_UpperCamelCase       (			self      :						Optional[Any]							):
													'''simple docstring'''





													A__  ,  A__						:  int		  =  self.model_tester.prepare_config_and_inputs_for_common()

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


							def 						_UpperCamelCase       (			self      :						int							):
													'''simple docstring'''





													A__  ,  A__						:  Optional[Any]		  =  self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																			A__						:  Union[str, Any]		  =  model_class(snake_case							)
																			A__						:  Union[str, Any]		  =  inspect.signature(model.forward							)
																			# signature.parameters is an OrderedDict => so arg_names order is deterministic
																			A__						:  List[Any]		  =  [*signature.parameters.keys()]

																			A__						:  Optional[int]		  =  ["""pixel_values"""]
																			self.assertListEqual(arg_names[:1]			,    snake_case							)


							def 						_UpperCamelCase       (			self      :						Dict							):
													'''simple docstring'''





													A__						:  Any		  =  self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*snake_case							)


							def 						_UpperCamelCase       (			self      :						Tuple							):
													'''simple docstring'''





													A__						:  List[Any]		  =  self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_pretraining(*snake_case							)


							@slow
							def 						_UpperCamelCase       (			self      :						Optional[Any]							):
													'''simple docstring'''





													for model_name in VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			A__						:  Optional[Any]		  =  VideoMAEModel.from_pretrained(snake_case							)
																			self.assertIsNotNone(snake_case							)


							def 						_UpperCamelCase       (			self      :						Tuple							):
													'''simple docstring'''





													if not self.has_attentions:
																			pass

													else:
																			A__  ,  A__						:  Any		  =  self.model_tester.prepare_config_and_inputs_for_common()
																			A__						:  Optional[Any]		  =  True

																			for model_class in self.all_model_classes:
																									A__						:  Union[str, Any]		  =  self.model_tester.seq_length - self.model_tester.num_masks
																									A__						:  Any		  =  (
																									    num_visible_patches if model_class == VideoMAEForPreTraining else self.model_tester.seq_length
																									)

																									A__						:  Tuple		  =  True
																									A__						:  Any		  =  False
																									A__						:  Any		  =  True
																									A__						:  Any		  =  model_class(snake_case							)
																									model.to(snake_case							)
																									model.eval()
																									with torch.no_grad():
																															A__						:  Any		  =  model(**self._prepare_for_class(snake_case			,    snake_case							)							)
																									A__						:  Optional[Any]		  =  outputs.attentions
																									self.assertEqual(len(snake_case							)			,    self.model_tester.num_hidden_layers							)

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

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

																									# Check attention is always last and order is fine
																									A__						:  str		  =  True
																									A__						:  List[Any]		  =  True
																									A__						:  List[Any]		  =  model_class(snake_case							)
																									model.to(snake_case							)
																									model.eval()
																									with torch.no_grad():
																															A__						:  int		  =  model(**self._prepare_for_class(snake_case			,    snake_case							)							)

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

																									A__						:  Union[str, Any]		  =  outputs.attentions

																									self.assertEqual(len(snake_case							)			,    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 						_UpperCamelCase       (			self      :						Optional[Any]							):
													'''simple docstring'''





													def check_hidden_states_output(snake_case      :						Optional[Any]			,    snake_case      :						List[str]			,    snake_case      :						str							):
																			A__						:  Dict		  =  model_class(snake_case							)
																			model.to(snake_case							)
																			model.eval()

																			with torch.no_grad():
																									A__						:  Any		  =  model(**self._prepare_for_class(snake_case			,    snake_case							)							)

																			A__						:  Optional[Any]		  =  outputs.hidden_states
																			A__						:  Any		  =  self.model_tester.num_hidden_layers + 1
																			self.assertEqual(len(snake_case							)			,    snake_case							)

																			A__						:  Any		  =  self.model_tester.seq_length - self.model_tester.num_masks
																			A__						:  Any		  =  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]			,    )

													A__  ,  A__						:  Any		  =  self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																			A__						:  str		  =  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"]
																			A__						:  str		  =  True

																			check_hidden_states_output(snake_case			,    snake_case			,    snake_case							)


							@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""							)
							def 						_UpperCamelCase       (			self      :						Optional[Any]							):
													'''simple docstring'''





													pass




def      _lowerCAmelCase		(							)		->int:
						A__						:  Tuple		  =  hf_hub_download(
						    repo_id="""hf-internal-testing/spaghetti-video""",						filename="""eating_spaghetti.npy""",						repo_type="""dataset"""      )
						A__						:  Tuple		  =  np.load(UpperCAmelCase__      )
						return list(UpperCAmelCase__      )






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


							@cached_property
							def 						_UpperCamelCase       (			self      :						List[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 						_UpperCamelCase       (			self      :						Optional[Any]							):
													'''simple docstring'''





													A__						:  int		  =  VideoMAEForVideoClassification.from_pretrained("""MCG-NJU/videomae-base-finetuned-kinetics"""							).to(
													    snake_case							)

													A__						:  int		  =  self.default_image_processor
													A__						:  int		  =  prepare_video()
													A__						:  List[Any]		  =  image_processor(snake_case			,    return_tensors="""pt"""							).to(snake_case							)

													# forward pass
													with torch.no_grad():
																			A__						:  str		  =  model(**snake_case							)

													# verify the logits
													A__						:  Tuple		  =  torch.Size((1, 400)							)
													self.assertEqual(outputs.logits.shape			,    snake_case							)

													A__						:  str		  =  torch.tensor([0.3669, -0.0688, -0.2421]							).to(snake_case							)

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


							@slow
							def 						_UpperCamelCase       (			self      :						Union[str, Any]							):
													'''simple docstring'''





													A__						:  Optional[Any]		  =  VideoMAEForPreTraining.from_pretrained("""MCG-NJU/videomae-base-short"""							).to(snake_case							)

													A__						:  Any		  =  self.default_image_processor
													A__						:  Union[str, Any]		  =  prepare_video()
													A__						:  List[str]		  =  image_processor(snake_case			,    return_tensors="""pt"""							).to(snake_case							)

													# add boolean mask, indicating which patches to mask
													A__						:  Union[str, Any]		  =  hf_hub_download(repo_id="""hf-internal-testing/bool-masked-pos"""			,    filename="""bool_masked_pos.pt"""							)
													A__						:  Dict		  =  torch.load(snake_case							)

													# forward pass
													with torch.no_grad():
																			A__						:  str		  =  model(**snake_case							)

													# verify the logits
													A__						:  Optional[int]		  =  torch.Size([1, 1408, 1536]							)
													A__						:  str		  =  torch.tensor(
													    [[0.7994, 0.9612, 0.8508], [0.7401, 0.8958, 0.8302], [0.5862, 0.7468, 0.7325]]			,    device=snake_case							)
													self.assertEqual(outputs.logits.shape			,    snake_case							)
													self.assertTrue(torch.allclose(outputs.logits[0, :3, :3]			,    snake_case			,    atol=1e-4							)							)

													# verify the loss (`config.norm_pix_loss` = `True`)
													A__						:  int		  =  torch.tensor([0.5142]			,    device=snake_case							)
													self.assertTrue(torch.allclose(outputs.loss			,    snake_case			,    atol=1e-4							)							)

													# verify the loss (`config.norm_pix_loss` = `False`)
													A__						:  List[Any]		  =  VideoMAEForPreTraining.from_pretrained("""MCG-NJU/videomae-base-short"""			,    norm_pix_loss=snake_case							).to(
													    snake_case							)

													with torch.no_grad():
																			A__						:  Optional[int]		  =  model(**snake_case							)

													A__						:  List[Any]		  =  torch.tensor(torch.tensor([0.6469]							)			,    device=snake_case							)
													self.assertTrue(torch.allclose(outputs.loss			,    snake_case			,    atol=1e-4							)							)














"""simple docstring"""



from typing import List, Optional, Tuple, Union

import torch

from ...models import UNetaDModel
from ...schedulers import KarrasVeScheduler
from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput






class 							__SCREAMING_SNAKE_CASE						(				UpperCamelCase		):
							snake_case_     =							42
							snake_case_     =							42


							def __init__(			self      :						Union[str, Any]			,    snake_case      :						UNetaDModel			,    snake_case      :						KarrasVeScheduler							):
													'''simple docstring'''





													super().__init__()
													self.register_modules(unet=snake_case			,    scheduler=snake_case							)


							@torch.no_grad()
							def __call__(			self      :						Any			,    snake_case      :						int = 1			,    snake_case      :						int = 50			,    snake_case      :						Optional[Union[torch.Generator, List[torch.Generator]]] = None			,    snake_case      :						Optional[str] = "pil"			,    snake_case      :						bool = True			,    **snake_case      :						Dict			,    ):
													'''simple docstring'''





													A__						:  Optional[int]		  =  self.unet.config.sample_size
													A__						:  Dict		  =  (batch_size, 3, img_size, img_size)

													A__						:  Tuple		  =  self.unet

													# sample x_0 ~ N(0, sigma_0^2 * I)
													A__						:  Union[str, Any]		  =  randn_tensor(snake_case			,    generator=snake_case			,    device=self.device							) * self.scheduler.init_noise_sigma

													self.scheduler.set_timesteps(snake_case							)

													for t in self.progress_bar(self.scheduler.timesteps							):
																			# here sigma_t == t_i from the paper
																			A__						:  Optional[Any]		  =  self.scheduler.schedule[t]
																			A__						:  Union[str, Any]		  =  self.scheduler.schedule[t - 1] if t > 0 else 0

																			# 1. Select temporarily increased noise level sigma_hat
																			# 2. Add new noise to move from sample_i to sample_hat
																			A__  ,  A__						:  Dict		  =  self.scheduler.add_noise_to_input(snake_case			,    snake_case			,    generator=snake_case							)

																			# 3. Predict the noise residual given the noise magnitude `sigma_hat`
																			# The model inputs and output are adjusted by following eq. (213) in [1].
																			A__						:  int		  =  (sigma_hat / 2) * model((sample_hat + 1) / 2			,    sigma_hat / 2							).sample

																			# 4. Evaluate dx/dt at sigma_hat
																			# 5. Take Euler step from sigma to sigma_prev
																			A__						:  Tuple		  =  self.scheduler.step(snake_case			,    snake_case			,    snake_case			,    snake_case							)

																			if sigma_prev != 0:
																									# 6. Apply 2nd order correction
																									# The model inputs and output are adjusted by following eq. (213) in [1].
																									A__						:  Dict		  =  (sigma_prev / 2) * model((step_output.prev_sample + 1) / 2			,    sigma_prev / 2							).sample
																									A__						:  Tuple		  =  self.scheduler.step_correct(
																									    snake_case			,    snake_case			,    snake_case			,    snake_case			,    step_output.prev_sample			,    step_output["""derivative"""]			,    )
																			A__						:  Tuple		  =  step_output.prev_sample

													A__						:  Dict		  =  (sample / 2 + 0.5).clamp(0			,    1							)
													A__						:  Dict		  =  sample.cpu().permute(0			,    2			,    3			,    1							).numpy()
													if output_type == "pil":
																			A__						:  Optional[Any]		  =  self.numpy_to_pil(snake_case							)

													if not return_dict:
																			return (image,)

													return ImagePipelineOutput(images=snake_case							)






