
'''simple docstring'''







from typing import Callable, List, Optional, Union

import torch

from ...models import UNetaDModel
from ...schedulers import CMStochasticIterativeScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput


_lowercase					:    Optional[int]       				=			logging.get_logger(__name__)  # pylint: disable=invalid-name


_lowercase					:    Dict       				=			"\n    Examples:\n        ```py\n        >>> import torch\n\n        >>> from diffusers import ConsistencyModelPipeline\n\n        >>> device = \"cuda\"\n        >>> # Load the cd_imagenet64_l2 checkpoint.\n        >>> model_id_or_path = \"openai/diffusers-cd_imagenet64_l2\"\n        >>> pipe = ConsistencyModelPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)\n        >>> pipe.to(device)\n\n        >>> # Onestep Sampling\n        >>> image = pipe(num_inference_steps=1).images[0]\n        >>> image.save(\"cd_imagenet64_l2_onestep_sample.png\")\n\n        >>> # Onestep sampling, class-conditional image generation\n        >>> # ImageNet-64 class label 145 corresponds to king penguins\n        >>> image = pipe(num_inference_steps=1, class_labels=145).images[0]\n        >>> image.save(\"cd_imagenet64_l2_onestep_sample_penguin.png\")\n\n        >>> # Multistep sampling, class-conditional image generation\n        >>> # Timesteps can be explicitly specified; the particular timesteps below are from the original Github repo:\n        >>> # https://github.com/openai/consistency_models/blob/main/scripts/launch.sh#L77\n        >>> image = pipe(num_inference_steps=None, timesteps=[22, 0], class_labels=145).images[0]\n        >>> image.save(\"cd_imagenet64_l2_multistep_sample_penguin.png\")\n        ```\n"






