

from typing import List, Optional, Union

import torch

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)


__magic_name__								=logging.get_logger(__name__)  # pylint: disable=invalid-name

__magic_name__								='''
    Examples:
        ```py
        >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline
        >>> import torch

        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained("kandinsky-community/kandinsky-2-2-prior")
        >>> pipe_prior.to("cuda")
        >>> prompt = "red cat, 4k photo"
        >>> out = pipe_prior(prompt)
        >>> image_emb = out.image_embeds
        >>> zero_image_emb = out.negative_image_embeds
        >>> pipe = KandinskyV22Pipeline.from_pretrained("kandinsky-community/kandinsky-2-2-decoder")
        >>> pipe.to("cuda")
        >>> image = pipe(
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=zero_image_emb,
        ...     height=768,
        ...     width=768,
        ...     num_inference_steps=50,
        ... ).images
        >>> image[0].save("cat.png")
        ```
'''







def  __UpperCamelCase    (						A ,							A ,							A=8							):
     UpperCamelCase__										=      height // scale_factor**2
     if height % scale_factor**2 != 0:
          new_height += 1
     UpperCamelCase__										=      width // scale_factor**2
     if width % scale_factor**2 != 0:
          new_width += 1
     return new_height * scale_factor, new_width * scale_factor



class     _A   (    __UpperCamelCase							):






  def __init__(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				)		->      List[str]:





       '''simple docstring'''





       super().__init__()

       self.register_modules(
           unet=SCREAMING_SNAKE_CASE_   ,				scheduler=SCREAMING_SNAKE_CASE_   ,				movq=SCREAMING_SNAKE_CASE_   ,				)
       UpperCamelCase__										=      2 ** (len(self.movq.config.block_out_channels	) - 1)






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)		->      int:





       '''simple docstring'''





       if latents is None:
            UpperCamelCase__										=      randn_tensor(SCREAMING_SNAKE_CASE_   ,				generator=SCREAMING_SNAKE_CASE_   ,				device=SCREAMING_SNAKE_CASE_   ,				dtype=SCREAMING_SNAKE_CASE_	)
       else:
            if latents.shape != shape:
                 raise ValueError(F"Unexpected latents shape, got {latents.shape}, expected {shape}"	)
            UpperCamelCase__										=      latents.to(SCREAMING_SNAKE_CASE_	)

       UpperCamelCase__										=      latents * scheduler.init_noise_sigma
       return latents






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_=0	)		->      str:





       '''simple docstring'''





       if is_accelerate_available():
            from accelerate import cpu_offload
       else:
            raise ImportError('''Please install accelerate via `pip install accelerate`'''	)

       UpperCamelCase__										=      torch.device(F"cuda:{gpu_id}"	)

       UpperCamelCase__										=      [
           self.unet,
           self.movq,
       ]
       for cpu_offloaded_model in models:
            if cpu_offloaded_model is not None:
                 cpu_offload(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_=0	)		->      Optional[Any]:





       '''simple docstring'''





       if is_accelerate_available() and is_accelerate_version('''>='''   ,				'''0.17.0.dev0'''	):
            from accelerate import cpu_offload_with_hook
       else:
            raise ImportError('''`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.'''	)

       UpperCamelCase__										=      torch.device(F"cuda:{gpu_id}"	)

       if self.device.type != "cpu":
            self.to('''cpu'''   ,				silence_dtype_warnings=SCREAMING_SNAKE_CASE_	)
            torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

       UpperCamelCase__										=      None
       for cpu_offloaded_model in [self.unet, self.movq]:
            UpperCamelCase__						,       UpperCamelCase__										=      cpu_offload_with_hook(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				prev_module_hook=SCREAMING_SNAKE_CASE_	)

       # We'll offload the last model manually.
       UpperCamelCase__										=      hook






  @property
  # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
  def 		_a					(self	)		->      List[str]:





       '''simple docstring'''





       if not hasattr(self.unet   ,				'''_hf_hook'''	):
            return self.device
       for module in self.unet.modules():
            if (
                hasattr(SCREAMING_SNAKE_CASE_   ,				'''_hf_hook'''	)
                and hasattr(module._hf_hook   ,				'''execution_device'''	)
                and module._hf_hook.execution_device is not None
            ):
                 return torch.device(module._hf_hook.execution_device	)
       return self.device






  @torch.no_grad()
  @replace_example_docstring(SCREAMING_SNAKE_CASE_	)
  def __call__(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_ = 512   ,				SCREAMING_SNAKE_CASE_ = 512   ,				SCREAMING_SNAKE_CASE_ = 100   ,				SCREAMING_SNAKE_CASE_ = 4.0   ,				SCREAMING_SNAKE_CASE_ = 1   ,				SCREAMING_SNAKE_CASE_ = None   ,				SCREAMING_SNAKE_CASE_ = None   ,				SCREAMING_SNAKE_CASE_ = "pil"   ,				SCREAMING_SNAKE_CASE_ = True   ,				)		->      Any:





       '''simple docstring'''





       UpperCamelCase__										=      self._execution_device

       UpperCamelCase__										=      guidance_scale > 1.0

       if isinstance(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	):
            UpperCamelCase__										=      torch.cat(SCREAMING_SNAKE_CASE_   ,				dim=0	)
       UpperCamelCase__										=      image_embeds.shape[0] * num_images_per_prompt
       if isinstance(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	):
            UpperCamelCase__										=      torch.cat(SCREAMING_SNAKE_CASE_   ,				dim=0	)

       if do_classifier_free_guidance:
            UpperCamelCase__										=      image_embeds.repeat_interleave(SCREAMING_SNAKE_CASE_   ,				dim=0	)
            UpperCamelCase__										=      negative_image_embeds.repeat_interleave(SCREAMING_SNAKE_CASE_   ,				dim=0	)

       UpperCamelCase__										=      torch.cat([negative_image_embeds, image_embeds]   ,				dim=0	).to(dtype=self.unet.dtype   ,				device=SCREAMING_SNAKE_CASE_	)

       self.scheduler.set_timesteps(SCREAMING_SNAKE_CASE_   ,				device=SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      self.scheduler.timesteps

       UpperCamelCase__										=      self.unet.config.in_channels

       UpperCamelCase__						,       UpperCamelCase__										=      downscale_height_and_width(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				self.movq_scale_factor	)

       # create initial latent
       UpperCamelCase__										=      self.prepare_latents(
           (batch_size, num_channels_latents, height, width)   ,				image_embeds.dtype   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				self.scheduler   ,				)

       for i, t in enumerate(self.progress_bar(SCREAMING_SNAKE_CASE_	)	):
            # expand the latents if we are doing classifier free guidance
            UpperCamelCase__										=      torch.cat([latents] * 2	) if do_classifier_free_guidance else latents

            UpperCamelCase__										=      {'''image_embeds''': image_embeds}
            UpperCamelCase__										=      self.unet(
                sample=SCREAMING_SNAKE_CASE_   ,				timestep=SCREAMING_SNAKE_CASE_   ,				encoder_hidden_states=SCREAMING_SNAKE_CASE_   ,				added_cond_kwargs=SCREAMING_SNAKE_CASE_   ,				return_dict=SCREAMING_SNAKE_CASE_   ,				)[0]

            if do_classifier_free_guidance:
                 UpperCamelCase__						,       UpperCamelCase__										=      noise_pred.split(latents.shape[1]   ,				dim=1	)
                 UpperCamelCase__						,       UpperCamelCase__										=      noise_pred.chunk(2	)
                 UpperCamelCase__						,       UpperCamelCase__										=      variance_pred.chunk(2	)
                 UpperCamelCase__										=      noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
                 UpperCamelCase__										=      torch.cat([noise_pred, variance_pred_text]   ,				dim=1	)

            if not (
                hasattr(self.scheduler.config   ,				'''variance_type'''	)
                and self.scheduler.config.variance_type in ["learned", "learned_range"]
            ):
                 UpperCamelCase__						,       UpperCamelCase__										=      noise_pred.split(latents.shape[1]   ,				dim=1	)

            # compute the previous noisy sample x_t -> x_t-1
            UpperCamelCase__										=      self.scheduler.step(
                SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				generator=SCREAMING_SNAKE_CASE_   ,				)[0]
       # post-processing
       UpperCamelCase__										=      self.movq.decode(SCREAMING_SNAKE_CASE_   ,				force_not_quantize=SCREAMING_SNAKE_CASE_	)['''sample''']

       if output_type not in ["pt", "np", "pil"]:
            raise ValueError(F"Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}"	)

       if output_type in ["np", "pil"]:
            UpperCamelCase__										=      image * 0.5 + 0.5
            UpperCamelCase__										=      image.clamp(0   ,				1	)
            UpperCamelCase__										=      image.cpu().permute(0   ,				2   ,				3   ,				1	).float().numpy()

       if output_type == "pil":
            UpperCamelCase__										=      self.numpy_to_pil(SCREAMING_SNAKE_CASE_	)

       if not return_dict:
            return (image,)

       return ImagePipelineOutput(images=SCREAMING_SNAKE_CASE_	)




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   (    __UpperCamelCase							):
  SCREAMING_SNAKE_CASE_	:	"DiagonalGaussianDistribution"



class     _A   (    __UpperCamelCase	, __UpperCamelCase							):
  SCREAMING_SNAKE_CASE_	:	List[str]						=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.18215   ,				)		->      Optional[int]:





       '''simple docstring'''





       super().__init__()

       # pass init params to Encoder
       UpperCamelCase__										=      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__										=      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__										=      nn.Convad(2 * latent_channels   ,				2 * latent_channels   ,				1	)
       UpperCamelCase__										=      nn.Convad(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				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   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_=False	)		->      Dict:





       '''simple docstring'''





       if isinstance(SCREAMING_SNAKE_CASE_   ,				(Encoder, Decoder)	):
            UpperCamelCase__										=      value






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_ = True	)		->      Union[str, Any]:





       '''simple docstring'''





       UpperCamelCase__										=      use_tiling






  def 		_a					(self	)		->      int:





       '''simple docstring'''





       self.enable_tiling(SCREAMING_SNAKE_CASE_	)






  def 		_a					(self	)		->      Any:





       '''simple docstring'''





       UpperCamelCase__										=      True






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





       '''simple docstring'''





       UpperCamelCase__										=      False






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





       '''simple docstring'''





       UpperCamelCase__										=      {}

       def fn_recursive_add_processors(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	):
            if hasattr(SCREAMING_SNAKE_CASE_   ,				'''set_processor'''	):
                 UpperCamelCase__										=      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_	)		->      List[Any]:





       '''simple docstring'''





       UpperCamelCase__										=      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:





       '''simple docstring'''





       self.set_attn_processor(AttnProcessor()	)






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

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

       UpperCamelCase__										=      self.quant_conv(SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      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]:





       '''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(SCREAMING_SNAKE_CASE_   ,				return_dict=SCREAMING_SNAKE_CASE_	)

       UpperCamelCase__										=      self.post_quant_conv(SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      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]:





       '''simple docstring'''





       if self.use_slicing and z.shape[0] > 1:
            UpperCamelCase__										=      [self._decode(SCREAMING_SNAKE_CASE_	).sample for z_slice in z.split(1	)]
            UpperCamelCase__										=      torch.cat(SCREAMING_SNAKE_CASE_	)
       else:
            UpperCamelCase__										=      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]:





       '''simple docstring'''





       UpperCamelCase__										=      min(a.shape[2]   ,				b.shape[2]   ,				SCREAMING_SNAKE_CASE_	)
       for y in range(SCREAMING_SNAKE_CASE_	):
            UpperCamelCase__										=      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_	)		->      List[Any]:





       '''simple docstring'''





       UpperCamelCase__										=      min(a.shape[3]   ,				b.shape[3]   ,				SCREAMING_SNAKE_CASE_	)
       for x in range(SCREAMING_SNAKE_CASE_	):
            UpperCamelCase__										=      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:





       '''simple docstring'''





       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]   ,				SCREAMING_SNAKE_CASE_	):
            UpperCamelCase__										=      []
            for j in range(0   ,				x.shape[3]   ,				SCREAMING_SNAKE_CASE_	):
                 UpperCamelCase__										=      x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
                 UpperCamelCase__										=      self.encoder(SCREAMING_SNAKE_CASE_	)
                 UpperCamelCase__										=      self.quant_conv(SCREAMING_SNAKE_CASE_	)
                 row.append(SCREAMING_SNAKE_CASE_	)
            rows.append(SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      []
       for i, row in enumerate(SCREAMING_SNAKE_CASE_	):
            UpperCamelCase__										=      []
            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__										=      self.blend_v(rows[i - 1][j]   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)
                 if j > 0:
                      UpperCamelCase__										=      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__										=      torch.cat(SCREAMING_SNAKE_CASE_   ,				dim=2	)
       UpperCamelCase__										=      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]:





       '''simple docstring'''





       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]   ,				SCREAMING_SNAKE_CASE_	):
            UpperCamelCase__										=      []
            for j in range(0   ,				z.shape[3]   ,				SCREAMING_SNAKE_CASE_	):
                 UpperCamelCase__										=      z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
                 UpperCamelCase__										=      self.post_quant_conv(SCREAMING_SNAKE_CASE_	)
                 UpperCamelCase__										=      self.decoder(SCREAMING_SNAKE_CASE_	)
                 row.append(SCREAMING_SNAKE_CASE_	)
            rows.append(SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      []
       for i, row in enumerate(SCREAMING_SNAKE_CASE_	):
            UpperCamelCase__										=      []
            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__										=      self.blend_v(rows[i - 1][j]   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)
                 if j > 0:
                      UpperCamelCase__										=      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__										=      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]:





       '''simple docstring'''





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

       if not return_dict:
            return (dec,)

       return DecoderOutput(sample=SCREAMING_SNAKE_CASE_	)



