

'''simple docstring'''







import gc
import math
import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, logging, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


__lowerCamelCase		       =    logging.get_logger(__name__)

enable_full_determinism()


class 					A__      (      _snake_case						,						_snake_case						,						unittest.TestCase			):
    lowercase						=		UNetaDModel
    lowercase						=		"sample"



    @property
    def 					snake_case_			(					self       )						->      Any:





           '''simple docstring'''
           A_						=      4
           A_						=      3
           A_						=      (32, 32)

           A_						=      floats_tensor((batch_size, num_channels) + sizes       ).to(UpperCamelCase__       )
           A_						=      torch.tensor([10]       ).to(UpperCamelCase__       )

           return {"sample": noise, "timestep": time_step}



    @property
    def 					snake_case_			(					self       )						->      int:





           '''simple docstring'''
           return (3, 32, 32)



    @property
    def 					snake_case_			(					self       )						->      List[str]:





           '''simple docstring'''
           return (3, 32, 32)







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





           '''simple docstring'''
           A_						=      {
               """block_out_channels""": (32, 64),
               """down_block_types""": ("""DownBlock2D""", """AttnDownBlock2D"""),
               """up_block_types""": ("""AttnUpBlock2D""", """UpBlock2D"""),
               """attention_head_dim""": 3,
               """out_channels""": 3,
               """in_channels""": 3,
               """layers_per_block""": 2,
               """sample_size""": 32,
           }
           A_						=      self.dummy_input
           return init_dict, inputs_dict


