





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__				(   lowerCAmelCase_     ):



     '''simple docstring'''

     __snake_case					:			"DiagonalGaussianDistribution"

class 	UpperCamelCase__				(   lowerCAmelCase_							,   lowerCAmelCase_     ):



     '''simple docstring'''

     __snake_case					:			Optional[int]				=    True






     @register_to_config
     def __init__(							self		:					str		,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.18215		,)  ->	Tuple:




           '''simple docstring'''


           super().__init__()

           # pass init params to Encoder
           SCREAMING_SNAKE_CASE	=		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
           SCREAMING_SNAKE_CASE	=		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__		,)

           SCREAMING_SNAKE_CASE	=		nn.Convad(2 * latent_channels		,2 * latent_channels		,1 )
           SCREAMING_SNAKE_CASE	=		nn.Convad(lowerCamelCase__		,lowerCamelCase__		,1 )

           SCREAMING_SNAKE_CASE	=		False
           SCREAMING_SNAKE_CASE	=		False

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






     def        SCREAMING_SNAKE_CASE__						(							self		:					Tuple		,lowerCamelCase__		:					Union[str, Any]		,lowerCamelCase__		:					Optional[int]=False )  ->	Tuple:




           '''simple docstring'''


           if isinstance(lowerCamelCase__		,(Encoder, Decoder) ):
                 SCREAMING_SNAKE_CASE	=		value






     def        SCREAMING_SNAKE_CASE__						(							self		:					int		,lowerCamelCase__		:					bool = True )  ->	Union[str, Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		use_tiling






     def        SCREAMING_SNAKE_CASE__						(							self		:					Any )  ->	Tuple:




           '''simple docstring'''


           self.enable_tiling(lowerCamelCase__ )






     def        SCREAMING_SNAKE_CASE__						(							self		:					int )  ->	Tuple:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		True






     def        SCREAMING_SNAKE_CASE__						(							self		:					Optional[int] )  ->	Optional[int]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		False






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




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		{}

           def fn_recursive_add_processors(lowerCamelCase__		:					str		,lowerCamelCase__		:					torch.nn.Module		,lowerCamelCase__		:					Dict[str, AttentionProcessor] ):
                 if hasattr(lowerCamelCase__		,"""set_processor""" ):
                       SCREAMING_SNAKE_CASE	=		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        SCREAMING_SNAKE_CASE__						(							self		:					List[str]		,lowerCamelCase__		:					Union[AttentionProcessor, Dict[str, AttentionProcessor]] )  ->	Any:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		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__		:					Union[str, Any] ):
                 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        SCREAMING_SNAKE_CASE__						(							self		:					Union[str, Any] )  ->	Optional[int]:




           '''simple docstring'''


           self.set_attn_processor(AttnProcessor() )






     @apply_forward_hook
     def        SCREAMING_SNAKE_CASE__						(							self		:					Any		,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:
                 SCREAMING_SNAKE_CASE	=		[self.encoder(lowerCamelCase__ ) for x_slice in x.split(1 )]
                 SCREAMING_SNAKE_CASE	=		torch.cat(lowerCamelCase__ )
           else:
                 SCREAMING_SNAKE_CASE	=		self.encoder(lowerCamelCase__ )

           SCREAMING_SNAKE_CASE	=		self.quant_conv(lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		DiagonalGaussianDistribution(lowerCamelCase__ )

           if not return_dict:
                 return (posterior,)

           return AutoencoderKLOutput(latent_dist=lowerCamelCase__ )






     def        SCREAMING_SNAKE_CASE__						(							self		:					List[str]		,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__ )

           SCREAMING_SNAKE_CASE	=		self.post_quant_conv(lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		self.decoder(lowerCamelCase__ )

           if not return_dict:
                 return (dec,)

           return DecoderOutput(sample=lowerCamelCase__ )






     @apply_forward_hook
     def        SCREAMING_SNAKE_CASE__						(							self		:					Any		,lowerCamelCase__		:					torch.FloatTensor		,lowerCamelCase__		:					bool = True )  ->	Union[DecoderOutput, torch.FloatTensor]:




           '''simple docstring'''


           if self.use_slicing and z.shape[0] > 1:
                 SCREAMING_SNAKE_CASE	=		[self._decode(lowerCamelCase__ ).sample for z_slice in z.split(1 )]
                 SCREAMING_SNAKE_CASE	=		torch.cat(lowerCamelCase__ )
           else:
                 SCREAMING_SNAKE_CASE	=		self._decode(lowerCamelCase__ ).sample

           if not return_dict:
                 return (decoded,)

           return DecoderOutput(sample=lowerCamelCase__ )






     def        SCREAMING_SNAKE_CASE__						(							self		:					List[str]		,lowerCamelCase__		:					Optional[int]		,lowerCamelCase__		:					List[str]		,lowerCamelCase__		:					Dict )  ->	Optional[int]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		min(a.shape[2]		,b.shape[2]		,lowerCamelCase__ )
           for y in range(lowerCamelCase__ ):
                 SCREAMING_SNAKE_CASE	=		a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
           return b






     def        SCREAMING_SNAKE_CASE__						(							self		:					Optional[int]		,lowerCamelCase__		:					Optional[Any]		,lowerCamelCase__		:					str		,lowerCamelCase__		:					int )  ->	Tuple:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		min(a.shape[3]		,b.shape[3]		,lowerCamelCase__ )
           for x in range(lowerCamelCase__ ):
                 SCREAMING_SNAKE_CASE	=		a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
           return b






     def        SCREAMING_SNAKE_CASE__						(							self		:					Optional[Any]		,lowerCamelCase__		:					torch.FloatTensor		,lowerCamelCase__		:					bool = True )  ->	AutoencoderKLOutput:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		int(self.tile_sample_min_size * (1 - self.tile_overlap_factor) )
           SCREAMING_SNAKE_CASE	=		int(self.tile_latent_min_size * self.tile_overlap_factor )
           SCREAMING_SNAKE_CASE	=		self.tile_latent_min_size - blend_extent

           # Split the image into 512x512 tiles and encode them separately.
           SCREAMING_SNAKE_CASE	=		[]
           for i in range(0		,x.shape[2]		,lowerCamelCase__ ):
                 SCREAMING_SNAKE_CASE	=		[]
                 for j in range(0		,x.shape[3]		,lowerCamelCase__ ):
                       SCREAMING_SNAKE_CASE	=		x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
                       SCREAMING_SNAKE_CASE	=		self.encoder(lowerCamelCase__ )
                       SCREAMING_SNAKE_CASE	=		self.quant_conv(lowerCamelCase__ )
                       row.append(lowerCamelCase__ )
                 rows.append(lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		[]
           for i, row in enumerate(lowerCamelCase__ ):
                 SCREAMING_SNAKE_CASE	=		[]
                 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:
                             SCREAMING_SNAKE_CASE	=		self.blend_v(rows[i - 1][j]		,lowerCamelCase__		,lowerCamelCase__ )
                       if j > 0:
                             SCREAMING_SNAKE_CASE	=		self.blend_h(row[j - 1]		,lowerCamelCase__		,lowerCamelCase__ )
                       result_row.append(tile[:, :, :row_limit, :row_limit] )
                 result_rows.append(torch.cat(lowerCamelCase__		,dim=3 ) )

           SCREAMING_SNAKE_CASE	=		torch.cat(lowerCamelCase__		,dim=2 )
           SCREAMING_SNAKE_CASE	=		DiagonalGaussianDistribution(lowerCamelCase__ )

           if not return_dict:
                 return (posterior,)

           return AutoencoderKLOutput(latent_dist=lowerCamelCase__ )






     def        SCREAMING_SNAKE_CASE__						(							self		:					Tuple		,lowerCamelCase__		:					torch.FloatTensor		,lowerCamelCase__		:					bool = True )  ->	Union[DecoderOutput, torch.FloatTensor]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		int(self.tile_latent_min_size * (1 - self.tile_overlap_factor) )
           SCREAMING_SNAKE_CASE	=		int(self.tile_sample_min_size * self.tile_overlap_factor )
           SCREAMING_SNAKE_CASE	=		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.
           SCREAMING_SNAKE_CASE	=		[]
           for i in range(0		,z.shape[2]		,lowerCamelCase__ ):
                 SCREAMING_SNAKE_CASE	=		[]
                 for j in range(0		,z.shape[3]		,lowerCamelCase__ ):
                       SCREAMING_SNAKE_CASE	=		z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
                       SCREAMING_SNAKE_CASE	=		self.post_quant_conv(lowerCamelCase__ )
                       SCREAMING_SNAKE_CASE	=		self.decoder(lowerCamelCase__ )
                       row.append(lowerCamelCase__ )
                 rows.append(lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		[]
           for i, row in enumerate(lowerCamelCase__ ):
                 SCREAMING_SNAKE_CASE	=		[]
                 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:
                             SCREAMING_SNAKE_CASE	=		self.blend_v(rows[i - 1][j]		,lowerCamelCase__		,lowerCamelCase__ )
                       if j > 0:
                             SCREAMING_SNAKE_CASE	=		self.blend_h(row[j - 1]		,lowerCamelCase__		,lowerCamelCase__ )
                       result_row.append(tile[:, :, :row_limit, :row_limit] )
                 result_rows.append(torch.cat(lowerCamelCase__		,dim=3 ) )

           SCREAMING_SNAKE_CASE	=		torch.cat(lowerCamelCase__		,dim=2 )
           if not return_dict:
                 return (dec,)

           return DecoderOutput(sample=lowerCamelCase__ )






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




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		sample
           SCREAMING_SNAKE_CASE	=		self.encode(lowerCamelCase__ ).latent_dist
           if sample_posterior:
                 SCREAMING_SNAKE_CASE	=		posterior.sample(generator=lowerCamelCase__ )
           else:
                 SCREAMING_SNAKE_CASE	=		posterior.mode()
           SCREAMING_SNAKE_CASE	=		self.decode(lowerCamelCase__ ).sample

           if not return_dict:
                 return (dec,)

           return DecoderOutput(sample=lowerCamelCase__ )









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


SCREAMING_SNAKE_CASE_		     = logging.get_logger(__name__)

SCREAMING_SNAKE_CASE_		     = {
    """google/pegasus-large""": """https://huggingface.co/google/pegasus-large/resolve/main/config.json""",
    # See all PEGASUS models at https://huggingface.co/models?filter=pegasus
}

class 	UpperCamelCase__				(   lowerCAmelCase_     ):



     '''simple docstring'''

     __snake_case					:			str				=    "pegasus"
     __snake_case					:			List[Any]				=    ["past_key_values"]
     __snake_case					:			Any				=    {"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}






     def __init__(							self		:					Dict		,lowerCamelCase__		:					List[str]=50265		,lowerCamelCase__		:					int=1024		,lowerCamelCase__		:					Optional[Any]=12		,lowerCamelCase__		:					Tuple=4096		,lowerCamelCase__		:					Dict=16		,lowerCamelCase__		:					Optional[Any]=12		,lowerCamelCase__		:					str=4096		,lowerCamelCase__		:					Dict=16		,lowerCamelCase__		:					Tuple=0.0		,lowerCamelCase__		:					Dict=0.0		,lowerCamelCase__		:					str=True		,lowerCamelCase__		:					Any=True		,lowerCamelCase__		:					List[Any]="gelu"		,lowerCamelCase__		:					Tuple=1024		,lowerCamelCase__		:					Optional[int]=0.1		,lowerCamelCase__		:					Tuple=0.0		,lowerCamelCase__		:					Dict=0.0		,lowerCamelCase__		:					Optional[int]=0.02		,lowerCamelCase__		:					Optional[int]=0		,lowerCamelCase__		:					List[Any]=False		,lowerCamelCase__		:					Union[str, Any]=0		,lowerCamelCase__		:					Optional[int]=1		,lowerCamelCase__		:					Any=1		,**lowerCamelCase__		:					Dict		,)  ->	str:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		vocab_size
           SCREAMING_SNAKE_CASE	=		max_position_embeddings
           SCREAMING_SNAKE_CASE	=		d_model
           SCREAMING_SNAKE_CASE	=		encoder_ffn_dim
           SCREAMING_SNAKE_CASE	=		encoder_layers
           SCREAMING_SNAKE_CASE	=		encoder_attention_heads
           SCREAMING_SNAKE_CASE	=		decoder_ffn_dim
           SCREAMING_SNAKE_CASE	=		decoder_layers
           SCREAMING_SNAKE_CASE	=		decoder_attention_heads
           SCREAMING_SNAKE_CASE	=		dropout
           SCREAMING_SNAKE_CASE	=		attention_dropout
           SCREAMING_SNAKE_CASE	=		activation_dropout
           SCREAMING_SNAKE_CASE	=		activation_function
           SCREAMING_SNAKE_CASE	=		init_std
           SCREAMING_SNAKE_CASE	=		encoder_layerdrop
           SCREAMING_SNAKE_CASE	=		decoder_layerdrop
           SCREAMING_SNAKE_CASE	=		use_cache
           SCREAMING_SNAKE_CASE	=		encoder_layers
           SCREAMING_SNAKE_CASE	=		scale_embedding  # scale factor will be sqrt(d_model) if True
           super().__init__(
               pad_token_id=lowerCamelCase__		,eos_token_id=lowerCamelCase__		,is_encoder_decoder=lowerCamelCase__		,decoder_start_token_id=lowerCamelCase__		,forced_eos_token_id=lowerCamelCase__		,**lowerCamelCase__		,)






     @property
     def        SCREAMING_SNAKE_CASE__						(							self		:					Tuple )  ->	int:




           '''simple docstring'''


           return self.encoder_attention_heads






     @property
     def        SCREAMING_SNAKE_CASE__						(							self		:					List[str] )  ->	int:




           '''simple docstring'''


           return self.d_model


