






'''simple docstring'''



import torch

from diffusers import UnCLIPScheduler

from .test_schedulers import SchedulerCommonTest
class 							__SCREAMING_SNAKE_CASE       (							_UpperCamelCase     ):
      __a  					=(UnCLIPScheduler,)




      def  __UpperCamelCase							(  self ,   **lowerCamelCase      )							->Tuple:



         '''simple docstring'''
         __a				     =    {
             "num_train_timesteps": 1000,
             "variance_type": "fixed_small_log",
             "clip_sample": True,
             "clip_sample_range": 1.0,
             "prediction_type": "epsilon",
         }

         config.update(**lowerCamelCase      )
         return config




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



         '''simple docstring'''
         for timesteps in [1, 5, 100, 1000]:
            self.check_over_configs(num_train_timesteps=lowerCamelCase      )




      def  __UpperCamelCase							(  self      )							->Dict:



         '''simple docstring'''
         for variance in ["fixed_small_log", "learned_range"]:
            self.check_over_configs(variance_type=lowerCamelCase      )




      def  __UpperCamelCase							(  self      )							->str:



         '''simple docstring'''
         for clip_sample in [True, False]:
            self.check_over_configs(clip_sample=lowerCamelCase      )




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



         '''simple docstring'''
         for clip_sample_range in [1, 5, 10, 20]:
            self.check_over_configs(clip_sample_range=lowerCamelCase      )




      def  __UpperCamelCase							(  self      )							->Dict:



         '''simple docstring'''
         for prediction_type in ["epsilon", "sample"]:
            self.check_over_configs(prediction_type=lowerCamelCase      )




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



         '''simple docstring'''
         for time_step in [0, 500, 999]:
            for prev_timestep in [None, 5, 100, 250, 500, 750]:
               if prev_timestep is not None and prev_timestep >= time_step:
                  continue

               self.check_over_forward(time_step=lowerCamelCase ,   prev_timestep=lowerCamelCase      )




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



         '''simple docstring'''
         __a				     =    self.scheduler_classes[0]
         __a				     =    self.get_scheduler_config(variance_type='fixed_small_log'      )
         __a				     =    scheduler_class(**lowerCamelCase      )

         assert torch.sum(torch.abs(scheduler._get_variance(0      ) - 1.0000e-10      )      ) < 1e-5
         assert torch.sum(torch.abs(scheduler._get_variance(487      ) - 0.0549625      )      ) < 1e-5
         assert torch.sum(torch.abs(scheduler._get_variance(999      ) - 0.9994987      )      ) < 1e-5




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



         '''simple docstring'''
         __a				     =    self.scheduler_classes[0]
         __a				     =    self.get_scheduler_config(variance_type='learned_range'      )
         __a				     =    scheduler_class(**lowerCamelCase      )

         __a				     =    0.5

         assert scheduler._get_variance(1 ,   predicted_variance=lowerCamelCase      ) - -10.1712790 < 1e-5
         assert scheduler._get_variance(487 ,   predicted_variance=lowerCamelCase      ) - -5.7998052 < 1e-5
         assert scheduler._get_variance(999 ,   predicted_variance=lowerCamelCase      ) - -0.0010011 < 1e-5




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



         '''simple docstring'''
         __a				     =    self.scheduler_classes[0]
         __a				     =    self.get_scheduler_config()
         __a				     =    scheduler_class(**lowerCamelCase      )

         __a				     =    scheduler.timesteps

         __a				     =    self.dummy_model()
         __a				     =    self.dummy_sample_deter
         __a				     =    torch.manual_seed(0      )

         for i, t in enumerate(lowerCamelCase      ):
            # 1. predict noise residual
            __a				     =    model(lowerCamelCase ,   lowerCamelCase      )

            # 2. predict previous mean of sample x_t-1
            __a				     =    scheduler.step(lowerCamelCase ,   lowerCamelCase ,   lowerCamelCase ,   generator=lowerCamelCase      ).prev_sample

            __a				     =    pred_prev_sample

         __a				     =    torch.sum(torch.abs(lowerCamelCase      )      )
         __a				     =    torch.mean(torch.abs(lowerCamelCase      )      )

         assert abs(result_sum.item() - 252.2682495      ) < 1e-2
         assert abs(result_mean.item() - 0.3284743      ) < 1e-3




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



         '''simple docstring'''
         __a				     =    self.scheduler_classes[0]
         __a				     =    self.get_scheduler_config()
         __a				     =    scheduler_class(**lowerCamelCase      )

         scheduler.set_timesteps(25      )

         __a				     =    scheduler.timesteps

         __a				     =    self.dummy_model()
         __a				     =    self.dummy_sample_deter
         __a				     =    torch.manual_seed(0      )

         for i, t in enumerate(lowerCamelCase      ):
            # 1. predict noise residual
            __a				     =    model(lowerCamelCase ,   lowerCamelCase      )

            if i + 1 == timesteps.shape[0]:
               __a				     =    None
            else:
               __a				     =    timesteps[i + 1]

            # 2. predict previous mean of sample x_t-1
            __a				     =    scheduler.step(
                lowerCamelCase ,   lowerCamelCase ,   lowerCamelCase ,   prev_timestep=lowerCamelCase ,   generator=lowerCamelCase      ).prev_sample

            __a				     =    pred_prev_sample

         __a				     =    torch.sum(torch.abs(lowerCamelCase      )      )
         __a				     =    torch.mean(torch.abs(lowerCamelCase      )      )

         assert abs(result_sum.item() - 258.2044983      ) < 1e-2
         assert abs(result_mean.item() - 0.3362038      ) < 1e-3




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



         '''simple docstring'''
         pass






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



         '''simple docstring'''
         pass







'''simple docstring'''



from ...configuration_utils import PretrainedConfig
from ...utils import logging


__UpperCamelCase      :       Dict      			=      logging.get_logger(__name__)

__UpperCamelCase      :       Any      			=      {
    """google/vivit-b-16x2-kinetics400""": (
        """https://huggingface.co/google/vivit-b-16x2-kinetics400/resolve/main/config.json"""
    ),
    # See all Vivit models at https://huggingface.co/models?filter=vivit
}
class 							__SCREAMING_SNAKE_CASE       (							_lowerCAmelCase     ):
      __a  					="vivit"




      def __init__(  self ,   lowerCamelCase=224 ,   lowerCamelCase=32 ,   lowerCamelCase=[2, 16, 16] ,   lowerCamelCase=3 ,   lowerCamelCase=768 ,   lowerCamelCase=12 ,   lowerCamelCase=12 ,   lowerCamelCase=3072 ,   lowerCamelCase="gelu_fast" ,   lowerCamelCase=0.0 ,   lowerCamelCase=0.0 ,   lowerCamelCase=0.02 ,   lowerCamelCase=1e-06 ,   lowerCamelCase=True ,   **lowerCamelCase ,   )							->Tuple:



         '''simple docstring'''
         __a				     =    hidden_size
         __a				     =    num_hidden_layers
         __a				     =    num_attention_heads
         __a				     =    intermediate_size
         __a				     =    hidden_act
         __a				     =    hidden_dropout_prob
         __a				     =    attention_probs_dropout_prob
         __a				     =    initializer_range
         __a				     =    layer_norm_eps

         __a				     =    image_size
         __a				     =    num_frames
         __a				     =    tubelet_size
         __a				     =    num_channels
         __a				     =    qkv_bias

         super().__init__(**lowerCamelCase      )