
"""simple docstring"""


import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion_safe import StableDiffusionPipelineSafe as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import require_torch_gpu







class 				lowerCAmelCase_				(      unittest.TestCase    ):



						'''simple docstring'''


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Any      )     ->   Any:
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()


						@property
						def 						_SCREAMING_SNAKE_CASE      ( self     :      Union[str, Any]      )     ->   Any:
							A						   =   1
							A						   =   3
							A						   =   (32, 32)

							A						   =   floats_tensor((batch_size, num_channels) + sizes  ,rng=random.Random(0      )      ).to(A_      )
							return image


						@property
						def 						_SCREAMING_SNAKE_CASE      ( self     :      Optional[int]      )     ->   Optional[Any]:
							torch.manual_seed(0      )
							A						   =   UNetaDConditionModel(
							    block_out_channels=(32, 64)  ,layers_per_block=2  ,sample_size=32  ,in_channels=4  ,out_channels=4  ,down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')  ,up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')  ,cross_attention_dim=32  ,)
							return model


						@property
						def 						_SCREAMING_SNAKE_CASE      ( self     :      Tuple      )     ->   Union[str, Any]:
							torch.manual_seed(0      )
							A						   =   AutoencoderKL(
							    block_out_channels=[32, 64]  ,in_channels=3  ,out_channels=3  ,down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']  ,up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']  ,latent_channels=4  ,)
							return model


						@property
						def 						_SCREAMING_SNAKE_CASE      ( self     :      Tuple      )     ->   Dict:
							torch.manual_seed(0      )
							A						   =   CLIPTextConfig(
							    bos_token_id=0  ,eos_token_id=2  ,hidden_size=32  ,intermediate_size=37  ,layer_norm_eps=1e-05  ,num_attention_heads=4  ,num_hidden_layers=5  ,pad_token_id=1  ,vocab_size=1000  ,)
							return CLIPTextModel(A_      )


						@property
						def 						_SCREAMING_SNAKE_CASE      ( self     :      List[Any]      )     ->   Optional[Any]:
							def extract(*A_     :      List[str]  ,**A_     :      int      ):







								class 				lowerCAmelCase_				:



														'''simple docstring'''


														def __init__( self     :      Optional[Any]      )     ->   Optional[Any]:
															A						   =   torch.ones([0]      )




														def 						_SCREAMING_SNAKE_CASE      ( self     :      str  ,A_     :      Union[str, Any]      )     ->   int:
															self.pixel_values.to(A_      )
															return self

								return Out()

							return extract


						def 						_SCREAMING_SNAKE_CASE      ( self     :      List[str]      )     ->   Dict:
							A						   =   'cpu'  # ensure determinism for the device-dependent torch.Generator
							A						   =   self.dummy_cond_unet
							A						   =   DDIMScheduler(
							    beta_start=0.0_00_85  ,beta_end=0.0_12  ,beta_schedule='scaled_linear'  ,clip_sample=A_  ,set_alpha_to_one=A_  ,)

							A						   =   self.dummy_vae
							A						   =   self.dummy_text_encoder
							A						   =   CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'      )

							# make sure here that pndm scheduler skips prk
							A						   =   StableDiffusionPipeline(
							    unet=A_  ,scheduler=A_  ,vae=A_  ,text_encoder=A_  ,tokenizer=A_  ,safety_checker=A_  ,feature_extractor=self.dummy_extractor  ,)
							A						   =   sd_pipe.to(A_      )
							sd_pipe.set_progress_bar_config(disable=A_      )

							A						   =   'A painting of a squirrel eating a burger'

							A						   =   torch.Generator(device=A_      ).manual_seed(0      )
							A						   =   sd_pipe([prompt]  ,generator=A_  ,guidance_scale=6.0  ,num_inference_steps=2  ,output_type='np'      )
							A						   =   output.images

							A						   =   torch.Generator(device=A_      ).manual_seed(0      )
							A						   =   sd_pipe(
							    [prompt]  ,generator=A_  ,guidance_scale=6.0  ,num_inference_steps=2  ,output_type='np'  ,return_dict=A_  ,)[0]

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

							assert image.shape == (1, 64, 64, 3)
							A						   =   np.array([0.57_56, 0.61_18, 0.50_05, 0.50_41, 0.54_71, 0.47_26, 0.49_76, 0.48_65, 0.48_64]      )

							assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-2
							assert np.abs(image_from_tuple_slice.flatten() - expected_slice      ).max() < 1e-2


						def 						_SCREAMING_SNAKE_CASE      ( self     :      List[Any]      )     ->   str:
							A						   =   'cpu'  # ensure determinism for the device-dependent torch.Generator
							A						   =   self.dummy_cond_unet
							A						   =   PNDMScheduler(skip_prk_steps=A_      )
							A						   =   self.dummy_vae
							A						   =   self.dummy_text_encoder
							A						   =   CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'      )

							# make sure here that pndm scheduler skips prk
							A						   =   StableDiffusionPipeline(
							    unet=A_  ,scheduler=A_  ,vae=A_  ,text_encoder=A_  ,tokenizer=A_  ,safety_checker=A_  ,feature_extractor=self.dummy_extractor  ,)
							A						   =   sd_pipe.to(A_      )
							sd_pipe.set_progress_bar_config(disable=A_      )

							A						   =   'A painting of a squirrel eating a burger'
							A						   =   torch.Generator(device=A_      ).manual_seed(0      )
							A						   =   sd_pipe([prompt]  ,generator=A_  ,guidance_scale=6.0  ,num_inference_steps=2  ,output_type='np'      )

							A						   =   output.images

							A						   =   torch.Generator(device=A_      ).manual_seed(0      )
							A						   =   sd_pipe(
							    [prompt]  ,generator=A_  ,guidance_scale=6.0  ,num_inference_steps=2  ,output_type='np'  ,return_dict=A_  ,)[0]

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

							assert image.shape == (1, 64, 64, 3)
							A						   =   np.array([0.51_25, 0.57_16, 0.48_28, 0.50_60, 0.56_50, 0.47_68, 0.51_85, 0.48_95, 0.49_93]      )

							assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-2
							assert np.abs(image_from_tuple_slice.flatten() - expected_slice      ).max() < 1e-2


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Dict      )     ->   Tuple:
							A						   =   StableDiffusionPipeline.from_pretrained(
							    'hf-internal-testing/tiny-stable-diffusion-lms-pipe'  ,safety_checker=A_      )
							assert isinstance(A_  ,A_      )
							assert isinstance(pipe.scheduler  ,A_      )
							assert pipe.safety_checker is None

							A						   =   pipe('example prompt'  ,num_inference_steps=2      ).images[0]
							assert image is not None

							# check that there's no error when saving a pipeline with one of the models being None
							with tempfile.TemporaryDirectory() as tmpdirname:
								pipe.save_pretrained(A_      )
								A						   =   StableDiffusionPipeline.from_pretrained(A_      )

							# sanity check that the pipeline still works
							assert pipe.safety_checker is None
							A						   =   pipe('example prompt'  ,num_inference_steps=2      ).images[0]
							assert image is not None




						@unittest.skipIf(torch_device != 'cuda'  ,'This test requires a GPU'      )
						def 						_SCREAMING_SNAKE_CASE      ( self     :      str      )     ->   List[str]:
							A						   =   self.dummy_cond_unet
							A						   =   PNDMScheduler(skip_prk_steps=A_      )
							A						   =   self.dummy_vae
							A						   =   self.dummy_text_encoder
							A						   =   CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'      )

							# put models in fp16
							A						   =   unet.half()
							A						   =   vae.half()
							A						   =   bert.half()

							# make sure here that pndm scheduler skips prk
							A						   =   StableDiffusionPipeline(
							    unet=A_  ,scheduler=A_  ,vae=A_  ,text_encoder=A_  ,tokenizer=A_  ,safety_checker=A_  ,feature_extractor=self.dummy_extractor  ,)
							A						   =   sd_pipe.to(A_      )
							sd_pipe.set_progress_bar_config(disable=A_      )

							A						   =   'A painting of a squirrel eating a burger'
							A						   =   sd_pipe([prompt]  ,num_inference_steps=2  ,output_type='np'      ).images

							assert image.shape == (1, 64, 64, 3)







