



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 				lowerCAmelCase_		(  a__							):
       UpperCAmelCase__				:					jnp.ndarray






@flax_register_to_config
class 				lowerCAmelCase_		(  nn.Module       ,    a__       ,    a__							):
       UpperCAmelCase__				:					int        =  32
       UpperCAmelCase__				:					int        =  4
       UpperCAmelCase__				:					int        =  4
       UpperCAmelCase__				:					Tuple[str]        =  (
           "CrossAttnDownBlock2D",
           "CrossAttnDownBlock2D",
           "CrossAttnDownBlock2D",
           "DownBlock2D",
       )
       UpperCAmelCase__				:					Tuple[str]        =  ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
       UpperCAmelCase__				:					Union[bool, Tuple[bool]]        =  False
       UpperCAmelCase__				:					Tuple[int]        =  (320, 640, 1280, 1280)
       UpperCAmelCase__				:					int        =  2
       UpperCAmelCase__				:					Union[int, Tuple[int]]        =  8
       UpperCAmelCase__				:					Optional[Union[int, Tuple[int]]]        =  None
       UpperCAmelCase__				:					int        =  1280
       UpperCAmelCase__				:					float        =  0.0
       UpperCAmelCase__				:					bool        =  False
       UpperCAmelCase__				:					jnp.dtype        =  jnp.floataa
       UpperCAmelCase__				:					bool        =  True
       UpperCAmelCase__				:					int        =  0
       UpperCAmelCase__				:					bool        =  False
       def 					snake_case_       (			self,			SCREAMING_SNAKE_CASE_       )     ->		FrozenDict:
          # init input tensors
          UpperCamelCase    :  Union[str, Any]							=	(1, self.in_channels, self.sample_size, self.sample_size)
          UpperCamelCase    :  Any							=	jnp.zeros(SCREAMING_SNAKE_CASE_,			dtype=jnp.floataa       )
          UpperCamelCase    :  Union[str, Any]							=	jnp.ones((1,),			dtype=jnp.intaa       )
          UpperCamelCase    :  Dict							=	jnp.zeros((1, 1, self.cross_attention_dim),			dtype=jnp.floataa       )

          UpperCamelCase	,     UpperCamelCase    :  Tuple							=	jax.random.split(SCREAMING_SNAKE_CASE_       )
          UpperCamelCase    :  int							=	{'params': params_rng, 'dropout': dropout_rng}

          return self.init(SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_       )["params"]
       def 					snake_case_       (			self       )     ->		Optional[int]:
          UpperCamelCase    :  str							=	self.block_out_channels
          UpperCamelCase    :  int							=	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    :  str							=	self.num_attention_heads or self.attention_head_dim

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

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

          UpperCamelCase    :  Union[str, Any]							=	self.only_cross_attention
          if isinstance(SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_       ):
             UpperCamelCase    :  str							=	(only_cross_attention,) * len(self.down_block_types       )

          if isinstance(SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_       ):
             UpperCamelCase    :  int							=	(num_attention_heads,) * len(self.down_block_types       )

          # down
          UpperCamelCase    :  Optional[int]							=	[]
          UpperCamelCase    :  Optional[Any]							=	block_out_channels[0]
          for i, down_block_type in enumerate(self.down_block_types       ):
             UpperCamelCase    :  Tuple							=	output_channel
             UpperCamelCase    :  str							=	block_out_channels[i]
             UpperCamelCase    :  List[Any]							=	i == len(SCREAMING_SNAKE_CASE_       ) - 1

             if down_block_type == "CrossAttnDownBlock2D":
                UpperCamelCase    :  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],			use_memory_efficient_attention=self.use_memory_efficient_attention,			dtype=self.dtype,			)
             else:
                UpperCamelCase    :  List[str]							=	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_       )
          UpperCamelCase    :  Dict							=	down_blocks

          # mid
          UpperCamelCase    :  Optional[int]							=	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    :  List[Any]							=	[]
          UpperCamelCase    :  int							=	list(reversed(SCREAMING_SNAKE_CASE_       )       )
          UpperCamelCase    :  Dict							=	list(reversed(SCREAMING_SNAKE_CASE_       )       )
          UpperCamelCase    :  List[Any]							=	list(reversed(SCREAMING_SNAKE_CASE_       )       )
          UpperCamelCase    :  List[Any]							=	reversed_block_out_channels[0]
          for i, up_block_type in enumerate(self.up_block_types       ):
             UpperCamelCase    :  str							=	output_channel
             UpperCamelCase    :  Union[str, Any]							=	reversed_block_out_channels[i]
             UpperCamelCase    :  Union[str, Any]							=	reversed_block_out_channels[min(i + 1,			len(SCREAMING_SNAKE_CASE_       ) - 1       )]

             UpperCamelCase    :  str							=	i == len(SCREAMING_SNAKE_CASE_       ) - 1

             if up_block_type == "CrossAttnUpBlock2D":
                UpperCamelCase    :  Dict							=	FlaxCrossAttnUpBlockaD(
                    in_channels=SCREAMING_SNAKE_CASE_,			out_channels=SCREAMING_SNAKE_CASE_,			prev_output_channel=SCREAMING_SNAKE_CASE_,			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    :  str							=	FlaxUpBlockaD(
                    in_channels=SCREAMING_SNAKE_CASE_,			out_channels=SCREAMING_SNAKE_CASE_,			prev_output_channel=SCREAMING_SNAKE_CASE_,			num_layers=self.layers_per_block + 1,			add_upsample=not is_final_block,			dropout=self.dropout,			dtype=self.dtype,			)

             up_blocks.append(SCREAMING_SNAKE_CASE_       )
             UpperCamelCase    :  Union[str, Any]							=	output_channel
          UpperCamelCase    :  str							=	up_blocks

          # out
          UpperCamelCase    :  Union[str, Any]							=	nn.GroupNorm(num_groups=32,			epsilon=1e-5       )
          UpperCamelCase    :  List[Any]							=	nn.Conv(
              self.out_channels,			kernel_size=(3, 3),			strides=(1, 1),			padding=((1, 1), (1, 1)),			dtype=self.dtype,			)






       def __call__(			self,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_=None,			SCREAMING_SNAKE_CASE_=None,			SCREAMING_SNAKE_CASE_ = True,			SCREAMING_SNAKE_CASE_ = False,			)     ->		Union[FlaxUNetaDConditionOutput, Tuple]:
          # 1. time
          if not isinstance(SCREAMING_SNAKE_CASE_,			jnp.ndarray       ):
             UpperCamelCase    :  List[str]							=	jnp.array([timesteps],			dtype=jnp.intaa       )
          elif isinstance(SCREAMING_SNAKE_CASE_,			jnp.ndarray       ) and len(timesteps.shape       ) == 0:
             UpperCamelCase    :  Optional[int]							=	timesteps.astype(dtype=jnp.floataa       )
             UpperCamelCase    :  Union[str, Any]							=	jnp.expand_dims(SCREAMING_SNAKE_CASE_,			0       )

          UpperCamelCase    :  List[str]							=	self.time_proj(SCREAMING_SNAKE_CASE_       )
          UpperCamelCase    :  List[Any]							=	self.time_embedding(SCREAMING_SNAKE_CASE_       )

          # 2. pre-process
          UpperCamelCase    :  Any							=	jnp.transpose(SCREAMING_SNAKE_CASE_,			(0, 2, 3, 1)       )
          UpperCamelCase    :  Any							=	self.conv_in(SCREAMING_SNAKE_CASE_       )

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

          if down_block_additional_residuals is not None:
             UpperCamelCase    :  List[Any]							=	()

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

             UpperCamelCase    :  Dict							=	new_down_block_res_samples

          # 4. mid
          UpperCamelCase    :  Optional[int]							=	self.mid_block(SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_,			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    :  Optional[int]							=	down_block_res_samples[-(self.layers_per_block + 1) :]
             UpperCamelCase    :  str							=	down_block_res_samples[: -(self.layers_per_block + 1)]
             if isinstance(SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_       ):
                UpperCamelCase    :  Union[str, Any]							=	up_block(
                    SCREAMING_SNAKE_CASE_,			temb=SCREAMING_SNAKE_CASE_,			encoder_hidden_states=SCREAMING_SNAKE_CASE_,			res_hidden_states_tuple=SCREAMING_SNAKE_CASE_,			deterministic=not train,			)
             else:
                UpperCamelCase    :  List[Any]							=	up_block(SCREAMING_SNAKE_CASE_,			temb=SCREAMING_SNAKE_CASE_,			res_hidden_states_tuple=SCREAMING_SNAKE_CASE_,			deterministic=not train       )

        # 6. post-process
          UpperCamelCase    :  Optional[int]							=	self.conv_norm_out(SCREAMING_SNAKE_CASE_       )
          UpperCamelCase    :  Optional[int]							=	nn.silu(SCREAMING_SNAKE_CASE_       )
          UpperCamelCase    :  Optional[Any]							=	self.conv_out(SCREAMING_SNAKE_CASE_       )
          UpperCamelCase    :  List[str]							=	jnp.transpose(SCREAMING_SNAKE_CASE_,			(0, 3, 1, 2)       )

          if not return_dict:
             return (sample,)

          return FlaxUNetaDConditionOutput(sample=SCREAMING_SNAKE_CASE_       )








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


__UpperCAmelCase          =       logging.get_logger(__name__)

__UpperCAmelCase          =       {
    '''facebook/s2t-wav2vec2-large-en-de''': (
        '''https://huggingface.co/facebook/s2t-wav2vec2-large-en-de/resolve/main/config.json'''
    ),
    # See all Speech2Text models at https://huggingface.co/models?filter=speech2text2
}



class 				lowerCAmelCase_		(  a__							):
       UpperCAmelCase__				:					Optional[Any]        =  "speech_to_text_2"
       UpperCAmelCase__				:					List[Any]        =  ["past_key_values"]
       UpperCAmelCase__				:					Any        =  {"num_attention_heads": "decoder_attention_heads", "hidden_size": "d_model"}
       def __init__(			self,			SCREAMING_SNAKE_CASE_=1_0000,			SCREAMING_SNAKE_CASE_=6,			SCREAMING_SNAKE_CASE_=2048,			SCREAMING_SNAKE_CASE_=4,			SCREAMING_SNAKE_CASE_=0.0,			SCREAMING_SNAKE_CASE_=True,			SCREAMING_SNAKE_CASE_="relu",			SCREAMING_SNAKE_CASE_=256,			SCREAMING_SNAKE_CASE_=0.1,			SCREAMING_SNAKE_CASE_=0.0,			SCREAMING_SNAKE_CASE_=0.0,			SCREAMING_SNAKE_CASE_=0.02,			SCREAMING_SNAKE_CASE_=2,			SCREAMING_SNAKE_CASE_=True,			SCREAMING_SNAKE_CASE_=1,			SCREAMING_SNAKE_CASE_=0,			SCREAMING_SNAKE_CASE_=2,			SCREAMING_SNAKE_CASE_=1024,			**SCREAMING_SNAKE_CASE_,			)     ->		int:
          UpperCamelCase    :  Optional[int]							=	vocab_size
          UpperCamelCase    :  List[str]							=	d_model
          UpperCamelCase    :  List[str]							=	decoder_ffn_dim
          UpperCamelCase    :  Optional[Any]							=	decoder_layers
          UpperCamelCase    :  Any							=	decoder_attention_heads
          UpperCamelCase    :  Tuple							=	dropout
          UpperCamelCase    :  str							=	attention_dropout
          UpperCamelCase    :  str							=	activation_dropout
          UpperCamelCase    :  Union[str, Any]							=	activation_function
          UpperCamelCase    :  Optional[int]							=	init_std
          UpperCamelCase    :  Tuple							=	decoder_layerdrop
          UpperCamelCase    :  Dict							=	use_cache
          UpperCamelCase    :  Any							=	decoder_layers
          UpperCamelCase    :  Union[str, Any]							=	scale_embedding  # scale factor will be sqrt(d_model) if True
          UpperCamelCase    :  Union[str, Any]							=	max_target_positions

          super().__init__(
              pad_token_id=SCREAMING_SNAKE_CASE_,			bos_token_id=SCREAMING_SNAKE_CASE_,			eos_token_id=SCREAMING_SNAKE_CASE_,			decoder_start_token_id=SCREAMING_SNAKE_CASE_,			**SCREAMING_SNAKE_CASE_,			)