class 					A__      (      _snake_case						,						_snake_case						,						unittest.TestCase			):
    lowercase						=		UNetaDModel
    lowercase						=		"sample"



    @property
    def 					snake_case_			(					self       )						->      Dict:





           '''simple docstring'''
           A_						=      4
           A_						=      4
           A_						=      (32, 32)

           A_						=      floats_tensor((batch_size, num_channels) + sizes       ).to(UpperCamelCase__       )
           A_						=      torch.tensor([10]       ).to(UpperCamelCase__       )

           return {"sample": noise, "timestep": time_step}



    @property
    def 					snake_case_			(					self       )						->      str:





           '''simple docstring'''
           return (4, 32, 32)



    @property
    def 					snake_case_			(					self       )						->      Tuple:





           '''simple docstring'''
           return (4, 32, 32)



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





           '''simple docstring'''
           A_						=      {
               """sample_size""": 32,
               """in_channels""": 4,
               """out_channels""": 4,
               """layers_per_block""": 2,
               """block_out_channels""": (32, 64),
               """attention_head_dim""": 32,
               """down_block_types""": ("""DownBlock2D""", """DownBlock2D"""),
               """up_block_types""": ("""UpBlock2D""", """UpBlock2D"""),
           }
           A_						=      self.dummy_input
           return init_dict, inputs_dict



    def 					snake_case_			(					self       )						->      str:





           '''simple docstring'''
           A_					,			A_						=      UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""		,				output_loading_info=UpperCamelCase__       )

           self.assertIsNotNone(UpperCamelCase__       )
           self.assertEqual(len(loading_info["""missing_keys"""]       )		,				0       )

           model.to(UpperCamelCase__       )
           A_						=      model(**self.dummy_input       ).sample

           assert image is not None, "Make sure output is not None"



    @unittest.skipIf(torch_device != """cuda"""		,				"""This test is supposed to run on GPU"""       )
    def 					snake_case_			(					self       )						->      int:





           '''simple docstring'''
           A_					,			A_						=      UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""		,				output_loading_info=UpperCamelCase__       )
           model.to(UpperCamelCase__       )
           A_						=      model(**self.dummy_input       ).sample

           assert image is not None, "Make sure output is not None"



    @unittest.skipIf(torch_device != """cuda"""		,				"""This test is supposed to run on GPU"""       )
    def 					snake_case_			(					self       )						->      List[str]:





           '''simple docstring'''
           # by defautl model loading will use accelerate as `low_cpu_mem_usage=True`
           A_					,			A_						=      UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""		,				output_loading_info=UpperCamelCase__       )
           model_accelerate.to(UpperCamelCase__       )
           model_accelerate.eval()

           A_						=      torch.randn(
               1		,				model_accelerate.config.in_channels		,				model_accelerate.config.sample_size		,				model_accelerate.config.sample_size		,				generator=torch.manual_seed(0       )		,				)
           A_						=      noise.to(UpperCamelCase__       )
           A_						=      torch.tensor([10] * noise.shape[0]       ).to(UpperCamelCase__       )

           A_						=      model_accelerate(UpperCamelCase__		,				UpperCamelCase__       )["""sample"""]

           # two models don't need to stay in the device at the same time
           del model_accelerate
           torch.cuda.empty_cache()
           gc.collect()

           A_					,			A_						=      UNetaDModel.from_pretrained(
               """fusing/unet-ldm-dummy-update"""		,				output_loading_info=UpperCamelCase__		,				low_cpu_mem_usage=UpperCamelCase__       )
           model_normal_load.to(UpperCamelCase__       )
           model_normal_load.eval()
           A_						=      model_normal_load(UpperCamelCase__		,				UpperCamelCase__       )["""sample"""]

           assert torch_all_close(UpperCamelCase__		,				UpperCamelCase__		,				rtol=1e-3       )







    def 					snake_case_			(					self       )						->      Dict:





           '''simple docstring'''
           A_						=      UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""       )
           model.eval()
           model.to(UpperCamelCase__       )

           A_						=      torch.randn(
               1		,				model.config.in_channels		,				model.config.sample_size		,				model.config.sample_size		,				generator=torch.manual_seed(0       )		,				)
           A_						=      noise.to(UpperCamelCase__       )
           A_						=      torch.tensor([10] * noise.shape[0]       ).to(UpperCamelCase__       )

           with torch.no_grad():
                  A_						=      model(UpperCamelCase__		,				UpperCamelCase__       ).sample

           A_						=      output[0, -1, -3:, -3:].flatten().cpu()
           # fmt: off
           A_						=      torch.tensor([-13.3258, -20.1100, -15.9873, -17.6617, -23.0596, -17.9419, -13.3675, -16.1889, -12.3800]       )
           # fmt: on

           self.assertTrue(torch_all_close(UpperCamelCase__		,				UpperCamelCase__		,				rtol=1e-3       )       )


