


'''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        lowercase__    (		_snake_case   ):





							'''simple docstring'''


							A_				: jnp.ndarray





@flax_register_to_config
class        lowercase__    (		nn.Module	,	_snake_case	,	_snake_case   ):





							'''simple docstring'''


							A_				: int       =    32
							A_				: int       =    4
							A_				: int       =    4
							A_				: Tuple[str]       =    (
							    "CrossAttnDownBlock2D",
							    "CrossAttnDownBlock2D",
							    "CrossAttnDownBlock2D",
							    "DownBlock2D",
							)
							A_				: Tuple[str]       =    ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
							A_				: Union[bool, Tuple[bool]]       =    False
							A_				: Tuple[int]       =    (320, 640, 1280, 1280)
							A_				: int       =    2
							A_				: Union[int, Tuple[int]]       =    8
							A_				: Optional[Union[int, Tuple[int]]]       =    None
							A_				: int       =    1280
							A_				: float       =    0.0
							A_				: bool       =    False
							A_				: jnp.dtype       =    jnp.floataa
							A_				: bool       =    True
							A_				: int       =    0
							A_				: bool       =    False






							def UpperCAmelCase_      (       self       ,   __snake_case    ):
											# init input tensors
											_SCREAMING_SNAKE_CASE					:			Optional[int]							= (1, self.in_channels, self.sample_size, self.sample_size)
											_SCREAMING_SNAKE_CASE					:			Any							= jnp.zeros(__snake_case       ,   dtype=jnp.floataa    )
											_SCREAMING_SNAKE_CASE					:			Any							= jnp.ones((1,)       ,   dtype=jnp.intaa    )
											_SCREAMING_SNAKE_CASE					:			Any							= jnp.zeros((1, 1, self.cross_attention_dim)       ,   dtype=jnp.floataa    )

											_SCREAMING_SNAKE_CASE						,			_SCREAMING_SNAKE_CASE					:			int							= jax.random.split(__snake_case    )
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= {"""params""": params_rng, """dropout""": dropout_rng}

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






							def UpperCAmelCase_      (       self    ):
											_SCREAMING_SNAKE_CASE					:			str							= self.block_out_channels
											_SCREAMING_SNAKE_CASE					:			List[Any]							= 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.
											_SCREAMING_SNAKE_CASE					:			Optional[int]							= self.num_attention_heads or self.attention_head_dim

											# input
											_SCREAMING_SNAKE_CASE					:			Tuple							= 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					:			int							= FlaxTimesteps(
											    block_out_channels[0]       ,   flip_sin_to_cos=self.flip_sin_to_cos       ,   freq_shift=self.config.freq_shift    )
											_SCREAMING_SNAKE_CASE					:			List[str]							= FlaxTimestepEmbedding(__snake_case       ,   dtype=self.dtype    )

											_SCREAMING_SNAKE_CASE					:			Tuple							= self.only_cross_attention
											if isinstance(__snake_case       ,   __snake_case    ):
															_SCREAMING_SNAKE_CASE					:			Tuple							= (only_cross_attention,) * len(self.down_block_types    )

											if isinstance(__snake_case       ,   __snake_case    ):
															_SCREAMING_SNAKE_CASE					:			Optional[int]							= (num_attention_heads,) * len(self.down_block_types    )

											# down
											_SCREAMING_SNAKE_CASE					:			Dict							= []
											_SCREAMING_SNAKE_CASE					:			Any							= block_out_channels[0]
											for i, down_block_type in enumerate(self.down_block_types    ):
															_SCREAMING_SNAKE_CASE					:			int							= output_channel
															_SCREAMING_SNAKE_CASE					:			Union[str, Any]							= block_out_channels[i]
															_SCREAMING_SNAKE_CASE					:			int							= i == len(__snake_case    ) - 1

															if down_block_type == "CrossAttnDownBlock2D":
																			_SCREAMING_SNAKE_CASE					:			Tuple							= FlaxCrossAttnDownBlockaD(
																			    in_channels=__snake_case       ,   out_channels=__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:
																			_SCREAMING_SNAKE_CASE					:			Optional[Any]							= FlaxDownBlockaD(
																			    in_channels=__snake_case       ,   out_channels=__snake_case       ,   dropout=self.dropout       ,   num_layers=self.layers_per_block       ,   add_downsample=not is_final_block       ,   dtype=self.dtype       ,   )

															down_blocks.append(__snake_case    )
											_SCREAMING_SNAKE_CASE					:			Any							= down_blocks

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

															_SCREAMING_SNAKE_CASE					:			int							= i == len(__snake_case    ) - 1

															if up_block_type == "CrossAttnUpBlock2D":
																			_SCREAMING_SNAKE_CASE					:			List[Any]							= FlaxCrossAttnUpBlockaD(
																			    in_channels=__snake_case       ,   out_channels=__snake_case       ,   prev_output_channel=__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:
																			_SCREAMING_SNAKE_CASE					:			List[str]							= FlaxUpBlockaD(
																			    in_channels=__snake_case       ,   out_channels=__snake_case       ,   prev_output_channel=__snake_case       ,   num_layers=self.layers_per_block + 1       ,   add_upsample=not is_final_block       ,   dropout=self.dropout       ,   dtype=self.dtype       ,   )

															up_blocks.append(__snake_case    )
															_SCREAMING_SNAKE_CASE					:			str							= output_channel
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= up_blocks

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






							def __call__(       self       ,   __snake_case       ,   __snake_case       ,   __snake_case       ,   __snake_case=None       ,   __snake_case=None       ,   __snake_case = True       ,   __snake_case = False       ,   ):
											# 1. time
											if not isinstance(__snake_case       ,   jnp.ndarray    ):
															_SCREAMING_SNAKE_CASE					:			Optional[Any]							= jnp.array([timesteps]       ,   dtype=jnp.intaa    )
											elif isinstance(__snake_case       ,   jnp.ndarray    ) and len(timesteps.shape    ) == 0:
															_SCREAMING_SNAKE_CASE					:			int							= timesteps.astype(dtype=jnp.floataa    )
															_SCREAMING_SNAKE_CASE					:			Dict							= jnp.expand_dims(__snake_case       ,   0    )

											_SCREAMING_SNAKE_CASE					:			List[Any]							= self.time_proj(__snake_case    )
											_SCREAMING_SNAKE_CASE					:			List[str]							= self.time_embedding(__snake_case    )

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

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

											if down_block_additional_residuals is not None:
															_SCREAMING_SNAKE_CASE					:			int							= ()

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

															_SCREAMING_SNAKE_CASE					:			Tuple							= new_down_block_res_samples

											# 4. mid
											_SCREAMING_SNAKE_CASE					:			Tuple							= self.mid_block(__snake_case       ,   __snake_case       ,   __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:
															_SCREAMING_SNAKE_CASE					:			Optional[int]							= down_block_res_samples[-(self.layers_per_block + 1) :]
															_SCREAMING_SNAKE_CASE					:			Any							= down_block_res_samples[: -(self.layers_per_block + 1)]
															if isinstance(__snake_case       ,   __snake_case    ):
																			_SCREAMING_SNAKE_CASE					:			Optional[Any]							= up_block(
																			    __snake_case       ,   temb=__snake_case       ,   encoder_hidden_states=__snake_case       ,   res_hidden_states_tuple=__snake_case       ,   deterministic=not train       ,   )
															else:
																			_SCREAMING_SNAKE_CASE					:			Dict							= up_block(__snake_case       ,   temb=__snake_case       ,   res_hidden_states_tuple=__snake_case       ,   deterministic=not train    )

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

											if not return_dict:
															return (sample,)

											return FlaxUNetaDConditionOutput(sample=__snake_case    )








'''simple docstring'''


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


UpperCAmelCase_						:  str				  =						logging.get_logger(__name__)

UpperCAmelCase_						:  Optional[int]				  =						{
    'SCUT-DLVCLab/lilt-roberta-en-base': (
        'https://huggingface.co/SCUT-DLVCLab/lilt-roberta-en-base/resolve/main/config.json'
    ),
}





class        lowercase__    (		_snake_case   ):





							'''simple docstring'''


							A_				: List[str]       =    """lilt"""






							def __init__(       self       ,   __snake_case=3_0522       ,   __snake_case=768       ,   __snake_case=12       ,   __snake_case=12       ,   __snake_case=3072       ,   __snake_case="gelu"       ,   __snake_case=0.1       ,   __snake_case=0.1       ,   __snake_case=512       ,   __snake_case=2       ,   __snake_case=0.02       ,   __snake_case=1e-12       ,   __snake_case=0       ,   __snake_case="absolute"       ,   __snake_case=None       ,   __snake_case=4       ,   __snake_case=1024       ,   **__snake_case       ,   ):
											super().__init__(pad_token_id=__snake_case       ,   **__snake_case    )

											_SCREAMING_SNAKE_CASE					:			str							= vocab_size
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= hidden_size
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= num_hidden_layers
											_SCREAMING_SNAKE_CASE					:			Optional[int]							= num_attention_heads
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= hidden_act
											_SCREAMING_SNAKE_CASE					:			Dict							= intermediate_size
											_SCREAMING_SNAKE_CASE					:			Tuple							= hidden_dropout_prob
											_SCREAMING_SNAKE_CASE					:			List[str]							= attention_probs_dropout_prob
											_SCREAMING_SNAKE_CASE					:			Union[str, Any]							= max_position_embeddings
											_SCREAMING_SNAKE_CASE					:			Dict							= type_vocab_size
											_SCREAMING_SNAKE_CASE					:			Union[str, Any]							= initializer_range
											_SCREAMING_SNAKE_CASE					:			Tuple							= layer_norm_eps
											_SCREAMING_SNAKE_CASE					:			str							= position_embedding_type
											_SCREAMING_SNAKE_CASE					:			List[str]							= classifier_dropout
											_SCREAMING_SNAKE_CASE					:			int							= channel_shrink_ratio
											_SCREAMING_SNAKE_CASE					:			str							= max_ad_position_embeddings




