
"""simple docstring"""





import copy
import inspect
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import TimesformerConfig
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,
				    TimesformerForVideoClassification,
				    TimesformerModel,
				)
				from transformers.models.timesformer.modeling_timesformer import TIMESFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
				from transformers import VideoMAEImageProcessor






class        _UpperCAmelCase:



		def __init__(   self							,						__a							,						__a=13							,						__a=10							,						__a=3							,						__a=2							,						__a=2							,						__a=True							,						__a=True							,						__a=32							,						__a=5							,						__a=4							,						__a=37							,						__a="gelu"							,						__a=0.1							,						__a=0.1							,						__a=10							,						__a=0.02							,						__a="divided_space_time"							,						__a=None							,						)			->			List[str]:






							'''simple docstring'''
							_UpperCamelCase									=    parent
							_UpperCamelCase									=    batch_size
							_UpperCamelCase									=    image_size
							_UpperCamelCase									=    num_channels
							_UpperCamelCase									=    patch_size
							_UpperCamelCase									=    num_frames
							_UpperCamelCase									=    is_training
							_UpperCamelCase									=    use_labels
							_UpperCamelCase									=    hidden_size
							_UpperCamelCase									=    num_hidden_layers
							_UpperCamelCase									=    num_attention_heads
							_UpperCamelCase									=    intermediate_size
							_UpperCamelCase									=    hidden_act
							_UpperCamelCase									=    hidden_dropout_prob
							_UpperCamelCase									=    attention_probs_dropout_prob
							_UpperCamelCase									=    attention_type
							_UpperCamelCase									=    initializer_range
							_UpperCamelCase									=    scope
							_UpperCamelCase									=    num_labels

							# in TimeSformer, the number of spatial tokens equals num_frames * num_patches per frame + 1 CLS token
							_UpperCamelCase									=    (image_size // patch_size) ** 2
							_UpperCamelCase									=    (num_frames) * self.num_patches_per_frame + 1



		def 		UpperCAmelCase				(   self)			->			Tuple:






							'''simple docstring'''
							_UpperCamelCase									=    floats_tensor(
							    [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size])

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

							_UpperCamelCase									=    self.get_config()

							return config, pixel_values, labels



		def 		UpperCAmelCase				(   self)			->			Tuple:






							'''simple docstring'''
							_UpperCamelCase									=    TimesformerConfig(
							    image_size=self.image_size							,						patch_size=self.patch_size							,						num_channels=self.num_channels							,						num_frames=self.num_frames							,						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							,						attention_type=self.attention_type							,						)
							_UpperCamelCase									=    self.num_labels
							return config



		def 		UpperCAmelCase				(   self							,						__a							,						__a							,						__a)			->			Optional[Any]:






							'''simple docstring'''
							_UpperCamelCase									=    TimesformerModel(config=__a)
							model.to(__a)
							model.eval()
							_UpperCamelCase									=    model(__a)
							self.parent.assertEqual(result.last_hidden_state.shape							,						(self.batch_size, self.seq_length, self.hidden_size))



		def 		UpperCAmelCase				(   self							,						__a							,						__a							,						__a)			->			Dict:






							'''simple docstring'''
							_UpperCamelCase									=    TimesformerForVideoClassification(__a)
							model.to(__a)
							model.eval()

							_UpperCamelCase									=    model(__a)

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




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






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






@require_torch
class        _UpperCAmelCase(       lowerCamelCase ,	lowerCamelCase ,	unittest.TestCase      ):
		lowercase__         = (TimesformerModel, TimesformerForVideoClassification) if is_torch_available() else ()
		lowercase__         = (
		    {'feature-extraction': TimesformerModel, 'video-classification': TimesformerForVideoClassification}
		    if is_torch_available()
		    else {}
		)

		lowercase__         = False
		lowercase__         = False
		lowercase__         = False
		lowercase__         = False



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






							'''simple docstring'''
							_UpperCamelCase									=    TimesformerModelTester(self)
							_UpperCamelCase									=    ConfigTester(
							    self							,						config_class=__a							,						has_text_modality=__a							,						hidden_size=37)



		def 		UpperCAmelCase				(   self							,						__a							,						__a							,						__a=False)			->			List[str]:






							'''simple docstring'''
							_UpperCamelCase									=    copy.deepcopy(__a)

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

							return inputs_dict



		def 		UpperCAmelCase				(   self)			->			Dict:






							'''simple docstring'''
							self.config_tester.run_common_tests()



		@unittest.skip(reason='''TimeSformer does not use inputs_embeds''')
		def 		UpperCAmelCase				(   self)			->			Any:






							'''simple docstring'''
							pass



		def 		UpperCAmelCase				(   self)			->			str:






							'''simple docstring'''
							_UpperCamelCase					,    _UpperCamelCase									=    self.model_tester.prepare_config_and_inputs_for_common()

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



		def 		UpperCAmelCase				(   self)			->			str:






							'''simple docstring'''
							_UpperCamelCase					,    _UpperCamelCase									=    self.model_tester.prepare_config_and_inputs_for_common()

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

												_UpperCamelCase									=    ['''pixel_values''']
												self.assertListEqual(arg_names[:1]							,						__a)



		def 		UpperCAmelCase				(   self)			->			Dict:






							'''simple docstring'''
							_UpperCamelCase									=    self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*__a)



		def 		UpperCAmelCase				(   self)			->			Optional[Any]:






							'''simple docstring'''
							_UpperCamelCase									=    self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_video_classification(*__a)



		@slow
		def 		UpperCAmelCase				(   self)			->			int:






							'''simple docstring'''
							for model_name in TIMESFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												_UpperCamelCase									=    TimesformerModel.from_pretrained(__a)
												self.assertIsNotNone(__a)



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






							'''simple docstring'''
							if not self.has_attentions:
												pass

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

												for model_class in self.all_model_classes:
																	_UpperCamelCase									=    self.model_tester.seq_length
																	_UpperCamelCase									=    self.model_tester.num_frames

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

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

																	# attentions has shape (batch_size x num_frames) x num_heads x (num_patches per frame + 1) x (num_patches per frame + 1)
																	self.assertListEqual(
																	    list(attentions[0].shape[-3:])							,						[self.model_tester.num_attention_heads, seq_len // num_frames + 1, seq_len // num_frames + 1]							,						)
																	_UpperCamelCase									=    len(__a)

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

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

																	_UpperCamelCase									=    outputs.attentions

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

																	# attentions has shape (batch_size x num_frames) x num_heads x (num_patches per frame + 1) x (num_patches per frame + 1)
																	self.assertListEqual(
																	    list(self_attentions[0].shape[-3:])							,						[self.model_tester.num_attention_heads, seq_len // num_frames + 1, seq_len // num_frames + 1]							,						)




		def 		UpperCAmelCase				(   self)			->			int:






							'''simple docstring'''
							def check_hidden_states_output(__a							,						__a							,						__a):
												_UpperCamelCase									=    model_class(__a)
												model.to(__a)
												model.eval()

												with torch.no_grad():
																	_UpperCamelCase									=    model(**self._prepare_for_class(__a							,						__a))

												_UpperCamelCase									=    outputs.hidden_states
												_UpperCamelCase									=    self.model_tester.num_hidden_layers + 1
												self.assertEqual(len(__a)							,						__a)

												_UpperCamelCase									=    self.model_tester.seq_length

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

							_UpperCamelCase					,    _UpperCamelCase									=    self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
												_UpperCamelCase									=    True
												check_hidden_states_output(__a							,						__a							,						__a)

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

												check_hidden_states_output(__a							,						__a							,						__a)





def  lowerCamelCase__						(     )				->						List[Any]:


					"""simple docstring"""





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






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



		@cached_property
		def 		UpperCAmelCase				(   self)			->			List[str]:






							'''simple docstring'''
							# logits were tested with a different mean and std, so we use the same here
							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)			->			List[str]:






							'''simple docstring'''
							_UpperCamelCase									=    TimesformerForVideoClassification.from_pretrained('''facebook/timesformer-base-finetuned-k400''').to(
							    __a)

							_UpperCamelCase									=    self.default_image_processor
							_UpperCamelCase									=    prepare_video()
							_UpperCamelCase									=    image_processor(video[:8]							,						return_tensors='''pt''').to(__a)

							# forward pass
							with torch.no_grad():
												_UpperCamelCase									=    model(**__a)

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

							_UpperCamelCase									=    torch.tensor([-0.3016, -0.7713, -0.4205]).to(__a)

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








"""simple docstring"""





from math import acos, sin
from typing import List, Tuple, Union

import numpy as np
import torch
from PIL import Image

from ...models import AutoencoderKL, UNetaDConditionModel
from ...schedulers import DDIMScheduler, DDPMScheduler
from ...utils import randn_tensor
from ..pipeline_utils import AudioPipelineOutput, BaseOutput, DiffusionPipeline, ImagePipelineOutput
from .mel import Mel






class        _UpperCAmelCase(       lowerCamelCase      ):
		lowercase__         = ['vqvae']



		def __init__(   self							,						__a							,						__a							,						__a							,						__a							,						)			->			List[str]:






							'''simple docstring'''
							super().__init__()
							self.register_modules(unet=__a							,						scheduler=__a							,						mel=__a							,						vqvae=__a)



		def 		UpperCAmelCase				(   self)			->			int:






							'''simple docstring'''
							return 50 if isinstance(self.scheduler							,						__a) else 10_00



		@torch.no_grad()
		def __call__(   self							,						__a = 1							,						__a = None							,						__a = None							,						__a = 0							,						__a = 0							,						__a = None							,						__a = None							,						__a = 0							,						__a = 0							,						__a = None							,						__a = 0							,						__a = None							,						__a = None							,						__a=True							,						)			->			Union[
		    Union[AudioPipelineOutput, ImagePipelineOutput],
		    Tuple[List[Image.Image], Tuple[int, List[np.ndarray]]],
		]:






							'''simple docstring'''
							_UpperCamelCase									=    steps or self.get_default_steps()
							self.scheduler.set_timesteps(__a)
							_UpperCamelCase									=    step_generator or generator
							# For backwards compatibility
							if type(self.unet.config.sample_size) == int:
												_UpperCamelCase									=    (self.unet.config.sample_size, self.unet.config.sample_size)
							if noise is None:
												_UpperCamelCase									=    randn_tensor(
												    (
												        batch_size,
												        self.unet.config.in_channels,
												        self.unet.config.sample_size[0],
												        self.unet.config.sample_size[1],
												    )							,						generator=__a							,						device=self.device							,						)
							_UpperCamelCase									=    noise
							_UpperCamelCase									=    None

							if audio_file is not None or raw_audio is not None:
												self.mel.load_audio(__a							,						__a)
												_UpperCamelCase									=    self.mel.audio_slice_to_image(__a)
												_UpperCamelCase									=    np.frombuffer(input_image.tobytes()							,						dtype='''uint8''').reshape(
												    (input_image.height, input_image.width))
												_UpperCamelCase									=    (input_image / 2_55) * 2 - 1
												_UpperCamelCase									=    torch.tensor(input_image[np.newaxis, :, :]							,						dtype=torch.float).to(self.device)

												if self.vqvae is not None:
																	_UpperCamelCase									=    self.vqvae.encode(torch.unsqueeze(__a							,						0)).latent_dist.sample(
																	    generator=__a)[0]
																	_UpperCamelCase									=    self.vqvae.config.scaling_factor * input_images

												if start_step > 0:
																	_UpperCamelCase									=    self.scheduler.add_noise(__a							,						__a							,						self.scheduler.timesteps[start_step - 1])

												_UpperCamelCase									=    (
												    self.unet.config.sample_size[1] * self.mel.get_sample_rate() / self.mel.x_res / self.mel.hop_length
												)
												_UpperCamelCase									=    int(mask_start_secs * pixels_per_second)
												_UpperCamelCase									=    int(mask_end_secs * pixels_per_second)
												_UpperCamelCase									=    self.scheduler.add_noise(__a							,						__a							,						torch.tensor(self.scheduler.timesteps[start_step:]))

							for step, t in enumerate(self.progress_bar(self.scheduler.timesteps[start_step:])):
												if isinstance(self.unet							,						__a):
																	_UpperCamelCase									=    self.unet(__a							,						__a							,						__a)['''sample''']
												else:
																	_UpperCamelCase									=    self.unet(__a							,						__a)['''sample''']

												if isinstance(self.scheduler							,						__a):
																	_UpperCamelCase									=    self.scheduler.step(
																	    model_output=__a							,						timestep=__a							,						sample=__a							,						eta=__a							,						generator=__a							,						)['''prev_sample''']
												else:
																	_UpperCamelCase									=    self.scheduler.step(
																	    model_output=__a							,						timestep=__a							,						sample=__a							,						generator=__a							,						)['''prev_sample''']

												if mask is not None:
																	if mask_start > 0:
																						_UpperCamelCase									=    mask[:, step, :, :mask_start]
																	if mask_end > 0:
																						_UpperCamelCase									=    mask[:, step, :, -mask_end:]

							if self.vqvae is not None:
												# 0.18215 was scaling factor used in training to ensure unit variance
												_UpperCamelCase									=    1 / self.vqvae.config.scaling_factor * images
												_UpperCamelCase									=    self.vqvae.decode(__a)['''sample''']

							_UpperCamelCase									=    (images / 2 + 0.5).clamp(0							,						1)
							_UpperCamelCase									=    images.cpu().permute(0							,						2							,						3							,						1).numpy()
							_UpperCamelCase									=    (images * 2_55).round().astype('''uint8''')
							_UpperCamelCase									=    list(
							    (Image.fromarray(_[:, :, 0]) for _ in images)
							    if images.shape[3] == 1
							    else (Image.fromarray(__a							,						mode='''RGB''').convert('''L''') for _ in images))

							_UpperCamelCase									=    [self.mel.image_to_audio(__a) for _ in images]
							if not return_dict:
												return images, (self.mel.get_sample_rate(), audios)

							return BaseOutput(**AudioPipelineOutput(np.array(__a)[:, np.newaxis, :])							,						**ImagePipelineOutput(__a))



		@torch.no_grad()
		def 		UpperCAmelCase				(   self							,						__a							,						__a = 50)			->			np.ndarray:






							'''simple docstring'''
							assert isinstance(self.scheduler							,						__a)
							self.scheduler.set_timesteps(__a)
							_UpperCamelCase									=    np.array(
							    [np.frombuffer(image.tobytes()							,						dtype='''uint8''').reshape((1, image.height, image.width)) for image in images])
							_UpperCamelCase									=    (sample / 2_55) * 2 - 1
							_UpperCamelCase									=    torch.Tensor(__a).to(self.device)

							for t in self.progress_bar(torch.flip(self.scheduler.timesteps							,						(0,))):
												_UpperCamelCase									=    t - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps
												_UpperCamelCase									=    self.scheduler.alphas_cumprod[t]
												_UpperCamelCase									=    (
												    self.scheduler.alphas_cumprod[prev_timestep]
												    if prev_timestep >= 0
												    else self.scheduler.final_alpha_cumprod
												)
												_UpperCamelCase									=    1 - alpha_prod_t
												_UpperCamelCase									=    self.unet(__a							,						__a)['''sample''']
												_UpperCamelCase									=    (1 - alpha_prod_t_prev) ** 0.5 * model_output
												_UpperCamelCase									=    (sample - pred_sample_direction) * alpha_prod_t_prev ** (-0.5)
												_UpperCamelCase									=    sample * alpha_prod_t ** 0.5 + beta_prod_t ** 0.5 * model_output

							return sample




		@staticmethod
		def 		UpperCAmelCase				(   __a							,						__a							,						__a)			->			torch.Tensor:






							'''simple docstring'''
							_UpperCamelCase									=    acos(torch.dot(torch.flatten(__a)							,						torch.flatten(__a)) / torch.norm(__a) / torch.norm(__a))
							return sin((1 - alpha) * theta) * xa / sin(__a) + sin(alpha * theta) * xa / sin(__a)






