




import argparse
import json

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import SegformerImageProcessor, SwinConfig, UperNetConfig, UperNetForSemanticSegmentation







def lowerCamelCase__       (      _lowercase				):

     '''simple docstring'''


     UpperCAmelCase_      :							str       =      384
     UpperCAmelCase_      :							Optional[int]       =      7
     if "tiny" in model_name:
          UpperCAmelCase_      :							Dict       =      96
          UpperCAmelCase_      :							Any       =      (2, 2, 6, 2)
          UpperCAmelCase_      :							str       =      (3, 6, 12, 24)
     elif "small" in model_name:
          UpperCAmelCase_      :							Optional[int]       =      96
          UpperCAmelCase_      :							Union[str, Any]       =      (2, 2, 18, 2)
          UpperCAmelCase_      :							Any       =      (3, 6, 12, 24)
     elif "base" in model_name:
          UpperCAmelCase_      :							List[Any]       =      128
          UpperCAmelCase_      :							List[Any]       =      (2, 2, 18, 2)
          UpperCAmelCase_      :							Optional[Any]       =      (4, 8, 16, 32)
          UpperCAmelCase_      :							Optional[Any]       =      12
          UpperCAmelCase_      :							int       =      512
     elif "large" in model_name:
          UpperCAmelCase_      :							List[Any]       =      192
          UpperCAmelCase_      :							int       =      (2, 2, 18, 2)
          UpperCAmelCase_      :							List[str]       =      (6, 12, 24, 48)
          UpperCAmelCase_      :							Tuple       =      12
          UpperCAmelCase_      :							int       =      768

     # set label information
     UpperCAmelCase_      :							Dict       =      150
     UpperCAmelCase_      :							Dict       =      '''huggingface/label-files'''
     UpperCAmelCase_      :							str       =      '''ade20k-id2label.json'''
     UpperCAmelCase_      :							List[str]       =      json.load(open(hf_hub_download(_lowercase			,						_lowercase			,						repo_type='''dataset'''				)			,						'''r'''				)				)
     UpperCAmelCase_      :							Any       =      {int(_lowercase				): v for k, v in idalabel.items()}
     UpperCAmelCase_      :							List[str]       =      {v: k for k, v in idalabel.items()}

     UpperCAmelCase_      :							List[Any]       =      SwinConfig(
         embed_dim=_lowercase			,						depths=_lowercase			,						num_heads=_lowercase			,						window_size=_lowercase			,						out_features=['''stage1''', '''stage2''', '''stage3''', '''stage4''']			,						)
     UpperCAmelCase_      :							str       =      UperNetConfig(
         backbone_config=_lowercase			,						auxiliary_in_channels=_lowercase			,						num_labels=_lowercase			,						idalabel=_lowercase			,						labelaid=_lowercase			,						)

     return config







