






from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


__A					=	logging.get_logger(__name__)

__A					=	{
    "microsoft/focalnet-tiny": "https://huggingface.co/microsoft/focalnet-tiny/resolve/main/config.json",
}



class 				_A			(     a__					,					a__				):
    """simple docstring"""

    lowerCamelCase	:							str				= "focalnet"



    def __init__(				self      :     Dict ,     __SCREAMING_SNAKE_CASE      :     Union[str, Any]=224 ,     __SCREAMING_SNAKE_CASE      :     str=4 ,     __SCREAMING_SNAKE_CASE      :     Any=3 ,     __SCREAMING_SNAKE_CASE      :     Dict=96 ,     __SCREAMING_SNAKE_CASE      :     Dict=False ,     __SCREAMING_SNAKE_CASE      :     Any=[192, 384, 768, 768] ,     __SCREAMING_SNAKE_CASE      :     Tuple=[2, 2, 6, 2] ,     __SCREAMING_SNAKE_CASE      :     Any=[2, 2, 2, 2] ,     __SCREAMING_SNAKE_CASE      :     Union[str, Any]=[3, 3, 3, 3] ,     __SCREAMING_SNAKE_CASE      :     Union[str, Any]="gelu" ,     __SCREAMING_SNAKE_CASE      :     Optional[Any]=4.0 ,     __SCREAMING_SNAKE_CASE      :     List[Any]=0.0 ,     __SCREAMING_SNAKE_CASE      :     List[Any]=0.1 ,     __SCREAMING_SNAKE_CASE      :     int=False ,     __SCREAMING_SNAKE_CASE      :     Dict=1e-4 ,     __SCREAMING_SNAKE_CASE      :     Any=False ,     __SCREAMING_SNAKE_CASE      :     Optional[Any]=False ,     __SCREAMING_SNAKE_CASE      :     Any=False ,     __SCREAMING_SNAKE_CASE      :     Dict=0.02 ,     __SCREAMING_SNAKE_CASE      :     Optional[Any]=1e-5 ,     __SCREAMING_SNAKE_CASE      :     Union[str, Any]=32 ,     __SCREAMING_SNAKE_CASE      :     Optional[Any]=None ,     __SCREAMING_SNAKE_CASE      :     Tuple=None ,     **__SCREAMING_SNAKE_CASE      :     List[str] ,     )      ->							str:
           super().__init__(**lowerCamelCase_      )

           __UpperCAmelCase     				=image_size
           __UpperCAmelCase     				=patch_size
           __UpperCAmelCase     				=num_channels
           __UpperCAmelCase     				=embed_dim
           __UpperCAmelCase     				=use_conv_embed
           __UpperCAmelCase     				=hidden_sizes
           __UpperCAmelCase     				=depths
           __UpperCAmelCase     				=focal_levels
           __UpperCAmelCase     				=focal_windows
           __UpperCAmelCase     				=hidden_act
           __UpperCAmelCase     				=mlp_ratio
           __UpperCAmelCase     				=hidden_dropout_prob
           __UpperCAmelCase     				=drop_path_rate
           __UpperCAmelCase     				=use_layerscale
           __UpperCAmelCase     				=layerscale_value
           __UpperCAmelCase     				=use_post_layernorm
           __UpperCAmelCase     				=use_post_layernorm_in_modulation
           __UpperCAmelCase     				=normalize_modulator
           __UpperCAmelCase     				=initializer_range
           __UpperCAmelCase     				=layer_norm_eps
           __UpperCAmelCase     				=encoder_stride
           __UpperCAmelCase     				=["""stem"""] + [f'''stage{idx}''' for idx in range(1 ,     len(self.depths      ) + 1      )]
           __UpperCAmelCase				, __UpperCAmelCase     				=get_aligned_output_features_output_indices(
               out_features=lowerCamelCase_ ,     out_indices=lowerCamelCase_ ,     stage_names=self.stage_names      )








