





'''simple docstring'''







import os
import unicodedata
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


snake_case_   :				List[Any]		=     logging.get_logger(__name__)
snake_case_   :				Tuple		=     {'vocab_file': 'spiece.model'}

snake_case_   :				int		=     {
    'vocab_file': {
        'albert-base-v1': 'https://huggingface.co/albert-base-v1/resolve/main/spiece.model',
        'albert-large-v1': 'https://huggingface.co/albert-large-v1/resolve/main/spiece.model',
        'albert-xlarge-v1': 'https://huggingface.co/albert-xlarge-v1/resolve/main/spiece.model',
        'albert-xxlarge-v1': 'https://huggingface.co/albert-xxlarge-v1/resolve/main/spiece.model',
        'albert-base-v2': 'https://huggingface.co/albert-base-v2/resolve/main/spiece.model',
        'albert-large-v2': 'https://huggingface.co/albert-large-v2/resolve/main/spiece.model',
        'albert-xlarge-v2': 'https://huggingface.co/albert-xlarge-v2/resolve/main/spiece.model',
        'albert-xxlarge-v2': 'https://huggingface.co/albert-xxlarge-v2/resolve/main/spiece.model',
    }
}

snake_case_   :				Any		=     {
    'albert-base-v1': 5_12,
    'albert-large-v1': 5_12,
    'albert-xlarge-v1': 5_12,
    'albert-xxlarge-v1': 5_12,
    'albert-base-v2': 5_12,
    'albert-large-v2': 5_12,
    'albert-xlarge-v2': 5_12,
    'albert-xxlarge-v2': 5_12,
}

snake_case_   :				Dict		=     '▁'


