




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


__lowerCAmelCase		:    List[Any]	      =		logging.get_logger(__name__)

__lowerCAmelCase		:    List[str]	      =		{
    'MCG-NJU/videomae-base': 'https://huggingface.co/MCG-NJU/videomae-base/resolve/main/config.json',
}




class    snake_case__							(_UpperCamelCase       ):
				"""simple docstring"""



				SCREAMING_SNAKE_CASE_	:							List[Any]										= """videomae"""





				def __init__(				self    :					str ,				__lowerCamelCase    :					str=2_24 ,				__lowerCamelCase    :					List[str]=16 ,				__lowerCamelCase    :					str=3 ,				__lowerCamelCase    :					Dict=16 ,				__lowerCamelCase    :					Tuple=2 ,				__lowerCamelCase    :					Union[str, Any]=7_68 ,				__lowerCamelCase    :					List[str]=12 ,				__lowerCamelCase    :					Union[str, Any]=12 ,				__lowerCamelCase    :					Dict=30_72 ,				__lowerCamelCase    :					Optional[Any]="gelu" ,				__lowerCamelCase    :					Optional[int]=0.0 ,				__lowerCamelCase    :					int=0.0 ,				__lowerCamelCase    :					Dict=0.02 ,				__lowerCamelCase    :					List[str]=1e-12 ,				__lowerCamelCase    :					List[Any]=True ,				__lowerCamelCase    :					List[str]=True ,				__lowerCamelCase    :					Union[str, Any]=6 ,				__lowerCamelCase    :					Dict=3_84 ,				__lowerCamelCase    :					Dict=4 ,				__lowerCamelCase    :					int=15_36 ,				__lowerCamelCase    :					Union[str, Any]=True ,				**__lowerCamelCase    :					Optional[Any] ,				) ->					Dict:
							super().__init__(**__lowerCamelCase							)

							a   =      image_size
							a   =      patch_size
							a   =      num_channels
							a   =      num_frames
							a   =      tubelet_size

							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   =      qkv_bias
							a   =      use_mean_pooling

							a   =      decoder_num_attention_heads
							a   =      decoder_hidden_size
							a   =      decoder_num_hidden_layers
							a   =      decoder_intermediate_size
							a   =      norm_pix_loss












import torch
import torch.nn.functional as F

from diffusers import VQDiffusionScheduler

from .test_schedulers import SchedulerCommonTest




class    snake_case__							(_UpperCamelCase       ):
				"""simple docstring"""



				SCREAMING_SNAKE_CASE_	:							Dict										= (VQDiffusionScheduler,)





				def 							__UpperCAmelCase							(				self    :					Tuple ,				**__lowerCamelCase    :					Optional[int]							) ->					Optional[Any]:
							a   =      {
							    "num_vec_classes": 40_97,
							    "num_train_timesteps": 1_00,
							}

							config.update(**__lowerCamelCase							)
							return config





				def 							__UpperCAmelCase							(				self    :					List[str] ,				__lowerCamelCase    :					Any							) ->					Union[str, Any]:
							a   =      4
							a   =      8
							a   =      8

							a   =      torch.randint(0 ,				__lowerCamelCase ,				(batch_size, height * width)							)

							return sample





				@property
				def 							__UpperCAmelCase							(				self    :					List[str]							) ->					Optional[Any]:
							assert False





				def 							__UpperCAmelCase							(				self    :					Any ,				__lowerCamelCase    :					Optional[int]							) ->					int:
							def model(__lowerCamelCase    :					Tuple ,				__lowerCamelCase    :					Any ,				*__lowerCamelCase    :					Optional[Any]							):
										a   ,		a   =      sample.shape
										a   =      torch.rand((batch_size, num_vec_classes - 1, num_latent_pixels)							)
										a   =      F.log_softmax(logits.double() ,				dim=1							).float()
										return return_value

							return model





				def 							__UpperCAmelCase							(				self    :					Optional[Any]							) ->					Any:
							for timesteps in [2, 5, 1_00, 10_00]:
										self.check_over_configs(num_train_timesteps=__lowerCamelCase							)





				def 							__UpperCAmelCase							(				self    :					Any							) ->					List[Any]:
							for num_vec_classes in [5, 1_00, 10_00, 40_00]:
										self.check_over_configs(num_vec_classes=__lowerCamelCase							)





				def 							__UpperCAmelCase							(				self    :					Optional[Any]							) ->					List[Any]:
							for t in [0, 50, 99]:
										self.check_over_forward(time_step=__lowerCamelCase							)





				def 							__UpperCAmelCase							(				self    :					List[Any]							) ->					Optional[Any]:
							pass






