





'''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,
    FlaxCrossAttnUpBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
    FlaxUpBlockaD,
)







@flax.struct.dataclass
class 		UpperCAmelCase (					lowercase_):


  """simple docstring"""







  lowerCAmelCase_											=					42







@flax_register_to_config
class 		UpperCAmelCase (					nn.Module ,  lowercase_ ,  lowercase_):


  """simple docstring"""







  lowerCAmelCase_											=					32
  lowerCAmelCase_											=					4
  lowerCAmelCase_											=					4
  lowerCAmelCase_											=					(
      """CrossAttnDownBlock2D""",
      """CrossAttnDownBlock2D""",
      """CrossAttnDownBlock2D""",
      """DownBlock2D""",
  )
  lowerCAmelCase_											=					("""UpBlock2D""", """CrossAttnUpBlock2D""", """CrossAttnUpBlock2D""", """CrossAttnUpBlock2D""")
  lowerCAmelCase_											=					False
  lowerCAmelCase_											=					(320, 640, 1_280, 1_280)
  lowerCAmelCase_											=					2
  lowerCAmelCase_											=					8
  lowerCAmelCase_											=					None
  lowerCAmelCase_											=					1_280
  lowerCAmelCase_											=					0.0
  lowerCAmelCase_											=					False
  lowerCAmelCase_											=					jnp.floataa
  lowerCAmelCase_											=					True
  lowerCAmelCase_											=					0
  lowerCAmelCase_											=					False


  def 	UpperCamelCase__   (					self     :       Tuple				,						UpperCamelCase__     :       Union[str, Any]	)  ->    FrozenDict:
    _UpperCamelCase							      =(1, self.in_channels, self.sample_size, self.sample_size)
    _UpperCamelCase							      =jnp.zeros(UpperCamelCase__				,						dtype=jnp.floataa	)
    _UpperCamelCase							      =jnp.ones((1,)				,						dtype=jnp.intaa	)
    _UpperCamelCase							      =jnp.zeros((1, 1, self.cross_attention_dim)				,						dtype=jnp.floataa	)

    _UpperCamelCase      ,  _UpperCamelCase							      =jax.random.split(UpperCamelCase__	)
    _UpperCamelCase							      ={'''params''': params_rng, '''dropout''': dropout_rng}

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


  def 	UpperCamelCase__   (					self     :       Union[str, Any]	)  ->    Optional[Any]:
    _UpperCamelCase							      =self.block_out_channels
    _UpperCamelCase							      =block_out_channels[0] * 4

    if self.num_attention_heads is not None:
      raise ValueError(
          '''At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19.'''	)

    # 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.
    _UpperCamelCase							      =self.num_attention_heads or self.attention_head_dim

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

    # time
    _UpperCamelCase							      =FlaxTimesteps(
        block_out_channels[0]				,						flip_sin_to_cos=self.flip_sin_to_cos				,						freq_shift=self.config.freq_shift	)
    _UpperCamelCase							      =FlaxTimestepEmbedding(UpperCamelCase__				,						dtype=self.dtype	)

    _UpperCamelCase							      =self.only_cross_attention
    if isinstance(UpperCamelCase__				,						UpperCamelCase__	):
      _UpperCamelCase							      =(only_cross_attention,) * len(self.down_block_types	)

    if isinstance(UpperCamelCase__				,						UpperCamelCase__	):
      _UpperCamelCase							      =(num_attention_heads,) * len(self.down_block_types	)

    # down
    _UpperCamelCase							      =[]
    _UpperCamelCase							      =block_out_channels[0]
    for i, down_block_type in enumerate(self.down_block_types	):
      _UpperCamelCase							      =output_channel
      _UpperCamelCase							      =block_out_channels[i]
      _UpperCamelCase							      =i == len(UpperCamelCase__	) - 1

      if down_block_type == "CrossAttnDownBlock2D":
        _UpperCamelCase							      =FlaxCrossAttnDownBlockaD(
            in_channels=UpperCamelCase__				,						out_channels=UpperCamelCase__				,						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]				,						use_memory_efficient_attention=self.use_memory_efficient_attention				,						dtype=self.dtype				,						)
      else:
        _UpperCamelCase							      =FlaxDownBlockaD(
            in_channels=UpperCamelCase__				,						out_channels=UpperCamelCase__				,						dropout=self.dropout				,						num_layers=self.layers_per_block				,						add_downsample=not is_final_block				,						dtype=self.dtype				,						)

      down_blocks.append(UpperCamelCase__	)
    _UpperCamelCase							      =down_blocks

    # mid
    _UpperCamelCase							      =FlaxUNetMidBlockaDCrossAttn(
        in_channels=block_out_channels[-1]				,						dropout=self.dropout				,						num_attention_heads=num_attention_heads[-1]				,						use_linear_projection=self.use_linear_projection				,						use_memory_efficient_attention=self.use_memory_efficient_attention				,						dtype=self.dtype				,						)

    # up
    _UpperCamelCase							      =[]
    _UpperCamelCase							      =list(reversed(UpperCamelCase__	)	)
    _UpperCamelCase							      =list(reversed(UpperCamelCase__	)	)
    _UpperCamelCase							      =list(reversed(UpperCamelCase__	)	)
    _UpperCamelCase							      =reversed_block_out_channels[0]
    for i, up_block_type in enumerate(self.up_block_types	):
      _UpperCamelCase							      =output_channel
      _UpperCamelCase							      =reversed_block_out_channels[i]
      _UpperCamelCase							      =reversed_block_out_channels[min(i + 1				,						len(UpperCamelCase__	) - 1	)]

      _UpperCamelCase							      =i == len(UpperCamelCase__	) - 1

      if up_block_type == "CrossAttnUpBlock2D":
        _UpperCamelCase							      =FlaxCrossAttnUpBlockaD(
            in_channels=UpperCamelCase__				,						out_channels=UpperCamelCase__				,						prev_output_channel=UpperCamelCase__				,						num_layers=self.layers_per_block + 1				,						num_attention_heads=reversed_num_attention_heads[i]				,						add_upsample=not is_final_block				,						dropout=self.dropout				,						use_linear_projection=self.use_linear_projection				,						only_cross_attention=only_cross_attention[i]				,						use_memory_efficient_attention=self.use_memory_efficient_attention				,						dtype=self.dtype				,						)
      else:
        _UpperCamelCase							      =FlaxUpBlockaD(
            in_channels=UpperCamelCase__				,						out_channels=UpperCamelCase__				,						prev_output_channel=UpperCamelCase__				,						num_layers=self.layers_per_block + 1				,						add_upsample=not is_final_block				,						dropout=self.dropout				,						dtype=self.dtype				,						)

      up_blocks.append(UpperCamelCase__	)
      _UpperCamelCase							      =output_channel
    _UpperCamelCase							      =up_blocks

    # out
    _UpperCamelCase							      =nn.GroupNorm(num_groups=32				,						epsilon=1E-5	)
    _UpperCamelCase							      =nn.Conv(
        self.out_channels				,						kernel_size=(3, 3)				,						strides=(1, 1)				,						padding=((1, 1), (1, 1))				,						dtype=self.dtype				,						)




  def __call__(					self     :       List[Any]				,						UpperCamelCase__     :       Optional[int]				,						UpperCamelCase__     :       str				,						UpperCamelCase__     :       Any				,						UpperCamelCase__     :       Tuple=None				,						UpperCamelCase__     :       List[str]=None				,						UpperCamelCase__     :       Tuple = True				,						UpperCamelCase__     :       str = False				,						)  ->    Union[FlaxUNetaDConditionOutput, Tuple]:






    if not isinstance(UpperCamelCase__				,						jnp.ndarray	):
      _UpperCamelCase							      =jnp.array([timesteps]				,						dtype=jnp.intaa	)
    elif isinstance(UpperCamelCase__				,						jnp.ndarray	) and len(timesteps.shape	) == 0:
      _UpperCamelCase							      =timesteps.astype(dtype=jnp.floataa	)
      _UpperCamelCase							      =jnp.expand_dims(UpperCamelCase__				,						0	)

    _UpperCamelCase							      =self.time_proj(UpperCamelCase__	)
    _UpperCamelCase							      =self.time_embedding(UpperCamelCase__	)

    # 2. pre-process
    _UpperCamelCase							      =jnp.transpose(UpperCamelCase__				,						(0, 2, 3, 1)	)
    _UpperCamelCase							      =self.conv_in(UpperCamelCase__	)

    # 3. down
    _UpperCamelCase							      =(sample,)
    for down_block in self.down_blocks:
      if isinstance(UpperCamelCase__				,						UpperCamelCase__	):
        _UpperCamelCase      ,  _UpperCamelCase							      =down_block(UpperCamelCase__				,						UpperCamelCase__				,						UpperCamelCase__				,						deterministic=not train	)
      else:
        _UpperCamelCase      ,  _UpperCamelCase							      =down_block(UpperCamelCase__				,						UpperCamelCase__				,						deterministic=not train	)
      down_block_res_samples += res_samples

    if down_block_additional_residuals is not None:
      _UpperCamelCase							      =()

      for down_block_res_sample, down_block_additional_residual in zip(
          UpperCamelCase__				,						UpperCamelCase__	):
        down_block_res_sample += down_block_additional_residual
        new_down_block_res_samples += (down_block_res_sample,)

      _UpperCamelCase							      =new_down_block_res_samples

    # 4. mid
    _UpperCamelCase							      =self.mid_block(UpperCamelCase__				,						UpperCamelCase__				,						UpperCamelCase__				,						deterministic=not train	)

    if mid_block_additional_residual is not None:
      sample += mid_block_additional_residual

    # 5. up
    for up_block in self.up_blocks:
      _UpperCamelCase							      =down_block_res_samples[-(self.layers_per_block + 1) :]
      _UpperCamelCase							      =down_block_res_samples[: -(self.layers_per_block + 1)]
      if isinstance(UpperCamelCase__				,						UpperCamelCase__	):
        _UpperCamelCase							      =up_block(
            UpperCamelCase__				,						temb=UpperCamelCase__				,						encoder_hidden_states=UpperCamelCase__				,						res_hidden_states_tuple=UpperCamelCase__				,						deterministic=not train				,						)
      else:
        _UpperCamelCase							      =up_block(UpperCamelCase__				,						temb=UpperCamelCase__				,						res_hidden_states_tuple=UpperCamelCase__				,						deterministic=not train	)

        # 6. post-process
    _UpperCamelCase							      =self.conv_norm_out(UpperCamelCase__	)
    _UpperCamelCase							      =nn.silu(UpperCamelCase__	)
    _UpperCamelCase							      =self.conv_out(UpperCamelCase__	)
    _UpperCamelCase							      =jnp.transpose(UpperCamelCase__				,						(0, 3, 1, 2)	)

    if not return_dict:
      return (sample,)

    return FlaxUNetaDConditionOutput(sample=UpperCamelCase__	)









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


lowercase_   							=  logging.get_logger(__name__)

lowercase_   							=  {
    'facebook/timesformer': 'https://huggingface.co/facebook/timesformer/resolve/main/config.json',
}





class 					__lowerCAmelCase	(   SCREAMING_SNAKE_CASE				):
  _a		     =    """timesformer"""




  def __init__(							self				,  lowerCAmelCase=224				,  lowerCAmelCase=16				,  lowerCAmelCase=3				,  lowerCAmelCase=8				,  lowerCAmelCase=768				,  lowerCAmelCase=12				,  lowerCAmelCase=12				,  lowerCAmelCase=3_072				,  lowerCAmelCase="gelu"				,  lowerCAmelCase=0.0				,  lowerCAmelCase=0.0				,  lowerCAmelCase=0.02				,  lowerCAmelCase=1e-6				,  lowerCAmelCase=True				,  lowerCAmelCase="divided_space_time"				,  lowerCAmelCase=0				,  **lowerCAmelCase				,  )    ->					Any:
        '''simple docstring'''






        super().__init__(**lowerCAmelCase   )

        _lowercase      			=image_size
        _lowercase      			=patch_size
        _lowercase      			=num_channels
        _lowercase      			=num_frames

        _lowercase      			=hidden_size
        _lowercase      			=num_hidden_layers
        _lowercase      			=num_attention_heads
        _lowercase      			=intermediate_size
        _lowercase      			=hidden_act
        _lowercase      			=hidden_dropout_prob
        _lowercase      			=attention_probs_dropout_prob
        _lowercase      			=initializer_range
        _lowercase      			=layer_norm_eps
        _lowercase      			=qkv_bias

        _lowercase      			=attention_type
        _lowercase      			=drop_path_rate