'''simple docstring'''






from dataclasses import dataclass
from typing import 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 .modeling_utils import ModelMixin
from .vae import Decoder, DecoderOutput, Encoder, VectorQuantizer




@dataclass
class 						a__							(				a__    ):





							'''simple docstring'''






							lowercase__     :  torch.FloatTensor




class 						a__							(				a__							,							a__    ):





							'''simple docstring'''






							@register_to_config
							def __init__(						self   ,       lowerCamelCase_ = 3   ,       lowerCamelCase_ = 3   ,       lowerCamelCase_ = ("DownEncoderBlock2D",)   ,       lowerCamelCase_ = ("UpDecoderBlock2D",)   ,       lowerCamelCase_ = (64,)   ,       lowerCamelCase_ = 1   ,       lowerCamelCase_ = "silu"   ,       lowerCamelCase_ = 3   ,       lowerCamelCase_ = 32   ,       lowerCamelCase_ = 2_56   ,       lowerCamelCase_ = 32   ,       lowerCamelCase_ = None   ,       lowerCamelCase_ = 0.18_215   ,       lowerCamelCase_ = "group"   ,       )			->				Union[str, Any]:
														super().__init__()

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

														lowerCAmelCase__			  =      vq_embed_dim if vq_embed_dim is not None else latent_channels

														lowerCAmelCase__			  =      nn.Convad(lowerCamelCase_   ,       lowerCamelCase_   ,       1  )
														lowerCAmelCase__			  =      VectorQuantizer(lowerCamelCase_   ,       lowerCamelCase_   ,       beta=0.25   ,       remap=lowerCamelCase_   ,       sane_index_shape=lowerCamelCase_  )
														lowerCAmelCase__			  =      nn.Convad(lowerCamelCase_   ,       lowerCamelCase_   ,       1  )

														# pass init params to Decoder
														lowerCAmelCase__			  =      Decoder(
														    in_channels=lowerCamelCase_   ,       out_channels=lowerCamelCase_   ,       up_block_types=lowerCamelCase_   ,       block_out_channels=lowerCamelCase_   ,       layers_per_block=lowerCamelCase_   ,       act_fn=lowerCamelCase_   ,       norm_num_groups=lowerCamelCase_   ,       norm_type=lowerCamelCase_   ,       )







							@apply_forward_hook
							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_ = True  )			->				VQEncoderOutput:
														lowerCAmelCase__			  =      self.encoder(lowerCamelCase_  )
														lowerCAmelCase__			  =      self.quant_conv(lowerCamelCase_  )

														if not return_dict:
																					return (h,)

														return VQEncoderOutput(latents=lowerCamelCase_  )







							@apply_forward_hook
							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_ = False   ,       lowerCamelCase_ = True  )			->				Union[DecoderOutput, torch.FloatTensor]:
														# also go through quantization layer
														if not force_not_quantize:
																					lowerCAmelCase__  ,   lowerCAmelCase__  ,   lowerCAmelCase__			  =      self.quantize(lowerCamelCase_  )
														else:
																					lowerCAmelCase__			  =      h
														lowerCAmelCase__			  =      self.post_quant_conv(lowerCamelCase_  )
														lowerCAmelCase__			  =      self.decoder(lowerCamelCase_   ,       quant if self.config.norm_type == '''spatial''' else None  )

														if not return_dict:
																					return (dec,)

														return DecoderOutput(sample=lowerCamelCase_  )







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_ = True  )			->				Union[DecoderOutput, torch.FloatTensor]:
														lowerCAmelCase__			  =      sample
														lowerCAmelCase__			  =      self.encode(lowerCamelCase_  ).latents
														lowerCAmelCase__			  =      self.decode(lowerCamelCase_  ).sample

														if not return_dict:
																					return (dec,)

														return DecoderOutput(sample=lowerCamelCase_  )