class 					lowerCAmelCase__  (    lowerCamelCase_	):






					def __init__(  self							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE						):


											"""simple docstring"""





											super().__init__()

											self.register_modules(
											    unet=__SCREAMING_SNAKE_CASE							,				scheduler=__SCREAMING_SNAKE_CASE							,				)

											lowercase_      :	Union[str, Any]						=			None






					def       _snake_case       (  self							,				__SCREAMING_SNAKE_CASE=0						):


											"""simple docstring"""





											if is_accelerate_available() and is_accelerate_version('''>='''							,				'''0.14.0'''						):
																	from accelerate import cpu_offload
											else:
																	raise ImportError('''`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher'''						)

											lowercase_      :	Any						=			torch.device(F'''cuda:{gpu_id}'''						)

											if self.device.type != "cpu":
																	self.to('''cpu'''							,				silence_dtype_warnings=__SCREAMING_SNAKE_CASE						)
																	torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

											for cpu_offloaded_model in [self.unet]:
																	cpu_offload(__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE						)

											if self.safety_checker is not None:
																	cpu_offload(self.safety_checker							,				execution_device=__SCREAMING_SNAKE_CASE							,				offload_buffers=__SCREAMING_SNAKE_CASE						)






					def       _snake_case       (  self							,				__SCREAMING_SNAKE_CASE=0						):


											"""simple docstring"""





											if is_accelerate_available() and is_accelerate_version('''>='''							,				'''0.17.0.dev0'''						):
																	from accelerate import cpu_offload_with_hook
											else:
																	raise ImportError('''`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.'''						)

											lowercase_      :	Dict						=			torch.device(F'''cuda:{gpu_id}'''						)

											if self.device.type != "cpu":
																	self.to('''cpu'''							,				silence_dtype_warnings=__SCREAMING_SNAKE_CASE						)
																	torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

											lowercase_      :	int						=			None
											for cpu_offloaded_model in [self.unet]:
																	lowercase_		,					lowercase_      :	int						=			cpu_offload_with_hook(__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				prev_module_hook=__SCREAMING_SNAKE_CASE						)

											if self.safety_checker is not None:
																	lowercase_		,					lowercase_      :	Union[str, Any]						=			cpu_offload_with_hook(self.safety_checker							,				__SCREAMING_SNAKE_CASE							,				prev_module_hook=__SCREAMING_SNAKE_CASE						)

											# We'll offload the last model manually.
											lowercase_      :	str						=			hook






					@property
					# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
					def       _snake_case       (  self						):


											"""simple docstring"""





											if not hasattr(self.unet							,				'''_hf_hook'''						):
																	return self.device
											for module in self.unet.modules():
																	if (
																	    hasattr(__SCREAMING_SNAKE_CASE							,				'''_hf_hook'''						)
																	    and hasattr(module._hf_hook							,				'''execution_device'''						)
																	    and module._hf_hook.execution_device is not None
																	):
																							return torch.device(module._hf_hook.execution_device						)
											return self.device






					def       _snake_case       (  self							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE=None						):


											"""simple docstring"""





											lowercase_      :	List[str]						=			(batch_size, num_channels, height, width)
											if isinstance(__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE						) and len(__SCREAMING_SNAKE_CASE						) != batch_size:
																	raise ValueError(
																	    F'''You have passed a list of generators of length {len(__SCREAMING_SNAKE_CASE						)}, but requested an effective batch'''
																	    F''' size of {batch_size}. Make sure the batch size matches the length of the generators.'''						)

											if latents is None:
																	lowercase_      :	List[Any]						=			randn_tensor(__SCREAMING_SNAKE_CASE							,				generator=__SCREAMING_SNAKE_CASE							,				device=__SCREAMING_SNAKE_CASE							,				dtype=__SCREAMING_SNAKE_CASE						)
											else:
																	lowercase_      :	str						=			latents.to(device=__SCREAMING_SNAKE_CASE							,				dtype=__SCREAMING_SNAKE_CASE						)

											# scale the initial noise by the standard deviation required by the scheduler
											lowercase_      :	Optional[int]						=			latents * self.scheduler.init_noise_sigma
											return latents






					def       _snake_case       (  self							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE = "pil"						):


											"""simple docstring"""





											if output_type not in ["pt", "np", "pil"]:
																	raise ValueError(
																	    F'''output_type={output_type} is not supported. Make sure to choose one of [\'pt\', \'np\', or \'pil\']'''						)

											# Equivalent to diffusers.VaeImageProcessor.denormalize
											lowercase_      :	Union[str, Any]						=			(sample / 2 + 0.5).clamp(0							,				1						)
											if output_type == "pt":
																	return sample

											# Equivalent to diffusers.VaeImageProcessor.pt_to_numpy
											lowercase_      :	Union[str, Any]						=			sample.cpu().permute(0							,				2							,				3							,				1						).numpy()
											if output_type == "np":
																	return sample

											# Output_type must be 'pil'
											lowercase_      :	Any						=			self.numpy_to_pil(__SCREAMING_SNAKE_CASE						)
											return sample






					def       _snake_case       (  self							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE=None						):


											"""simple docstring"""





											if self.unet.config.num_class_embeds is not None:
																	if isinstance(__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE						):
																							lowercase_      :	Tuple						=			torch.tensor(__SCREAMING_SNAKE_CASE							,				dtype=torch.int						)
																	elif isinstance(__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE						):
																							assert batch_size == 1, "Batch size must be 1 if classes is an int"
																							lowercase_      :	Optional[int]						=			torch.tensor([class_labels]							,				dtype=torch.int						)
																	elif class_labels is None:
																							# Randomly generate batch_size class labels
																							# TODO: should use generator here? int analogue of randn_tensor is not exposed in ...utils
																							lowercase_      :	Dict						=			torch.randint(0							,				self.unet.config.num_class_embeds							,				size=(batch_size,)						)
																	lowercase_      :	Dict						=			class_labels.to(__SCREAMING_SNAKE_CASE						)
											else:
																	lowercase_      :	Tuple						=			None
											return class_labels






					def       _snake_case       (  self							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE						):


											"""simple docstring"""





											if num_inference_steps is None and timesteps is None:
																	raise ValueError('''Exactly one of `num_inference_steps` or `timesteps` must be supplied.'''						)

											if num_inference_steps is not None and timesteps is not None:
																	logger.warning(
																	    F'''Both `num_inference_steps`: {num_inference_steps} and `timesteps`: {timesteps} are supplied;'''
																	    ''' `timesteps` will be used over `num_inference_steps`.'''						)

											if latents is not None:
																	lowercase_      :	Any						=			(batch_size, 3, img_size, img_size)
																	if latents.shape != expected_shape:
																							raise ValueError(F'''The shape of latents is {latents.shape} but is expected to be {expected_shape}.'''						)

											if (callback_steps is None) or (
											    callback_steps is not None and (not isinstance(__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE						) or callback_steps <= 0)
											):
																	raise ValueError(
																	    F'''`callback_steps` has to be a positive integer but is {callback_steps} of type'''
																	    F''' {type(__SCREAMING_SNAKE_CASE						)}.'''						)






					@torch.no_grad()
					@replace_example_docstring(__SCREAMING_SNAKE_CASE						)
					def __call__(  self							,				__SCREAMING_SNAKE_CASE = 1							,				__SCREAMING_SNAKE_CASE = None							,				__SCREAMING_SNAKE_CASE = 1							,				__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							,				):


											"""simple docstring"""





											lowercase_      :	Optional[Any]						=			self.unet.config.sample_size
											lowercase_      :	Union[str, Any]						=			self._execution_device

											# 1. Check inputs
											self.check_inputs(__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE						)

											# 2. Prepare image latents
											# Sample image latents x_0 ~ N(0, sigma_0^2 * I)
											lowercase_      :	Tuple						=			self.prepare_latents(
											    batch_size=__SCREAMING_SNAKE_CASE							,				num_channels=self.unet.config.in_channels							,				height=__SCREAMING_SNAKE_CASE							,				width=__SCREAMING_SNAKE_CASE							,				dtype=self.unet.dtype							,				device=__SCREAMING_SNAKE_CASE							,				generator=__SCREAMING_SNAKE_CASE							,				latents=__SCREAMING_SNAKE_CASE							,				)

											# 3. Handle class_labels for class-conditional models
											lowercase_      :	Tuple						=			self.prepare_class_labels(__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				class_labels=__SCREAMING_SNAKE_CASE						)

											# 4. Prepare timesteps
											if timesteps is not None:
																	self.scheduler.set_timesteps(timesteps=__SCREAMING_SNAKE_CASE							,				device=__SCREAMING_SNAKE_CASE						)
																	lowercase_      :	Tuple						=			self.scheduler.timesteps
																	lowercase_      :	Dict						=			len(__SCREAMING_SNAKE_CASE						)
											else:
																	self.scheduler.set_timesteps(__SCREAMING_SNAKE_CASE						)
																	lowercase_      :	str						=			self.scheduler.timesteps

											# 5. Denoising loop
											# Multistep sampling: implements Algorithm 1 in the paper
											with self.progress_bar(total=__SCREAMING_SNAKE_CASE						) as progress_bar:
																	for i, t in enumerate(__SCREAMING_SNAKE_CASE						):
																							lowercase_      :	Dict						=			self.scheduler.scale_model_input(__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE						)
																							lowercase_      :	Dict						=			self.unet(__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				class_labels=__SCREAMING_SNAKE_CASE							,				return_dict=__SCREAMING_SNAKE_CASE						)[0]

																							lowercase_      :	Dict						=			self.scheduler.step(__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				generator=__SCREAMING_SNAKE_CASE						)[0]

																							# call the callback, if provided
																							progress_bar.update()
																							if callback is not None and i % callback_steps == 0:
																													callback(__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE						)

        # 6. Post-process image sample
											lowercase_      :	List[Any]						=			self.postprocess_image(__SCREAMING_SNAKE_CASE							,				output_type=__SCREAMING_SNAKE_CASE						)

											# Offload last model to CPU
											if hasattr(self							,				'''final_offload_hook'''						) and self.final_offload_hook is not None:
																	self.final_offload_hook.offload()

											if not return_dict:
																	return (image,)

											return ImagePipelineOutput(images=__SCREAMING_SNAKE_CASE						)





