

import itertools
import os
import random
import tempfile
import unittest

import numpy as np

from transformers import TvltFeatureExtractor, is_datasets_available
from transformers.testing_utils import check_json_file_has_correct_format, require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_torch_available():
     import torch

if is_datasets_available():
     from datasets import load_dataset

__magic_name__								=random.Random()







def  __UpperCamelCase    (						A ,							A=1.0 ,							A=None ,							A=None							):
     if rng is None:
          UpperCamelCase__										=      global_rng

     UpperCamelCase__										=      []
     for batch_idx in range(shape[0]							):
          values.append([]							)
          for _ in range(shape[1]							):
               values[-1].append(rng.random() * scale							)

     return values



class     _A   (    unittest.TestCase							):






  def __init__(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_=7   ,				SCREAMING_SNAKE_CASE_=400   ,				SCREAMING_SNAKE_CASE_=2000   ,				SCREAMING_SNAKE_CASE_=2048   ,				SCREAMING_SNAKE_CASE_=128   ,				SCREAMING_SNAKE_CASE_=1   ,				SCREAMING_SNAKE_CASE_=512   ,				SCREAMING_SNAKE_CASE_=30   ,				SCREAMING_SNAKE_CASE_=4_4100   ,				)		->      Tuple:





       '''simple docstring'''





       UpperCamelCase__										=      parent
       UpperCamelCase__										=      batch_size
       UpperCamelCase__										=      min_seq_length
       UpperCamelCase__										=      max_seq_length
       UpperCamelCase__										=      (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
       UpperCamelCase__										=      spectrogram_length
       UpperCamelCase__										=      feature_size
       UpperCamelCase__										=      num_audio_channels
       UpperCamelCase__										=      hop_length
       UpperCamelCase__										=      chunk_length
       UpperCamelCase__										=      sampling_rate






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





       '''simple docstring'''





       return {
           "spectrogram_length": self.spectrogram_length,
           "feature_size": self.feature_size,
           "num_audio_channels": self.num_audio_channels,
           "hop_length": self.hop_length,
           "chunk_length": self.chunk_length,
           "sampling_rate": self.sampling_rate,
       }






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_=False   ,				SCREAMING_SNAKE_CASE_=False	)		->      List[Any]:





       '''simple docstring'''





       def _flatten(SCREAMING_SNAKE_CASE_	):
            return list(itertools.chain(*SCREAMING_SNAKE_CASE_	)	)

       if equal_length:
            UpperCamelCase__										=      [floats_list((self.max_seq_length, self.feature_size)	) for _ in range(self.batch_size	)]
       else:
            # make sure that inputs increase in size
            UpperCamelCase__										=      [
                floats_list((x, self.feature_size)	)
                for x in range(self.min_seq_length   ,				self.max_seq_length   ,				self.seq_length_diff	)
            ]
       if numpify:
            UpperCamelCase__										=      [np.asarray(SCREAMING_SNAKE_CASE_	) for x in speech_inputs]
       return speech_inputs



