



from dataclasses import dataclass
from typing import Dict, Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, apply_forward_hook
from .attention_processor import AttentionProcessor, AttnProcessor
from .modeling_utils import ModelMixin
from .vae import Decoder, DecoderOutput, DiagonalGaussianDistribution, Encoder







@dataclass
class 		snake_case						(	SCREAMING_SNAKE_CASE_  ):



   '''simple docstring'''






   snake_case_					: "DiagonalGaussianDistribution"







class 		snake_case						(	SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_  ):



   '''simple docstring'''






   snake_case_					: List[str]          =		True

   @register_to_config
   def __init__(					self							:			Optional[int]				,		lowerCAmelCase							:			int = 3				,		lowerCAmelCase							:			int = 3				,		lowerCAmelCase							:			Tuple[str] = ("DownEncoderBlock2D",)				,		lowerCAmelCase							:			Tuple[str] = ("UpDecoderBlock2D",)				,		lowerCAmelCase							:			Tuple[int] = (64,)				,		lowerCAmelCase							:			int = 1				,		lowerCAmelCase							:			str = "silu"				,		lowerCAmelCase							:			int = 4				,		lowerCAmelCase							:			int = 32				,		lowerCAmelCase							:			int = 32				,		lowerCAmelCase							:			float = 0.18_215				,		)				->				List[str]:

    """simple docstring"""







    super().__init__()

    # pass init params to Encoder
    _snake_case    :  int					=					Encoder(
        in_channels=lowerCAmelCase				,		out_channels=lowerCAmelCase				,		down_block_types=lowerCAmelCase				,		block_out_channels=lowerCAmelCase				,		layers_per_block=lowerCAmelCase				,		act_fn=lowerCAmelCase				,		norm_num_groups=lowerCAmelCase				,		double_z=lowerCAmelCase				,		)

    # pass init params to Decoder
    _snake_case    :  List[Any]					=					Decoder(
        in_channels=lowerCAmelCase				,		out_channels=lowerCAmelCase				,		up_block_types=lowerCAmelCase				,		block_out_channels=lowerCAmelCase				,		layers_per_block=lowerCAmelCase				,		norm_num_groups=lowerCAmelCase				,		act_fn=lowerCAmelCase				,		)

    _snake_case    :  Optional[Any]					=					nn.Convad(2 * latent_channels				,		2 * latent_channels				,		1)
    _snake_case    :  Tuple					=					nn.Convad(lowerCAmelCase				,		lowerCAmelCase				,		1)

    _snake_case    :  Optional[int]					=					False
    _snake_case    :  Union[str, Any]					=					False

    # only relevant if vae tiling is enabled
    _snake_case    :  Optional[int]					=					self.config.sample_size
    _snake_case    :  str					=					(
        self.config.sample_size[0]
        if isinstance(self.config.sample_size				,		(list, tuple))
        else self.config.sample_size
    )
    _snake_case    :  Dict					=					int(sample_size / (2 ** (len(self.config.block_out_channels) - 1)))
    _snake_case    :  Dict					=					0.25

   def      UpperCamelCase_  (					self							:			Union[str, Any]				,		lowerCAmelCase							:			str				,		lowerCAmelCase							:			Optional[Any]=False)				->				Tuple:

    """simple docstring"""







    if isinstance(lowerCAmelCase				,		(Encoder, Decoder)):
     _snake_case    :  Dict					=					value

   def      UpperCamelCase_  (					self							:			Dict				,		lowerCAmelCase							:			bool = True)				->				Optional[Any]:

    """simple docstring"""







    _snake_case    :  Union[str, Any]					=					use_tiling

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

    """simple docstring"""







    self.enable_tiling(lowerCAmelCase)

   def      UpperCamelCase_  (					self							:			str)				->				Any:

    """simple docstring"""







    _snake_case    :  Optional[Any]					=					True

   def      UpperCamelCase_  (					self							:			Any)				->				int:

    """simple docstring"""







    _snake_case    :  List[Any]					=					False

   @property
   # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
   def      UpperCamelCase_  (					self							:			Optional[Any])				->				Dict[str, AttentionProcessor]:

    """simple docstring"""







    _snake_case    :  Optional[Any]					=					{}

    def fn_recursive_add_processors(lowerCAmelCase							:			str				,		lowerCAmelCase							:			torch.nn.Module				,		lowerCAmelCase							:			Dict[str, AttentionProcessor]):
     if hasattr(lowerCAmelCase				,		"""set_processor"""):
      _snake_case    :  List[Any]					=					module.processor

     for sub_name, child in module.named_children():
      fn_recursive_add_processors(F'''{name}.{sub_name}'''				,		lowerCAmelCase				,		lowerCAmelCase)

     return processors

    for name, module in self.named_children():
     fn_recursive_add_processors(lowerCAmelCase				,		lowerCAmelCase				,		lowerCAmelCase)

    return processors

   def      UpperCamelCase_  (					self							:			Tuple				,		lowerCAmelCase							:			Union[AttentionProcessor, Dict[str, AttentionProcessor]])				->				Optional[Any]:

    """simple docstring"""







    _snake_case    :  Dict					=					len(self.attn_processors.keys())

    if isinstance(lowerCAmelCase				,		lowerCAmelCase) and len(lowerCAmelCase) != count:
     raise ValueError(
         F'''A dict of processors was passed, but the number of processors {len(lowerCAmelCase)} does not match the'''
         F''' number of attention layers: {count}. Please make sure to pass {count} processor classes.''')

    def fn_recursive_attn_processor(lowerCAmelCase							:			str				,		lowerCAmelCase							:			torch.nn.Module				,		lowerCAmelCase							:			Dict):
     if hasattr(lowerCAmelCase				,		"""set_processor"""):
      if not isinstance(lowerCAmelCase				,		lowerCAmelCase):
       module.set_processor(lowerCAmelCase)
      else:
       module.set_processor(processor.pop(F'''{name}.processor'''))

     for sub_name, child in module.named_children():
      fn_recursive_attn_processor(F'''{name}.{sub_name}'''				,		lowerCAmelCase				,		lowerCAmelCase)

    for name, module in self.named_children():
     fn_recursive_attn_processor(lowerCAmelCase				,		lowerCAmelCase				,		lowerCAmelCase)

   def      UpperCamelCase_  (					self							:			List[Any])				->				List[Any]:

    """simple docstring"""







    self.set_attn_processor(AttnProcessor())

   @apply_forward_hook
   def      UpperCamelCase_  (					self							:			str				,		lowerCAmelCase							:			torch.FloatTensor				,		lowerCAmelCase							:			bool = True)				->				AutoencoderKLOutput:

    """simple docstring"""







    if self.use_tiling and (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size):
     return self.tiled_encode(lowerCAmelCase				,		return_dict=lowerCAmelCase)

    if self.use_slicing and x.shape[0] > 1:
     _snake_case    :  Dict					=					[self.encoder(lowerCAmelCase) for x_slice in x.split(1)]
     _snake_case    :  List[Any]					=					torch.cat(lowerCAmelCase)
    else:
     _snake_case    :  str					=					self.encoder(lowerCAmelCase)

    _snake_case    :  List[Any]					=					self.quant_conv(lowerCAmelCase)
    _snake_case    :  Tuple					=					DiagonalGaussianDistribution(lowerCAmelCase)

    if not return_dict:
     return (posterior,)

    return AutoencoderKLOutput(latent_dist=lowerCAmelCase)

   def      UpperCamelCase_  (					self							:			Optional[int]				,		lowerCAmelCase							:			torch.FloatTensor				,		lowerCAmelCase							:			bool = True)				->				Union[DecoderOutput, torch.FloatTensor]:

    """simple docstring"""







    if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size):
     return self.tiled_decode(lowerCAmelCase				,		return_dict=lowerCAmelCase)

    _snake_case    :  int					=					self.post_quant_conv(lowerCAmelCase)
    _snake_case    :  Optional[int]					=					self.decoder(lowerCAmelCase)

    if not return_dict:
     return (dec,)

    return DecoderOutput(sample=lowerCAmelCase)

   @apply_forward_hook
   def      UpperCamelCase_  (					self							:			Tuple				,		lowerCAmelCase							:			torch.FloatTensor				,		lowerCAmelCase							:			bool = True)				->				Union[DecoderOutput, torch.FloatTensor]:

    """simple docstring"""







    if self.use_slicing and z.shape[0] > 1:
     _snake_case    :  Optional[int]					=					[self._decode(lowerCAmelCase).sample for z_slice in z.split(1)]
     _snake_case    :  Any					=					torch.cat(lowerCAmelCase)
    else:
     _snake_case    :  Dict					=					self._decode(lowerCAmelCase).sample

    if not return_dict:
     return (decoded,)

    return DecoderOutput(sample=lowerCAmelCase)

   def      UpperCamelCase_  (					self							:			Dict				,		lowerCAmelCase							:			Any				,		lowerCAmelCase							:			Optional[int]				,		lowerCAmelCase							:			Tuple)				->				Dict:

    """simple docstring"""







    _snake_case    :  Union[str, Any]					=					min(a.shape[2]				,		b.shape[2]				,		lowerCAmelCase)
    for y in range(lowerCAmelCase):
     _snake_case    :  List[str]					=					a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
    return b

   def      UpperCamelCase_  (					self							:			List[str]				,		lowerCAmelCase							:			List[str]				,		lowerCAmelCase							:			str				,		lowerCAmelCase							:			List[Any])				->				Tuple:

    """simple docstring"""







    _snake_case    :  List[Any]					=					min(a.shape[3]				,		b.shape[3]				,		lowerCAmelCase)
    for x in range(lowerCAmelCase):
     _snake_case    :  Tuple					=					a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
    return b

   def      UpperCamelCase_  (					self							:			Dict				,		lowerCAmelCase							:			torch.FloatTensor				,		lowerCAmelCase							:			bool = True)				->				AutoencoderKLOutput:

    """simple docstring"""







    _snake_case    :  List[str]					=					int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
    _snake_case    :  Dict					=					int(self.tile_latent_min_size * self.tile_overlap_factor)
    _snake_case    :  Union[str, Any]					=					self.tile_latent_min_size - blend_extent

    # Split the image into 512x512 tiles and encode them separately.
    _snake_case    :  Optional[int]					=					[]
    for i in range(0				,		x.shape[2]				,		lowerCAmelCase):
     _snake_case    :  Optional[int]					=					[]
     for j in range(0				,		x.shape[3]				,		lowerCAmelCase):
      _snake_case    :  List[Any]					=					x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
      _snake_case    :  str					=					self.encoder(lowerCAmelCase)
      _snake_case    :  Tuple					=					self.quant_conv(lowerCAmelCase)
      row.append(lowerCAmelCase)
     rows.append(lowerCAmelCase)
    _snake_case    :  Tuple					=					[]
    for i, row in enumerate(lowerCAmelCase):
     _snake_case    :  List[Any]					=					[]
     for j, tile in enumerate(lowerCAmelCase):
      # blend the above tile and the left tile
      # to the current tile and add the current tile to the result row
      if i > 0:
       _snake_case    :  List[Any]					=					self.blend_v(rows[i - 1][j]				,		lowerCAmelCase				,		lowerCAmelCase)
      if j > 0:
       _snake_case    :  Any					=					self.blend_h(row[j - 1]				,		lowerCAmelCase				,		lowerCAmelCase)
      result_row.append(tile[:, :, :row_limit, :row_limit])
     result_rows.append(torch.cat(lowerCAmelCase				,		dim=3))

    _snake_case    :  Union[str, Any]					=					torch.cat(lowerCAmelCase				,		dim=2)
    _snake_case    :  Optional[int]					=					DiagonalGaussianDistribution(lowerCAmelCase)

    if not return_dict:
     return (posterior,)

    return AutoencoderKLOutput(latent_dist=lowerCAmelCase)

   def      UpperCamelCase_  (					self							:			Any				,		lowerCAmelCase							:			torch.FloatTensor				,		lowerCAmelCase							:			bool = True)				->				Union[DecoderOutput, torch.FloatTensor]:

    """simple docstring"""







    _snake_case    :  Any					=					int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
    _snake_case    :  Optional[int]					=					int(self.tile_sample_min_size * self.tile_overlap_factor)
    _snake_case    :  Tuple					=					self.tile_sample_min_size - blend_extent

    # Split z into overlapping 64x64 tiles and decode them separately.
    # The tiles have an overlap to avoid seams between tiles.
    _snake_case    :  Optional[Any]					=					[]
    for i in range(0				,		z.shape[2]				,		lowerCAmelCase):
     _snake_case    :  Optional[int]					=					[]
     for j in range(0				,		z.shape[3]				,		lowerCAmelCase):
      _snake_case    :  Tuple					=					z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
      _snake_case    :  Tuple					=					self.post_quant_conv(lowerCAmelCase)
      _snake_case    :  Tuple					=					self.decoder(lowerCAmelCase)
      row.append(lowerCAmelCase)
     rows.append(lowerCAmelCase)
    _snake_case    :  Optional[int]					=					[]
    for i, row in enumerate(lowerCAmelCase):
     _snake_case    :  List[Any]					=					[]
     for j, tile in enumerate(lowerCAmelCase):
      # blend the above tile and the left tile
      # to the current tile and add the current tile to the result row
      if i > 0:
       _snake_case    :  int					=					self.blend_v(rows[i - 1][j]				,		lowerCAmelCase				,		lowerCAmelCase)
      if j > 0:
       _snake_case    :  List[Any]					=					self.blend_h(row[j - 1]				,		lowerCAmelCase				,		lowerCAmelCase)
      result_row.append(tile[:, :, :row_limit, :row_limit])
     result_rows.append(torch.cat(lowerCAmelCase				,		dim=3))

    _snake_case    :  Optional[int]					=					torch.cat(lowerCAmelCase				,		dim=2)
    if not return_dict:
     return (dec,)

    return DecoderOutput(sample=lowerCAmelCase)


   def      UpperCamelCase_  (					self							:			Union[str, Any]				,		lowerCAmelCase							:			torch.FloatTensor				,		lowerCAmelCase							:			bool = False				,		lowerCAmelCase							:			bool = True				,		lowerCAmelCase							:			Optional[torch.Generator] = None				,		)				->				Union[DecoderOutput, torch.FloatTensor]:

    """simple docstring"""







    _snake_case    :  str					=					sample
    _snake_case    :  Union[str, Any]					=					self.encode(lowerCAmelCase).latent_dist
    if sample_posterior:
     _snake_case    :  str					=					posterior.sample(generator=lowerCAmelCase)
    else:
     _snake_case    :  Dict					=					posterior.mode()
    _snake_case    :  List[str]					=					self.decode(lowerCAmelCase).sample

    if not return_dict:
     return (dec,)

    return DecoderOutput(sample=lowerCAmelCase)