'''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__  :






					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-0_6							,				__SCREAMING_SNAKE_CASE=True							,				__SCREAMING_SNAKE_CASE=None							,				):


											"""simple docstring"""





											lowercase_      :	Any						=			parent
											lowercase_      :	Union[str, Any]						=			batch_size
											lowercase_      :	Union[str, Any]						=			is_training
											lowercase_      :	Tuple						=			use_labels
											lowercase_      :	Any						=			num_labels
											lowercase_      :	str						=			image_size
											lowercase_      :	Union[str, Any]						=			num_frames
											lowercase_      :	Any						=			tubelet_size
											lowercase_      :	Tuple						=			num_channels
											lowercase_      :	Union[str, Any]						=			hidden_size
											lowercase_      :	Dict						=			num_hidden_layers
											lowercase_      :	Optional[int]						=			num_attention_heads
											lowercase_      :	str						=			intermediate_size
											lowercase_      :	Dict						=			hidden_act
											lowercase_      :	Dict						=			hidden_dropout_prob
											lowercase_      :	Optional[Any]						=			attention_probs_dropout_prob
											lowercase_      :	Dict						=			initializer_range
											lowercase_      :	Optional[int]						=			layer_norm_eps
											lowercase_      :	str						=			qkv_bias
											lowercase_      :	Dict						=			scope

											lowercase_      :	Dict						=			(
											    (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       _snake_case       (  self						):


											"""simple docstring"""





											lowercase_      :	Tuple						=			floats_tensor(
											    [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]						)

											lowercase_      :	Optional[int]						=			None
											if self.use_labels:
																	lowercase_      :	Tuple						=			ids_tensor([self.batch_size]							,				self.num_labels						)

											lowercase_      :	Union[str, Any]						=			self.get_config()

											return config, pixel_values, labels






					def       _snake_case       (  self						):


											"""simple docstring"""





											lowercase_      :	Optional[int]						=			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							,				)
											lowercase_      :	Optional[Any]						=			self.num_labels
											return config






					def       _snake_case       (  self							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE						):


											"""simple docstring"""





											lowercase_      :	str						=			VivitModel(config=__SCREAMING_SNAKE_CASE						)
											model.to(__SCREAMING_SNAKE_CASE						)
											model.eval()
											lowercase_      :	Union[str, Any]						=			model(__SCREAMING_SNAKE_CASE						)
											self.parent.assertEqual(result.last_hidden_state.shape							,				(self.batch_size, self.seq_length, self.hidden_size)						)






					def       _snake_case       (  self							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE						):


											"""simple docstring"""





											lowercase_      :	List[Any]						=			VivitForVideoClassification(__SCREAMING_SNAKE_CASE						)
											model.to(__SCREAMING_SNAKE_CASE						)
											model.eval()

											lowercase_      :	Union[str, Any]						=			model(__SCREAMING_SNAKE_CASE						)

											# verify the logits shape
											lowercase_      :	Union[str, Any]						=			torch.Size((self.batch_size, self.num_labels)						)
											self.parent.assertEqual(result.logits.shape							,				__SCREAMING_SNAKE_CASE						)






					def       _snake_case       (  self						):


											"""simple docstring"""





											lowercase_      :	int						=			self.prepare_config_and_inputs()
											lowercase_		,					lowercase_		,					lowercase_      :	Dict						=			config_and_inputs
											lowercase_      :	List[str]						=			{'''pixel_values''': pixel_values}
											return config, inputs_dict







@require_torch
class 					lowerCAmelCase__  (    lowerCamelCase_       ,		unittest.TestCase	):
					lowerCAmelCase_	  =	(VivitModel, VivitForVideoClassification) if is_torch_available() else ()

					lowerCAmelCase_	  =	False
					lowerCAmelCase_	  =	False
					lowerCAmelCase_	  =	False
					lowerCAmelCase_	  =	False






					def       _snake_case       (  self						):


											"""simple docstring"""





											lowercase_      :	int						=			VivitModelTester(self						)
											lowercase_      :	Optional[int]						=			ConfigTester(self							,				config_class=__SCREAMING_SNAKE_CASE							,				has_text_modality=__SCREAMING_SNAKE_CASE							,				hidden_size=37						)






					def       _snake_case       (  self							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE=False						):


											"""simple docstring"""





											lowercase_      :	str						=			copy.deepcopy(__SCREAMING_SNAKE_CASE						)

											if return_labels:
																	if model_class in get_values(__SCREAMING_SNAKE_CASE						):
																							lowercase_      :	List[str]						=			torch.zeros(
																							    self.model_tester.batch_size							,				dtype=torch.long							,				device=__SCREAMING_SNAKE_CASE						)

											return inputs_dict






					def       _snake_case       (  self						):


											"""simple docstring"""





											self.config_tester.run_common_tests()






					@unittest.skip(reason='''Vivit does not use inputs_embeds'''						)
					def       _snake_case       (  self						):


											"""simple docstring"""





											pass






					def       _snake_case       (  self						):


											"""simple docstring"""





											lowercase_		,					lowercase_      :	str						=			self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	lowercase_      :	Optional[Any]						=			model_class(__SCREAMING_SNAKE_CASE						)
																	self.assertIsInstance(model.get_input_embeddings()							,				(nn.Module)						)
																	lowercase_      :	int						=			model.get_output_embeddings()
																	self.assertTrue(x is None or isinstance(__SCREAMING_SNAKE_CASE							,				nn.Linear						)						)






					def       _snake_case       (  self						):


											"""simple docstring"""





											lowercase_		,					lowercase_      :	str						=			self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	lowercase_      :	Optional[int]						=			model_class(__SCREAMING_SNAKE_CASE						)
																	lowercase_      :	Dict						=			inspect.signature(model.forward						)
																	# signature.parameters is an OrderedDict => so arg_names order is deterministic
																	lowercase_      :	Any						=			[*signature.parameters.keys()]

																	lowercase_      :	str						=			['''pixel_values''', '''head_mask''']
																	self.assertListEqual(arg_names[:2]							,				__SCREAMING_SNAKE_CASE						)






					def       _snake_case       (  self						):


											"""simple docstring"""





											lowercase_      :	Union[str, Any]						=			self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE						)






					def       _snake_case       (  self						):


											"""simple docstring"""





											lowercase_      :	int						=			self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_video_classification(*__SCREAMING_SNAKE_CASE						)






					@slow
					def       _snake_case       (  self						):


											"""simple docstring"""





											for model_name in VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	lowercase_      :	Union[str, Any]						=			VivitModel.from_pretrained(__SCREAMING_SNAKE_CASE						)
																	self.assertIsNotNone(__SCREAMING_SNAKE_CASE						)






					def       _snake_case       (  self						):


											"""simple docstring"""





											lowercase_		,					lowercase_      :	Union[str, Any]						=			self.model_tester.prepare_config_and_inputs_for_common()
											lowercase_      :	Union[str, Any]						=			True

											for model_class in self.all_model_classes:
																	lowercase_      :	int						=			self.model_tester.seq_length

																	lowercase_      :	Optional[Any]						=			True
																	lowercase_      :	int						=			False
																	lowercase_      :	Union[str, Any]						=			True
																	lowercase_      :	str						=			model_class(__SCREAMING_SNAKE_CASE						)
																	model.to(__SCREAMING_SNAKE_CASE						)
																	model.eval()
																	with torch.no_grad():
																							lowercase_      :	List[str]						=			model(**self._prepare_for_class(__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE						)						)
																	lowercase_      :	Optional[Any]						=			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"]
																	lowercase_      :	List[Any]						=			True
																	lowercase_      :	Union[str, Any]						=			model_class(__SCREAMING_SNAKE_CASE						)
																	model.to(__SCREAMING_SNAKE_CASE						)
																	model.eval()
																	with torch.no_grad():
																							lowercase_      :	str						=			model(**self._prepare_for_class(__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE						)						)
																	lowercase_      :	Dict						=			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]							,				)
																	lowercase_      :	Dict						=			len(__SCREAMING_SNAKE_CASE						)

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

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

																	lowercase_      :	List[str]						=			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       _snake_case       (  self						):


											"""simple docstring"""





											def check_hidden_states_output(__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE						):
																	lowercase_      :	List[str]						=			model_class(__SCREAMING_SNAKE_CASE						)
																	model.to(__SCREAMING_SNAKE_CASE						)
																	model.eval()

																	with torch.no_grad():
																							lowercase_      :	List[Any]						=			model(**self._prepare_for_class(__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE						)						)

																	lowercase_      :	List[str]						=			outputs.hidden_states
																	lowercase_      :	Any						=			self.model_tester.num_hidden_layers + 1
																	self.assertEqual(len(__SCREAMING_SNAKE_CASE						)							,				__SCREAMING_SNAKE_CASE						)

																	lowercase_      :	List[Any]						=			self.model_tester.seq_length

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

											lowercase_		,					lowercase_      :	List[str]						=			self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	lowercase_      :	int						=			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"]
																	lowercase_      :	List[Any]						=			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       _snake_case       (  self						):


											"""simple docstring"""





											pass







def    snake_case_						(							):



						"""simple docstring"""





						lowercase_      :	List[str]						=			hf_hub_download(
						    repo_id='''hf-internal-testing/spaghetti-video''' , filename='''eating_spaghetti_32_frames.npy''' , repo_type='''dataset'''				)
						lowercase_      :	List[str]						=			np.load(__SCREAMING_SNAKE_CASE				)
						return list(__SCREAMING_SNAKE_CASE				)






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






					@cached_property
					def       _snake_case       (  self						):


											"""simple docstring"""





											return VivitImageProcessor() if is_vision_available() else None






					@slow
					def       _snake_case       (  self						):


											"""simple docstring"""





											lowercase_      :	List[str]						=			VivitForVideoClassification.from_pretrained('''google/vivit-b-16x2-kinetics400'''						).to(__SCREAMING_SNAKE_CASE						)

											lowercase_      :	Optional[int]						=			self.default_image_processor
											lowercase_      :	Union[str, Any]						=			prepare_video()
											lowercase_      :	str						=			image_processor(__SCREAMING_SNAKE_CASE							,				return_tensors='''pt'''						).to(__SCREAMING_SNAKE_CASE						)

											# forward pass
											with torch.no_grad():
																	lowercase_      :	Any						=			model(**__SCREAMING_SNAKE_CASE						)

											# verify the logits
											lowercase_      :	Any						=			torch.Size((1, 4_00)						)
											self.assertEqual(outputs.logits.shape							,				__SCREAMING_SNAKE_CASE						)

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

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