class 					A__      (      _snake_case						,						_snake_case						,						unittest.TestCase			):
    lowercase						=		UNetaDModel
    lowercase						=		"sample"



    @property
    def 					snake_case_			(					self		,				UpperCamelCase__=(32, 32)       )						->      Union[str, Any]:





           '''simple docstring'''
           A_						=      4
           A_						=      3

           A_						=      floats_tensor((batch_size, num_channels) + sizes       ).to(UpperCamelCase__       )
           A_						=      torch.tensor(batch_size * [10]       ).to(dtype=torch.intaa		,				device=UpperCamelCase__       )

           return {"sample": noise, "timestep": time_step}



    @property
    def 					snake_case_			(					self       )						->      int:





           '''simple docstring'''
           return (3, 32, 32)



    @property
    def 					snake_case_			(					self       )						->      List[Any]:





           '''simple docstring'''
           return (3, 32, 32)



    def 					snake_case_			(					self       )						->      int:





           '''simple docstring'''
           A_						=      {
               """block_out_channels""": [32, 64, 64, 64],
               """in_channels""": 3,
               """layers_per_block""": 1,
               """out_channels""": 3,
               """time_embedding_type""": """fourier""",
               """norm_eps""": 1e-6,
               """mid_block_scale_factor""": math.sqrt(2.0       ),
               """norm_num_groups""": None,
               """down_block_types""": [
                   """SkipDownBlock2D""",
                   """AttnSkipDownBlock2D""",
                   """SkipDownBlock2D""",
                   """SkipDownBlock2D""",
               ],
               """up_block_types""": [
                   """SkipUpBlock2D""",
                   """SkipUpBlock2D""",
                   """AttnSkipUpBlock2D""",
                   """SkipUpBlock2D""",
               ],
           }
           A_						=      self.dummy_input
           return init_dict, inputs_dict



    @slow
    def 					snake_case_			(					self       )						->      Optional[int]:





           '''simple docstring'''
           A_					,			A_						=      UNetaDModel.from_pretrained("""google/ncsnpp-celebahq-256"""		,				output_loading_info=UpperCamelCase__       )
           self.assertIsNotNone(UpperCamelCase__       )
           self.assertEqual(len(loading_info["""missing_keys"""]       )		,				0       )

           model.to(UpperCamelCase__       )
           A_						=      self.dummy_input
           A_						=      floats_tensor((4, 3) + (256, 256)       ).to(UpperCamelCase__       )
           A_						=      noise
           A_						=      model(**UpperCamelCase__       )

           assert image is not None, "Make sure output is not None"



    @slow
    def 					snake_case_			(					self       )						->      List[str]:





           '''simple docstring'''
           A_						=      UNetaDModel.from_pretrained("""google/ncsnpp-celebahq-256"""       )
           model.to(UpperCamelCase__       )

           A_						=      4
           A_						=      3
           A_						=      (256, 256)

           A_						=      torch.ones((batch_size, num_channels) + sizes       ).to(UpperCamelCase__       )
           A_						=      torch.tensor(batch_size * [1e-4]       ).to(UpperCamelCase__       )

           with torch.no_grad():
                  A_						=      model(UpperCamelCase__		,				UpperCamelCase__       ).sample

           A_						=      output[0, -3:, -3:, -1].flatten().cpu()
           # fmt: off
           A_						=      torch.tensor([-4842.8691, -6499.6631, -3800.1953, -7978.2686, -10980.7129, -20028.8535, 8148.2822, 2342.2905, 567.7608]       )
           # fmt: on

           self.assertTrue(torch_all_close(UpperCamelCase__		,				UpperCamelCase__		,				rtol=1e-2       )       )



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





           '''simple docstring'''
           A_						=      UNetaDModel.from_pretrained("""fusing/ncsnpp-ffhq-ve-dummy-update"""       )
           model.to(UpperCamelCase__       )

           A_						=      4
           A_						=      3
           A_						=      (32, 32)

           A_						=      torch.ones((batch_size, num_channels) + sizes       ).to(UpperCamelCase__       )
           A_						=      torch.tensor(batch_size * [1e-4]       ).to(UpperCamelCase__       )

           with torch.no_grad():
                  A_						=      model(UpperCamelCase__		,				UpperCamelCase__       ).sample

           A_						=      output[0, -3:, -3:, -1].flatten().cpu()
           # fmt: off
           A_						=      torch.tensor([-0.0325, -0.0900, -0.0869, -0.0332, -0.0725, -0.0270, -0.0101, 0.0227, 0.0256]       )
           # fmt: on

           self.assertTrue(torch_all_close(UpperCamelCase__		,				UpperCamelCase__		,				rtol=1e-2       )       )







    def 					snake_case_			(					self       )						->      int:





           '''simple docstring'''
           # not required for this model
           pass



'''simple docstring'''







from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block


@dataclass
class 					A__      (      _snake_case			):
    lowercase						=		42


