






"""simple docstring"""



from logging import getLogger
from typing import Any, Callable, List, Optional, Union

import numpy as np
import PIL
import torch

from ...schedulers import DDPMScheduler
from ..onnx_utils import ORT_TO_NP_TYPE, OnnxRuntimeModel
from ..pipeline_utils import ImagePipelineOutput
from . import StableDiffusionUpscalePipeline


_snake_case       :      List[Any]       =					getLogger(__name__)


_snake_case       :      Any       =					4
_snake_case       :      Optional[int]       =					7

_snake_case       :      List[str]       =					{
    "float16": torch.floataa,
    "float32": torch.floataa,
}




def 						lowerCamelCase_				(  __A  :		Dict			):
					if isinstance(__A			,	torch.Tensor			):
										return image
					elif isinstance(__A			,	PIL.Image.Image			):
										lowerCamelCase__				=[image]

					if isinstance(image[0]			,	PIL.Image.Image			):
										lowerCamelCase__				,     lowerCamelCase__				=image[0].size
										lowerCamelCase__				,     lowerCamelCase__				=(x - x % 64 for x in (w, h))  # resize to integer multiple of 32

										lowerCamelCase__				=[np.array(i.resize((w, h)			)			)[None, :] for i in image]
										lowerCamelCase__				=np.concatenate(__A			,	axis=0			)
										lowerCamelCase__				=np.array(__A			).astype(np.floataa			) / 2_5_5.0
										lowerCamelCase__				=image.transpose(0			,	3			,	1			,	2			)
										lowerCamelCase__				=2.0 * image - 1.0
										lowerCamelCase__				=torch.from_numpy(__A			)
					elif isinstance(image[0]			,	torch.Tensor			):
										lowerCamelCase__				=torch.cat(__A			,	dim=0			)

					return image