@nightly
@require_torch_gpu
class 				lowerCAmelCase_				(      unittest.TestCase    ):



						'''simple docstring'''


						def 						_SCREAMING_SNAKE_CASE      ( self     :      str      )     ->   Tuple:
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Tuple      )     ->   int:
							A						   =   StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'  ,safety_checker=A_      )
							A						   =   LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config      )
							A						   =   sd_pipe.to(A_      )
							sd_pipe.set_progress_bar_config(disable=A_      )

							A						   =   (
							    'portrait of girl with smokey eyes makeup in abandoned hotel, grange clothes, redshift, wide high angle'
							    ' coloured polaroid photograph with flash, kodak film, hyper real, stunning moody cinematography, with'
							    ' anamorphic lenses, by maripol, fallen angels by wong kar - wai, style of suspiria and neon demon and'
							    ' children from bahnhof zoo, detailed '
							)
							A						   =   40_0366_0346
							A						   =   7

							# without safety guidance (sld_guidance_scale = 0)
							A						   =   torch.manual_seed(A_      )
							A						   =   sd_pipe(
							    [prompt]  ,generator=A_  ,guidance_scale=A_  ,num_inference_steps=50  ,output_type='np'  ,width=512  ,height=512  ,sld_guidance_scale=0  ,)

							A						   =   output.images
							A						   =   image[0, -3:, -3:, -1]
							A						   =   [0.22_78, 0.22_31, 0.22_49, 0.23_33, 0.23_03, 0.18_85, 0.22_73, 0.21_44, 0.21_76]

							assert image.shape == (1, 512, 512, 3)

							assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-2

							# without safety guidance (strong configuration)
							A						   =   torch.manual_seed(A_      )
							A						   =   sd_pipe(
							    [prompt]  ,generator=A_  ,guidance_scale=A_  ,num_inference_steps=50  ,output_type='np'  ,width=512  ,height=512  ,sld_guidance_scale=2000  ,sld_warmup_steps=7  ,sld_threshold=0.0_25  ,sld_momentum_scale=0.5  ,sld_mom_beta=0.7  ,)

							A						   =   output.images
							A						   =   image[0, -3:, -3:, -1]
							A						   =   [0.23_83, 0.22_76, 0.2_36, 0.21_92, 0.21_86, 0.20_53, 0.19_71, 0.19_01, 0.17_19]

							assert image.shape == (1, 512, 512, 3)

							assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-2


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Dict      )     ->   List[str]:
							A						   =   StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'  ,safety_checker=A_      )
							A						   =   LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config      )
							A						   =   sd_pipe.to(A_      )
							sd_pipe.set_progress_bar_config(disable=A_      )

							A						   =   'padme amidala taking a bath artwork, safe for work, no nudity'
							A						   =   27_3497_1755
							A						   =   7

							A						   =   torch.manual_seed(A_      )
							A						   =   sd_pipe(
							    [prompt]  ,generator=A_  ,guidance_scale=A_  ,num_inference_steps=50  ,output_type='np'  ,width=512  ,height=512  ,sld_guidance_scale=0  ,)

							A						   =   output.images
							A						   =   image[0, -3:, -3:, -1]
							A						   =   [0.35_02, 0.36_22, 0.33_96, 0.36_42, 0.34_78, 0.33_18, 0.35, 0.33_48, 0.32_97]

							assert image.shape == (1, 512, 512, 3)

							assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-2

							A						   =   torch.manual_seed(A_      )
							A						   =   sd_pipe(
							    [prompt]  ,generator=A_  ,guidance_scale=A_  ,num_inference_steps=50  ,output_type='np'  ,width=512  ,height=512  ,sld_guidance_scale=2000  ,sld_warmup_steps=7  ,sld_threshold=0.0_25  ,sld_momentum_scale=0.5  ,sld_mom_beta=0.7  ,)

							A						   =   output.images
							A						   =   image[0, -3:, -3:, -1]
							A						   =   [0.55_31, 0.52_06, 0.48_95, 0.51_56, 0.51_82, 0.47_51, 0.48_02, 0.48_03, 0.44_43]

							assert image.shape == (1, 512, 512, 3)

							assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-2




						def 						_SCREAMING_SNAKE_CASE      ( self     :      int      )     ->   Dict:
							A						   =   StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'      )
							A						   =   sd_pipe.to(A_      )
							sd_pipe.set_progress_bar_config(disable=A_      )

							A						   =   (
							    'the four horsewomen of the apocalypse, painting by tom of finland, gaston bussiere, craig mullins, j. c.'
							    ' leyendecker'
							)
							A						   =   10_4435_5234
							A						   =   12

							A						   =   torch.manual_seed(A_      )
							A						   =   sd_pipe(
							    [prompt]  ,generator=A_  ,guidance_scale=A_  ,num_inference_steps=50  ,output_type='np'  ,width=512  ,height=512  ,sld_guidance_scale=0  ,)

							A						   =   output.images
							A						   =   image[0, -3:, -3:, -1]
							A						   =   np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]      )

							assert image.shape == (1, 512, 512, 3)

							assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-7

							A						   =   torch.manual_seed(A_      )
							A						   =   sd_pipe(
							    [prompt]  ,generator=A_  ,guidance_scale=A_  ,num_inference_steps=50  ,output_type='np'  ,width=512  ,height=512  ,sld_guidance_scale=2000  ,sld_warmup_steps=7  ,sld_threshold=0.0_25  ,sld_momentum_scale=0.5  ,sld_mom_beta=0.7  ,)

							A						   =   output.images
							A						   =   image[0, -3:, -3:, -1]
							A						   =   np.array([0.58_18, 0.62_85, 0.68_35, 0.60_19, 0.6_25, 0.67_54, 0.60_96, 0.63_34, 0.65_61]      )
							assert image.shape == (1, 512, 512, 3)

							assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-2

