


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     UpperCAmelCase__       (					A_		):



    """simple docstring"""






    UpperCAmelCase__						:		"DiagonalGaussianDistribution"






class     UpperCAmelCase__       (					A_   ,  A_		):



    """simple docstring"""






    UpperCAmelCase__						:		List[Any]      =  True






    @register_to_config
    def __init__(			self					,   A_ = 3					,   A_ = 3					,   A_ = ("DownEncoderBlock2D",)					,   A_ = ("UpDecoderBlock2D",)					,   A_ = (64,)					,   A_ = 1					,   A_ = "silu"					,   A_ = 4					,   A_ = 32					,   A_ = 32					,   A_ = 0.1_8215					,   )  ->       Any:
         super().__init__()

         # pass init params to Encoder
         __UpperCamelCase    							=Encoder(
             in_channels=A_					,   out_channels=A_					,   down_block_types=A_					,   block_out_channels=A_					,   layers_per_block=A_					,   act_fn=A_					,   norm_num_groups=A_					,   double_z=A_					,   )

         # pass init params to Decoder
         __UpperCamelCase    							=Decoder(
             in_channels=A_					,   out_channels=A_					,   up_block_types=A_					,   block_out_channels=A_					,   layers_per_block=A_					,   norm_num_groups=A_					,   act_fn=A_					,   )

         __UpperCamelCase    							=nn.Convad(2 * latent_channels					,   2 * latent_channels					,   1    )
         __UpperCamelCase    							=nn.Convad(A_					,   A_					,   1    )

         __UpperCamelCase    							=False
         __UpperCamelCase    							=False

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






    def _a					(			self					,   A_					,   A_=False    )  ->       Any:
         if isinstance(A_					,   (Encoder, Decoder)    ):
              __UpperCamelCase    							=value






    def _a					(			self					,   A_ = True    )  ->       List[Any]:
         __UpperCamelCase    							=use_tiling






    def _a					(			self    )  ->       Union[str, Any]:
         self.enable_tiling(A_    )






    def _a					(			self    )  ->       List[Any]:
         __UpperCamelCase    							=True






    def _a					(			self    )  ->       Any:
         __UpperCamelCase    							=False






    @property
    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
    def _a					(			self    )  ->       Dict[str, AttentionProcessor]:
         __UpperCamelCase    							={}

         def fn_recursive_add_processors(A_					,   A_					,   A_    ):
              if hasattr(A_					,   'set_processor'    ):
                   __UpperCamelCase    							=module.processor

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

              return processors

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

         return processors






    def _a					(			self					,   A_    )  ->       Tuple:
         __UpperCamelCase    							=len(self.attn_processors.keys()    )

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

         def fn_recursive_attn_processor(A_					,   A_					,   A_    ):
              if hasattr(A_					,   'set_processor'    ):
                   if not isinstance(A_					,   A_    ):
                        module.set_processor(A_    )
                   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}'					,   A_					,   A_    )

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






    def _a					(			self    )  ->       List[Any]:
         self.set_attn_processor(AttnProcessor()    )






    @apply_forward_hook
    def _a					(			self					,   A_					,   A_ = True    )  ->       AutoencoderKLOutput:
         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(A_					,   return_dict=A_    )

         if self.use_slicing and x.shape[0] > 1:
              __UpperCamelCase    							=[self.encoder(A_    ) for x_slice in x.split(1    )]
              __UpperCamelCase    							=torch.cat(A_    )
         else:
              __UpperCamelCase    							=self.encoder(A_    )

         __UpperCamelCase    							=self.quant_conv(A_    )
         __UpperCamelCase    							=DiagonalGaussianDistribution(A_    )

         if not return_dict:
              return (posterior,)

         return AutoencoderKLOutput(latent_dist=A_    )






    def _a					(			self					,   A_					,   A_ = True    )  ->       Union[DecoderOutput, torch.FloatTensor]:
         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(A_					,   return_dict=A_    )

         __UpperCamelCase    							=self.post_quant_conv(A_    )
         __UpperCamelCase    							=self.decoder(A_    )

         if not return_dict:
              return (dec,)

         return DecoderOutput(sample=A_    )






    @apply_forward_hook
    def _a					(			self					,   A_					,   A_ = True    )  ->       Union[DecoderOutput, torch.FloatTensor]:
         if self.use_slicing and z.shape[0] > 1:
              __UpperCamelCase    							=[self._decode(A_    ).sample for z_slice in z.split(1    )]
              __UpperCamelCase    							=torch.cat(A_    )
         else:
              __UpperCamelCase    							=self._decode(A_    ).sample

         if not return_dict:
              return (decoded,)

         return DecoderOutput(sample=A_    )






    def _a					(			self					,   A_					,   A_					,   A_    )  ->       Optional[Any]:
         __UpperCamelCase    							=min(a.shape[2]					,   b.shape[2]					,   A_    )
         for y in range(A_    ):
              __UpperCamelCase    							=a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
         return b






    def _a					(			self					,   A_					,   A_					,   A_    )  ->       Dict:
         __UpperCamelCase    							=min(a.shape[3]					,   b.shape[3]					,   A_    )
         for x in range(A_    ):
              __UpperCamelCase    							=a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
         return b






    def _a					(			self					,   A_					,   A_ = True    )  ->       AutoencoderKLOutput:
         __UpperCamelCase    							=int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)    )
         __UpperCamelCase    							=int(self.tile_latent_min_size * self.tile_overlap_factor    )
         __UpperCamelCase    							=self.tile_latent_min_size - blend_extent

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

         __UpperCamelCase    							=torch.cat(A_					,   dim=2    )
         __UpperCamelCase    							=DiagonalGaussianDistribution(A_    )

         if not return_dict:
              return (posterior,)

         return AutoencoderKLOutput(latent_dist=A_    )






    def _a					(			self					,   A_					,   A_ = True    )  ->       Union[DecoderOutput, torch.FloatTensor]:
         __UpperCamelCase    							=int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)    )
         __UpperCamelCase    							=int(self.tile_sample_min_size * self.tile_overlap_factor    )
         __UpperCamelCase    							=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.
         __UpperCamelCase    							=[]
         for i in range(0					,   z.shape[2]					,   A_    ):
              __UpperCamelCase    							=[]
              for j in range(0					,   z.shape[3]					,   A_    ):
                   __UpperCamelCase    							=z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
                   __UpperCamelCase    							=self.post_quant_conv(A_    )
                   __UpperCamelCase    							=self.decoder(A_    )
                   row.append(A_    )
              rows.append(A_    )
         __UpperCamelCase    							=[]
         for i, row in enumerate(A_    ):
              __UpperCamelCase    							=[]
              for j, tile in enumerate(A_    ):
                   # blend the above tile and the left tile
                   # to the current tile and add the current tile to the result row
                   if i > 0:
                        __UpperCamelCase    							=self.blend_v(rows[i - 1][j]					,   A_					,   A_    )
                   if j > 0:
                        __UpperCamelCase    							=self.blend_h(row[j - 1]					,   A_					,   A_    )
                   result_row.append(tile[:, :, :row_limit, :row_limit]    )
              result_rows.append(torch.cat(A_					,   dim=3    )    )

         __UpperCamelCase    							=torch.cat(A_					,   dim=2    )
         if not return_dict:
              return (dec,)

         return DecoderOutput(sample=A_    )






    def _a					(			self					,   A_					,   A_ = False					,   A_ = True					,   A_ = None					,   )  ->       Union[DecoderOutput, torch.FloatTensor]:
         __UpperCamelCase    							=sample
         __UpperCamelCase    							=self.encode(A_    ).latent_dist
         if sample_posterior:
              __UpperCamelCase    							=posterior.sample(generator=A_    )
         else:
              __UpperCamelCase    							=posterior.mode()
         __UpperCamelCase    							=self.decode(A_    ).sample

         if not return_dict:
              return (dec,)

         return DecoderOutput(sample=A_    )