class        __lowercase(       _UpperCamelCase  ):






						def __init__(					self    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE = 3_50    ,						)			->  Dict:




											"""simple docstring"""





											super().__init__(
											    vae=_SCREAMING_SNAKE_CASE    ,						text_encoder=_SCREAMING_SNAKE_CASE    ,						tokenizer=_SCREAMING_SNAKE_CASE    ,						unet=_SCREAMING_SNAKE_CASE    ,						low_res_scheduler=_SCREAMING_SNAKE_CASE    ,						scheduler=_SCREAMING_SNAKE_CASE    ,						safety_checker=_SCREAMING_SNAKE_CASE    ,						feature_extractor=_SCREAMING_SNAKE_CASE    ,						watermarker=_SCREAMING_SNAKE_CASE    ,						max_noise_level=_SCREAMING_SNAKE_CASE    ,						)






						def __call__(					self    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE = 75    ,						_SCREAMING_SNAKE_CASE = 9.0    ,						_SCREAMING_SNAKE_CASE = 20    ,						_SCREAMING_SNAKE_CASE = None    ,						_SCREAMING_SNAKE_CASE = 1    ,						_SCREAMING_SNAKE_CASE = 0.0    ,						_SCREAMING_SNAKE_CASE = None    ,						_SCREAMING_SNAKE_CASE = None    ,						_SCREAMING_SNAKE_CASE = None    ,						_SCREAMING_SNAKE_CASE = None    ,						_SCREAMING_SNAKE_CASE = "pil"    ,						_SCREAMING_SNAKE_CASE = True    ,						_SCREAMING_SNAKE_CASE = None    ,						_SCREAMING_SNAKE_CASE = 1    ,						)			->  Tuple:




											"""simple docstring"""





											self.check_inputs(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE)

											# 2. Define call parameters
											if prompt is not None and isinstance(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE):
																lowerCamelCase__				=1
											elif prompt is not None and isinstance(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE):
																lowerCamelCase__				=len(_SCREAMING_SNAKE_CASE)
											else:
																lowerCamelCase__				=prompt_embeds.shape[0]

											lowerCamelCase__				=self._execution_device
											# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
											# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
											# corresponds to doing no classifier free guidance.
											lowerCamelCase__				=guidance_scale > 1.0

											# 3. Encode input prompt
											lowerCamelCase__				=self._encode_prompt(
											    _SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						prompt_embeds=_SCREAMING_SNAKE_CASE    ,						negative_prompt_embeds=_SCREAMING_SNAKE_CASE    ,						)

											lowerCamelCase__				=ORT_TO_PT_TYPE[str(text_embeddings.dtype)]

											# 4. Preprocess image
											lowerCamelCase__				=preprocess(_SCREAMING_SNAKE_CASE)
											lowerCamelCase__				=image.cpu()

											# 5. set timesteps
											self.scheduler.set_timesteps(_SCREAMING_SNAKE_CASE    ,						device=_SCREAMING_SNAKE_CASE)
											lowerCamelCase__				=self.scheduler.timesteps

											# 5. Add noise to image
											lowerCamelCase__				=torch.tensor([noise_level]    ,						dtype=torch.long    ,						device=_SCREAMING_SNAKE_CASE)
											lowerCamelCase__				=torch.randn(image.shape    ,						generator=_SCREAMING_SNAKE_CASE    ,						device=_SCREAMING_SNAKE_CASE    ,						dtype=_SCREAMING_SNAKE_CASE)
											lowerCamelCase__				=self.low_res_scheduler.add_noise(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE)

											lowerCamelCase__				=2 if do_classifier_free_guidance else 1
											lowerCamelCase__				=np.concatenate([image] * batch_multiplier * num_images_per_prompt)
											lowerCamelCase__				=np.concatenate([noise_level] * image.shape[0])

											# 6. Prepare latent variables
											lowerCamelCase__				,     lowerCamelCase__				=image.shape[2:]
											lowerCamelCase__				=self.prepare_latents(
											    batch_size * num_images_per_prompt    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						)

											# 7. Check that sizes of image and latents match
											lowerCamelCase__				=image.shape[1]
											if NUM_LATENT_CHANNELS + num_channels_image != NUM_UNET_INPUT_CHANNELS:
																raise ValueError(
																    'Incorrect configuration settings! The config of `pipeline.unet` expects'
																    F" {NUM_UNET_INPUT_CHANNELS} but received `num_channels_latents`: {NUM_LATENT_CHANNELS} +"
																    F" `num_channels_image`: {num_channels_image} "
																    F" = {NUM_LATENT_CHANNELS+num_channels_image}. Please verify the config of"
																    ' `pipeline.unet` or your `image` input.')

											# 8. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
											lowerCamelCase__				=self.prepare_extra_step_kwargs(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE)

											lowerCamelCase__				=next(
											    (input.type for input in self.unet.model.get_inputs() if input.name == 'timestep')    ,						'tensor(float)')
											lowerCamelCase__				=ORT_TO_NP_TYPE[timestep_dtype]

											# 9. Denoising loop
											lowerCamelCase__				=len(_SCREAMING_SNAKE_CASE) - num_inference_steps * self.scheduler.order
											with self.progress_bar(total=_SCREAMING_SNAKE_CASE) as progress_bar:
																for i, t in enumerate(_SCREAMING_SNAKE_CASE):
																					# expand the latents if we are doing classifier free guidance
																					lowerCamelCase__				=np.concatenate([latents] * 2) if do_classifier_free_guidance else latents

																					# concat latents, mask, masked_image_latents in the channel dimension
																					lowerCamelCase__				=self.scheduler.scale_model_input(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE)
																					lowerCamelCase__				=np.concatenate([latent_model_input, image]    ,						axis=1)

																					# timestep to tensor
																					lowerCamelCase__				=np.array([t]    ,						dtype=_SCREAMING_SNAKE_CASE)

																					# predict the noise residual
																					lowerCamelCase__				=self.unet(
																					    sample=_SCREAMING_SNAKE_CASE    ,						timestep=_SCREAMING_SNAKE_CASE    ,						encoder_hidden_states=_SCREAMING_SNAKE_CASE    ,						class_labels=noise_level.astype(np.intaa)    ,						)[0]

																					# perform guidance
																					if do_classifier_free_guidance:
																										lowerCamelCase__				,     lowerCamelCase__				=np.split(_SCREAMING_SNAKE_CASE    ,						2)
																					lowerCamelCase__				=noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

																					# compute the previous noisy sample x_t -> x_t-1
																					lowerCamelCase__				=self.scheduler.step(
																					    torch.from_numpy(_SCREAMING_SNAKE_CASE)    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						**_SCREAMING_SNAKE_CASE).prev_sample

																					# call the callback, if provided
																					if i == len(_SCREAMING_SNAKE_CASE) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
																										progress_bar.update()
																										if callback is not None and i % callback_steps == 0:
																															callback(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE)

        # 10. Post-processing
											lowerCamelCase__				=self.decode_latents(latents.float())

											# 11. Convert to PIL
											if output_type == "pil":
																lowerCamelCase__				=self.numpy_to_pil(_SCREAMING_SNAKE_CASE)

											if not return_dict:
																return (image,)

											return ImagePipelineOutput(images=_SCREAMING_SNAKE_CASE)






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




											"""simple docstring"""





											lowerCamelCase__				=1 / 0.08333 * latents
											lowerCamelCase__				=self.vae(latent_sample=_SCREAMING_SNAKE_CASE)[0]
											lowerCamelCase__				=np.clip(image / 2 + 0.5    ,						0    ,						1)
											lowerCamelCase__				=image.transpose((0, 2, 3, 1))
											return image






						def 		lowerCAmelCase   (					self    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE = None    ,						_SCREAMING_SNAKE_CASE = None    ,						)			->  Dict:




											"""simple docstring"""





											if prompt is not None and isinstance(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE):
																lowerCamelCase__				=1
											elif prompt is not None and isinstance(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE):
																lowerCamelCase__				=len(_SCREAMING_SNAKE_CASE)
											else:
																lowerCamelCase__				=prompt_embeds.shape[0]

											if prompt_embeds is None:
																lowerCamelCase__				=self.tokenizer(
																    _SCREAMING_SNAKE_CASE    ,						padding='max_length'    ,						max_length=self.tokenizer.model_max_length    ,						truncation=_SCREAMING_SNAKE_CASE    ,						return_tensors='pt'    ,						)
																lowerCamelCase__				=text_inputs.input_ids
																lowerCamelCase__				=self.tokenizer(_SCREAMING_SNAKE_CASE    ,						padding='longest'    ,						return_tensors='pt').input_ids

																if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
																    _SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE):
																					lowerCamelCase__				=self.tokenizer.batch_decode(
																					    untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1])
																					logger.warning(
																					    'The following part of your input was truncated because CLIP can only handle sequences up to'
																					    F" {self.tokenizer.model_max_length} tokens: {removed_text}")

																# no positional arguments to text_encoder
																lowerCamelCase__				=self.text_encoder(
																    input_ids=text_input_ids.int().to(_SCREAMING_SNAKE_CASE)    ,						)
																lowerCamelCase__				=prompt_embeds[0]

											lowerCamelCase__				,     lowerCamelCase__				,     lowerCamelCase__				=prompt_embeds.shape
											# duplicate text embeddings for each generation per prompt, using mps friendly method
											lowerCamelCase__				=prompt_embeds.repeat(1    ,						_SCREAMING_SNAKE_CASE)
											lowerCamelCase__				=prompt_embeds.reshape(bs_embed * num_images_per_prompt    ,						_SCREAMING_SNAKE_CASE    ,						-1)

											# get unconditional embeddings for classifier free guidance
											if do_classifier_free_guidance and negative_prompt_embeds is None:
																lowerCamelCase__				=42
																if negative_prompt is None:
																					lowerCamelCase__				=[''] * batch_size
																elif type(_SCREAMING_SNAKE_CASE) is not type(_SCREAMING_SNAKE_CASE):
																					raise TypeError(
																					    F"`negative_prompt` should be the same type to `prompt`, but got {type(_SCREAMING_SNAKE_CASE)} !="
																					    F" {type(_SCREAMING_SNAKE_CASE)}.")
																elif isinstance(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE):
																					lowerCamelCase__				=[negative_prompt]
																elif batch_size != len(_SCREAMING_SNAKE_CASE):
																					raise ValueError(
																					    F"`negative_prompt`: {negative_prompt} has batch size {len(_SCREAMING_SNAKE_CASE)}, but `prompt`:"
																					    F" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
																					    ' the batch size of `prompt`.')
																else:
																					lowerCamelCase__				=negative_prompt

																lowerCamelCase__				=text_input_ids.shape[-1]
																lowerCamelCase__				=self.tokenizer(
																    _SCREAMING_SNAKE_CASE    ,						padding='max_length'    ,						max_length=_SCREAMING_SNAKE_CASE    ,						truncation=_SCREAMING_SNAKE_CASE    ,						return_tensors='pt'    ,						)

																# if hasattr(uncond_input, "attention_mask"):
																#     attention_mask = uncond_input.attention_mask.to(device)
																# else:
																#     attention_mask = None

																lowerCamelCase__				=self.text_encoder(
																    input_ids=uncond_input.input_ids.int().to(_SCREAMING_SNAKE_CASE)    ,						)
																lowerCamelCase__				=uncond_embeddings[0]

											if do_classifier_free_guidance:
																lowerCamelCase__				=uncond_embeddings.shape[1]
																# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
																lowerCamelCase__				=uncond_embeddings.repeat(1    ,						_SCREAMING_SNAKE_CASE)
																lowerCamelCase__				=uncond_embeddings.reshape(batch_size * num_images_per_prompt    ,						_SCREAMING_SNAKE_CASE    ,						-1)

																# For classifier free guidance, we need to do two forward passes.
																# Here we concatenate the unconditional and text embeddings into a single batch
																# to avoid doing two forward passes
																lowerCamelCase__				=np.concatenate([uncond_embeddings, prompt_embeds])

											return prompt_embeds







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






						def __init__(					self    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE=2    ,						_SCREAMING_SNAKE_CASE=True    ,						_SCREAMING_SNAKE_CASE=True    ,						_SCREAMING_SNAKE_CASE=10    ,						_SCREAMING_SNAKE_CASE=10    ,						_SCREAMING_SNAKE_CASE=8    ,						_SCREAMING_SNAKE_CASE=[2, 4, 4]    ,						_SCREAMING_SNAKE_CASE=3    ,						_SCREAMING_SNAKE_CASE=7_68    ,						_SCREAMING_SNAKE_CASE=5    ,						_SCREAMING_SNAKE_CASE=4    ,						_SCREAMING_SNAKE_CASE=37    ,						_SCREAMING_SNAKE_CASE="gelu_fast"    ,						_SCREAMING_SNAKE_CASE=0.0    ,						_SCREAMING_SNAKE_CASE=0.0    ,						_SCREAMING_SNAKE_CASE=0.02    ,						_SCREAMING_SNAKE_CASE=1E-06    ,						_SCREAMING_SNAKE_CASE=True    ,						_SCREAMING_SNAKE_CASE=None    ,						)			->  Dict:




											"""simple docstring"""





											lowerCamelCase__				=parent
											lowerCamelCase__				=batch_size
											lowerCamelCase__				=is_training
											lowerCamelCase__				=use_labels
											lowerCamelCase__				=num_labels
											lowerCamelCase__				=image_size
											lowerCamelCase__				=num_frames
											lowerCamelCase__				=tubelet_size
											lowerCamelCase__				=num_channels
											lowerCamelCase__				=hidden_size
											lowerCamelCase__				=num_hidden_layers
											lowerCamelCase__				=num_attention_heads
											lowerCamelCase__				=intermediate_size
											lowerCamelCase__				=hidden_act
											lowerCamelCase__				=hidden_dropout_prob
											lowerCamelCase__				=attention_probs_dropout_prob
											lowerCamelCase__				=initializer_range
											lowerCamelCase__				=layer_norm_eps
											lowerCamelCase__				=qkv_bias
											lowerCamelCase__				=scope

											lowerCamelCase__				=(
											    (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)			->  Union[str, Any]:




											"""simple docstring"""





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

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

											lowerCamelCase__				=self.get_config()

											return config, pixel_values, labels






						def 		lowerCAmelCase   (					self)			->  Dict:




											"""simple docstring"""





											lowerCamelCase__				=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    ,						)
											lowerCamelCase__				=self.num_labels
											return config






						def 		lowerCAmelCase   (					self    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE)			->  int:




											"""simple docstring"""





											lowerCamelCase__				=VivitModel(config=_SCREAMING_SNAKE_CASE)
											model.to(_SCREAMING_SNAKE_CASE)
											model.eval()
											lowerCamelCase__				=model(_SCREAMING_SNAKE_CASE)
											self.parent.assertEqual(result.last_hidden_state.shape    ,						(self.batch_size, self.seq_length, self.hidden_size))






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




											"""simple docstring"""





											lowerCamelCase__				=VivitForVideoClassification(_SCREAMING_SNAKE_CASE)
											model.to(_SCREAMING_SNAKE_CASE)
											model.eval()

											lowerCamelCase__				=model(_SCREAMING_SNAKE_CASE)

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






						def 		lowerCAmelCase   (					self)			->  Dict:




											"""simple docstring"""





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