class 					A__      (      nn.Module			):



    def __init__(					self		,				UpperCamelCase__=3		,				UpperCamelCase__=3		,				UpperCamelCase__=("DownEncoderBlock2D",)		,				UpperCamelCase__=(64,)		,				UpperCamelCase__=2		,				UpperCamelCase__=32		,				UpperCamelCase__="silu"		,				UpperCamelCase__=True		,				)						->      Union[str, Any]:





           '''simple docstring'''
           super().__init__()
           A_						=      layers_per_block

           A_						=      torch.nn.Convad(
               UpperCamelCase__		,				block_out_channels[0]		,				kernel_size=3		,				stride=1		,				padding=1		,				)

           A_						=      None
           A_						=      nn.ModuleList([]       )

           # down
           A_						=      block_out_channels[0]
           for i, down_block_type in enumerate(UpperCamelCase__       ):
                  A_						=      output_channel
                  A_						=      block_out_channels[i]
                  A_						=      i == len(UpperCamelCase__       ) - 1

                  A_						=      get_down_block(
                      UpperCamelCase__		,				num_layers=self.layers_per_block		,				in_channels=UpperCamelCase__		,				out_channels=UpperCamelCase__		,				add_downsample=not is_final_block		,				resnet_eps=1e-6		,				downsample_padding=0		,				resnet_act_fn=UpperCamelCase__		,				resnet_groups=UpperCamelCase__		,				attention_head_dim=UpperCamelCase__		,				temb_channels=UpperCamelCase__		,				)
                  self.down_blocks.append(UpperCamelCase__       )

           # mid
           A_						=      UNetMidBlockaD(
               in_channels=block_out_channels[-1]		,				resnet_eps=1e-6		,				resnet_act_fn=UpperCamelCase__		,				output_scale_factor=1		,				resnet_time_scale_shift="""default"""		,				attention_head_dim=block_out_channels[-1]		,				resnet_groups=UpperCamelCase__		,				temb_channels=UpperCamelCase__		,				)

           # out
           A_						=      nn.GroupNorm(num_channels=block_out_channels[-1]		,				num_groups=UpperCamelCase__		,				eps=1e-6       )
           A_						=      nn.SiLU()

           A_						=      2 * out_channels if double_z else out_channels
           A_						=      nn.Convad(block_out_channels[-1]		,				UpperCamelCase__		,				3		,				padding=1       )

           A_						=      False







    def 					snake_case_			(					self		,				UpperCamelCase__       )						->      str:





           '''simple docstring'''
           A_						=      x
           A_						=      self.conv_in(UpperCamelCase__       )

           if self.training and self.gradient_checkpointing:

                  def create_custom_forward(UpperCamelCase__       ):
                         def custom_forward(*UpperCamelCase__       ):
                                return module(*UpperCamelCase__       )

                         return custom_forward

                  # down
                  if is_torch_version(""">="""		,				"""1.11.0"""       ):
                         for down_block in self.down_blocks:
                                A_						=      torch.utils.checkpoint.checkpoint(
                                    create_custom_forward(UpperCamelCase__       )		,				UpperCamelCase__		,				use_reentrant=UpperCamelCase__       )
                         # middle
                         A_						=      torch.utils.checkpoint.checkpoint(
                             create_custom_forward(self.mid_block       )		,				UpperCamelCase__		,				use_reentrant=UpperCamelCase__       )
                  else:
                         for down_block in self.down_blocks:
                                A_						=      torch.utils.checkpoint.checkpoint(create_custom_forward(UpperCamelCase__       )		,				UpperCamelCase__       )
                         # middle
                         A_						=      torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block       )		,				UpperCamelCase__       )

           else:
                  # down
                  for down_block in self.down_blocks:
                         A_						=      down_block(UpperCamelCase__       )

                  # middle
                  A_						=      self.mid_block(UpperCamelCase__       )

           # post-process
           A_						=      self.conv_norm_out(UpperCamelCase__       )
           A_						=      self.conv_act(UpperCamelCase__       )
           A_						=      self.conv_out(UpperCamelCase__       )

           return sample


