



'''simple docstring'''
import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    PNDMScheduler,
    StableDiffusionLDMaDPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS


enable_full_determinism()






class 							lowercase__	(							unittest.TestCase   ):
	lowercase__  				=			StableDiffusionLDMaDPipeline
	lowercase__  				=			TEXT_TO_IMAGE_PARAMS
	lowercase__  				=			TEXT_TO_IMAGE_BATCH_PARAMS
	lowercase__  				=			TEXT_TO_IMAGE_IMAGE_PARAMS





	def 	UpperCamelCase_  ( self					:				int ):







							'''simple docstring'''
							torch.manual_seed(0 )
							_UpperCamelCase					:		List[str] 				=   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 ,)
							_UpperCamelCase					:		List[Any] 				=   DDIMScheduler(
							    beta_start=0.0_0_0_8_5 ,beta_end=0.0_1_2 ,beta_schedule='scaled_linear' ,clip_sample=lowerCAmelCase_ ,set_alpha_to_one=lowerCAmelCase_ ,)
							torch.manual_seed(0 )
							_UpperCamelCase					:		str 				=   AutoencoderKL(
							    block_out_channels=[32, 64] ,in_channels=6 ,out_channels=6 ,down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D'] ,up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D'] ,latent_channels=4 ,)
							torch.manual_seed(0 )
							_UpperCamelCase					:		Dict 				=   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 ,)
							_UpperCamelCase					:		Any 				=   CLIPTextModel(lowerCAmelCase_ )
							_UpperCamelCase					:		Dict 				=   CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip' )

							_UpperCamelCase					:		str 				=   {
							    'unet': unet,
							    'scheduler': scheduler,
							    'vae': vae,
							    'text_encoder': text_encoder,
							    'tokenizer': tokenizer,
							    'safety_checker': None,
							    'feature_extractor': None,
							}
							return components





	def 	UpperCamelCase_  ( self					:				Optional[int] ,lowerCamelCase__					:				Union[str, Any] ,lowerCamelCase__					:				Optional[int]=0 ):







							'''simple docstring'''






							if str(lowerCAmelCase_ ).startswith('mps' ):
													_UpperCamelCase					:		Optional[int] 				=   torch.manual_seed(lowerCAmelCase_ )
							else:
													_UpperCamelCase					:		List[str] 				=   torch.Generator(device=lowerCAmelCase_ ).manual_seed(lowerCAmelCase_ )
							_UpperCamelCase					:		List[str] 				=   {
							    'prompt': 'A painting of a squirrel eating a burger',
							    'generator': generator,
							    'num_inference_steps': 2,
							    'guidance_scale': 6.0,
							    'output_type': 'numpy',
							}
							return inputs





	def 	UpperCamelCase_  ( self					:				str ):







							'''simple docstring'''
							_UpperCamelCase					:		Tuple 				=   'cpu'  # ensure determinism for the device-dependent torch.Generator

							_UpperCamelCase					:		Any 				=   self.get_dummy_components()
							_UpperCamelCase					:		Dict 				=   StableDiffusionLDMaDPipeline(**lowerCAmelCase_ )
							_UpperCamelCase					:		Union[str, Any] 				=   ldmad_pipe.to(lowerCAmelCase_ )
							ldmad_pipe.set_progress_bar_config(disable=lowerCAmelCase_ )

							_UpperCamelCase					:		List[str] 				=   self.get_dummy_inputs(lowerCAmelCase_ )
							_UpperCamelCase					:		Union[str, Any] 				=   ldmad_pipe(**lowerCAmelCase_ )
							_UpperCamelCase				,			_UpperCamelCase					:		Dict 				=   output.rgb, output.depth

							_UpperCamelCase					:		int 				=   rgb[0, -3:, -3:, -1]
							_UpperCamelCase					:		Optional[int] 				=   depth[0, -3:, -1]

							assert rgb.shape == (1, 64, 64, 3)
							assert depth.shape == (1, 64, 64)

							_UpperCamelCase					:		int 				=   np.array(
							    [0.3_7_3_3_8_1_7_6, 0.7_0_2_4_7, 0.7_4_2_0_3_1_9_3, 0.5_1_6_4_3_6_0_4, 0.5_8_2_5_6_7_9_3, 0.6_0_9_3_2_1_3_6, 0.4_1_8_1_0_9_5, 0.4_8_3_5_5_8_7_7, 0.4_6_5_3_5_2_6_2] )
							_UpperCamelCase					:		Union[str, Any] 				=   np.array([1_0_3.4_6_7_2_7, 8_5.8_1_2_0_0_4, 8_7.8_4_9_2_3_6] )

							assert np.abs(image_slice_rgb.flatten() - expected_slice_rgb ).max() < 1E-2
							assert np.abs(image_slice_depth.flatten() - expected_slice_depth ).max() < 1E-2





	def 	UpperCamelCase_  ( self					:				str ):







							'''simple docstring'''
							_UpperCamelCase					:		int 				=   self.get_dummy_components()
							_UpperCamelCase					:		List[str] 				=   StableDiffusionLDMaDPipeline(**lowerCAmelCase_ )
							_UpperCamelCase					:		Tuple 				=   ldmad_pipe.to(lowerCAmelCase_ )
							ldmad_pipe.set_progress_bar_config(disable=lowerCAmelCase_ )

							_UpperCamelCase					:		Dict 				=   self.get_dummy_inputs(lowerCAmelCase_ )
							_UpperCamelCase					:		int 				=   3 * [inputs['prompt']]

							# forward
							_UpperCamelCase					:		Tuple 				=   ldmad_pipe(**lowerCAmelCase_ )
							_UpperCamelCase				,			_UpperCamelCase					:		Optional[Any] 				=   output.rgb, output.depth
							_UpperCamelCase					:		Any 				=   rgb_slice_a[0, -3:, -3:, -1]
							_UpperCamelCase					:		Optional[Any] 				=   depth_slice_a[0, -3:, -1]

							_UpperCamelCase					:		int 				=   self.get_dummy_inputs(lowerCAmelCase_ )
							_UpperCamelCase					:		Union[str, Any] 				=   3 * [inputs.pop('prompt' )]

							_UpperCamelCase					:		Tuple 				=   ldmad_pipe.tokenizer(
							    lowerCAmelCase_ ,padding='max_length' ,max_length=ldmad_pipe.tokenizer.model_max_length ,truncation=lowerCAmelCase_ ,return_tensors='pt' ,)
							_UpperCamelCase					:		Union[str, Any] 				=   text_inputs['input_ids'].to(lowerCAmelCase_ )

							_UpperCamelCase					:		Dict 				=   ldmad_pipe.text_encoder(lowerCAmelCase_ )[0]

							_UpperCamelCase					:		Any 				=   prompt_embeds

							# forward
							_UpperCamelCase					:		List[str] 				=   ldmad_pipe(**lowerCAmelCase_ )
							_UpperCamelCase				,			_UpperCamelCase					:		Dict 				=   output.rgb, output.depth
							_UpperCamelCase					:		int 				=   rgb_slice_a[0, -3:, -3:, -1]
							_UpperCamelCase					:		List[str] 				=   depth_slice_a[0, -3:, -1]

							assert np.abs(rgb_slice_a.flatten() - rgb_slice_a.flatten() ).max() < 1E-4
							assert np.abs(depth_slice_a.flatten() - depth_slice_a.flatten() ).max() < 1E-4





	def 	UpperCamelCase_  ( self					:				Optional[int] ):







							'''simple docstring'''
							_UpperCamelCase					:		int 				=   'cpu'  # ensure determinism for the device-dependent torch.Generator
							_UpperCamelCase					:		Tuple 				=   self.get_dummy_components()
							_UpperCamelCase					:		List[str] 				=   PNDMScheduler(skip_prk_steps=lowerCAmelCase_ )
							_UpperCamelCase					:		str 				=   StableDiffusionLDMaDPipeline(**lowerCAmelCase_ )
							_UpperCamelCase					:		List[str] 				=   ldmad_pipe.to(lowerCAmelCase_ )
							ldmad_pipe.set_progress_bar_config(disable=lowerCAmelCase_ )

							_UpperCamelCase					:		Optional[Any] 				=   self.get_dummy_inputs(lowerCAmelCase_ )
							_UpperCamelCase					:		Optional[int] 				=   'french fries'
							_UpperCamelCase					:		str 				=   ldmad_pipe(**lowerCAmelCase_ ,negative_prompt=lowerCAmelCase_ )

							_UpperCamelCase				,			_UpperCamelCase					:		List[str] 				=   output.rgb, output.depth
							_UpperCamelCase					:		Optional[Any] 				=   rgb[0, -3:, -3:, -1]
							_UpperCamelCase					:		List[str] 				=   depth[0, -3:, -1]

							assert rgb.shape == (1, 64, 64, 3)
							assert depth.shape == (1, 64, 64)

							_UpperCamelCase					:		List[str] 				=   np.array(
							    [0.3_7_0_4_4, 0.7_1_8_1_1_5_0_3, 0.7_2_2_3_2_5_1, 0.4_8_6_0_3_6_7_5, 0.5_6_3_8_3_9_1, 0.6_3_6_4_9_4_8, 0.4_2_8_3_3_7_0_4, 0.4_9_0_1_3_1_5, 0.4_7_9_2_6_2_1_7] )
							_UpperCamelCase					:		Tuple 				=   np.array([1_0_7.8_4_7_3_8, 8_4.6_2_8_0_2, 8_9.9_6_2_1_3_5] )
							assert np.abs(rgb_slice.flatten() - expected_slice_rgb ).max() < 1E-2
							assert np.abs(depth_slice.flatten() - expected_slice_depth ).max() < 1E-2






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





	def 	UpperCamelCase_  ( self					:				Union[str, Any] ):







							'''simple docstring'''
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()





	def 	UpperCamelCase_  ( self					:				Dict ,lowerCamelCase__					:				int ,lowerCamelCase__					:				Dict="cpu" ,lowerCamelCase__					:				Optional[int]=torch.floataa ,lowerCamelCase__					:				Union[str, Any]=0 ):







							'''simple docstring'''
							_UpperCamelCase					:		Dict 				=   torch.Generator(device=lowerCAmelCase_ ).manual_seed(lowerCAmelCase_ )
							_UpperCamelCase					:		List[str] 				=   np.random.RandomState(lowerCAmelCase_ ).standard_normal((1, 4, 64, 64) )
							_UpperCamelCase					:		List[str] 				=   torch.from_numpy(lowerCAmelCase_ ).to(device=lowerCAmelCase_ ,dtype=lowerCAmelCase_ )
							_UpperCamelCase					:		str 				=   {
							    'prompt': 'a photograph of an astronaut riding a horse',
							    'latents': latents,
							    'generator': generator,
							    'num_inference_steps': 3,
							    'guidance_scale': 7.5,
							    'output_type': 'numpy',
							}
							return inputs





	def 	UpperCamelCase_  ( self					:				Tuple ):







							'''simple docstring'''
							_UpperCamelCase					:		Optional[int] 				=   StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d' )
							_UpperCamelCase					:		List[Any] 				=   ldmad_pipe.to(lowerCAmelCase_ )
							ldmad_pipe.set_progress_bar_config(disable=lowerCAmelCase_ )

							_UpperCamelCase					:		str 				=   self.get_inputs(lowerCAmelCase_ )
							_UpperCamelCase					:		List[str] 				=   ldmad_pipe(**lowerCAmelCase_ )
							_UpperCamelCase				,			_UpperCamelCase					:		Dict 				=   output.rgb, output.depth
							_UpperCamelCase					:		List[str] 				=   rgb[0, -3:, -3:, -1].flatten()
							_UpperCamelCase					:		Union[str, Any] 				=   rgb[0, -3:, -1].flatten()

							assert rgb.shape == (1, 512, 512, 3)
							assert depth.shape == (1, 512, 512)

							_UpperCamelCase					:		Tuple 				=   np.array(
							    [0.5_3_8_0_5_4_6_5, 0.5_6_7_0_7_3_0_5, 0.5_4_8_6_5_1_5, 0.5_7_0_1_2_2_3_6, 0.5_8_1_4_5_1_1, 0.5_6_2_5_3_4_8_7, 0.5_4_8_4_3_0_1_4, 0.5_5_0_9_2_2_6_3, 0.6_4_5_9_7_0_6] )
							_UpperCamelCase					:		int 				=   np.array(
							    [0.9_2_6_3_7_8_1, 0.6_6_7_8_6_7_2, 0.5_4_8_6_5_1_5, 0.9_2_2_0_2_1_4_5, 0.6_7_8_3_1_1_3_5, 0.5_6_2_5_3_4_8_7, 0.9_2_4_1_6_9_4, 0.7_5_5_1_4_7_8, 0.6_4_5_9_7_0_6] )
							assert np.abs(rgb_slice - expected_slice_rgb ).max() < 3E-3
							assert np.abs(depth_slice - expected_slice_depth ).max() < 3E-3






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





	def 	UpperCamelCase_  ( self					:				List[str] ):







							'''simple docstring'''
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()





	def 	UpperCamelCase_  ( self					:				List[str] ,lowerCamelCase__					:				int ,lowerCamelCase__					:				Optional[Any]="cpu" ,lowerCamelCase__					:				List[Any]=torch.floataa ,lowerCamelCase__					:				Dict=0 ):







							'''simple docstring'''
							_UpperCamelCase					:		Optional[int] 				=   torch.Generator(device=lowerCAmelCase_ ).manual_seed(lowerCAmelCase_ )
							_UpperCamelCase					:		Dict 				=   np.random.RandomState(lowerCAmelCase_ ).standard_normal((1, 4, 64, 64) )
							_UpperCamelCase					:		Dict 				=   torch.from_numpy(lowerCAmelCase_ ).to(device=lowerCAmelCase_ ,dtype=lowerCAmelCase_ )
							_UpperCamelCase					:		List[str] 				=   {
							    'prompt': 'a photograph of an astronaut riding a horse',
							    'latents': latents,
							    'generator': generator,
							    'num_inference_steps': 50,
							    'guidance_scale': 7.5,
							    'output_type': 'numpy',
							}
							return inputs





	def 	UpperCamelCase_  ( self					:				Dict ):







							'''simple docstring'''
							_UpperCamelCase					:		str 				=   StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d' ).to(lowerCAmelCase_ )
							ldmad_pipe.set_progress_bar_config(disable=lowerCAmelCase_ )

							_UpperCamelCase					:		Dict 				=   self.get_inputs(lowerCAmelCase_ )
							_UpperCamelCase					:		Union[str, Any] 				=   ldmad_pipe(**lowerCAmelCase_ )
							_UpperCamelCase				,			_UpperCamelCase					:		Union[str, Any] 				=   output.rgb, output.depth

							_UpperCamelCase					:		Dict 				=   0.4_9_5_5_8_6
							_UpperCamelCase					:		Tuple 				=   0.3_3_7_9_5_5_1_5
							_UpperCamelCase					:		Tuple 				=   1_1_2.4_8_5_1_8
							_UpperCamelCase					:		Union[str, Any] 				=   9_8.4_8_9_7_4_6
							assert np.abs(expected_rgb_mean - rgb.mean() ) < 1E-3
							assert np.abs(expected_rgb_std - rgb.std() ) < 1E-3
							assert np.abs(expected_depth_mean - depth.mean() ) < 1E-3
							assert np.abs(expected_depth_std - depth.std() ) < 1E-3





	def 	UpperCamelCase_  ( self					:				Dict ):







							'''simple docstring'''
							_UpperCamelCase					:		Optional[int] 				=   StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d-4c' ).to(lowerCAmelCase_ )
							ldmad_pipe.set_progress_bar_config(disable=lowerCAmelCase_ )

							_UpperCamelCase					:		Union[str, Any] 				=   self.get_inputs(lowerCAmelCase_ )
							_UpperCamelCase					:		Optional[Any] 				=   ldmad_pipe(**lowerCAmelCase_ )
							_UpperCamelCase				,			_UpperCamelCase					:		Union[str, Any] 				=   output.rgb, output.depth

							_UpperCamelCase					:		Optional[Any] 				=   0.4_1_9_4_1_2_7
							_UpperCamelCase					:		Union[str, Any] 				=   0.3_5_3_7_5_5_8_6
							_UpperCamelCase					:		Any 				=   0.5_6_3_8_5_0_2
							_UpperCamelCase					:		Tuple 				=   0.3_4_6_8_6_1_0_3

							assert rgb.shape == (1, 512, 512, 3)
							assert depth.shape == (1, 512, 512, 1)
							assert np.abs(expected_rgb_mean - rgb.mean() ) < 1E-3
							assert np.abs(expected_rgb_std - rgb.std() ) < 1E-3
							assert np.abs(expected_depth_mean - depth.mean() ) < 1E-3
							assert np.abs(expected_depth_std - depth.std() ) < 1E-3