class     __a					(_lowerCamelCase   ):
							__a    :      Tuple  	=				VOCAB_FILES_NAMES
							__a    :      Tuple  	=				PRETRAINED_VOCAB_FILES_MAP
							__a    :      Dict  	=				PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES

							def __init__(							self							:   Dict				,   __magic_name__							:   Optional[int]				,   __magic_name__							:   Optional[int]=True				,   __magic_name__							:   Optional[int]=True				,   __magic_name__							:   Dict=False				,   __magic_name__							:   Optional[int]="[CLS]"				,   __magic_name__							:   Dict="[SEP]"				,   __magic_name__							:   str="<unk>"				,   __magic_name__							:   Dict="[SEP]"				,   __magic_name__							:   Optional[Any]="<pad>"				,   __magic_name__							:   Union[str, Any]="[CLS]"				,   __magic_name__							:   Tuple="[MASK]"				,   __magic_name__							:   Optional[Dict[str, Any]] = None				,   **__magic_name__							:   Dict				,   )    ->				None:







													"""simple docstring"""
													# Mask token behave like a normal word, i.e. include the space before it and
													# is included in the raw text, there should be a match in a non-normalized sentence.
													UpperCAmelCase_	:					List[Any]     =     (
													    AddedToken(_UpperCamelCase				,   lstrip=_UpperCamelCase				,   rstrip=_UpperCamelCase				,   normalized=_UpperCamelCase      )
													    if isinstance(_UpperCamelCase				,   _UpperCamelCase      )
													    else mask_token
													)

													UpperCAmelCase_	:					str     =     {} if sp_model_kwargs is None else sp_model_kwargs

													super().__init__(
													    do_lower_case=_UpperCamelCase				,   remove_space=_UpperCamelCase				,   keep_accents=_UpperCamelCase				,   bos_token=_UpperCamelCase				,   eos_token=_UpperCamelCase				,   unk_token=_UpperCamelCase				,   sep_token=_UpperCamelCase				,   pad_token=_UpperCamelCase				,   cls_token=_UpperCamelCase				,   mask_token=_UpperCamelCase				,   sp_model_kwargs=self.sp_model_kwargs				,   **_UpperCamelCase				,   )

													UpperCAmelCase_	:					List[Any]     =     do_lower_case
													UpperCAmelCase_	:					Union[str, Any]     =     remove_space
													UpperCAmelCase_	:					Dict     =     keep_accents
													UpperCAmelCase_	:					Dict     =     vocab_file

													UpperCAmelCase_	:					Union[str, Any]     =     spm.SentencePieceProcessor(**self.sp_model_kwargs      )
													self.sp_model.Load(_UpperCamelCase      )

							@property
							def 			UpperCAmelCase__		(							self							:   str      )    ->				Tuple:







													"""simple docstring"""
													return len(self.sp_model      )

							def 			UpperCAmelCase__		(							self							:   Optional[Any]      )    ->				List[Any]:







													"""simple docstring"""
													UpperCAmelCase_	:					Optional[Any]     =     {self.convert_ids_to_tokens(_UpperCamelCase      ): i for i in range(self.vocab_size      )}
													vocab.update(self.added_tokens_encoder      )
													return vocab

							def __getstate__(							self							:   Union[str, Any]      )    ->				List[Any]:







													"""simple docstring"""
													UpperCAmelCase_	:					Union[str, Any]     =     self.__dict__.copy()
													UpperCAmelCase_	:					Dict     =     None
													return state

							def __setstate__(							self							:   Dict				,   __magic_name__							:   str      )    ->				Any:







													"""simple docstring"""
													UpperCAmelCase_	:					Optional[Any]     =     d

													# for backward compatibility
													if not hasattr(self				,   '''sp_model_kwargs'''      ):
																			UpperCAmelCase_	:					int     =     {}

													UpperCAmelCase_	:					List[str]     =     spm.SentencePieceProcessor(**self.sp_model_kwargs      )
													self.sp_model.Load(self.vocab_file      )

							def 			UpperCAmelCase__		(							self							:   Optional[Any]				,   __magic_name__							:   Any      )    ->				List[Any]:







													"""simple docstring"""
													if self.remove_space:
																			UpperCAmelCase_	:					int     =     ''' '''.join(inputs.strip().split()      )
													else:
																			UpperCAmelCase_	:					Union[str, Any]     =     inputs
													UpperCAmelCase_	:					Union[str, Any]     =     outputs.replace('''``'''				,   '''"'''      ).replace('''\'\''''				,   '''"'''      )

													if not self.keep_accents:
																			UpperCAmelCase_	:					List[str]     =     unicodedata.normalize('''NFKD'''				,   _UpperCamelCase      )
																			UpperCAmelCase_	:					str     =     ''''''.join([c for c in outputs if not unicodedata.combining(_UpperCamelCase      )]      )
													if self.do_lower_case:
																			UpperCAmelCase_	:					Union[str, Any]     =     outputs.lower()

													return outputs

							def 			UpperCAmelCase__		(							self							:   Any				,   __magic_name__							:   str      )    ->				List[str]:







													"""simple docstring"""
													UpperCAmelCase_	:					Optional[Any]     =     self.preprocess_text(_UpperCamelCase      )
													UpperCAmelCase_	:					Tuple     =     self.sp_model.encode(_UpperCamelCase				,   out_type=_UpperCamelCase      )
													UpperCAmelCase_	:					Any     =     []
													for piece in pieces:
																			if len(_UpperCamelCase      ) > 1 and piece[-1] == str(''','''      ) and piece[-2].isdigit():
																									UpperCAmelCase_	:					List[Any]     =     self.sp_model.EncodeAsPieces(piece[:-1].replace(_UpperCamelCase				,   ''''''      )      )
																									if piece[0] != SPIECE_UNDERLINE and cur_pieces[0][0] == SPIECE_UNDERLINE:
																															if len(cur_pieces[0]      ) == 1:
																																					UpperCAmelCase_	:					int     =     cur_pieces[1:]
																															else:
																																					UpperCAmelCase_	:					Any     =     cur_pieces[0][1:]
																									cur_pieces.append(piece[-1]      )
																									new_pieces.extend(_UpperCamelCase      )
																			else:
																									new_pieces.append(_UpperCamelCase      )

													return new_pieces

							def 			UpperCAmelCase__		(							self							:   Any				,   __magic_name__							:   int      )    ->				Optional[Any]:







													"""simple docstring"""
													return self.sp_model.PieceToId(_UpperCamelCase      )

							def 			UpperCAmelCase__		(							self							:   Tuple				,   __magic_name__							:   int      )    ->				Union[str, Any]:







													"""simple docstring"""
													return self.sp_model.IdToPiece(_UpperCamelCase      )

							def 			UpperCAmelCase__		(							self							:   List[Any]				,   __magic_name__							:   Tuple      )    ->				Any:







													"""simple docstring"""
													UpperCAmelCase_	:					int     =     []
													UpperCAmelCase_	:					Dict     =     ''''''
													UpperCAmelCase_	:					int     =     False
													for token in tokens:
																			# make sure that special tokens are not decoded using sentencepiece model
																			if token in self.all_special_tokens:
																									if not prev_is_special:
																															out_string += " "
																									out_string += self.sp_model.decode(_UpperCamelCase      ) + token
																									UpperCAmelCase_	:					Any     =     True
																									UpperCAmelCase_	:					int     =     []
																			else:
																									current_sub_tokens.append(_UpperCamelCase      )
																									UpperCAmelCase_	:					int     =     False
													out_string += self.sp_model.decode(_UpperCamelCase      )
													return out_string.strip()

							def 			UpperCAmelCase__		(							self							:   Tuple				,   __magic_name__							:   List[int]				,   __magic_name__							:   Optional[List[int]] = None      )    ->				List[int]:







													"""simple docstring"""
													UpperCAmelCase_	:					Optional[Any]     =     [self.sep_token_id]
													UpperCAmelCase_	:					Optional[Any]     =     [self.cls_token_id]
													if token_ids_a is None:
																			return cls + token_ids_a + sep
													return cls + token_ids_a + sep + token_ids_a + sep

							def 			UpperCAmelCase__		(							self							:   List[Any]				,   __magic_name__							:   List[int]				,   __magic_name__							:   Optional[List[int]] = None				,   __magic_name__							:   bool = False      )    ->				List[int]:







													"""simple docstring"""

													if already_has_special_tokens:
																			return super().get_special_tokens_mask(
																			    token_ids_a=_UpperCamelCase				,   token_ids_a=_UpperCamelCase				,   already_has_special_tokens=_UpperCamelCase      )

													if token_ids_a is not None:
																			return [1] + ([0] * len(_UpperCamelCase      )) + [1] + ([0] * len(_UpperCamelCase      )) + [1]
													return [1] + ([0] * len(_UpperCamelCase      )) + [1]

							def 			UpperCAmelCase__		(							self							:   Optional[Any]				,   __magic_name__							:   List[int]				,   __magic_name__							:   Optional[List[int]] = None      )    ->				List[int]:







													"""simple docstring"""
													UpperCAmelCase_	:					List[Any]     =     [self.sep_token_id]
													UpperCAmelCase_	:					Optional[int]     =     [self.cls_token_id]

													if token_ids_a is None:
																			return len(cls + token_ids_a + sep      ) * [0]
													return len(cls + token_ids_a + sep      ) * [0] + len(token_ids_a + sep      ) * [1]



							def 			UpperCAmelCase__		(							self							:   Union[str, Any]				,   __magic_name__							:   str				,   __magic_name__							:   Optional[str] = None      )    ->				Tuple[str]:







													"""simple docstring"""
													if not os.path.isdir(_UpperCamelCase      ):
																			logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""      )
																			return
													UpperCAmelCase_	:					Union[str, Any]     =     os.path.join(
													    _UpperCamelCase				,   (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']      )

													if os.path.abspath(self.vocab_file      ) != os.path.abspath(_UpperCamelCase      ) and os.path.isfile(self.vocab_file      ):
																			copyfile(self.vocab_file				,   _UpperCamelCase      )
													elif not os.path.isfile(self.vocab_file      ):
																			with open(_UpperCamelCase				,   '''wb'''      ) as fi:
																									UpperCAmelCase_	:					Optional[int]     =     self.sp_model.serialized_model_proto()
																									fi.write(_UpperCamelCase      )

													return (out_vocab_file,)












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 A_					(			_lowerCamelCase   ):
   lowerCAmelCase__        = 42





@flax_register_to_config
class A_					(			nn.Module ,      _lowerCamelCase ,      _lowerCamelCase   ):
   lowerCAmelCase__        = 32
   lowerCAmelCase__        = 4
   lowerCAmelCase__        = 4
   lowerCAmelCase__        = (
       "CrossAttnDownBlock2D",
       "CrossAttnDownBlock2D",
       "CrossAttnDownBlock2D",
       "DownBlock2D",
   )
   lowerCAmelCase__        = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
   lowerCAmelCase__        = False
   lowerCAmelCase__        = (320, 640, 1280, 1280)
   lowerCAmelCase__        = 2
   lowerCAmelCase__        = 8
   lowerCAmelCase__        = None
   lowerCAmelCase__        = 1280
   lowerCAmelCase__        = 0.0
   lowerCAmelCase__        = False
   lowerCAmelCase__        = jnp.floataa
   lowerCAmelCase__        = True
   lowerCAmelCase__        = 0
   lowerCAmelCase__        = False




   def _lowerCAmelCase      (self	:Tuple    ,   _UpperCamelCase	:jax.random.KeyArray )->					FrozenDict:
         # init input tensors
         __A          =				(1, self.in_channels, self.sample_size, self.sample_size)
         __A          =				jnp.zeros(_UpperCamelCase    ,   dtype=jnp.floataa )
         __A          =				jnp.ones((1,)    ,   dtype=jnp.intaa )
         __A          =				jnp.zeros((1, 1, self.cross_attention_dim)    ,   dtype=jnp.floataa )

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

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




   def _lowerCAmelCase      (self	:Tuple )->					Optional[int]:
         __A          =				self.block_out_channels
         __A          =				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.
         __A          =				self.num_attention_heads or self.attention_head_dim

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

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

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

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

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

               if down_block_type == "CrossAttnDownBlock2D":
                     __A          =				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:
                     __A          =				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 )
         __A          =				down_blocks

         # mid
         __A          =				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
         __A          =				[]
         __A          =				list(reversed(_UpperCamelCase ) )
         __A          =				list(reversed(_UpperCamelCase ) )
         __A          =				list(reversed(_UpperCamelCase ) )
         __A          =				reversed_block_out_channels[0]
         for i, up_block_type in enumerate(self.up_block_types ):
               __A          =				output_channel
               __A          =				reversed_block_out_channels[i]
               __A          =				reversed_block_out_channels[min(i + 1    ,   len(_UpperCamelCase ) - 1 )]

               __A          =				i == len(_UpperCamelCase ) - 1

               if up_block_type == "CrossAttnUpBlock2D":
                     __A          =				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:
                     __A          =				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 )
               __A          =				output_channel
         __A          =				up_blocks

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






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

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

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

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

         if down_block_additional_residuals is not None:
               __A          =				()

               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,)

               __A          =				new_down_block_res_samples

         # 4. mid
         __A          =				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:
               __A          =				down_block_res_samples[-(self.layers_per_block + 1) :]
               __A          =				down_block_res_samples[: -(self.layers_per_block + 1)]
               if isinstance(_UpperCamelCase    ,   _UpperCamelCase ):
                     __A          =				up_block(
                         _UpperCamelCase    ,   temb=_UpperCamelCase    ,   encoder_hidden_states=_UpperCamelCase    ,   res_hidden_states_tuple=_UpperCamelCase    ,   deterministic=not train    ,   )
               else:
                     __A          =				up_block(_UpperCamelCase    ,   temb=_UpperCamelCase    ,   res_hidden_states_tuple=_UpperCamelCase    ,   deterministic=not train )

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

         if not return_dict:
               return (sample,)

         return FlaxUNetaDConditionOutput(sample=_UpperCamelCase )





