



from dataclasses import dataclass
from typing import Optional

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .modeling_utils import ModelMixin







@dataclass
class 				_a				(   _lowerCAmelCase					):
			A      		=	42







class 				_a				(   _lowerCAmelCase ,		_lowerCAmelCase					):

			@register_to_config
			def __init__(self,   SCREAMING_SNAKE_CASE_ = 16,   SCREAMING_SNAKE_CASE_ = 88,   SCREAMING_SNAKE_CASE_ = None,   SCREAMING_SNAKE_CASE_ = None,   SCREAMING_SNAKE_CASE_ = 1,   SCREAMING_SNAKE_CASE_ = 0.0,   SCREAMING_SNAKE_CASE_ = 32,   SCREAMING_SNAKE_CASE_ = None,   SCREAMING_SNAKE_CASE_ = False,   SCREAMING_SNAKE_CASE_ = None,   SCREAMING_SNAKE_CASE_ = "geglu",   SCREAMING_SNAKE_CASE_ = True,   SCREAMING_SNAKE_CASE_ = True,   )       ->   int:
						super().__init__()
						UpperCAmelCase_:		List[Any]							      =     num_attention_heads
						UpperCAmelCase_:		Union[str, Any]							      =     attention_head_dim
						UpperCAmelCase_:		List[str]							      =     num_attention_heads * attention_head_dim

						UpperCAmelCase_:		Optional[Any]							      =     in_channels

						UpperCAmelCase_:		Optional[Any]							      =     torch.nn.GroupNorm(num_groups=SCREAMING_SNAKE_CASE_,   num_channels=SCREAMING_SNAKE_CASE_,   eps=1E-6,   affine=SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		int							      =     nn.Linear(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)

						# 3. Define transformers blocks
						UpperCAmelCase_:		Optional[int]							      =     nn.ModuleList(
						    [
						        BasicTransformerBlock(
						            SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   dropout=SCREAMING_SNAKE_CASE_,   cross_attention_dim=SCREAMING_SNAKE_CASE_,   activation_fn=SCREAMING_SNAKE_CASE_,   attention_bias=SCREAMING_SNAKE_CASE_,   double_self_attention=SCREAMING_SNAKE_CASE_,   norm_elementwise_affine=SCREAMING_SNAKE_CASE_,   )
						        for d in range(SCREAMING_SNAKE_CASE_					)
						    ]					)

						UpperCAmelCase_:		str							      =     nn.Linear(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_=None,   SCREAMING_SNAKE_CASE_=None,   SCREAMING_SNAKE_CASE_=None,   SCREAMING_SNAKE_CASE_=1,   SCREAMING_SNAKE_CASE_=None,   SCREAMING_SNAKE_CASE_ = True,   )       ->   Tuple:
						UpperCAmelCase_      ,							UpperCAmelCase_      ,							UpperCAmelCase_      ,							UpperCAmelCase_:		Tuple							      =     hidden_states.shape
						UpperCAmelCase_:		Tuple							      =     batch_frames // num_frames

						UpperCAmelCase_:		Optional[int]							      =     hidden_states

						UpperCAmelCase_:		List[str]							      =     hidden_states[None, :].reshape(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		int							      =     hidden_states.permute(0,   2,   1,   3,   4					)

						UpperCAmelCase_:		int							      =     self.norm(SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Dict							      =     hidden_states.permute(0,   3,   4,   2,   1					).reshape(batch_size * height * width,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		str							      =     self.proj_in(SCREAMING_SNAKE_CASE_					)

						# 2. Blocks
						for block in self.transformer_blocks:
									UpperCAmelCase_:		Optional[Any]							      =     block(
									    SCREAMING_SNAKE_CASE_,   encoder_hidden_states=SCREAMING_SNAKE_CASE_,   timestep=SCREAMING_SNAKE_CASE_,   cross_attention_kwargs=SCREAMING_SNAKE_CASE_,   class_labels=SCREAMING_SNAKE_CASE_,   )

						# 3. Output
						UpperCAmelCase_:		Union[str, Any]							      =     self.proj_out(SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Any							      =     (
						    hidden_states[None, None, :]
						    .reshape(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)
						    .permute(0,   3,   4,   1,   2					)
						    .contiguous()
						)
						UpperCAmelCase_:		Union[str, Any]							      =     hidden_states.reshape(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		Any							      =     hidden_states + residual

						if not return_dict:
									return (output,)

						return TransformerTemporalModelOutput(sample=SCREAMING_SNAKE_CASE_					)











import numpy as np







def  lowerCAmelCase_	(lowerCAmelCase__:		np.ndarray		,						lowerCAmelCase__:		float  ):





			"""simple docstring"""


			return np.where(vector > 0		,						lowerCAmelCase__		,						(alpha * (np.exp(lowerCAmelCase__  ) - 1))  )


if __name__ == "__main__":
	import doctest

	doctest.testmod()