'''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      (					_lowerCamelCase						):
							__lowercase      		=       42


class 						__UpperCAmelCase      (					nn.Module						):


							def __init__(       self    ,   lowerCAmelCase_=3    ,   lowerCAmelCase_=3    ,   lowerCAmelCase_=("DownEncoderBlock2D",)    ,   lowerCAmelCase_=(64,)    ,   lowerCAmelCase_=2    ,   lowerCAmelCase_=32    ,   lowerCAmelCase_="silu"    ,   lowerCAmelCase_=True    ,   ):





								"""simple docstring"""






								super().__init__()
								_snake_case   						=		layers_per_block

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

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

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

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

								# mid
								_snake_case   						=		UNetMidBlockaD(
								    in_channels=block_out_channels[-1]    ,   resnet_eps=1E-6    ,   resnet_act_fn=lowerCAmelCase_    ,   output_scale_factor=1    ,   resnet_time_scale_shift='default'    ,   attention_head_dim=block_out_channels[-1]    ,   resnet_groups=lowerCAmelCase_    ,   temb_channels=lowerCAmelCase_    ,   )

								# out
								_snake_case   						=		nn.GroupNorm(num_channels=block_out_channels[-1]    ,   num_groups=lowerCAmelCase_    ,   eps=1E-6  )
								_snake_case   						=		nn.SiLU()

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

								_snake_case   						=		False




							def  lowerCamelCase				(       self    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								_snake_case   						=		x
								_snake_case   						=		self.conv_in(lowerCAmelCase_  )

								if self.training and self.gradient_checkpointing:

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

										return custom_forward

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

								else:
									# down
									for down_block in self.down_blocks:
										_snake_case   						=		down_block(lowerCAmelCase_  )

									# middle
									_snake_case   						=		self.mid_block(lowerCAmelCase_  )

								# post-process
								_snake_case   						=		self.conv_norm_out(lowerCAmelCase_  )
								_snake_case   						=		self.conv_act(lowerCAmelCase_  )
								_snake_case   						=		self.conv_out(lowerCAmelCase_  )

								return sample


class 						__UpperCAmelCase      (					nn.Module						):


							def __init__(       self    ,   lowerCAmelCase_=3    ,   lowerCAmelCase_=3    ,   lowerCAmelCase_=("UpDecoderBlock2D",)    ,   lowerCAmelCase_=(64,)    ,   lowerCAmelCase_=2    ,   lowerCAmelCase_=32    ,   lowerCAmelCase_="silu"    ,   lowerCAmelCase_="group"    ,   ):





								"""simple docstring"""






								super().__init__()
								_snake_case   						=		layers_per_block

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

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

								_snake_case   						=		in_channels if norm_type == 'spatial' else None

								# mid
								_snake_case   						=		UNetMidBlockaD(
								    in_channels=block_out_channels[-1]    ,   resnet_eps=1E-6    ,   resnet_act_fn=lowerCAmelCase_    ,   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=lowerCAmelCase_    ,   temb_channels=lowerCAmelCase_    ,   )

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

									_snake_case   						=		i == len(lowerCAmelCase_  ) - 1

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

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

								_snake_case   						=		False




							def  lowerCamelCase				(       self    ,   lowerCAmelCase_    ,   lowerCAmelCase_=None  ):





								"""simple docstring"""






								_snake_case   						=		z
								_snake_case   						=		self.conv_in(lowerCAmelCase_  )

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

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

										return custom_forward

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

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

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

									# up
									for up_block in self.up_blocks:
										_snake_case   						=		up_block(lowerCAmelCase_    ,   lowerCAmelCase_  )

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

								return sample


class 						__UpperCAmelCase      (					nn.Module						):


							def __init__(       self    ,   lowerCAmelCase_    ,   lowerCAmelCase_    ,   lowerCAmelCase_    ,   lowerCAmelCase_=None    ,   lowerCAmelCase_="random"    ,   lowerCAmelCase_=False    ,   lowerCAmelCase_=True  ):





								"""simple docstring"""






								super().__init__()
								_snake_case   						=		n_e
								_snake_case   						=		vq_embed_dim
								_snake_case   						=		beta
								_snake_case   						=		legacy

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

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

								_snake_case   						=		sane_index_shape


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_  ):





								"""simple docstring"""






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


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_  ):





								"""simple docstring"""






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


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								_snake_case   						=		z.permute(0    ,   2    ,   3    ,   1  ).contiguous()
								_snake_case   						=		z.view(-1    ,   self.vq_embed_dim  )

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

								_snake_case   						=		self.embedding(lowerCAmelCase_  ).view(z.shape  )
								_snake_case   						=		None
								_snake_case   						=		None

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

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

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

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

								if self.sane_index_shape:
									_snake_case   						=		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  lowerCamelCase				(       self    ,   lowerCAmelCase_    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								if self.remap is not None:
									_snake_case   						=		indices.reshape(shape[0]    ,   -1  )  # add batch axis
									_snake_case   						=		self.unmap_to_all(lowerCAmelCase_  )
									_snake_case   						=		indices.reshape(-1  )  # flatten again

								# get quantized latent vectors
								_snake_case   						=		self.embedding(lowerCAmelCase_  )

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

								return z_q


class 						__UpperCAmelCase      (					_lowerCamelCase						):


							def __init__(       self    ,   lowerCAmelCase_    ,   lowerCAmelCase_=False  ):





								"""simple docstring"""






								_snake_case   						=		parameters
								_snake_case						, _snake_case   						=		torch.chunk(lowerCAmelCase_    ,   2    ,   dim=1  )
								_snake_case   						=		torch.clamp(self.logvar    ,   -30.0    ,   20.0  )
								_snake_case   						=		deterministic
								_snake_case   						=		torch.exp(0.5 * self.logvar  )
								_snake_case   						=		torch.exp(self.logvar  )
								if self.deterministic:
									_snake_case   						=		_snake_case   						=		torch.zeros_like(
									    self.mean    ,   device=self.parameters.device    ,   dtype=self.parameters.dtype  )


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_ = None  ):





								"""simple docstring"""






								_snake_case   						=		randn_tensor(
								    self.mean.shape    ,   generator=lowerCAmelCase_    ,   device=self.parameters.device    ,   dtype=self.parameters.dtype  )
								_snake_case   						=		self.mean + self.std * sample
								return x


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_=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  lowerCamelCase				(       self    ,   lowerCAmelCase_    ,   lowerCAmelCase_=[1, 2, 3]  ):





								"""simple docstring"""






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




							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								return self.mean

