



import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyInpaintPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()



class        _SCREAMING_SNAKE_CASE		(       _lowerCAmelCase			,			unittest.TestCase ):
   a_							:      Union[str, Any]						= KandinskyInpaintPipeline
   a_							:      Tuple						= ["prompt", "image_embeds", "negative_image_embeds", "image", "mask_image"]
   a_							:      Dict						= [
       "prompt",
       "negative_prompt",
       "image_embeds",
       "negative_image_embeds",
       "image",
       "mask_image",
   ]
   a_							:      Optional[Any]						= [
       "generator",
       "height",
       "width",
       "latents",
       "guidance_scale",
       "negative_prompt",
       "num_inference_steps",
       "return_dict",
       "guidance_scale",
       "num_images_per_prompt",
       "output_type",
       "return_dict",
   ]
   a_							:      Optional[Any]						= False





   @property
   def 	A__      (self):
       '''simple docstring'''







       return 3_2





   @property
   def 	A__      (self):
       '''simple docstring'''







       return 3_2





   @property
   def 	A__      (self):
       '''simple docstring'''







       return self.time_input_dim





   @property
   def 	A__      (self):
       '''simple docstring'''







       return self.time_input_dim * 4





   @property
   def 	A__      (self):
       '''simple docstring'''







       return 1_0_0





   @property
   def 	A__      (self):
       '''simple docstring'''







       __UpperCAmelCase			 =XLMRobertaTokenizerFast.from_pretrained('''YiYiXu/tiny-random-mclip-base''')
       return tokenizer





   @property
   def 	A__      (self):
       '''simple docstring'''







       torch.manual_seed(0)
       __UpperCAmelCase			 =MCLIPConfig(
           numDims=self.cross_attention_dim    ,  transformerDimensions=self.text_embedder_hidden_size    ,  hidden_size=self.text_embedder_hidden_size    ,  intermediate_size=3_7    ,  num_attention_heads=4    ,  num_hidden_layers=5    ,  vocab_size=1_0_0_5    ,  )

       __UpperCAmelCase			 =MultilingualCLIP(snake_case__)
       __UpperCAmelCase			 =text_encoder.eval()

       return text_encoder





   @property
   def 	A__      (self):
       '''simple docstring'''







       torch.manual_seed(0)

       __UpperCAmelCase			 ={
           "in_channels": 9,
           # Out channels is double in channels because predicts mean and variance
           "out_channels": 8,
           "addition_embed_type": "text_image",
           "down_block_types": ("ResnetDownsampleBlock2D", "SimpleCrossAttnDownBlock2D"),
           "up_block_types": ("SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"),
           "mid_block_type": "UNetMidBlock2DSimpleCrossAttn",
           "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
           "layers_per_block": 1,
           "encoder_hid_dim": self.text_embedder_hidden_size,
           "encoder_hid_dim_type": "text_image_proj",
           "cross_attention_dim": self.cross_attention_dim,
           "attention_head_dim": 4,
           "resnet_time_scale_shift": "scale_shift",
           "class_embed_type": None,
       }

       __UpperCAmelCase			 =UNetaDConditionModel(**snake_case__)
       return model





   @property
   def 	A__      (self):
       '''simple docstring'''







       return {
           "block_out_channels": [3_2, 6_4],
           "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
           "in_channels": 3,
           "latent_channels": 4,
           "layers_per_block": 1,
           "norm_num_groups": 8,
           "norm_type": "spatial",
           "num_vq_embeddings": 1_2,
           "out_channels": 3,
           "up_block_types": [
               "AttnUpDecoderBlock2D",
               "UpDecoderBlock2D",
           ],
           "vq_embed_dim": 4,
       }





   @property
   def 	A__      (self):
       '''simple docstring'''







       torch.manual_seed(0)
       __UpperCAmelCase			 =VQModel(**self.dummy_movq_kwargs)
       return model





   def 	A__      (self):
       '''simple docstring'''







       __UpperCAmelCase			 =self.dummy_text_encoder
       __UpperCAmelCase			 =self.dummy_tokenizer
       __UpperCAmelCase			 =self.dummy_unet
       __UpperCAmelCase			 =self.dummy_movq

       __UpperCAmelCase			 =DDIMScheduler(
           num_train_timesteps=1_0_0_0    ,  beta_schedule='''linear'''    ,  beta_start=0.0_0085    ,  beta_end=0.012    ,  clip_sample=snake_case__    ,  set_alpha_to_one=snake_case__    ,  steps_offset=1    ,  prediction_type='''epsilon'''    ,  thresholding=snake_case__    ,  )

       __UpperCAmelCase			 ={
           "text_encoder": text_encoder,
           "tokenizer": tokenizer,
           "unet": unet,
           "scheduler": scheduler,
           "movq": movq,
       }

       return components





   def 	A__      (self    ,  UpperCAmelCase    ,  UpperCAmelCase=0):
       '''simple docstring'''







       __UpperCAmelCase			 =floats_tensor((1, self.cross_attention_dim)    ,  rng=random.Random(snake_case__)).to(snake_case__)
       __UpperCAmelCase			 =floats_tensor((1, self.cross_attention_dim)    ,  rng=random.Random(seed + 1)).to(snake_case__)
       # create init_image
       __UpperCAmelCase			 =floats_tensor((1, 3, 6_4, 6_4)    ,  rng=random.Random(snake_case__)).to(snake_case__)
       __UpperCAmelCase			 =image.cpu().permute(0    ,  2    ,  3    ,  1)[0]
       __UpperCAmelCase			 =Image.fromarray(np.uinta(snake_case__)).convert('''RGB''').resize((2_5_6, 2_5_6))
       # create mask
       __UpperCAmelCase			 =np.ones((6_4, 6_4)    ,  dtype=np.floataa)
       __UpperCAmelCase			 =0

       if str(snake_case__).startswith('''mps'''):
           __UpperCAmelCase			 =torch.manual_seed(snake_case__)
       else:
           __UpperCAmelCase			 =torch.Generator(device=snake_case__).manual_seed(snake_case__)
       __UpperCAmelCase			 ={
           "prompt": "horse",
           "image": init_image,
           "mask_image": mask,
           "image_embeds": image_embeds,
           "negative_image_embeds": negative_image_embeds,
           "generator": generator,
           "height": 6_4,
           "width": 6_4,
           "num_inference_steps": 2,
           "guidance_scale": 4.0,
           "output_type": "np",
       }
       return inputs





   def 	A__      (self):
       '''simple docstring'''







       __UpperCAmelCase			 ="cpu"

       __UpperCAmelCase			 =self.get_dummy_components()

       __UpperCAmelCase			 =self.pipeline_class(**snake_case__)
       __UpperCAmelCase			 =pipe.to(snake_case__)

       pipe.set_progress_bar_config(disable=snake_case__)

       __UpperCAmelCase			 =pipe(**self.get_dummy_inputs(snake_case__))
       __UpperCAmelCase			 =output.images

       __UpperCAmelCase			 =pipe(
           **self.get_dummy_inputs(snake_case__)    ,  return_dict=snake_case__    ,  )[0]

       __UpperCAmelCase			 =image[0, -3:, -3:, -1]
       __UpperCAmelCase			 =image_from_tuple[0, -3:, -3:, -1]

       print(f"""image.shape {image.shape}""")

       assert image.shape == (1, 6_4, 6_4, 3)

       __UpperCAmelCase			 =np.array(
           [0.832_6919, 0.7379_0467, 0.2091_8581, 0.930_9612, 0.551_1791, 0.4371_3328, 0.551_3321, 0.4992_2934, 0.5949_7786])

       assert (
           np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
       ), f""" expected_slice {expected_slice}, but got {image_slice.flatten()}"""
       assert (
           np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1e-2
       ), f""" expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}"""






   def 	A__      (self):
       '''simple docstring'''







       super().test_inference_batch_single_identical(expected_max_diff=3e-3)