class 					A__      (      nn.Module			):



    def __init__(					self		,				UpperCamelCase__=3		,				UpperCamelCase__=3		,				UpperCamelCase__=("UpDecoderBlock2D",)		,				UpperCamelCase__=(64,)		,				UpperCamelCase__=2		,				UpperCamelCase__=32		,				UpperCamelCase__="silu"		,				UpperCamelCase__="group"		,				)						->      List[Any]:





           '''simple docstring'''
           super().__init__()
           A_						=      layers_per_block

           A_						=      nn.Convad(
               UpperCamelCase__		,				block_out_channels[-1]		,				kernel_size=3		,				stride=1		,				padding=1		,				)

           A_						=      None
           A_						=      nn.ModuleList([]       )

           A_						=      in_channels if norm_type == """spatial""" else None

           # mid
           A_						=      UNetMidBlockaD(
               in_channels=block_out_channels[-1]		,				resnet_eps=1e-6		,				resnet_act_fn=UpperCamelCase__		,				output_scale_factor=1		,				resnet_time_scale_shift="""default""" if norm_type == """group""" else norm_type		,				attention_head_dim=block_out_channels[-1]		,				resnet_groups=UpperCamelCase__		,				temb_channels=UpperCamelCase__		,				)

           # up
           A_						=      list(reversed(UpperCamelCase__       )       )
           A_						=      reversed_block_out_channels[0]
           for i, up_block_type in enumerate(UpperCamelCase__       ):
                  A_						=      output_channel
                  A_						=      reversed_block_out_channels[i]

                  A_						=      i == len(UpperCamelCase__       ) - 1

                  A_						=      get_up_block(
                      UpperCamelCase__		,				num_layers=self.layers_per_block + 1		,				in_channels=UpperCamelCase__		,				out_channels=UpperCamelCase__		,				prev_output_channel=UpperCamelCase__		,				add_upsample=not is_final_block		,				resnet_eps=1e-6		,				resnet_act_fn=UpperCamelCase__		,				resnet_groups=UpperCamelCase__		,				attention_head_dim=UpperCamelCase__		,				temb_channels=UpperCamelCase__		,				resnet_time_scale_shift=UpperCamelCase__		,				)
                  self.up_blocks.append(UpperCamelCase__       )
                  A_						=      output_channel

           # out
           if norm_type == "spatial":
                  A_						=      SpatialNorm(block_out_channels[0]		,				UpperCamelCase__       )
           else:
                  A_						=      nn.GroupNorm(num_channels=block_out_channels[0]		,				num_groups=UpperCamelCase__		,				eps=1e-6       )
           A_						=      nn.SiLU()
           A_						=      nn.Convad(block_out_channels[0]		,				UpperCamelCase__		,				3		,				padding=1       )

           A_						=      False







    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__=None       )						->      Optional[Any]:





           '''simple docstring'''
           A_						=      z
           A_						=      self.conv_in(UpperCamelCase__       )

           A_						=      next(iter(self.up_blocks.parameters()       )       ).dtype
           if self.training and self.gradient_checkpointing:

                  def create_custom_forward(UpperCamelCase__       ):
                         def custom_forward(*UpperCamelCase__       ):
                                return module(*UpperCamelCase__       )

                         return custom_forward

                  if is_torch_version(""">="""		,				"""1.11.0"""       ):
                         # middle
                         A_						=      torch.utils.checkpoint.checkpoint(
                             create_custom_forward(self.mid_block       )		,				UpperCamelCase__		,				UpperCamelCase__		,				use_reentrant=UpperCamelCase__       )
                         A_						=      sample.to(UpperCamelCase__       )

                         # up
                         for up_block in self.up_blocks:
                                A_						=      torch.utils.checkpoint.checkpoint(
                                    create_custom_forward(UpperCamelCase__       )		,				UpperCamelCase__		,				UpperCamelCase__		,				use_reentrant=UpperCamelCase__       )
                  else:
                         # middle
                         A_						=      torch.utils.checkpoint.checkpoint(
                             create_custom_forward(self.mid_block       )		,				UpperCamelCase__		,				UpperCamelCase__       )
                         A_						=      sample.to(UpperCamelCase__       )

                         # up
                         for up_block in self.up_blocks:
                                A_						=      torch.utils.checkpoint.checkpoint(create_custom_forward(UpperCamelCase__       )		,				UpperCamelCase__		,				UpperCamelCase__       )
           else:
                  # middle
                  A_						=      self.mid_block(UpperCamelCase__		,				UpperCamelCase__       )
                  A_						=      sample.to(UpperCamelCase__       )

                  # up
                  for up_block in self.up_blocks:
                         A_						=      up_block(UpperCamelCase__		,				UpperCamelCase__       )

        # post-process
           if latent_embeds is None:
                  A_						=      self.conv_norm_out(UpperCamelCase__       )
           else:
                  A_						=      self.conv_norm_out(UpperCamelCase__		,				UpperCamelCase__       )
           A_						=      self.conv_act(UpperCamelCase__       )
           A_						=      self.conv_out(UpperCamelCase__       )

           return sample