def lowerCamelCase__       (      _lowercase				):

     '''simple docstring'''


     UpperCAmelCase_      :							Optional[Any]       =      []

     # fmt: off
     # stem
     rename_keys.append(('''backbone.patch_embed.projection.weight''', '''backbone.embeddings.patch_embeddings.projection.weight''')				)
     rename_keys.append(('''backbone.patch_embed.projection.bias''', '''backbone.embeddings.patch_embeddings.projection.bias''')				)
     rename_keys.append(('''backbone.patch_embed.norm.weight''', '''backbone.embeddings.norm.weight''')				)
     rename_keys.append(('''backbone.patch_embed.norm.bias''', '''backbone.embeddings.norm.bias''')				)
     # stages
     for i in range(len(config.backbone_config.depths				)				):
          for j in range(config.backbone_config.depths[i]				):
               rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.norm1.weight''', f'''backbone.encoder.layers.{i}.blocks.{j}.layernorm_before.weight''')				)
               rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.norm1.bias''', f'''backbone.encoder.layers.{i}.blocks.{j}.layernorm_before.bias''')				)
               rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.attn.w_msa.relative_position_bias_table''', f'''backbone.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_bias_table''')				)
               rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.attn.w_msa.relative_position_index''', f'''backbone.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_index''')				)
               rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.attn.w_msa.proj.weight''', f'''backbone.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight''')				)
               rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.attn.w_msa.proj.bias''', f'''backbone.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias''')				)
               rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.norm2.weight''', f'''backbone.encoder.layers.{i}.blocks.{j}.layernorm_after.weight''')				)
               rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.norm2.bias''', f'''backbone.encoder.layers.{i}.blocks.{j}.layernorm_after.bias''')				)
               rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.ffn.layers.0.0.weight''', f'''backbone.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight''')				)
               rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.ffn.layers.0.0.bias''', f'''backbone.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias''')				)
               rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.ffn.layers.1.weight''', f'''backbone.encoder.layers.{i}.blocks.{j}.output.dense.weight''')				)
               rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.ffn.layers.1.bias''', f'''backbone.encoder.layers.{i}.blocks.{j}.output.dense.bias''')				)

          if i < 3:
               rename_keys.append((f'''backbone.stages.{i}.downsample.reduction.weight''', f'''backbone.encoder.layers.{i}.downsample.reduction.weight''')				)
               rename_keys.append((f'''backbone.stages.{i}.downsample.norm.weight''', f'''backbone.encoder.layers.{i}.downsample.norm.weight''')				)
               rename_keys.append((f'''backbone.stages.{i}.downsample.norm.bias''', f'''backbone.encoder.layers.{i}.downsample.norm.bias''')				)
          rename_keys.append((f'''backbone.norm{i}.weight''', f'''backbone.hidden_states_norms.stage{i+1}.weight''')				)
          rename_keys.append((f'''backbone.norm{i}.bias''', f'''backbone.hidden_states_norms.stage{i+1}.bias''')				)

     # decode head
     rename_keys.extend(
         [
             ('''decode_head.conv_seg.weight''', '''decode_head.classifier.weight'''),
             ('''decode_head.conv_seg.bias''', '''decode_head.classifier.bias'''),
             ('''auxiliary_head.conv_seg.weight''', '''auxiliary_head.classifier.weight'''),
             ('''auxiliary_head.conv_seg.bias''', '''auxiliary_head.classifier.bias'''),
         ]				)
     # fmt: on

     return rename_keys







def lowerCamelCase__       (      _lowercase			,						_lowercase			,						_lowercase				):

     '''simple docstring'''


     UpperCAmelCase_      :							Any       =      dct.pop(_lowercase				)
     UpperCAmelCase_      :							Union[str, Any]       =      val







def lowerCamelCase__       (      _lowercase			,						_lowercase				):

     '''simple docstring'''


     UpperCAmelCase_      :							Optional[int]       =      [int(backbone_config.embed_dim * 2**i				) for i in range(len(backbone_config.depths				)				)]
     for i in range(len(backbone_config.depths				)				):
          UpperCAmelCase_      :							Optional[int]       =      num_features[i]
          for j in range(backbone_config.depths[i]				):
               # fmt: off
               # read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias)
               UpperCAmelCase_      :							List[str]       =      state_dict.pop(f'''backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.weight'''				)
               UpperCAmelCase_      :							Any       =      state_dict.pop(f'''backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.bias'''				)
               # next, add query, keys and values (in that order) to the state dict
               UpperCAmelCase_      :							int       =      in_proj_weight[:dim, :]
               UpperCAmelCase_      :							Optional[int]       =      in_proj_bias[: dim]
               UpperCAmelCase_      :							int       =      in_proj_weight[
                   dim : dim * 2, :
               ]
               UpperCAmelCase_      :							str       =      in_proj_bias[
                   dim : dim * 2
               ]
               UpperCAmelCase_      :							Any       =      in_proj_weight[
                   -dim :, :
               ]
               UpperCAmelCase_      :							int       =      in_proj_bias[-dim :]
               # fmt: on