@slow
@require_torch_gpu
class        _SCREAMING_SNAKE_CASE		(       unittest.TestCase ):





   def 	A__      (self):
       '''simple docstring'''







       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()






   def 	A__      (self):
       '''simple docstring'''







       __UpperCAmelCase			 =load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/kandinsky/kandinsky_inpaint_cat_with_hat_fp16.npy''')

       __UpperCAmelCase			 =load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/kandinsky/cat.png''')
       __UpperCAmelCase			 =np.ones((7_6_8, 7_6_8)    ,  dtype=np.floataa)
       __UpperCAmelCase			 =0

       __UpperCAmelCase			 ="a hat"

       __UpperCAmelCase			 =KandinskyPriorPipeline.from_pretrained(
           '''kandinsky-community/kandinsky-2-1-prior'''    ,  torch_dtype=torch.floataa)
       pipe_prior.to(snake_case__)

       __UpperCAmelCase			 =KandinskyInpaintPipeline.from_pretrained(
           '''kandinsky-community/kandinsky-2-1-inpaint'''    ,  torch_dtype=torch.floataa)
       __UpperCAmelCase			 =pipeline.to(snake_case__)
       pipeline.set_progress_bar_config(disable=snake_case__)

       __UpperCAmelCase			 =torch.Generator(device='''cpu''').manual_seed(0)
       __UpperCAmelCase			 =pipe_prior(
           snake_case__    ,  generator=snake_case__    ,  num_inference_steps=5    ,  negative_prompt=''''''    ,  ).to_tuple()

       __UpperCAmelCase			 =pipeline(
           snake_case__    ,  image=snake_case__    ,  mask_image=snake_case__    ,  image_embeds=snake_case__    ,  negative_image_embeds=snake_case__    ,  generator=snake_case__    ,  num_inference_steps=1_0_0    ,  height=7_6_8    ,  width=7_6_8    ,  output_type='''np'''    ,  )

       __UpperCAmelCase			 =output.images[0]

       assert image.shape == (7_6_8, 7_6_8, 3)

       assert_mean_pixel_difference(snake_case__    ,  snake_case__)