@require_torch
@require_torchaudio
class     _A   (    __UpperCAmelCase	, unittest.TestCase							):
  SCREAMING_SNAKE_CASE_	:	List[Any]						=TvltFeatureExtractor






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





       '''simple docstring'''





       UpperCamelCase__										=      TvltFeatureExtractionTester(self	)






  def 		_a					(self	)		->      Dict:





       '''simple docstring'''





       UpperCamelCase__										=      self.feature_extraction_class(**self.feat_extract_dict	)
       self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_   ,				'''spectrogram_length'''	)	)
       self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_   ,				'''feature_size'''	)	)
       self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_   ,				'''num_audio_channels'''	)	)
       self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_   ,				'''hop_length'''	)	)
       self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_   ,				'''chunk_length'''	)	)
       self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_   ,				'''sampling_rate'''	)	)






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





       '''simple docstring'''





       UpperCamelCase__										=      self.feature_extraction_class(**self.feat_extract_dict	)

       with tempfile.TemporaryDirectory() as tmpdirname:
            UpperCamelCase__										=      feat_extract_first.save_pretrained(SCREAMING_SNAKE_CASE_	)[0]
            check_json_file_has_correct_format(SCREAMING_SNAKE_CASE_	)
            UpperCamelCase__										=      self.feature_extraction_class.from_pretrained(SCREAMING_SNAKE_CASE_	)

       UpperCamelCase__										=      feat_extract_first.to_dict()
       UpperCamelCase__										=      feat_extract_second.to_dict()
       UpperCamelCase__										=      dict_first.pop('''mel_filters'''	)
       UpperCamelCase__										=      dict_second.pop('''mel_filters'''	)
       self.assertTrue(np.allclose(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)	)
       self.assertEqual(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)






  def 		_a					(self	)		->      str:





       '''simple docstring'''





       UpperCamelCase__										=      self.feature_extraction_class(**self.feat_extract_dict	)

       with tempfile.TemporaryDirectory() as tmpdirname:
            UpperCamelCase__										=      os.path.join(SCREAMING_SNAKE_CASE_   ,				'''feat_extract.json'''	)
            feat_extract_first.to_json_file(SCREAMING_SNAKE_CASE_	)
            UpperCamelCase__										=      self.feature_extraction_class.from_json_file(SCREAMING_SNAKE_CASE_	)

       UpperCamelCase__										=      feat_extract_first.to_dict()
       UpperCamelCase__										=      feat_extract_second.to_dict()
       UpperCamelCase__										=      dict_first.pop('''mel_filters'''	)
       UpperCamelCase__										=      dict_second.pop('''mel_filters'''	)
       self.assertTrue(np.allclose(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)	)
       self.assertEqual(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)






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





       '''simple docstring'''





       UpperCamelCase__										=      self.feature_extraction_class(**self.feat_extract_dict	)

       # create three inputs of length 800, 1000, and 1200
       UpperCamelCase__										=      [floats_list((1, x)	)[0] for x in range(800   ,				1400   ,				200	)]
       UpperCamelCase__										=      [np.asarray(SCREAMING_SNAKE_CASE_	) for speech_input in speech_inputs]

       # Test not batched input
       UpperCamelCase__										=      feature_extractor(np_speech_inputs[0]   ,				return_tensors='''np'''   ,				sampling_rate=4_4100	).audio_values

       self.assertTrue(encoded_audios.ndim == 4	)
       self.assertTrue(encoded_audios.shape[-1] == feature_extractor.feature_size	)
       self.assertTrue(encoded_audios.shape[-2] <= feature_extractor.spectrogram_length	)
       self.assertTrue(encoded_audios.shape[-3] == feature_extractor.num_channels	)

       # Test batched
       UpperCamelCase__										=      feature_extractor(SCREAMING_SNAKE_CASE_   ,				return_tensors='''np'''   ,				sampling_rate=4_4100	).audio_values

       self.assertTrue(encoded_audios.ndim == 4	)
       self.assertTrue(encoded_audios.shape[-1] == feature_extractor.feature_size	)
       self.assertTrue(encoded_audios.shape[-2] <= feature_extractor.spectrogram_length	)
       self.assertTrue(encoded_audios.shape[-3] == feature_extractor.num_channels	)

       # Test audio masking
       UpperCamelCase__										=      feature_extractor(
           SCREAMING_SNAKE_CASE_   ,				return_tensors='''np'''   ,				sampling_rate=4_4100   ,				mask_audio=SCREAMING_SNAKE_CASE_	).audio_values

       self.assertTrue(encoded_audios.ndim == 4	)
       self.assertTrue(encoded_audios.shape[-1] == feature_extractor.feature_size	)
       self.assertTrue(encoded_audios.shape[-2] <= feature_extractor.spectrogram_length	)
       self.assertTrue(encoded_audios.shape[-3] == feature_extractor.num_channels	)

       # Test 2-D numpy arrays are batched.
       UpperCamelCase__										=      [floats_list((1, x)	)[0] for x in (800, 800, 800)]
       UpperCamelCase__										=      np.asarray(SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      feature_extractor(SCREAMING_SNAKE_CASE_   ,				return_tensors='''np'''   ,				sampling_rate=4_4100	).audio_values
       self.assertTrue(encoded_audios.ndim == 4	)
       self.assertTrue(encoded_audios.shape[-1] == feature_extractor.feature_size	)
       self.assertTrue(encoded_audios.shape[-2] <= feature_extractor.spectrogram_length	)
       self.assertTrue(encoded_audios.shape[-3] == feature_extractor.num_channels	)






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_	)		->      Optional[int]:





       '''simple docstring'''





       UpperCamelCase__										=      load_dataset('''hf-internal-testing/librispeech_asr_dummy'''   ,				'''clean'''   ,				split='''validation'''	)
       # automatic decoding with librispeech
       UpperCamelCase__										=      ds.sort('''id'''	).select(range(SCREAMING_SNAKE_CASE_	)	)[:num_samples]['''audio''']

       return [x["array"] for x in speech_samples]






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





       '''simple docstring'''





       UpperCamelCase__										=      self._load_datasamples(1	)
       UpperCamelCase__										=      TvltFeatureExtractor()
       UpperCamelCase__										=      feature_extractor(SCREAMING_SNAKE_CASE_   ,				return_tensors='''pt'''	).audio_values

       self.assertEquals(audio_values.shape   ,				(1, 1, 192, 128)	)

       UpperCamelCase__										=      torch.tensor([[-0.3032, -0.2708], [-0.4434, -0.4007]]	)
       self.assertTrue(torch.allclose(audio_values[0, 0, :2, :2]   ,				SCREAMING_SNAKE_CASE_   ,				atol=1E-4	)	)




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_	)