import json
import os
import unittest

from transformers.models.gptsan_japanese.tokenization_gptsan_japanese import (
    VOCAB_FILES_NAMES,
    GPTSanJapaneseTokenizer,
)
from transformers.testing_utils import require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin







@require_tokenizers
class     __snake_case       (							__lowerCamelCase   ,			unittest.TestCase			):





       '''simple docstring'''







       lowerCAmelCase__      		=     GPTSanJapaneseTokenizer
       lowerCAmelCase__      		=     False
       lowerCAmelCase__      		=     {"""do_clean_text""": False, """add_prefix_space""": False}




       def 							UpperCAmelCase__ (  self :    Union[str, Any]		):
         super().setUp()

         # fmt: off
         __snake_case:						str   =						["""こん""", """こんに""", """にちは""", """ばんは""", """世界,㔺界""", """、""", """。""", """<BR>""", """<SP>""", """<TAB>""", """<URL>""", """<EMAIL>""", """<TEL>""", """<DATE>""", """<PRICE>""", """<BLOCK>""", """<KIGOU>""", """<U2000U2BFF>""", """<|emoji1|>""", """<unk>""", """<|bagoftoken|>""", """<|endoftext|>"""]
         # fmt: on
         __snake_case:						List[Any]   =						{"""emoji""": {"""\ud83d\ude00""": """<|emoji1|>"""}, """emoji_inv""": {"""<|emoji1|>""": """\ud83d\ude00"""}}  # 😀
         __snake_case:						Optional[int]   =						{"""unk_token""": """<unk>"""}

         __snake_case:						Any   =						os.path.join(self.tmpdirname							,				VOCAB_FILES_NAMES["""vocab_file"""]		)
         __snake_case:						int   =						os.path.join(self.tmpdirname							,				VOCAB_FILES_NAMES["""emoji_file"""]		)
         with open(self.vocab_file							,				"""w"""							,				encoding="""utf-8"""		) as vocab_writer:
           vocab_writer.write("""""".join([x + """\n""" for x in vocab_tokens]		)		)
         with open(self.emoji_file							,				"""w"""		) as emoji_writer:
           emoji_writer.write(json.dumps(A		)		)




       def 							UpperCAmelCase__ (  self :    Optional[int]							,				**A :    int		):
         kwargs.update(self.special_tokens_map		)
         return GPTSanJapaneseTokenizer.from_pretrained(self.tmpdirname							,				**A		)




       def 							UpperCAmelCase__ (  self :    Optional[Any]							,				A :    Dict		):
         __snake_case:						Tuple   =						"""こんにちは、世界。 \nこんばんは、㔺界。😀"""
         __snake_case:						str   =						"""こんにちは、世界。 \nこんばんは、世界。😀"""
         return input_text, output_text




       def 							UpperCAmelCase__ (  self :    Optional[int]							,				A :    Optional[int]		):
         __snake_case  ,   __snake_case:						Optional[int]   =						self.get_input_output_texts(A		)
         __snake_case:						Optional[Any]   =						tokenizer.encode(A							,				add_special_tokens=A		)
         __snake_case:						int   =						tokenizer.decode(A							,				clean_up_tokenization_spaces=A		)
         return text, ids




       def 							UpperCAmelCase__ (  self :    int		):
         pass  # TODO add if relevant




       def 							UpperCAmelCase__ (  self :    Dict		):
         pass  # TODO add if relevant




       def 							UpperCAmelCase__ (  self :    Optional[int]		):
         pass  # TODO add if relevant




       def 							UpperCAmelCase__ (  self :    Union[str, Any]		):
         __snake_case:						Dict   =						self.get_tokenizer()

         # Testing tokenization
         __snake_case:						List[str]   =						"""こんにちは、世界。　こんばんは、㔺界。"""
         __snake_case:						Optional[int]   =						["""こん""", """にちは""", """、""", """世界""", """。""", """<SP>""", """こん""", """ばんは""", """、""", """㔺界""", """。"""]
         __snake_case:						Any   =						tokenizer.tokenize(A		)
         self.assertListEqual(A							,				A		)

         # Testing conversion to ids without special tokens
         __snake_case:						str   =						[0, 2, 5, 4, 6, 8, 0, 3, 5, 4, 6]
         __snake_case:						str   =						tokenizer.convert_tokens_to_ids(A		)
         self.assertListEqual(A							,				A		)

         # Testing conversion to ids with special tokens
         __snake_case:						Optional[int]   =						tokens + [tokenizer.unk_token]
         __snake_case:						List[str]   =						[0, 2, 5, 4, 6, 8, 0, 3, 5, 4, 6, 19]
         __snake_case:						str   =						tokenizer.convert_tokens_to_ids(A		)
         self.assertListEqual(A							,				A		)




       def 							UpperCAmelCase__ (  self :    str		):
         __snake_case:						Union[str, Any]   =						self.get_tokenizer()

         # Testing tokenization
         __snake_case:						Optional[Any]   =						"""こんにちは、<|bagoftoken|>世界。こんばんは、<|bagoftoken|>㔺界。"""
         __snake_case:						Optional[Any]   =						"""こんにちは、、、、世界。こんばんは、、、、世界。"""
         __snake_case:						Union[str, Any]   =						tokenizer.encode(A		)
         __snake_case:						List[Any]   =						tokenizer.decode(A		)
         self.assertEqual(A							,				A		)




       @slow
       def 							UpperCAmelCase__ (  self :    Any		):
         __snake_case:						Dict   =						self.tokenizer_class.from_pretrained("""Tanrei/GPTSAN-japanese"""		)

         # Testing tokenization
         __snake_case:						Any   =						"""こんにちは、世界。"""
         __snake_case:						Tuple   =						"""こんばんは、㔺界。😀"""
         __snake_case:						Optional[Any]   =						"""こんにちは、世界。こんばんは、世界。😀"""
         __snake_case:						int   =						tokenizer.encode(prefix_text + input_text		)
         __snake_case:						Union[str, Any]   =						tokenizer.encode(""""""							,				prefix_text=prefix_text + input_text		)
         __snake_case:						Tuple   =						tokenizer.encode(A							,				prefix_text=A		)
         __snake_case:						Union[str, Any]   =						tokenizer.decode(A		)
         __snake_case:						Dict   =						tokenizer.decode(A		)
         __snake_case:						Optional[Any]   =						tokenizer.decode(A		)
         self.assertEqual(A							,				A		)
         self.assertEqual(A							,				A		)
         self.assertEqual(A							,				A		)




       @slow
       def 							UpperCAmelCase__ (  self :    List[str]		):
         __snake_case:						int   =						self.tokenizer_class.from_pretrained("""Tanrei/GPTSAN-japanese"""		)

         # Testing tokenization
         __snake_case:						Optional[int]   =						"""こんにちは、世界。"""
         __snake_case:						Any   =						"""こんばんは、㔺界。😀"""

         __snake_case:						Optional[int]   =						len(tokenizer.encode(A		)		) - 2
         __snake_case:						str   =						len(tokenizer.encode(A		)		) - 2

         __snake_case:						Dict   =						[1] + [0] * (len_prefix + len_text + 1)
         __snake_case:						str   =						[1] * (len_prefix + len_text + 1) + [0]
         __snake_case:						List[Any]   =						[1] + [1] * (len_prefix) + [0] * (len_text + 1)

         __snake_case:						int   =						tokenizer(prefix_text + input_text		).token_type_ids
         __snake_case:						Optional[int]   =						tokenizer(""""""							,				prefix_text=prefix_text + input_text		).token_type_ids
         __snake_case:						Tuple   =						tokenizer(A							,				prefix_text=A		).token_type_ids
         self.assertListEqual(A							,				A		)
         self.assertListEqual(A							,				A		)
         self.assertListEqual(A							,				A		)




       @slow
       def 							UpperCAmelCase__ (  self :    Dict		):
         __snake_case:						Dict   =						self.tokenizer_class.from_pretrained("""Tanrei/GPTSAN-japanese"""		)

         __snake_case:						int   =						tokenizer.encode("""あンいワ"""		)
         __snake_case:						Optional[int]   =						tokenizer.encode(""""""							,				prefix_text="""あンいワ"""		)
         __snake_case:						List[Any]   =						tokenizer.encode("""いワ"""							,				prefix_text="""あン"""		)

         self.assertEqual(tokenizer.decode(A		)							,				tokenizer.decode(A		)		)
         self.assertEqual(tokenizer.decode(A		)							,				tokenizer.decode(A		)		)
         self.assertNotEqual(A							,				A		)
         self.assertNotEqual(A							,				A		)
         self.assertEqual(x_token_a[1]							,				x_token_a[-1]		)  # SEG token
         self.assertEqual(x_token_a[1]							,				x_token_a[3]		)  # SEG token




       @slow
       def 							UpperCAmelCase__ (  self :    int		):
         __snake_case:						List[str]   =						self.tokenizer_class.from_pretrained("""Tanrei/GPTSAN-japanese"""		)

         __snake_case:						Union[str, Any]   =						[["""武田信玄""", """は、"""], ["""織田信長""", """の配下の、"""]]
         __snake_case:						Optional[int]   =						tokenizer(A							,				padding=A		)
         __snake_case:						int   =						tokenizer.batch_encode_plus(A							,				padding=A		)

         # fmt: off
         __snake_case:						List[str]   =						[[35_993, 8_640, 25_948, 35_998, 30_647, 35_675, 35_999, 35_999], [35_993, 10_382, 9_868, 35_998, 30_646, 9_459, 30_646, 35_675]]
         __snake_case:						int   =						[[1, 1, 1, 0, 0, 0, 0, 0], [1, 1, 1, 0, 0, 0, 0, 0]]
         __snake_case:						Dict   =						[[1, 1, 1, 1, 1, 1, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1]]
         # fmt: on
         self.assertListEqual(x_token.input_ids							,				A		)
         self.assertListEqual(x_token.token_type_ids							,				A		)
         self.assertListEqual(x_token.attention_mask							,				A		)
         self.assertListEqual(x_token_a.input_ids							,				A		)
         self.assertListEqual(x_token_a.token_type_ids							,				A		)
         self.assertListEqual(x_token_a.attention_mask							,				A		)




       def 							UpperCAmelCase__ (  self :    Union[str, Any]		):
         # Intentionally convert some words to accommodate character fluctuations unique to Japanese
         pass






       def 							UpperCAmelCase__ (  self :    List[str]		):
         # tokenizer has no padding token
         pass