'''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 							UpperCAmelCase_							(    A				):

 '''simple docstring'''
 lowercase_					:    torch.FloatTensor






class 							UpperCAmelCase_							(    nn.Module				):

 '''simple docstring'''







 def __init__(					self		:  Union[str, Any]			,				snake_case__		:  Tuple=3			,				snake_case__		:  Dict=3			,				snake_case__		:  Dict=("DownEncoderBlock2D",)			,				snake_case__		:  Optional[Any]=(64,)			,				snake_case__		:  List[Any]=2			,				snake_case__		:  Any=32			,				snake_case__		:  Tuple="silu"			,				snake_case__		:  Tuple=True			,				):



   '''simple docstring'''




   super().__init__()
   UpperCAmelCase__      :					Tuple				       =							layers_per_block

   UpperCAmelCase__      :					Optional[int]				       =							torch.nn.Convad(
       snake_case__			,				block_out_channels[0]			,				kernel_size=3			,				stride=1			,				padding=1			,				)

   UpperCAmelCase__      :					Tuple				       =							None
   UpperCAmelCase__      :					Union[str, Any]				       =							nn.ModuleList([] )

   # down
   UpperCAmelCase__      :					Any				       =							block_out_channels[0]
   for i, down_block_type in enumerate(snake_case__ ):
     UpperCAmelCase__      :					List[str]				       =							output_channel
     UpperCAmelCase__      :					Dict				       =							block_out_channels[i]
     UpperCAmelCase__      :					Tuple				       =							i == len(snake_case__ ) - 1

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

   # mid
   UpperCAmelCase__      :					Optional[Any]				       =							UNetMidBlockaD(
       in_channels=block_out_channels[-1]			,				resnet_eps=1e-6			,				resnet_act_fn=snake_case__			,				output_scale_factor=1			,				resnet_time_scale_shift="default"			,				attention_head_dim=block_out_channels[-1]			,				resnet_groups=snake_case__			,				temb_channels=snake_case__			,				)

   # out
   UpperCAmelCase__      :					Tuple				       =							nn.GroupNorm(num_channels=block_out_channels[-1]			,				num_groups=snake_case__			,				eps=1e-6 )
   UpperCAmelCase__      :					Union[str, Any]				       =							nn.SiLU()

   UpperCAmelCase__      :					Dict				       =							2 * out_channels if double_z else out_channels
   UpperCAmelCase__      :					Union[str, Any]				       =							nn.Convad(block_out_channels[-1]			,				snake_case__			,				3			,				padding=1 )

   UpperCAmelCase__      :					Union[str, Any]				       =							False







 def 							UpperCamelCase  (					self		:  int			,				snake_case__		:  Tuple ):



   '''simple docstring'''




   UpperCAmelCase__      :					List[str]				       =							x
   UpperCAmelCase__      :					Dict				       =							self.conv_in(snake_case__ )

   if self.training and self.gradient_checkpointing:

     def create_custom_forward(snake_case__		:  Dict ):
       def custom_forward(*snake_case__		:  List[str] ):
         return module(*snake_case__ )

       return custom_forward

     # down
     if is_torch_version(">="			,				"1.11.0" ):
       for down_block in self.down_blocks:
         UpperCAmelCase__      :					List[str]				       =							torch.utils.checkpoint.checkpoint(
             create_custom_forward(snake_case__ )			,				snake_case__			,				use_reentrant=snake_case__ )
       # middle
       UpperCAmelCase__      :					int				       =							torch.utils.checkpoint.checkpoint(
           create_custom_forward(self.mid_block )			,				snake_case__			,				use_reentrant=snake_case__ )
     else:
       for down_block in self.down_blocks:
         UpperCAmelCase__      :					Optional[Any]				       =							torch.utils.checkpoint.checkpoint(create_custom_forward(snake_case__ )			,				snake_case__ )
       # middle
       UpperCAmelCase__      :					int				       =							torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block )			,				snake_case__ )

   else:
     # down
     for down_block in self.down_blocks:
       UpperCAmelCase__      :					Dict				       =							down_block(snake_case__ )

     # middle
     UpperCAmelCase__      :					Optional[Any]				       =							self.mid_block(snake_case__ )

   # post-process
   UpperCAmelCase__      :					Tuple				       =							self.conv_norm_out(snake_case__ )
   UpperCAmelCase__      :					List[Any]				       =							self.conv_act(snake_case__ )
   UpperCAmelCase__      :					Union[str, Any]				       =							self.conv_out(snake_case__ )

   return sample






class 							UpperCAmelCase_							(    nn.Module				):

 '''simple docstring'''







 def __init__(					self		:  str			,				snake_case__		:  int=3			,				snake_case__		:  str=3			,				snake_case__		:  Union[str, Any]=("UpDecoderBlock2D",)			,				snake_case__		:  Dict=(64,)			,				snake_case__		:  Optional[Any]=2			,				snake_case__		:  Dict=32			,				snake_case__		:  str="silu"			,				snake_case__		:  Any="group"			,				):



   '''simple docstring'''




   super().__init__()
   UpperCAmelCase__      :					Any				       =							layers_per_block

   UpperCAmelCase__      :					Any				       =							nn.Convad(
       snake_case__			,				block_out_channels[-1]			,				kernel_size=3			,				stride=1			,				padding=1			,				)

   UpperCAmelCase__      :					str				       =							None
   UpperCAmelCase__      :					Optional[int]				       =							nn.ModuleList([] )

   UpperCAmelCase__      :					str				       =							in_channels if norm_type == "spatial" else None

   # mid
   UpperCAmelCase__      :					Optional[Any]				       =							UNetMidBlockaD(
       in_channels=block_out_channels[-1]			,				resnet_eps=1e-6			,				resnet_act_fn=snake_case__			,				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=snake_case__			,				temb_channels=snake_case__			,				)

   # up
   UpperCAmelCase__      :					Tuple				       =							list(reversed(snake_case__ ) )
   UpperCAmelCase__      :					Optional[Any]				       =							reversed_block_out_channels[0]
   for i, up_block_type in enumerate(snake_case__ ):
     UpperCAmelCase__      :					Dict				       =							output_channel
     UpperCAmelCase__      :					List[Any]				       =							reversed_block_out_channels[i]

     UpperCAmelCase__      :					List[Any]				       =							i == len(snake_case__ ) - 1

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

   # out
   if norm_type == "spatial":
     UpperCAmelCase__      :					Optional[Any]				       =							SpatialNorm(block_out_channels[0]			,				snake_case__ )
   else:
     UpperCAmelCase__      :					Optional[int]				       =							nn.GroupNorm(num_channels=block_out_channels[0]			,				num_groups=snake_case__			,				eps=1e-6 )
   UpperCAmelCase__      :					Dict				       =							nn.SiLU()
   UpperCAmelCase__      :					Union[str, Any]				       =							nn.Convad(block_out_channels[0]			,				snake_case__			,				3			,				padding=1 )

   UpperCAmelCase__      :					Union[str, Any]				       =							False







 def 							UpperCamelCase  (					self		:  List[str]			,				snake_case__		:  List[str]			,				snake_case__		:  Tuple=None ):



   '''simple docstring'''




   UpperCAmelCase__      :					str				       =							z
   UpperCAmelCase__      :					List[str]				       =							self.conv_in(snake_case__ )

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

     def create_custom_forward(snake_case__		:  Dict ):
       def custom_forward(*snake_case__		:  List[Any] ):
         return module(*snake_case__ )

       return custom_forward

     if is_torch_version(">="			,				"1.11.0" ):
       # middle
       UpperCAmelCase__      :					str				       =							torch.utils.checkpoint.checkpoint(
           create_custom_forward(self.mid_block )			,				snake_case__			,				snake_case__			,				use_reentrant=snake_case__ )
       UpperCAmelCase__      :					List[Any]				       =							sample.to(snake_case__ )

       # up
       for up_block in self.up_blocks:
         UpperCAmelCase__      :					Tuple				       =							torch.utils.checkpoint.checkpoint(
             create_custom_forward(snake_case__ )			,				snake_case__			,				snake_case__			,				use_reentrant=snake_case__ )
     else:
       # middle
       UpperCAmelCase__      :					List[Any]				       =							torch.utils.checkpoint.checkpoint(
           create_custom_forward(self.mid_block )			,				snake_case__			,				snake_case__ )
       UpperCAmelCase__      :					int				       =							sample.to(snake_case__ )

       # up
       for up_block in self.up_blocks:
         UpperCAmelCase__      :					str				       =							torch.utils.checkpoint.checkpoint(create_custom_forward(snake_case__ )			,				snake_case__			,				snake_case__ )
   else:
     # middle
     UpperCAmelCase__      :					Union[str, Any]				       =							self.mid_block(snake_case__			,				snake_case__ )
     UpperCAmelCase__      :					Optional[Any]				       =							sample.to(snake_case__ )

     # up
     for up_block in self.up_blocks:
       UpperCAmelCase__      :					int				       =							up_block(snake_case__			,				snake_case__ )

        # post-process
   if latent_embeds is None:
     UpperCAmelCase__      :					List[Any]				       =							self.conv_norm_out(snake_case__ )
   else:
     UpperCAmelCase__      :					Any				       =							self.conv_norm_out(snake_case__			,				snake_case__ )
   UpperCAmelCase__      :					List[Any]				       =							self.conv_act(snake_case__ )
   UpperCAmelCase__      :					Optional[Any]				       =							self.conv_out(snake_case__ )

   return sample






class 							UpperCAmelCase_							(    nn.Module				):

 '''simple docstring'''







 def __init__(					self		:  str			,				snake_case__		:  str			,				snake_case__		:  int			,				snake_case__		:  Tuple			,				snake_case__		:  Union[str, Any]=None			,				snake_case__		:  Optional[Any]="random"			,				snake_case__		:  Any=False			,				snake_case__		:  Any=True ):



   '''simple docstring'''




   super().__init__()
   UpperCAmelCase__      :					Any				       =							n_e
   UpperCAmelCase__      :					str				       =							vq_embed_dim
   UpperCAmelCase__      :					List[Any]				       =							beta
   UpperCAmelCase__      :					List[Any]				       =							legacy

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

   UpperCAmelCase__      :					Optional[Any]				       =							remap
   if self.remap is not None:
     self.register_buffer("used"			,				torch.tensor(np.load(self.remap ) ) )
     UpperCAmelCase__      :					Optional[Any]				       =							self.used.shape[0]
     UpperCAmelCase__      :					Any				       =							unknown_index  # "random" or "extra" or integer
     if self.unknown_index == "extra":
       UpperCAmelCase__      :					Union[str, Any]				       =							self.re_embed
       UpperCAmelCase__      :					Optional[Any]				       =							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:
     UpperCAmelCase__      :					int				       =							n_e

   UpperCAmelCase__      :					Dict				       =							sane_index_shape







 def 							UpperCamelCase  (					self		:  Tuple			,				snake_case__		:  int ):



   '''simple docstring'''




   UpperCAmelCase__      :					Any				       =							inds.shape
   assert len(snake_case__ ) > 1
   UpperCAmelCase__      :					Tuple				       =							inds.reshape(ishape[0]			,				-1 )
   UpperCAmelCase__      :					Optional[Any]				       =							self.used.to(snake_case__ )
   UpperCAmelCase__      :					Union[str, Any]				       =							(inds[:, :, None] == used[None, None, ...]).long()
   UpperCAmelCase__      :					Optional[int]				       =							match.argmax(-1 )
   UpperCAmelCase__      :					str				       =							match.sum(2 ) < 1
   if self.unknown_index == "random":
     UpperCAmelCase__      :					Union[str, Any]				       =							torch.randint(0			,				self.re_embed			,				size=new[unknown].shape ).to(device=new.device )
   else:
     UpperCAmelCase__      :					Tuple				       =							self.unknown_index
   return new.reshape(snake_case__ )







 def 							UpperCamelCase  (					self		:  int			,				snake_case__		:  str ):



   '''simple docstring'''




   UpperCAmelCase__      :					List[str]				       =							inds.shape
   assert len(snake_case__ ) > 1
   UpperCAmelCase__      :					List[Any]				       =							inds.reshape(ishape[0]			,				-1 )
   UpperCAmelCase__      :					int				       =							self.used.to(snake_case__ )
   if self.re_embed > self.used.shape[0]:  # extra token
     UpperCAmelCase__      :					Any				       =							0  # simply set to zero
   UpperCAmelCase__      :					List[Any]				       =							torch.gather(used[None, :][inds.shape[0] * [0], :]			,				1			,				snake_case__ )
   return back.reshape(snake_case__ )







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



   '''simple docstring'''




   UpperCAmelCase__      :					List[str]				       =							z.permute(0			,				2			,				3			,				1 ).contiguous()
   UpperCAmelCase__      :					str				       =							z.view(-1			,				self.vq_embed_dim )

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

   UpperCAmelCase__      :					Any				       =							self.embedding(snake_case__ ).view(z.shape )
   UpperCAmelCase__      :					Union[str, Any]				       =							None
   UpperCAmelCase__      :					Any				       =							None

   # compute loss for embedding
   if not self.legacy:
     UpperCAmelCase__      :					Optional[int]				       =							self.beta * torch.mean((z_q.detach() - z) ** 2 ) + torch.mean((z_q - z.detach()) ** 2 )
   else:
     UpperCAmelCase__      :					Tuple				       =							torch.mean((z_q.detach() - z) ** 2 ) + self.beta * torch.mean((z_q - z.detach()) ** 2 )

   # preserve gradients
   UpperCAmelCase__      :					Any				       =							z + (z_q - z).detach()

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

   if self.remap is not None:
     UpperCAmelCase__      :					Dict				       =							min_encoding_indices.reshape(z.shape[0]			,				-1 )  # add batch axis
     UpperCAmelCase__      :					Optional[int]				       =							self.remap_to_used(snake_case__ )
     UpperCAmelCase__      :					Optional[Any]				       =							min_encoding_indices.reshape(-1			,				1 )  # flatten

   if self.sane_index_shape:
     UpperCAmelCase__      :					Union[str, Any]				       =							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 							UpperCamelCase  (					self		:  int			,				snake_case__		:  str			,				snake_case__		:  int ):



   '''simple docstring'''




   if self.remap is not None:
     UpperCAmelCase__      :					List[Any]				       =							indices.reshape(shape[0]			,				-1 )  # add batch axis
     UpperCAmelCase__      :					Dict				       =							self.unmap_to_all(snake_case__ )
     UpperCAmelCase__      :					Dict				       =							indices.reshape(-1 )  # flatten again

   # get quantized latent vectors
   UpperCAmelCase__      :					Optional[Any]				       =							self.embedding(snake_case__ )

   if shape is not None:
     UpperCAmelCase__      :					List[str]				       =							z_q.view(snake_case__ )
     # reshape back to match original input shape
     UpperCAmelCase__      :					List[Any]				       =							z_q.permute(0			,				3			,				1			,				2 ).contiguous()

   return z_q






class 							UpperCAmelCase_							(    A				):

 '''simple docstring'''







 def __init__(					self		:  int			,				snake_case__		:  Optional[Any]			,				snake_case__		:  List[str]=False ):



   '''simple docstring'''




   UpperCAmelCase__      :					Dict				       =							parameters
   UpperCAmelCase__       ,			UpperCAmelCase__      :					Optional[int]				       =							torch.chunk(snake_case__			,				2			,				dim=1 )
   UpperCAmelCase__      :					int				       =							torch.clamp(self.logvar			,				-30.0			,				20.0 )
   UpperCAmelCase__      :					Optional[Any]				       =							deterministic
   UpperCAmelCase__      :					Dict				       =							torch.exp(0.5 * self.logvar )
   UpperCAmelCase__      :					List[str]				       =							torch.exp(self.logvar )
   if self.deterministic:
     UpperCAmelCase__      :					Dict				       =							torch.zeros_like(
         self.mean			,				device=self.parameters.device			,				dtype=self.parameters.dtype )







 def 							UpperCamelCase  (					self		:  List[Any]			,				snake_case__		:  Optional[torch.Generator] = None ):



   '''simple docstring'''




   UpperCAmelCase__      :					int				       =							randn_tensor(
       self.mean.shape			,				generator=snake_case__			,				device=self.parameters.device			,				dtype=self.parameters.dtype )
   UpperCAmelCase__      :					List[str]				       =							self.mean + self.std * sample
   return x







 def 							UpperCamelCase  (					self		:  Optional[Any]			,				snake_case__		:  Any=None ):



   '''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 							UpperCamelCase  (					self		:  List[Any]			,				snake_case__		:  Tuple			,				snake_case__		:  Tuple=[1, 2, 3] ):



   '''simple docstring'''




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







 def 							UpperCamelCase  (					self		:  List[str] ):



   '''simple docstring'''




   return self.mean