def lowerCamelCase__       (      _lowercase				):

     '''simple docstring'''


     UpperCAmelCase_      :							str       =      x.shape
     UpperCAmelCase_      :							Dict       =      x.reshape(_lowercase			,						4			,						in_channel // 4				)
     UpperCAmelCase_      :							Dict       =      x[:, [0, 2, 1, 3], :].transpose(1			,						2				).reshape(_lowercase			,						_lowercase				)
     return x







def lowerCamelCase__       (      _lowercase				):

     '''simple docstring'''


     UpperCAmelCase_      :							int       =      x.shape
     UpperCAmelCase_      :							Optional[int]       =      x.reshape(_lowercase			,						in_channel // 4			,						4				)
     UpperCAmelCase_      :							Dict       =      x[:, :, [0, 2, 1, 3]].transpose(1			,						2				).reshape(_lowercase			,						_lowercase				)

     return x







def lowerCamelCase__       (      _lowercase				):

     '''simple docstring'''


     UpperCAmelCase_      :							str       =      x.shape[0]
     UpperCAmelCase_      :							Optional[int]       =      x.reshape(4			,						in_channel // 4				)
     UpperCAmelCase_      :							Dict       =      x[[0, 2, 1, 3], :].transpose(0			,						1				).reshape(_lowercase				)
     return x







def lowerCamelCase__       (      _lowercase				):

     '''simple docstring'''


     UpperCAmelCase_      :							Tuple       =      x.shape[0]
     UpperCAmelCase_      :							int       =      x.reshape(in_channel // 4			,						4				)
     UpperCAmelCase_      :							Any       =      x[:, [0, 2, 1, 3]].transpose(0			,						1				).reshape(_lowercase				)
     return x







def lowerCamelCase__       (      _lowercase			,						_lowercase			,						_lowercase				):

     '''simple docstring'''


     UpperCAmelCase_      :							int       =      {
         '''upernet-swin-tiny''': '''https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K_20210531_112542-e380ad3e.pth''',
         '''upernet-swin-small''': '''https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_small_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K/upernet_swin_small_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K_20210526_192015-ee2fff1c.pth''',
         '''upernet-swin-base''': '''https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_base_patch4_window12_512x512_160k_ade20k_pretrain_384x384_22K/upernet_swin_base_patch4_window12_512x512_160k_ade20k_pretrain_384x384_22K_20210531_125459-429057bf.pth''',
         '''upernet-swin-large''': '''https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_large_patch4_window12_512x512_pretrain_384x384_22K_160k_ade20k/upernet_swin_large_patch4_window12_512x512_pretrain_384x384_22K_160k_ade20k_20220318_091743-9ba68901.pth''',
     }
     UpperCAmelCase_      :							List[Any]       =      model_name_to_url[model_name]
     UpperCAmelCase_      :							List[str]       =      torch.hub.load_state_dict_from_url(_lowercase			,						map_location='''cpu'''			,						file_name=_lowercase				)[
         '''state_dict'''
     ]

     for name, param in state_dict.items():
          print(_lowercase			,						param.shape				)

     UpperCAmelCase_      :							List[Any]       =      get_upernet_config(_lowercase				)
     UpperCAmelCase_      :							Union[str, Any]       =      UperNetForSemanticSegmentation(_lowercase				)
     model.eval()

     # replace "bn" => "batch_norm"
     for key in state_dict.copy().keys():
          UpperCAmelCase_      :							Dict       =      state_dict.pop(_lowercase				)
          if "bn" in key:
               UpperCAmelCase_      :							Optional[int]       =      key.replace('''bn'''			,						'''batch_norm'''				)
          UpperCAmelCase_      :							Dict       =      val

     # rename keys
     UpperCAmelCase_      :							Dict       =      create_rename_keys(_lowercase				)
     for src, dest in rename_keys:
          rename_key(_lowercase			,						_lowercase			,						_lowercase				)
     read_in_q_k_v(_lowercase			,						config.backbone_config				)

     # fix downsample parameters
     for key, value in state_dict.items():
          if "downsample" in key:
               if "reduction" in key:
                    UpperCAmelCase_      :							int       =      reverse_correct_unfold_reduction_order(_lowercase				)
               if "norm" in key:
                    UpperCAmelCase_      :							Any       =      reverse_correct_unfold_norm_order(_lowercase				)

     model.load_state_dict(_lowercase				)

     # verify on image
     UpperCAmelCase_      :							List[str]       =      '''https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg'''
     UpperCAmelCase_      :							Any       =      Image.open(requests.get(_lowercase			,						stream=_lowercase				).raw				).convert('''RGB'''				)

     UpperCAmelCase_      :							Union[str, Any]       =      SegformerImageProcessor()
     UpperCAmelCase_      :							Optional[int]       =      processor(_lowercase			,						return_tensors='''pt'''				).pixel_values

     with torch.no_grad():
          UpperCAmelCase_      :							Dict       =      model(_lowercase				)
          UpperCAmelCase_      :							Any       =      outputs.logits

     print(logits.shape				)
     print('''First values of logits:'''			,						logits[0, 0, :3, :3]				)
     # assert values
     if model_name == "upernet-swin-tiny":
          UpperCAmelCase_      :							Any       =      torch.tensor(
              [[-7.5958, -7.5958, -7.4302], [-7.5958, -7.5958, -7.4302], [-7.4797, -7.4797, -7.3068]]				)
     elif model_name == "upernet-swin-small":
          UpperCAmelCase_      :							str       =      torch.tensor(
              [[-7.1921, -7.1921, -6.9532], [-7.1921, -7.1921, -6.9532], [-7.0908, -7.0908, -6.8534]]				)
     elif model_name == "upernet-swin-base":
          UpperCAmelCase_      :							Any       =      torch.tensor(
              [[-6.5851, -6.5851, -6.4330], [-6.5851, -6.5851, -6.4330], [-6.4763, -6.4763, -6.3254]]				)
     elif model_name == "upernet-swin-large":
          UpperCAmelCase_      :							List[Any]       =      torch.tensor(
              [[-7.5297, -7.5297, -7.3802], [-7.5297, -7.5297, -7.3802], [-7.4044, -7.4044, -7.2586]]				)
     print('''Logits:'''			,						outputs.logits[0, 0, :3, :3]				)
     assert torch.allclose(outputs.logits[0, 0, :3, :3]			,						_lowercase			,						atol=1E-4				)
     print('''Looks ok!'''				)

     if pytorch_dump_folder_path is not None:
          print(f'''Saving model {model_name} to {pytorch_dump_folder_path}'''				)
          model.save_pretrained(_lowercase				)
          print(f'''Saving processor to {pytorch_dump_folder_path}'''				)
          processor.save_pretrained(_lowercase				)

     if push_to_hub:
          print(f'''Pushing model and processor for {model_name} to hub'''				)
          model.push_to_hub(f'''openmmlab/{model_name}'''				)
          processor.push_to_hub(f'''openmmlab/{model_name}'''				)


if __name__ == "__main__":
     __a          =			argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '--model_name',
         default='upernet-swin-tiny',
         type=str,
         choices=[F"""upernet-swin-{size}""" for size in ['tiny', 'small', 'base', 'large']],
         help='Name of the Swin + UperNet model you\'d like to convert.',
     )
     parser.add_argument(
         '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
     )
     parser.add_argument(
         '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
     )

     __a          =			parser.parse_args()
     convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)










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   __a(					_a       ):






    """simple docstring"""





    lowerCAmelCase        =				42


class   __a(					_a    ,						_a       ):






    """simple docstring"""





    lowerCAmelCase        =				True





    @register_to_config
    def __init__(				self       ,_SCREAMING_SNAKE_CASE = 3       ,_SCREAMING_SNAKE_CASE = 3       ,_SCREAMING_SNAKE_CASE = ("DownEncoderBlock2D",)       ,_SCREAMING_SNAKE_CASE = ("UpDecoderBlock2D",)       ,_SCREAMING_SNAKE_CASE = (64,)       ,_SCREAMING_SNAKE_CASE = 1       ,_SCREAMING_SNAKE_CASE = "silu"       ,_SCREAMING_SNAKE_CASE = 4       ,_SCREAMING_SNAKE_CASE = 32       ,_SCREAMING_SNAKE_CASE = 32       ,_SCREAMING_SNAKE_CASE = 0.1_82_15       ,)			->		Optional[int]:
         super().__init__()

         # pass init params to Encoder
         UpperCAmelCase_      :							List[Any]       =      Encoder(
             in_channels=_SCREAMING_SNAKE_CASE       ,out_channels=_SCREAMING_SNAKE_CASE       ,down_block_types=_SCREAMING_SNAKE_CASE       ,block_out_channels=_SCREAMING_SNAKE_CASE       ,layers_per_block=_SCREAMING_SNAKE_CASE       ,act_fn=_SCREAMING_SNAKE_CASE       ,norm_num_groups=_SCREAMING_SNAKE_CASE       ,double_z=_SCREAMING_SNAKE_CASE       ,)

         # pass init params to Decoder
         UpperCAmelCase_      :							List[str]       =      Decoder(
             in_channels=_SCREAMING_SNAKE_CASE       ,out_channels=_SCREAMING_SNAKE_CASE       ,up_block_types=_SCREAMING_SNAKE_CASE       ,block_out_channels=_SCREAMING_SNAKE_CASE       ,layers_per_block=_SCREAMING_SNAKE_CASE       ,norm_num_groups=_SCREAMING_SNAKE_CASE       ,act_fn=_SCREAMING_SNAKE_CASE       ,)

         UpperCAmelCase_      :							int       =      nn.Convad(2 * latent_channels       ,2 * latent_channels       ,1  )
         UpperCAmelCase_      :							Union[str, Any]       =      nn.Convad(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,1  )

         UpperCAmelCase_      :							Optional[Any]       =      False
         UpperCAmelCase_      :							Dict       =      False

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





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE=False  )			->		int:
         if isinstance(_SCREAMING_SNAKE_CASE       ,(Encoder, Decoder)  ):
              UpperCAmelCase_      :							Union[str, Any]       =      value





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE = True  )			->		Optional[Any]:
         UpperCAmelCase_      :							Dict       =      use_tiling





    def 							a__				(				self  )			->		Optional[Any]:
         self.enable_tiling(_SCREAMING_SNAKE_CASE  )





    def 							a__				(				self  )			->		List[str]:
         UpperCAmelCase_      :							str       =      True





    def 							a__				(				self  )			->		Any:
         UpperCAmelCase_      :							Any       =      False





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

         def fn_recursive_add_processors(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  ):
              if hasattr(_SCREAMING_SNAKE_CASE       ,'''set_processor'''  ):
                   UpperCAmelCase_      :							Tuple       =      module.processor

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

              return processors

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

         return processors





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE  )			->		Union[str, Any]:
         UpperCAmelCase_      :							Optional[int]       =      len(self.attn_processors.keys()  )

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

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

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





    def 							a__				(				self  )			->		Tuple:
         self.set_attn_processor(AttnProcessor()  )





    @apply_forward_hook
    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE = 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(_SCREAMING_SNAKE_CASE       ,return_dict=_SCREAMING_SNAKE_CASE  )

         if self.use_slicing and x.shape[0] > 1:
              UpperCAmelCase_      :							int       =      [self.encoder(_SCREAMING_SNAKE_CASE  ) for x_slice in x.split(1  )]
              UpperCAmelCase_      :							List[str]       =      torch.cat(_SCREAMING_SNAKE_CASE  )
         else:
              UpperCAmelCase_      :							Optional[Any]       =      self.encoder(_SCREAMING_SNAKE_CASE  )

         UpperCAmelCase_      :							Any       =      self.quant_conv(_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Union[str, Any]       =      DiagonalGaussianDistribution(_SCREAMING_SNAKE_CASE  )

         if not return_dict:
              return (posterior,)

         return AutoencoderKLOutput(latent_dist=_SCREAMING_SNAKE_CASE  )





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE = 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(_SCREAMING_SNAKE_CASE       ,return_dict=_SCREAMING_SNAKE_CASE  )

         UpperCAmelCase_      :							Tuple       =      self.post_quant_conv(_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Tuple       =      self.decoder(_SCREAMING_SNAKE_CASE  )

         if not return_dict:
              return (dec,)

         return DecoderOutput(sample=_SCREAMING_SNAKE_CASE  )





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

         if not return_dict:
              return (decoded,)

         return DecoderOutput(sample=_SCREAMING_SNAKE_CASE  )





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )			->		Union[str, Any]:
         UpperCAmelCase_      :							str       =      min(a.shape[2]       ,b.shape[2]       ,_SCREAMING_SNAKE_CASE  )
         for y in range(_SCREAMING_SNAKE_CASE  ):
              UpperCAmelCase_      :							str       =      a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
         return b





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )			->		int:
         UpperCAmelCase_      :							Any       =      min(a.shape[3]       ,b.shape[3]       ,_SCREAMING_SNAKE_CASE  )
         for x in range(_SCREAMING_SNAKE_CASE  ):
              UpperCAmelCase_      :							List[str]       =      a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
         return b





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE = True  )			->		AutoencoderKLOutput:
         UpperCAmelCase_      :							Optional[int]       =      int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)  )
         UpperCAmelCase_      :							Any       =      int(self.tile_latent_min_size * self.tile_overlap_factor  )
         UpperCAmelCase_      :							str       =      self.tile_latent_min_size - blend_extent

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

         UpperCAmelCase_      :							Optional[Any]       =      torch.cat(_SCREAMING_SNAKE_CASE       ,dim=2  )
         UpperCAmelCase_      :							Union[str, Any]       =      DiagonalGaussianDistribution(_SCREAMING_SNAKE_CASE  )

         if not return_dict:
              return (posterior,)

         return AutoencoderKLOutput(latent_dist=_SCREAMING_SNAKE_CASE  )





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE = True  )			->		Union[DecoderOutput, torch.FloatTensor]:
         UpperCAmelCase_      :							Optional[Any]       =      int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)  )
         UpperCAmelCase_      :							str       =      int(self.tile_sample_min_size * self.tile_overlap_factor  )
         UpperCAmelCase_      :							Any       =      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_      :							Tuple       =      []
         for i in range(0       ,z.shape[2]       ,_SCREAMING_SNAKE_CASE  ):
              UpperCAmelCase_      :							Tuple       =      []
              for j in range(0       ,z.shape[3]       ,_SCREAMING_SNAKE_CASE  ):
                   UpperCAmelCase_      :							Optional[int]       =      z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
                   UpperCAmelCase_      :							List[str]       =      self.post_quant_conv(_SCREAMING_SNAKE_CASE  )
                   UpperCAmelCase_      :							Dict       =      self.decoder(_SCREAMING_SNAKE_CASE  )
                   row.append(_SCREAMING_SNAKE_CASE  )
              rows.append(_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Optional[Any]       =      []
         for i, row in enumerate(_SCREAMING_SNAKE_CASE  ):
              UpperCAmelCase_      :							int       =      []
              for j, tile in enumerate(_SCREAMING_SNAKE_CASE  ):
                   # 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_      :							Optional[int]       =      self.blend_v(rows[i - 1][j]       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )
                   if j > 0:
                        UpperCAmelCase_      :							Optional[int]       =      self.blend_h(row[j - 1]       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )
                   result_row.append(tile[:, :, :row_limit, :row_limit]  )
              result_rows.append(torch.cat(_SCREAMING_SNAKE_CASE       ,dim=3  )  )

         UpperCAmelCase_      :							Dict       =      torch.cat(_SCREAMING_SNAKE_CASE       ,dim=2  )
         if not return_dict:
              return (dec,)

         return DecoderOutput(sample=_SCREAMING_SNAKE_CASE  )





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE = False       ,_SCREAMING_SNAKE_CASE = True       ,_SCREAMING_SNAKE_CASE = None       ,)			->		Union[DecoderOutput, torch.FloatTensor]:
         UpperCAmelCase_      :							str       =      sample
         UpperCAmelCase_      :							Optional[Any]       =      self.encode(_SCREAMING_SNAKE_CASE  ).latent_dist
         if sample_posterior:
              UpperCAmelCase_      :							List[Any]       =      posterior.sample(generator=_SCREAMING_SNAKE_CASE  )
         else:
              UpperCAmelCase_      :							Union[str, Any]       =      posterior.mode()
         UpperCAmelCase_      :							List[Any]       =      self.decode(_SCREAMING_SNAKE_CASE  ).sample

         if not return_dict:
              return (dec,)

         return DecoderOutput(sample=_SCREAMING_SNAKE_CASE  )