


import unittest

import torch

from diffusers import VQModel
from diffusers.utils import floats_tensor, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()





class  lowerCAmelCase__(    snake_case__    ,    snake_case__    ,    unittest.TestCase						):






    '''simple docstring'''





    A_  :     Optional[int]			= VQModel
    A_  :     Optional[int]			= 'sample'





    @property
    def 					_lowerCamelCase (	self				:   List[Any]       ,   __snake_case				:   Union[str, Any]=(32, 32)	):





       '''simple docstring'''
       UpperCAmelCase_							:      int						=   4
       UpperCAmelCase_							:      List[str]						=   3

       UpperCAmelCase_							:      Tuple						=   floats_tensor((batch_size, num_channels) + sizes	).to(__snake_case	)

       return {"sample": image}





    @property
    def 					_lowerCamelCase (	self				:   Optional[int]	):





       '''simple docstring'''
       return (3, 32, 32)





    @property
    def 					_lowerCamelCase (	self				:   Optional[int]	):





       '''simple docstring'''
       return (3, 32, 32)





    def 					_lowerCamelCase (	self				:   Dict	):





       '''simple docstring'''
       UpperCAmelCase_							:      Union[str, Any]						=   {
           '''block_out_channels''': [32, 64],
           '''in_channels''': 3,
           '''out_channels''': 3,
           '''down_block_types''': ['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],
           '''up_block_types''': ['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],
           '''latent_channels''': 3,
       }
       UpperCAmelCase_							:      Union[str, Any]						=   self.dummy_input
       return init_dict, inputs_dict





    def 					_lowerCamelCase (	self				:   int	):





       '''simple docstring'''
       pass





    def 					_lowerCamelCase (	self				:   Tuple	):





       '''simple docstring'''
       pass





    def 					_lowerCamelCase (	self				:   List[Any]	):





       '''simple docstring'''
       UpperCAmelCase_							:      Union[str, Any]						=   VQModel.from_pretrained('''fusing/vqgan-dummy'''       ,   output_loading_info=__snake_case	)
       self.assertIsNotNone(__snake_case	)
       self.assertEqual(len(loading_info['''missing_keys''']	)       ,   0	)

       model.to(__snake_case	)
       UpperCAmelCase_							:      Dict						=   model(**self.dummy_input	)

       assert image is not None, "Make sure output is not None"





    def 					_lowerCamelCase (	self				:   Dict	):





       '''simple docstring'''
       UpperCAmelCase_							:      int						=   VQModel.from_pretrained('''fusing/vqgan-dummy'''	)
       model.to(__snake_case	).eval()

       torch.manual_seed(0	)
       if torch.cuda.is_available():
          torch.cuda.manual_seed_all(0	)

       UpperCAmelCase_							:      Union[str, Any]						=   torch.randn(1       ,   model.config.in_channels       ,   model.config.sample_size       ,   model.config.sample_size	)
       UpperCAmelCase_							:      List[Any]						=   image.to(__snake_case	)
       with torch.no_grad():
          UpperCAmelCase_							:      Union[str, Any]						=   model(__snake_case	).sample

       UpperCAmelCase_							:      Tuple						=   output[0, -1, -3:, -3:].flatten().cpu()
       # fmt: off
       UpperCAmelCase_							:      Dict						=   torch.tensor([-0.0_153, -0.4_044, -0.1_880, -0.5_161, -0.2_418, -0.4_072, -0.1_612, -0.0_633, -0.0_143]	)
       # fmt: on
       self.assertTrue(torch.allclose(__snake_case       ,   __snake_case       ,   atol=1E-3	)	)



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


__UpperCamelCase				:  int      					=							logging.get_logger(__name__)

__UpperCamelCase				:  Union[str, Any]      					=							{
    'facebook/timesformer': 'https://huggingface.co/facebook/timesformer/resolve/main/config.json',
}





class  lowerCAmelCase__(    snake_case__						):






    '''simple docstring'''





    A_  :     str			= 'timesformer'





    def __init__(	self				:   int       ,   __snake_case				:   Any=224       ,   __snake_case				:   str=16       ,   __snake_case				:   Any=3       ,   __snake_case				:   List[Any]=8       ,   __snake_case				:   Dict=768       ,   __snake_case				:   Dict=12       ,   __snake_case				:   Tuple=12       ,   __snake_case				:   Dict=3_072       ,   __snake_case				:   str="gelu"       ,   __snake_case				:   Union[str, Any]=0.0       ,   __snake_case				:   List[Any]=0.0       ,   __snake_case				:   Tuple=0.02       ,   __snake_case				:   Optional[Any]=1E-6       ,   __snake_case				:   List[Any]=True       ,   __snake_case				:   List[str]="divided_space_time"       ,   __snake_case				:   Optional[int]=0       ,   **__snake_case				:   Dict       ,   ):





       '''simple docstring'''
       super().__init__(**__snake_case	)

       UpperCAmelCase_							:      Optional[int]						=   image_size
       UpperCAmelCase_							:      Optional[Any]						=   patch_size
       UpperCAmelCase_							:      Any						=   num_channels
       UpperCAmelCase_							:      int						=   num_frames

       UpperCAmelCase_							:      List[Any]						=   hidden_size
       UpperCAmelCase_							:      Any						=   num_hidden_layers
       UpperCAmelCase_							:      Dict						=   num_attention_heads
       UpperCAmelCase_							:      List[Any]						=   intermediate_size
       UpperCAmelCase_							:      str						=   hidden_act
       UpperCAmelCase_							:      List[str]						=   hidden_dropout_prob
       UpperCAmelCase_							:      List[str]						=   attention_probs_dropout_prob
       UpperCAmelCase_							:      Optional[int]						=   initializer_range
       UpperCAmelCase_							:      List[str]						=   layer_norm_eps
       UpperCAmelCase_							:      str						=   qkv_bias

       UpperCAmelCase_							:      Dict						=   attention_type
       UpperCAmelCase_							:      str						=   drop_path_rate