@require_torch
class        __lowercase(       _UpperCamelCase						,							unittest.TestCase  ):
						lowerCAmelCase_		:							str=       (VivitModel, VivitForVideoClassification) if is_torch_available() else ()

						lowerCAmelCase_		:							str=       False
						lowerCAmelCase_		:							Tuple=       False
						lowerCAmelCase_		:							Union[str, Any]=       False
						lowerCAmelCase_		:							Dict=       False






						def 		lowerCAmelCase   (					self)			->  List[str]:




											"""simple docstring"""





											lowerCamelCase__				=VivitModelTester(self)
											lowerCamelCase__				=ConfigTester(self    ,						config_class=_SCREAMING_SNAKE_CASE    ,						has_text_modality=_SCREAMING_SNAKE_CASE    ,						hidden_size=37)






						def 		lowerCAmelCase   (					self    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE=False)			->  Optional[int]:




											"""simple docstring"""





											lowerCamelCase__				=copy.deepcopy(_SCREAMING_SNAKE_CASE)

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

											return inputs_dict






						def 		lowerCAmelCase   (					self)			->  Tuple:




											"""simple docstring"""





											self.config_tester.run_common_tests()






						@unittest.skip(reason='Vivit does not use inputs_embeds')
						def 		lowerCAmelCase   (					self)			->  Dict:




											"""simple docstring"""





											pass






						def 		lowerCAmelCase   (					self)			->  str:




											"""simple docstring"""





											lowerCamelCase__				,     lowerCamelCase__				=self.model_tester.prepare_config_and_inputs_for_common()

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






						def 		lowerCAmelCase   (					self)			->  Dict:




											"""simple docstring"""





											lowerCamelCase__				,     lowerCamelCase__				=self.model_tester.prepare_config_and_inputs_for_common()

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

																lowerCamelCase__				=['pixel_values', 'head_mask']
																self.assertListEqual(arg_names[:2]    ,						_SCREAMING_SNAKE_CASE)






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




											"""simple docstring"""





											lowerCamelCase__				=self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE)






						def 		lowerCAmelCase   (					self)			->  str:




											"""simple docstring"""





											lowerCamelCase__				=self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_video_classification(*_SCREAMING_SNAKE_CASE)






						@slow
						def 		lowerCAmelCase   (					self)			->  Tuple:




											"""simple docstring"""





											for model_name in VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																lowerCamelCase__				=VivitModel.from_pretrained(_SCREAMING_SNAKE_CASE)
																self.assertIsNotNone(_SCREAMING_SNAKE_CASE)






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




											"""simple docstring"""





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

											for model_class in self.all_model_classes:
																lowerCamelCase__				=self.model_tester.seq_length

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

																# check that output_attentions also work using config
																del inputs_dict["output_attentions"]
																lowerCamelCase__				=True
																lowerCamelCase__				=model_class(_SCREAMING_SNAKE_CASE)
																model.to(_SCREAMING_SNAKE_CASE)
																model.eval()
																with torch.no_grad():
																					lowerCamelCase__				=model(**self._prepare_for_class(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE))
																lowerCamelCase__				=outputs.attentions
																self.assertEqual(len(_SCREAMING_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]    ,						)
																lowerCamelCase__				=len(_SCREAMING_SNAKE_CASE)

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

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

																lowerCamelCase__				=outputs.attentions

																self.assertEqual(len(_SCREAMING_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 		lowerCAmelCase   (					self)			->  Dict:




											"""simple docstring"""





											def check_hidden_states_output(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE):
																lowerCamelCase__				=model_class(_SCREAMING_SNAKE_CASE)
																model.to(_SCREAMING_SNAKE_CASE)
																model.eval()

																with torch.no_grad():
																					lowerCamelCase__				=model(**self._prepare_for_class(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE))

																lowerCamelCase__				=outputs.hidden_states
																lowerCamelCase__				=self.model_tester.num_hidden_layers + 1
																self.assertEqual(len(_SCREAMING_SNAKE_CASE)    ,						_SCREAMING_SNAKE_CASE)

																lowerCamelCase__				=self.model_tester.seq_length

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

											lowerCamelCase__				,     lowerCamelCase__				=self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																lowerCamelCase__				=True
																check_hidden_states_output(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE)

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

																check_hidden_states_output(_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE)






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




											"""simple docstring"""





											pass




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




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






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




											"""simple docstring"""





											return VivitImageProcessor() if is_vision_available() else None






						@slow
						def 		lowerCAmelCase   (					self)			->  str:




											"""simple docstring"""





											lowerCamelCase__				=VivitForVideoClassification.from_pretrained('google/vivit-b-16x2-kinetics400').to(_SCREAMING_SNAKE_CASE)

											lowerCamelCase__				=self.default_image_processor
											lowerCamelCase__				=prepare_video()
											lowerCamelCase__				=image_processor(_SCREAMING_SNAKE_CASE    ,						return_tensors='pt').to(_SCREAMING_SNAKE_CASE)

											# forward pass
											with torch.no_grad():
																lowerCamelCase__				=model(**_SCREAMING_SNAKE_CASE)

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

											# taken from original model
											lowerCamelCase__				=torch.tensor([-1.0543, 2.0764, -0.2104, 0.4439, -0.9658]).to(_SCREAMING_SNAKE_CASE)

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