

"""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.0_2    ,_A="divided_space_time"    ,_A=None    ,):
      '''simple docstring'''

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

      # in TimeSformer, the number of spatial tokens equals num_frames * num_patches per frame + 1 CLS token
      _lowerCAmelCase       :   str  =    (image_size // patch_size) ** 2
      _lowerCAmelCase       :   List[Any]  =    (num_frames) * self.num_patches_per_frame + 1

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[str]  =    floats_tensor(
          [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]      )

      _lowerCAmelCase       :   Any  =    None
      if self.use_labels:
          _lowerCAmelCase       :   List[Any]  =    ids_tensor([self.batch_size]    ,self.num_labels      )

      _lowerCAmelCase       :   Optional[Any]  =    self.get_config()

      return config, pixel_values, labels

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   int  =    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    ,)
      _lowerCAmelCase       :   List[str]  =    self.num_labels
      return config

  def __lowerCamelCase		(		self    ,_A    ,_A    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[str]  =    TimesformerModel(config=_A      )
      model.to(_A      )
      model.eval()
      _lowerCAmelCase       :   int  =    model(_A      )
      self.parent.assertEqual(result.last_hidden_state.shape    ,(self.batch_size, self.seq_length, self.hidden_size)      )

  def __lowerCamelCase		(		self    ,_A    ,_A    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Dict  =    TimesformerForVideoClassification(_A      )
      model.to(_A      )
      model.eval()

      _lowerCAmelCase       :   Optional[Any]  =    model(_A      )

      # verify the logits shape
      _lowerCAmelCase       :   List[Any]  =    torch.Size((self.batch_size, self.num_labels)      )
      self.parent.assertEqual(result.logits.shape    ,_A      )







  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Optional[Any]  =    self.prepare_config_and_inputs()
      _lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase       :   int  =    config_and_inputs
      _lowerCAmelCase       :   Optional[int]  =    {'pixel_values': pixel_values}
      return config, inputs_dict







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

  _UpperCAmelCase          =		False
  _UpperCAmelCase          =		False
  _UpperCAmelCase          =		False
  _UpperCAmelCase          =		False

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[str]  =    TimesformerModelTester(self      )
      _lowerCAmelCase       :   Optional[int]  =    ConfigTester(
          self    ,config_class=_A    ,has_text_modality=_A    ,hidden_size=37      )

  def __lowerCamelCase		(		self    ,_A    ,_A    ,_A=False      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Union[str, Any]  =    copy.deepcopy(_A      )

      if return_labels:
          if model_class in get_values(_A      ):
              _lowerCAmelCase       :   int  =    torch.zeros(
                  self.model_tester.batch_size    ,dtype=torch.long    ,device=_A      )

      return inputs_dict

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      self.config_tester.run_common_tests()

  @unittest.skip(reason='TimeSformer does not use inputs_embeds'      )
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      pass

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

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

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

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase,					_lowerCAmelCase       :   Dict  =    self.model_tester.prepare_config_and_inputs_for_common()

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

          _lowerCAmelCase       :   Any  =    ['pixel_values']
          self.assertListEqual(arg_names[:1]    ,_A      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[str]  =    self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*_A      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Optional[int]  =    self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_video_classification(*_A      )

  @slow
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      for model_name in TIMESFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          _lowerCAmelCase       :   Any  =    TimesformerModel.from_pretrained(_A      )
          self.assertIsNotNone(_A      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      if not self.has_attentions:
          pass

      else:
          _lowerCAmelCase,					_lowerCAmelCase       :   Dict  =    self.model_tester.prepare_config_and_inputs_for_common()
          _lowerCAmelCase       :   List[str]  =    True

          for model_class in self.all_model_classes:
              _lowerCAmelCase       :   List[str]  =    self.model_tester.seq_length
              _lowerCAmelCase       :   Union[str, Any]  =    self.model_tester.num_frames

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

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

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

              _lowerCAmelCase       :   Tuple  =    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 __lowerCamelCase		(		self      ):
      '''simple docstring'''

      def check_hidden_states_output(_A    ,_A    ,_A      ):
          _lowerCAmelCase       :   Union[str, Any]  =    model_class(_A      )
          model.to(_A      )
          model.eval()

          with torch.no_grad():
              _lowerCAmelCase       :   Union[str, Any]  =    model(**self._prepare_for_class(_A    ,_A      )      )

          _lowerCAmelCase       :   List[Any]  =    outputs.hidden_states
          _lowerCAmelCase       :   int  =    self.model_tester.num_hidden_layers + 1
          self.assertEqual(len(_A      )    ,_A      )

          _lowerCAmelCase       :   Any  =    self.model_tester.seq_length

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

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

      for model_class in self.all_model_classes:
          _lowerCAmelCase       :   Optional[Any]  =    True
          check_hidden_states_output(_A    ,_A    ,_A      )

          # check that output_hidden_states also work using config
          del inputs_dict["output_hidden_states"]
          _lowerCAmelCase       :   str  =    True

          check_hidden_states_output(_A    ,_A    ,_A      )





def 		lowerCamelCase__		(					):

    '''simple docstring'''



    _lowerCAmelCase       :   Union[str, Any]  =    hf_hub_download(
        repo_id='hf-internal-testing/spaghetti-video' ,   filename='eating_spaghetti.npy' ,   repo_type='dataset'   )
    _lowerCAmelCase       :   Any  =    np.load(_lowerCamelCase   )
    return list(_lowerCamelCase   )







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

  @cached_property
  def __lowerCamelCase		(		self      ):
      '''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 __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Dict  =    TimesformerForVideoClassification.from_pretrained('facebook/timesformer-base-finetuned-k400'      ).to(
          _A      )

      _lowerCAmelCase       :   List[Any]  =    self.default_image_processor
      _lowerCAmelCase       :   Tuple  =    prepare_video()
      _lowerCAmelCase       :   int  =    image_processor(video[:8]    ,return_tensors='pt'      ).to(_A      )

      # forward pass
      with torch.no_grad():
          _lowerCAmelCase       :   Optional[int]  =    model(**_A      )

      # verify the logits
      _lowerCAmelCase       :   Optional[int]  =    torch.Size((1, 400)      )
      self.assertEqual(outputs.logits.shape    ,_A      )

      _lowerCAmelCase       :   Optional[Any]  =    torch.tensor([-0.3_0_1_6, -0.7_7_1_3, -0.4_2_0_5]      ).to(_A      )

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








"""simple docstring"""






import gc
import unittest

import numpy as np
import torch

from diffusers import DanceDiffusionPipeline, IPNDMScheduler, UNetaDModel
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS, UNCONDITIONAL_AUDIO_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()







class 	__UpperCamelCase    (    a__							,							unittest.TestCase						):
  _UpperCAmelCase          =		DanceDiffusionPipeline
  _UpperCAmelCase          =		UNCONDITIONAL_AUDIO_GENERATION_PARAMS
  _UpperCAmelCase          =		PipelineTesterMixin.required_optional_params - {
      "callback",
      "latents",
      "callback_steps",
      "output_type",
      "num_images_per_prompt",
  }
  _UpperCAmelCase          =		UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
  _UpperCAmelCase          =		False
  _UpperCAmelCase          =		False

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      torch.manual_seed(0      )
      _lowerCAmelCase       :   List[Any]  =    UNetaDModel(
          block_out_channels=(32, 32, 64)    ,extra_in_channels=16    ,sample_size=512    ,sample_rate=1_6000    ,in_channels=2    ,out_channels=2    ,flip_sin_to_cos=_A    ,use_timestep_embedding=_A    ,time_embedding_type='fourier'    ,mid_block_type='UNetMidBlock1D'    ,down_block_types=('DownBlock1DNoSkip', 'DownBlock1D', 'AttnDownBlock1D')    ,up_block_types=('AttnUpBlock1D', 'UpBlock1D', 'UpBlock1DNoSkip')    ,)
      _lowerCAmelCase       :   int  =    IPNDMScheduler()

      _lowerCAmelCase       :   Union[str, Any]  =    {
          'unet': unet,
          'scheduler': scheduler,
      }
      return components

  def __lowerCamelCase		(		self    ,_A    ,_A=0      ):
      '''simple docstring'''

      if str(_A      ).startswith('mps'      ):
          _lowerCAmelCase       :   str  =    torch.manual_seed(_A      )
      else:
          _lowerCAmelCase       :   Optional[Any]  =    torch.Generator(device=_A      ).manual_seed(_A      )
      _lowerCAmelCase       :   int  =    {
          'batch_size': 1,
          'generator': generator,
          'num_inference_steps': 4,
      }
      return inputs

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[str]  =    'cpu'  # ensure determinism for the device-dependent torch.Generator
      _lowerCAmelCase       :   int  =    self.get_dummy_components()
      _lowerCAmelCase       :   Optional[Any]  =    DanceDiffusionPipeline(**_A      )
      _lowerCAmelCase       :   int  =    pipe.to(_A      )
      pipe.set_progress_bar_config(disable=_A      )

      _lowerCAmelCase       :   Union[str, Any]  =    self.get_dummy_inputs(_A      )
      _lowerCAmelCase       :   List[str]  =    pipe(**_A      )
      _lowerCAmelCase       :   List[Any]  =    output.audios

      _lowerCAmelCase       :   List[str]  =    audio[0, -3:, -3:]

      assert audio.shape == (1, 2, components["unet"].sample_size)
      _lowerCAmelCase       :   Optional[Any]  =    np.array([-0.7_2_6_5, 1.0_0_0_0, -0.8_3_8_8, 0.1_1_7_5, 0.9_4_9_8, -1.0_0_0_0]      )
      assert np.abs(audio_slice.flatten() - expected_slice      ).max() < 1E-2

  @skip_mps
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      return super().test_save_load_local()

  @skip_mps
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3      )

  @skip_mps
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      return super().test_save_load_optional_components()

  @skip_mps
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      return super().test_attention_slicing_forward_pass()







  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      super().test_inference_batch_single_identical(expected_max_diff=3E-3      )







@slow
@require_torch_gpu
class 	__UpperCamelCase    (    unittest.TestCase						):

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[Any]  =    torch_device

      _lowerCAmelCase       :   int  =    DanceDiffusionPipeline.from_pretrained('harmonai/maestro-150k'      )
      _lowerCAmelCase       :   int  =    pipe.to(_A      )
      pipe.set_progress_bar_config(disable=_A      )

      _lowerCAmelCase       :   Optional[int]  =    torch.manual_seed(0      )
      _lowerCAmelCase       :   str  =    pipe(generator=_A    ,num_inference_steps=100    ,audio_length_in_s=4.0_9_6      )
      _lowerCAmelCase       :   str  =    output.audios

      _lowerCAmelCase       :   List[str]  =    audio[0, -3:, -3:]

      assert audio.shape == (1, 2, pipe.unet.sample_size)
      _lowerCAmelCase       :   Union[str, Any]  =    np.array([-0.0_1_9_2, -0.0_2_3_1, -0.0_3_1_8, -0.0_0_5_9, 0.0_0_0_2, -0.0_0_2_0]      )

      assert np.abs(audio_slice.flatten() - expected_slice      ).max() < 1E-2







  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   str  =    torch_device

      _lowerCAmelCase       :   Tuple  =    DanceDiffusionPipeline.from_pretrained('harmonai/maestro-150k'    ,torch_dtype=torch.floataa      )
      _lowerCAmelCase       :   Optional[int]  =    pipe.to(_A      )
      pipe.set_progress_bar_config(disable=_A      )

      _lowerCAmelCase       :   Union[str, Any]  =    torch.manual_seed(0      )
      _lowerCAmelCase       :   Optional[int]  =    pipe(generator=_A    ,num_inference_steps=100    ,audio_length_in_s=4.0_9_6      )
      _lowerCAmelCase       :   Union[str, Any]  =    output.audios

      _lowerCAmelCase       :   int  =    audio[0, -3:, -3:]

      assert audio.shape == (1, 2, pipe.unet.sample_size)
      _lowerCAmelCase       :   List[str]  =    np.array([-0.0_3_6_7, -0.0_4_8_8, -0.0_7_7_1, -0.0_5_2_5, -0.0_4_4_4, -0.0_3_4_1]      )

      assert np.abs(audio_slice.flatten() - expected_slice      ).max() < 1E-2