import os
from shutil import copyfile
from typing import List, Optional, Tuple

from ...tokenization_utils import AddedToken
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
  from .tokenization_camembert import CamembertTokenizer
else:
  a__    =	None


a__    =	logging.get_logger(__name__)

a__    =	{"""vocab_file""": """sentencepiece.bpe.model""", """tokenizer_file""": """tokenizer.json"""}

a__    =	{
    """vocab_file""": {
        """camembert-base""": """https://huggingface.co/camembert-base/resolve/main/sentencepiece.bpe.model""",
    },
    """tokenizer_file""": {
        """camembert-base""": """https://huggingface.co/camembert-base/resolve/main/tokenizer.json""",
    },
}

a__    =	{
    """camembert-base""": 5_12,
}

a__    =	"""▁"""







class 		snake_case						(	SCREAMING_SNAKE_CASE_  ):



   '''simple docstring'''






   snake_case_					: List[str]          =		VOCAB_FILES_NAMES
   snake_case_					: str          =		PRETRAINED_VOCAB_FILES_MAP
   snake_case_					: Optional[Any]          =		PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
   snake_case_					: List[Any]          =		["""input_ids""", """attention_mask"""]
   snake_case_					: int          =		CamembertTokenizer

   def __init__(					self							:			Optional[int]				,		lowerCAmelCase							:			Optional[Any]=None				,		lowerCAmelCase							:			int=None				,		lowerCAmelCase							:			List[Any]="<s>"				,		lowerCAmelCase							:			List[str]="</s>"				,		lowerCAmelCase							:			Union[str, Any]="</s>"				,		lowerCAmelCase							:			Dict="<s>"				,		lowerCAmelCase							:			Dict="<unk>"				,		lowerCAmelCase							:			Any="<pad>"				,		lowerCAmelCase							:			List[str]="<mask>"				,		lowerCAmelCase							:			List[str]=["<s>NOTUSED", "</s>NOTUSED"]				,		**lowerCAmelCase							:			Dict				,		)				->				Tuple:

    """simple docstring"""







    _snake_case    :  Dict					=					AddedToken(lowerCAmelCase				,		lstrip=lowerCAmelCase				,		rstrip=lowerCAmelCase) if isinstance(lowerCAmelCase				,		lowerCAmelCase) else mask_token

    super().__init__(
        lowerCAmelCase				,		tokenizer_file=lowerCAmelCase				,		bos_token=lowerCAmelCase				,		eos_token=lowerCAmelCase				,		sep_token=lowerCAmelCase				,		cls_token=lowerCAmelCase				,		unk_token=lowerCAmelCase				,		pad_token=lowerCAmelCase				,		mask_token=lowerCAmelCase				,		additional_special_tokens=lowerCAmelCase				,		**lowerCAmelCase				,		)

    _snake_case    :  List[str]					=					vocab_file
    _snake_case    :  Optional[Any]					=					False if not self.vocab_file else True

   def      UpperCamelCase_  (					self							:			Optional[int]				,		lowerCAmelCase							:			List[int]				,		lowerCAmelCase							:			Optional[List[int]] = None)				->				List[int]:

    """simple docstring"""







    if token_ids_a is None:
     return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
    _snake_case    :  Optional[Any]					=					[self.cls_token_id]
    _snake_case    :  List[Any]					=					[self.sep_token_id]
    return cls + token_ids_a + sep + sep + token_ids_a + sep

   def      UpperCamelCase_  (					self							:			str				,		lowerCAmelCase							:			List[int]				,		lowerCAmelCase							:			Optional[List[int]] = None)				->				List[int]:

    """simple docstring"""







    _snake_case    :  str					=					[self.sep_token_id]
    _snake_case    :  List[str]					=					[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 + sep + token_ids_a + sep) * [0]


   def      UpperCamelCase_  (					self							:			str				,		lowerCAmelCase							:			str				,		lowerCAmelCase							:			Optional[str] = None)				->				Tuple[str]:

    """simple docstring"""







    if not self.can_save_slow_tokenizer:
     raise ValueError(
         """Your fast tokenizer does not have the necessary information to save the vocabulary for a slow """
         """tokenizer.""")

    if not os.path.isdir(lowerCAmelCase):
     logger.error(F'''Vocabulary path ({save_directory}) should be a directory''')
     return
    _snake_case    :  Optional[int]					=					os.path.join(
        lowerCAmelCase				,		(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""])

    if os.path.abspath(self.vocab_file) != os.path.abspath(lowerCAmelCase):
     copyfile(self.vocab_file				,		lowerCAmelCase)

    return (out_vocab_file,)