class 					A__      (      nn.Module			):



    def __init__(					self		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__=None		,				UpperCamelCase__="random"		,				UpperCamelCase__=False		,				UpperCamelCase__=True       )						->      str:





           '''simple docstring'''
           super().__init__()
           A_						=      n_e
           A_						=      vq_embed_dim
           A_						=      beta
           A_						=      legacy

           A_						=      nn.Embedding(self.n_e		,				self.vq_embed_dim       )
           self.embedding.weight.data.uniform_(-1.0 / self.n_e		,				1.0 / self.n_e       )

           A_						=      remap
           if self.remap is not None:
                  self.register_buffer("""used"""		,				torch.tensor(np.load(self.remap       )       )       )
                  A_						=      self.used.shape[0]
                  A_						=      unknown_index  # "random" or "extra" or integer
                  if self.unknown_index == "extra":
                         A_						=      self.re_embed
                         A_						=      self.re_embed + 1
                  print(
                      f'''Remapping {self.n_e} indices to {self.re_embed} indices. '''
                      f'''Using {self.unknown_index} for unknown indices.'''       )
           else:
                  A_						=      n_e

           A_						=      sane_index_shape



    def 					snake_case_			(					self		,				UpperCamelCase__       )						->      Optional[Any]:





           '''simple docstring'''
           A_						=      inds.shape
           assert len(UpperCamelCase__       ) > 1
           A_						=      inds.reshape(ishape[0]		,				-1       )
           A_						=      self.used.to(UpperCamelCase__       )
           A_						=      (inds[:, :, None] == used[None, None, ...]).long()
           A_						=      match.argmax(-1       )
           A_						=      match.sum(2       ) < 1
           if self.unknown_index == "random":
                  A_						=      torch.randint(0		,				self.re_embed		,				size=new[unknown].shape       ).to(device=new.device       )
           else:
                  A_						=      self.unknown_index
           return new.reshape(UpperCamelCase__       )



    def 					snake_case_			(					self		,				UpperCamelCase__       )						->      Dict:





           '''simple docstring'''
           A_						=      inds.shape
           assert len(UpperCamelCase__       ) > 1
           A_						=      inds.reshape(ishape[0]		,				-1       )
           A_						=      self.used.to(UpperCamelCase__       )
           if self.re_embed > self.used.shape[0]:  # extra token
                  A_						=      0  # simply set to zero
           A_						=      torch.gather(used[None, :][inds.shape[0] * [0], :]		,				1		,				UpperCamelCase__       )
           return back.reshape(UpperCamelCase__       )



    def 					snake_case_			(					self		,				UpperCamelCase__       )						->      str:





           '''simple docstring'''
           # reshape z -> (batch, height, width, channel) and flatten
           A_						=      z.permute(0		,				2		,				3		,				1       ).contiguous()
           A_						=      z.view(-1		,				self.vq_embed_dim       )

           # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
           A_						=      torch.argmin(torch.cdist(UpperCamelCase__		,				self.embedding.weight       )		,				dim=1       )

           A_						=      self.embedding(UpperCamelCase__       ).view(z.shape       )
           A_						=      None
           A_						=      None

           # compute loss for embedding
           if not self.legacy:
                  A_						=      self.beta * torch.mean((z_q.detach() - z) ** 2       ) + torch.mean((z_q - z.detach()) ** 2       )
           else:
                  A_						=      torch.mean((z_q.detach() - z) ** 2       ) + self.beta * torch.mean((z_q - z.detach()) ** 2       )

           # preserve gradients
           A_						=      z + (z_q - z).detach()

           # reshape back to match original input shape
           A_						=      z_q.permute(0		,				3		,				1		,				2       ).contiguous()

           if self.remap is not None:
                  A_						=      min_encoding_indices.reshape(z.shape[0]		,				-1       )  # add batch axis
                  A_						=      self.remap_to_used(UpperCamelCase__       )
                  A_						=      min_encoding_indices.reshape(-1		,				1       )  # flatten

           if self.sane_index_shape:
                  A_						=      min_encoding_indices.reshape(z_q.shape[0]		,				z_q.shape[2]		,				z_q.shape[3]       )

           return z_q, loss, (perplexity, min_encodings, min_encoding_indices)







    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__       )						->      Optional[int]:





           '''simple docstring'''
           # shape specifying (batch, height, width, channel)
           if self.remap is not None:
                  A_						=      indices.reshape(shape[0]		,				-1       )  # add batch axis
                  A_						=      self.unmap_to_all(UpperCamelCase__       )
                  A_						=      indices.reshape(-1       )  # flatten again

           # get quantized latent vectors
           A_						=      self.embedding(UpperCamelCase__       )

           if shape is not None:
                  A_						=      z_q.view(UpperCamelCase__       )
                  # reshape back to match original input shape
                  A_						=      z_q.permute(0		,				3		,				1		,				2       ).contiguous()

           return z_q