"""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 				lowerCAmelCase_				(      _lowercase    ):



						'''simple docstring'''


						_lowerCamelCase:    torch.FloatTensor







class 				lowerCAmelCase_				(      nn.Module    ):



						'''simple docstring'''


						def __init__( self     :      List[str]  ,A_     :      Dict=3  ,A_     :      int=3  ,A_     :      str=("DownEncoderBlock2D",)  ,A_     :      Dict=(64,)  ,A_     :      str=2  ,A_     :      Union[str, Any]=32  ,A_     :      Optional[int]="silu"  ,A_     :      str=True  ,)     ->   Union[str, Any]:
							super().__init__()
							A						   =   layers_per_block

							A						   =   torch.nn.Convad(
							    A_  ,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(A_      ):
								A						   =   output_channel
								A						   =   block_out_channels[i]
								A						   =   i == len(A_      ) - 1

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

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

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

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

							A						   =   False




						def 						_SCREAMING_SNAKE_CASE      ( self     :      Union[str, Any]  ,A_     :      Optional[int]      )     ->   Union[str, Any]:
							A						   =   x
							A						   =   self.conv_in(A_      )

							if self.training and self.gradient_checkpointing:

								def create_custom_forward(A_     :      Dict      ):
									def custom_forward(*A_     :      Tuple      ):
										return module(*A_      )

									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(A_      )  ,A_  ,use_reentrant=A_      )
									# middle
									A						   =   torch.utils.checkpoint.checkpoint(
									    create_custom_forward(self.mid_block      )  ,A_  ,use_reentrant=A_      )
								else:
									for down_block in self.down_blocks:
										A						   =   torch.utils.checkpoint.checkpoint(create_custom_forward(A_      )  ,A_      )
									# middle
									A						   =   torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block      )  ,A_      )

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

								# middle
								A						   =   self.mid_block(A_      )

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

							return sample







class 				lowerCAmelCase_				(      nn.Module    ):



						'''simple docstring'''


						def __init__( self     :      List[Any]  ,A_     :      Optional[Any]=3  ,A_     :      Optional[int]=3  ,A_     :      str=("UpDecoderBlock2D",)  ,A_     :      Any=(64,)  ,A_     :      Optional[int]=2  ,A_     :      Optional[int]=32  ,A_     :      Tuple="silu"  ,A_     :      Optional[int]="group"  ,)     ->   Any:
							super().__init__()
							A						   =   layers_per_block

							A						   =   nn.Convad(
							    A_  ,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=A_  ,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=A_  ,temb_channels=A_  ,)

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

								A						   =   i == len(A_      ) - 1

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

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

							A						   =   False




						def 						_SCREAMING_SNAKE_CASE      ( self     :      Union[str, Any]  ,A_     :      str  ,A_     :      Union[str, Any]=None      )     ->   Any:
							A						   =   z
							A						   =   self.conv_in(A_      )

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

								def create_custom_forward(A_     :      List[Any]      ):
									def custom_forward(*A_     :      Tuple      ):
										return module(*A_      )

									return custom_forward

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

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

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

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

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

							return sample







class 				lowerCAmelCase_				(      nn.Module    ):



						'''simple docstring'''


						def __init__( self     :      Optional[int]  ,A_     :      Optional[int]  ,A_     :      Any  ,A_     :      str  ,A_     :      Dict=None  ,A_     :      List[Any]="random"  ,A_     :      Optional[int]=False  ,A_     :      str=True      )     ->   List[str]:
							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 						_SCREAMING_SNAKE_CASE      ( self     :      List[str]  ,A_     :      Optional[int]      )     ->   Any:
							A						   =   inds.shape
							assert len(A_      ) > 1
							A						   =   inds.reshape(ishape[0]  ,-1      )
							A						   =   self.used.to(A_      )
							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(A_      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      List[str]  ,A_     :      Optional[Any]      )     ->   List[Any]:
							A						   =   inds.shape
							assert len(A_      ) > 1
							A						   =   inds.reshape(ishape[0]  ,-1      )
							A						   =   self.used.to(A_      )
							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  ,A_      )
							return back.reshape(A_      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Optional[int]  ,A_     :      List[Any]      )     ->   str:
							# 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(A_  ,self.embedding.weight      )  ,dim=1      )

							A						   =   self.embedding(A_      ).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(A_      )
								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 						_SCREAMING_SNAKE_CASE      ( self     :      Dict  ,A_     :      Dict  ,A_     :      str      )     ->   Union[str, Any]:
							# 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(A_      )
								A						   =   indices.reshape(-1      )  # flatten again

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

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

							return z_q







class 				lowerCAmelCase_				(      _lowercase    ):



						'''simple docstring'''


						def __init__( self     :      str  ,A_     :      Tuple  ,A_     :      Dict=False      )     ->   List[str]:
							A						   =   parameters
							A   ,      A						   =   torch.chunk(A_  ,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 						_SCREAMING_SNAKE_CASE      ( self     :      Any  ,A_     :      Optional[torch.Generator] = None      )     ->   torch.FloatTensor:
							# make sure sample is on the same device as the parameters and has same dtype
							A						   =   randn_tensor(
							    self.mean.shape  ,generator=A_  ,device=self.parameters.device  ,dtype=self.parameters.dtype      )
							A						   =   self.mean + self.std * sample
							return x


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Any  ,A_     :      Tuple=None      )     ->   int:
							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 						_SCREAMING_SNAKE_CASE      ( self     :      List[str]  ,A_     :      List[str]  ,A_     :      Union[str, Any]=[1, 2, 3]      )     ->   List[str]:
							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=A_      )




						def 						_SCREAMING_SNAKE_CASE      ( self     :      Dict      )     ->   Optional[int]:
							return self.mean