





"""simple docstring"""
from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
)






@flax.struct.dataclass
class 					_SCREAMING_SNAKE_CASE(   A  ):
       SCREAMING_SNAKE_CASE_   :      jnp.ndarray
       SCREAMING_SNAKE_CASE_   :      jnp.ndarray






class 					_SCREAMING_SNAKE_CASE(   nn.Module  ):
       SCREAMING_SNAKE_CASE_   :      int
       SCREAMING_SNAKE_CASE_   :      Tuple[int]							=  (16, 32, 96, 256)
       SCREAMING_SNAKE_CASE_   :      jnp.dtype							=  jnp.floataa




       def     _UpperCamelCase     (						self							)			->      int:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  nn.Conv(
                 self.block_out_channels[0]		,kernel_size=(3, 3)		,padding=((1, 1), (1, 1))		,dtype=self.dtype		,)

             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  []
             for i in range(len(self.block_out_channels							) - 1							):
                   __SCREAMING_SNAKE_CASE   :Optional[int]						=  self.block_out_channels[i]
                   __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  self.block_out_channels[i + 1]
                   __SCREAMING_SNAKE_CASE   :Any						=  nn.Conv(
                       SCREAMING_SNAKE_CASE__		,kernel_size=(3, 3)		,padding=((1, 1), (1, 1))		,dtype=self.dtype		,)
                   blocks.append(SCREAMING_SNAKE_CASE__							)
                   __SCREAMING_SNAKE_CASE   :List[Any]						=  nn.Conv(
                       SCREAMING_SNAKE_CASE__		,kernel_size=(3, 3)		,strides=(2, 2)		,padding=((1, 1), (1, 1))		,dtype=self.dtype		,)
                   blocks.append(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :Tuple						=  blocks

             __SCREAMING_SNAKE_CASE   :Optional[int]						=  nn.Conv(
                 self.conditioning_embedding_channels		,kernel_size=(3, 3)		,padding=((1, 1), (1, 1))		,kernel_init=nn.initializers.zeros_init()		,bias_init=nn.initializers.zeros_init()		,dtype=self.dtype		,)




       def __call__(						self		,SCREAMING_SNAKE_CASE__							)			->      Dict:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Optional[int]						=  self.conv_in(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :Dict						=  nn.silu(SCREAMING_SNAKE_CASE__							)

             for block in self.blocks:
                   __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  block(SCREAMING_SNAKE_CASE__							)
                   __SCREAMING_SNAKE_CASE   :Optional[int]						=  nn.silu(SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  self.conv_out(SCREAMING_SNAKE_CASE__							)

             return embedding






@flax_register_to_config
class 					_SCREAMING_SNAKE_CASE(   nn.Module    ,	A    ,	A  ):
       SCREAMING_SNAKE_CASE_   :      int							=  32
       SCREAMING_SNAKE_CASE_   :      int							=  4
       SCREAMING_SNAKE_CASE_   :      Tuple[str]							=  (
           "CrossAttnDownBlock2D",
           "CrossAttnDownBlock2D",
           "CrossAttnDownBlock2D",
           "DownBlock2D",
       )
       SCREAMING_SNAKE_CASE_   :      Union[bool, Tuple[bool]]							=  False
       SCREAMING_SNAKE_CASE_   :      Tuple[int]							=  (320, 640, 1280, 1280)
       SCREAMING_SNAKE_CASE_   :      int							=  2
       SCREAMING_SNAKE_CASE_   :      Union[int, Tuple[int]]							=  8
       SCREAMING_SNAKE_CASE_   :      Optional[Union[int, Tuple[int]]]							=  None
       SCREAMING_SNAKE_CASE_   :      int							=  1280
       SCREAMING_SNAKE_CASE_   :      float							=  0.0
       SCREAMING_SNAKE_CASE_   :      bool							=  False
       SCREAMING_SNAKE_CASE_   :      jnp.dtype							=  jnp.floataa
       SCREAMING_SNAKE_CASE_   :      bool							=  True
       SCREAMING_SNAKE_CASE_   :      int							=  0
       SCREAMING_SNAKE_CASE_   :      str							=  "rgb"
       SCREAMING_SNAKE_CASE_   :      Tuple[int]							=  (16, 32, 96, 256)




       def     _UpperCamelCase     (						self		,SCREAMING_SNAKE_CASE__							)			->      FrozenDict:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  (1, self.in_channels, self.sample_size, self.sample_size)
             __SCREAMING_SNAKE_CASE   :str						=  jnp.zeros(SCREAMING_SNAKE_CASE__		,dtype=jnp.floataa							)
             __SCREAMING_SNAKE_CASE   :int						=  jnp.ones((1,)		,dtype=jnp.intaa							)
             __SCREAMING_SNAKE_CASE   :int						=  jnp.zeros((1, 1, self.cross_attention_dim)		,dtype=jnp.floataa							)
             __SCREAMING_SNAKE_CASE   :Any						=  (1, 3, self.sample_size * 8, self.sample_size * 8)
             __SCREAMING_SNAKE_CASE   :Tuple						=  jnp.zeros(SCREAMING_SNAKE_CASE__		,dtype=jnp.floataa							)

             __SCREAMING_SNAKE_CASE				,__SCREAMING_SNAKE_CASE   :Any						=  jax.random.split(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  {'''params''': params_rng, '''dropout''': dropout_rng}

             return self.init(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							)["params"]




       def     _UpperCamelCase     (						self							)			->      str:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :str						=  self.block_out_channels
             __SCREAMING_SNAKE_CASE   :List[str]						=  block_out_channels[0] * 4

             # If `num_attention_heads` is not defined (which is the case for most models)
             # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
             # The reason for this behavior is to correct for incorrectly named variables that were introduced
             # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
             # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
             # which is why we correct for the naming here.
             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  self.num_attention_heads or self.attention_head_dim

             # input
             __SCREAMING_SNAKE_CASE   :List[str]						=  nn.Conv(
                 block_out_channels[0]		,kernel_size=(3, 3)		,strides=(1, 1)		,padding=((1, 1), (1, 1))		,dtype=self.dtype		,)

             # time
             __SCREAMING_SNAKE_CASE   :List[Any]						=  FlaxTimesteps(
                 block_out_channels[0]		,flip_sin_to_cos=self.flip_sin_to_cos		,freq_shift=self.config.freq_shift							)
             __SCREAMING_SNAKE_CASE   :Dict						=  FlaxTimestepEmbedding(SCREAMING_SNAKE_CASE__		,dtype=self.dtype							)

             __SCREAMING_SNAKE_CASE   :Dict						=  FlaxControlNetConditioningEmbedding(
                 conditioning_embedding_channels=block_out_channels[0]		,block_out_channels=self.conditioning_embedding_out_channels		,)

             __SCREAMING_SNAKE_CASE   :List[str]						=  self.only_cross_attention
             if isinstance(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							):
                   __SCREAMING_SNAKE_CASE   :Optional[int]						=  (only_cross_attention,) * len(self.down_block_types							)

             if isinstance(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							):
                   __SCREAMING_SNAKE_CASE   :Dict						=  (num_attention_heads,) * len(self.down_block_types							)

             # down
             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  []
             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  []

             __SCREAMING_SNAKE_CASE   :str						=  block_out_channels[0]

             __SCREAMING_SNAKE_CASE   :int						=  nn.Conv(
                 SCREAMING_SNAKE_CASE__		,kernel_size=(1, 1)		,padding='''VALID'''		,kernel_init=nn.initializers.zeros_init()		,bias_init=nn.initializers.zeros_init()		,dtype=self.dtype		,)
             controlnet_down_blocks.append(SCREAMING_SNAKE_CASE__							)

             for i, down_block_type in enumerate(self.down_block_types							):
                   __SCREAMING_SNAKE_CASE   :List[str]						=  output_channel
                   __SCREAMING_SNAKE_CASE   :List[str]						=  block_out_channels[i]
                   __SCREAMING_SNAKE_CASE   :str						=  i == len(SCREAMING_SNAKE_CASE__							) - 1

                   if down_block_type == "CrossAttnDownBlock2D":
                         __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  FlaxCrossAttnDownBlockaD(
                             in_channels=SCREAMING_SNAKE_CASE__		,out_channels=SCREAMING_SNAKE_CASE__		,dropout=self.dropout		,num_layers=self.layers_per_block		,num_attention_heads=num_attention_heads[i]		,add_downsample=not is_final_block		,use_linear_projection=self.use_linear_projection		,only_cross_attention=only_cross_attention[i]		,dtype=self.dtype		,)
                   else:
                         __SCREAMING_SNAKE_CASE   :Any						=  FlaxDownBlockaD(
                             in_channels=SCREAMING_SNAKE_CASE__		,out_channels=SCREAMING_SNAKE_CASE__		,dropout=self.dropout		,num_layers=self.layers_per_block		,add_downsample=not is_final_block		,dtype=self.dtype		,)

                   down_blocks.append(SCREAMING_SNAKE_CASE__							)

                   for _ in range(self.layers_per_block							):
                         __SCREAMING_SNAKE_CASE   :Optional[int]						=  nn.Conv(
                             SCREAMING_SNAKE_CASE__		,kernel_size=(1, 1)		,padding='''VALID'''		,kernel_init=nn.initializers.zeros_init()		,bias_init=nn.initializers.zeros_init()		,dtype=self.dtype		,)
                         controlnet_down_blocks.append(SCREAMING_SNAKE_CASE__							)

                   if not is_final_block:
                         __SCREAMING_SNAKE_CASE   :int						=  nn.Conv(
                             SCREAMING_SNAKE_CASE__		,kernel_size=(1, 1)		,padding='''VALID'''		,kernel_init=nn.initializers.zeros_init()		,bias_init=nn.initializers.zeros_init()		,dtype=self.dtype		,)
                         controlnet_down_blocks.append(SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :str						=  down_blocks
             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  controlnet_down_blocks

             # mid
             __SCREAMING_SNAKE_CASE   :str						=  block_out_channels[-1]
             __SCREAMING_SNAKE_CASE   :int						=  FlaxUNetMidBlockaDCrossAttn(
                 in_channels=SCREAMING_SNAKE_CASE__		,dropout=self.dropout		,num_attention_heads=num_attention_heads[-1]		,use_linear_projection=self.use_linear_projection		,dtype=self.dtype		,)

             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  nn.Conv(
                 SCREAMING_SNAKE_CASE__		,kernel_size=(1, 1)		,padding='''VALID'''		,kernel_init=nn.initializers.zeros_init()		,bias_init=nn.initializers.zeros_init()		,dtype=self.dtype		,)




       def __call__(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__ = 1.0		,SCREAMING_SNAKE_CASE__ = True		,SCREAMING_SNAKE_CASE__ = False		,)			->      Union[FlaxControlNetOutput, Tuple]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Tuple						=  self.controlnet_conditioning_channel_order
             if channel_order == "bgr":
                   __SCREAMING_SNAKE_CASE   :Tuple						=  jnp.flip(SCREAMING_SNAKE_CASE__		,axis=1							)

             # 1. time
             if not isinstance(SCREAMING_SNAKE_CASE__		,jnp.ndarray							):
                   __SCREAMING_SNAKE_CASE   :List[Any]						=  jnp.array([timesteps]		,dtype=jnp.intaa							)
             elif isinstance(SCREAMING_SNAKE_CASE__		,jnp.ndarray							) and len(timesteps.shape							) == 0:
                   __SCREAMING_SNAKE_CASE   :List[str]						=  timesteps.astype(dtype=jnp.floataa							)
                   __SCREAMING_SNAKE_CASE   :List[Any]						=  jnp.expand_dims(SCREAMING_SNAKE_CASE__		,0							)

             __SCREAMING_SNAKE_CASE   :Dict						=  self.time_proj(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :Tuple						=  self.time_embedding(SCREAMING_SNAKE_CASE__							)

             # 2. pre-process
             __SCREAMING_SNAKE_CASE   :List[str]						=  jnp.transpose(SCREAMING_SNAKE_CASE__		,(0, 2, 3, 1)							)
             __SCREAMING_SNAKE_CASE   :List[str]						=  self.conv_in(SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :str						=  jnp.transpose(SCREAMING_SNAKE_CASE__		,(0, 2, 3, 1)							)
             __SCREAMING_SNAKE_CASE   :Any						=  self.controlnet_cond_embedding(SCREAMING_SNAKE_CASE__							)
             sample += controlnet_cond

             # 3. down
             __SCREAMING_SNAKE_CASE   :Dict						=  (sample,)
             for down_block in self.down_blocks:
                   if isinstance(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							):
                         __SCREAMING_SNAKE_CASE				,__SCREAMING_SNAKE_CASE   :str						=  down_block(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,deterministic=not train							)
                   else:
                         __SCREAMING_SNAKE_CASE				,__SCREAMING_SNAKE_CASE   :Optional[Any]						=  down_block(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,deterministic=not train							)
                   down_block_res_samples += res_samples

             # 4. mid
             __SCREAMING_SNAKE_CASE   :Optional[int]						=  self.mid_block(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,deterministic=not train							)

             # 5. contronet blocks
             __SCREAMING_SNAKE_CASE   :List[str]						=  ()
             for down_block_res_sample, controlnet_block in zip(SCREAMING_SNAKE_CASE__		,self.controlnet_down_blocks							):
                   __SCREAMING_SNAKE_CASE   :Tuple						=  controlnet_block(SCREAMING_SNAKE_CASE__							)
                   controlnet_down_block_res_samples += (down_block_res_sample,)

             __SCREAMING_SNAKE_CASE   :List[Any]						=  controlnet_down_block_res_samples

             __SCREAMING_SNAKE_CASE   :Dict						=  self.controlnet_mid_block(SCREAMING_SNAKE_CASE__							)

             # 6. scaling
             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  [sample * conditioning_scale for sample in down_block_res_samples]
             mid_block_res_sample *= conditioning_scale

             if not return_dict:
                   return (down_block_res_samples, mid_block_res_sample)

             return FlaxControlNetOutput(
                 down_block_res_samples=SCREAMING_SNAKE_CASE__		,mid_block_res_sample=SCREAMING_SNAKE_CASE__							)






"""simple docstring"""
from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCamelCase_								=		logging.get_logger(__name__)

lowerCamelCase_								=		{
    "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(   A  ):
       SCREAMING_SNAKE_CASE_   :      List[str]							=  '''vivit'''




       def __init__(						self		,SCREAMING_SNAKE_CASE__=2_24		,SCREAMING_SNAKE_CASE__=32		,SCREAMING_SNAKE_CASE__=[2, 16, 16]		,SCREAMING_SNAKE_CASE__=3		,SCREAMING_SNAKE_CASE__=7_68		,SCREAMING_SNAKE_CASE__=12		,SCREAMING_SNAKE_CASE__=12		,SCREAMING_SNAKE_CASE__=30_72		,SCREAMING_SNAKE_CASE__="gelu_fast"		,SCREAMING_SNAKE_CASE__=0.0		,SCREAMING_SNAKE_CASE__=0.0		,SCREAMING_SNAKE_CASE__=0.0_2		,SCREAMING_SNAKE_CASE__=1E-06		,SCREAMING_SNAKE_CASE__=True		,**SCREAMING_SNAKE_CASE__		,)			->      Any:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :int						=  hidden_size
             __SCREAMING_SNAKE_CASE   :List[Any]						=  num_hidden_layers
             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  num_attention_heads
             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  intermediate_size
             __SCREAMING_SNAKE_CASE   :Any						=  hidden_act
             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  hidden_dropout_prob
             __SCREAMING_SNAKE_CASE   :str						=  attention_probs_dropout_prob
             __SCREAMING_SNAKE_CASE   :Any						=  initializer_range
             __SCREAMING_SNAKE_CASE   :Optional[int]						=  layer_norm_eps

             __SCREAMING_SNAKE_CASE   :Optional[int]						=  image_size
             __SCREAMING_SNAKE_CASE   :List[str]						=  num_frames
             __SCREAMING_SNAKE_CASE   :Any						=  tubelet_size
             __SCREAMING_SNAKE_CASE   :str						=  num_channels
             __SCREAMING_SNAKE_CASE   :Any						=  qkv_bias

             super().__init__(**SCREAMING_SNAKE_CASE__							)