class 					A__      (      _snake_case			):



    def __init__(					self		,				UpperCamelCase__		,				UpperCamelCase__=False       )						->      Dict:





           '''simple docstring'''
           A_						=      parameters
           A_					,			A_						=      torch.chunk(UpperCamelCase__		,				2		,				dim=1       )
           A_						=      torch.clamp(self.logvar		,				-30.0		,				20.0       )
           A_						=      deterministic
           A_						=      torch.exp(0.5 * self.logvar       )
           A_						=      torch.exp(self.logvar       )
           if self.deterministic:
                  A_						=      A_						=      torch.zeros_like(
                      self.mean		,				device=self.parameters.device		,				dtype=self.parameters.dtype       )



    def 					snake_case_			(					self		,				UpperCamelCase__ = None       )						->      torch.FloatTensor:





           '''simple docstring'''
           # make sure sample is on the same device as the parameters and has same dtype
           A_						=      randn_tensor(
               self.mean.shape		,				generator=UpperCamelCase__		,				device=self.parameters.device		,				dtype=self.parameters.dtype       )
           A_						=      self.mean + self.std * sample
           return x



    def 					snake_case_			(					self		,				UpperCamelCase__=None       )						->      int:





           '''simple docstring'''
           if self.deterministic:
                  return torch.Tensor([0.0]       )
           else:
                  if other is None:
                         return 0.5 * torch.sum(torch.pow(self.mean		,				2       ) + self.var - 1.0 - self.logvar		,				dim=[1, 2, 3]       )
                  else:
                         return 0.5 * torch.sum(
                             torch.pow(self.mean - other.mean		,				2       ) / other.var
                             + self.var / other.var
                             - 1.0
                             - self.logvar
                             + other.logvar		,				dim=[1, 2, 3]		,				)



    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__=[1, 2, 3]       )						->      Optional[Any]:





           '''simple docstring'''
           if self.deterministic:
                  return torch.Tensor([0.0]       )
           A_						=      np.log(2.0 * np.pi       )
           return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean		,				2       ) / self.var		,				dim=UpperCamelCase__       )







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





           '''simple docstring'''
           return self.mean
