





'''simple docstring'''






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  (					lowercase_							):






   '''simple docstring'''





   lowerCamelCase__    							=      42



class 					__snake_case  (					lowercase_		,				lowercase_							):






   '''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 = (6_4,)	,		__SCREAMING_SNAKE_CASE = 1	,		__SCREAMING_SNAKE_CASE = "silu"	,		__SCREAMING_SNAKE_CASE = 4	,		__SCREAMING_SNAKE_CASE = 3_2	,		__SCREAMING_SNAKE_CASE = 3_2	,		__SCREAMING_SNAKE_CASE = 0.1_8215	,		):
    super().__init__()

    # pass init params to Encoder
    snake_case__     :					int    =  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
    snake_case__     :					Union[str, Any]    =  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	,		)

    snake_case__     :					Any    =  nn.Convad(2 * latent_channels	,		2 * latent_channels	,		1       )
    snake_case__     :					List[str]    =  nn.Convad(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		1       )

    snake_case__     :					List[Any]    =  False
    snake_case__     :					List[str]    =  False

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

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE=False       ):







    if isinstance(__SCREAMING_SNAKE_CASE	,		(Encoder, Decoder)       ):
     snake_case__     :					str    =  value

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE = True       ):
    snake_case__     :					Union[str, Any]    =  use_tiling

   def 		__UpperCamelCase						(							self       ):
    self.enable_tiling(__SCREAMING_SNAKE_CASE       )

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					List[Any]    =  True

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Any    =  False

   @property
   # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					str    =  {}

    def fn_recursive_add_processors(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       ):
     if hasattr(__SCREAMING_SNAKE_CASE	,		"""set_processor"""       ):
      snake_case__     :					int    =  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 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE       ):
    snake_case__     :					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 		__UpperCamelCase						(							self       ):
    self.set_attn_processor(AttnProcessor()       )

   @apply_forward_hook
   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE = True       ):







    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:
     snake_case__     :					Optional[int]    =  [self.encoder(__SCREAMING_SNAKE_CASE       ) for x_slice in x.split(1       )]
     snake_case__     :					Any    =  torch.cat(__SCREAMING_SNAKE_CASE       )
    else:
     snake_case__     :					int    =  self.encoder(__SCREAMING_SNAKE_CASE       )

    snake_case__     :					str    =  self.quant_conv(__SCREAMING_SNAKE_CASE       )
    snake_case__     :					Optional[int]    =  DiagonalGaussianDistribution(__SCREAMING_SNAKE_CASE       )

    if not return_dict:
     return (posterior,)

    return AutoencoderKLOutput(latent_dist=__SCREAMING_SNAKE_CASE       )

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE = True       ):







    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       )

    snake_case__     :					Optional[Any]    =  self.post_quant_conv(__SCREAMING_SNAKE_CASE       )
    snake_case__     :					str    =  self.decoder(__SCREAMING_SNAKE_CASE       )

    if not return_dict:
     return (dec,)

    return DecoderOutput(sample=__SCREAMING_SNAKE_CASE       )

   @apply_forward_hook
   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE = True       ):







    if self.use_slicing and z.shape[0] > 1:
     snake_case__     :					Any    =  [self._decode(__SCREAMING_SNAKE_CASE       ).sample for z_slice in z.split(1       )]
     snake_case__     :					List[str]    =  torch.cat(__SCREAMING_SNAKE_CASE       )
    else:
     snake_case__     :					Union[str, Any]    =  self._decode(__SCREAMING_SNAKE_CASE       ).sample

    if not return_dict:
     return (decoded,)

    return DecoderOutput(sample=__SCREAMING_SNAKE_CASE       )

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       ):
    snake_case__     :					Optional[int]    =  min(a.shape[2]	,		b.shape[2]	,		__SCREAMING_SNAKE_CASE       )
    for y in range(__SCREAMING_SNAKE_CASE       ):
     snake_case__     :					str    =  a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
    return b

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       ):
    snake_case__     :					List[Any]    =  min(a.shape[3]	,		b.shape[3]	,		__SCREAMING_SNAKE_CASE       )
    for x in range(__SCREAMING_SNAKE_CASE       ):
     snake_case__     :					Tuple    =  a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
    return b

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE = True       ):
    snake_case__     :					Dict    =  int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)       )
    snake_case__     :					List[Any]    =  int(self.tile_latent_min_size * self.tile_overlap_factor       )
    snake_case__     :					List[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]	,		__SCREAMING_SNAKE_CASE       ):
     snake_case__     :					Optional[Any]    =  []
     for j in range(0	,		x.shape[3]	,		__SCREAMING_SNAKE_CASE       ):
      snake_case__     :					int    =  x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
      snake_case__     :					List[str]    =  self.encoder(__SCREAMING_SNAKE_CASE       )
      snake_case__     :					Dict    =  self.quant_conv(__SCREAMING_SNAKE_CASE       )
      row.append(__SCREAMING_SNAKE_CASE       )
     rows.append(__SCREAMING_SNAKE_CASE       )
    snake_case__     :					List[str]    =  []
    for i, row in enumerate(__SCREAMING_SNAKE_CASE       ):
     snake_case__     :					Optional[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:
       snake_case__     :					List[Any]    =  self.blend_v(rows[i - 1][j]	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )
      if j > 0:
       snake_case__     :					Optional[Any]    =  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       )       )

    snake_case__     :					Dict    =  torch.cat(__SCREAMING_SNAKE_CASE	,		dim=2       )
    snake_case__     :					Tuple    =  DiagonalGaussianDistribution(__SCREAMING_SNAKE_CASE       )

    if not return_dict:
     return (posterior,)

    return AutoencoderKLOutput(latent_dist=__SCREAMING_SNAKE_CASE       )

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE = True       ):
    snake_case__     :					Optional[Any]    =  int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)       )
    snake_case__     :					Optional[Any]    =  int(self.tile_sample_min_size * self.tile_overlap_factor       )
    snake_case__     :					Dict    =  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__     :					Dict    =  []
    for i in range(0	,		z.shape[2]	,		__SCREAMING_SNAKE_CASE       ):
     snake_case__     :					Optional[Any]    =  []
     for j in range(0	,		z.shape[3]	,		__SCREAMING_SNAKE_CASE       ):
      snake_case__     :					List[str]    =  z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
      snake_case__     :					int    =  self.post_quant_conv(__SCREAMING_SNAKE_CASE       )
      snake_case__     :					List[Any]    =  self.decoder(__SCREAMING_SNAKE_CASE       )
      row.append(__SCREAMING_SNAKE_CASE       )
     rows.append(__SCREAMING_SNAKE_CASE       )
    snake_case__     :					List[Any]    =  []
    for i, row in enumerate(__SCREAMING_SNAKE_CASE       ):
     snake_case__     :					Union[str, Any]    =  []
     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:
       snake_case__     :					Optional[Any]    =  self.blend_v(rows[i - 1][j]	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )
      if j > 0:
       snake_case__     :					List[Any]    =  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       )       )

    snake_case__     :					Any    =  torch.cat(__SCREAMING_SNAKE_CASE	,		dim=2       )
    if not return_dict:
     return (dec,)

    return DecoderOutput(sample=__SCREAMING_SNAKE_CASE       )


   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE = False	,		__SCREAMING_SNAKE_CASE = True	,		__SCREAMING_SNAKE_CASE = None	,		):
    snake_case__     :					Dict    =  sample
    snake_case__     :					Optional[Any]    =  self.encode(__SCREAMING_SNAKE_CASE       ).latent_dist
    if sample_posterior:
     snake_case__     :					Dict    =  posterior.sample(generator=__SCREAMING_SNAKE_CASE       )
    else:
     snake_case__     :					Optional[Any]    =  posterior.mode()
    snake_case__     :					Optional[int]    =  self.decode(__SCREAMING_SNAKE_CASE       ).sample

    if not return_dict:
     return (dec,)

    return DecoderOutput(sample=__SCREAMING_SNAKE_CASE       )








'''simple docstring'''

from __future__ import annotations

from typing import Any
class     __lowerCAmelCase		:





      '''simple docstring'''






      def __init__(self       :       Optional[Any]  ,		UpperCamelCase       :       int = 6					):

             '''simple docstring'''







             lowercase__      =					None
             lowercase__      =					None
             self.create_linked_list(UpperCamelCase					)



      def  UpperCamelCase__						(self       :       List[str]  ,		UpperCamelCase       :       int					):

             '''simple docstring'''







             lowercase__      =					Node()
             lowercase__      =					current_node
             lowercase__      =					current_node
             lowercase__      =					current_node
             for _ in range(1  ,		UpperCamelCase					):
                    lowercase__      =					Node()
                    lowercase__      =					current_node
                    lowercase__      =					previous_node
                    lowercase__      =					current_node
             lowercase__      =					self.front
             lowercase__      =					previous_node



      def  UpperCamelCase__						(self       :       Union[str, Any]					):

             '''simple docstring'''







             return (
                 self.front == self.rear
                 and self.front is not None
                 and self.front.data is None
             )



      def  UpperCamelCase__						(self       :       Dict					):

             '''simple docstring'''







             self.check_can_perform_operation()
             return self.front.data if self.front else None



      def  UpperCamelCase__						(self       :       Any  ,		UpperCamelCase       :       Any					):

             '''simple docstring'''







             if self.rear is None:
                    return

             self.check_is_full()
             if not self.is_empty():
                    lowercase__      =					self.rear.next
             if self.rear:
                    lowercase__      =					data



      def  UpperCamelCase__						(self       :       List[Any]					):

             '''simple docstring'''







             self.check_can_perform_operation()
             if self.rear is None or self.front is None:
                    return None
             if self.front == self.rear:
                    lowercase__      =					self.front.data
                    lowercase__      =					None
                    return data

             lowercase__      =					self.front
             lowercase__      =					old_front.next
             lowercase__      =					old_front.data
             lowercase__      =					None
             return data



      def  UpperCamelCase__						(self       :       Union[str, Any]					):

             '''simple docstring'''







             if self.is_empty():
                    raise Exception('''Empty Queue'''					)






      def  UpperCamelCase__						(self       :       int					):

             '''simple docstring'''







             if self.rear and self.rear.next == self.front:
                    raise Exception('''Full Queue'''					)






class     __lowerCAmelCase		:





      '''simple docstring'''






      def __init__(self       :       Any					):

             '''simple docstring'''







             lowercase__      =					None
             lowercase__      =					None
             lowercase__      =					None


if __name__ == "__main__":
 import doctest

 doctest.testmod()




