



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_):
							lowerCamelCase__   			=		42



class   UpperCAmelCase_    (       nn.Module):







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

														'''simple docstring'''

														super().__init__()
														_lowerCAmelCase     :							Optional[int]    =    layers_per_block

														_lowerCAmelCase     :							Tuple    =    torch.nn.Convad(
														    lowerCAmelCase__,			block_out_channels[0],			kernel_size=3,			stride=1,			padding=1,			)

														_lowerCAmelCase     :							Tuple    =    None
														_lowerCAmelCase     :							Dict    =    nn.ModuleList([])

														# down
														_lowerCAmelCase     :							Dict    =    block_out_channels[0]
														for i, down_block_type in enumerate(lowerCAmelCase__):
																					_lowerCAmelCase     :							Tuple    =    output_channel
																					_lowerCAmelCase     :							int    =    block_out_channels[i]
																					_lowerCAmelCase     :							str    =    i == len(lowerCAmelCase__) - 1

																					_lowerCAmelCase     :							Optional[int]    =    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
														_lowerCAmelCase     :							str    =    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
														_lowerCAmelCase     :							str    =    nn.GroupNorm(num_channels=block_out_channels[-1],			num_groups=lowerCAmelCase__,			eps=1E-6)
														_lowerCAmelCase     :							Union[str, Any]    =    nn.SiLU()

														_lowerCAmelCase     :							List[str]    =    2 * out_channels if double_z else out_channels
														_lowerCAmelCase     :							List[Any]    =    nn.Convad(block_out_channels[-1],			lowerCAmelCase__,			3,			padding=1)

														_lowerCAmelCase     :							Optional[Any]    =    False







							def    snake_case__      (   self,			__a):

														'''simple docstring'''

														_lowerCAmelCase     :							Optional[int]    =    x
														_lowerCAmelCase     :							Dict    =    self.conv_in(lowerCAmelCase__)

														if self.training and self.gradient_checkpointing:

																					def create_custom_forward(__a):
																												def custom_forward(*__a):
																																			return module(*lowerCAmelCase__)

																												return custom_forward

																					# down
																					if is_torch_version(">=",			"1.11.0"):
																												for down_block in self.down_blocks:
																																			_lowerCAmelCase     :							Union[str, Any]    =    torch.utils.checkpoint.checkpoint(
																																			    create_custom_forward(lowerCAmelCase__),			lowerCAmelCase__,			use_reentrant=lowerCAmelCase__)
																												# middle
																												_lowerCAmelCase     :							List[str]    =    torch.utils.checkpoint.checkpoint(
																												    create_custom_forward(self.mid_block),			lowerCAmelCase__,			use_reentrant=lowerCAmelCase__)
																					else:
																												for down_block in self.down_blocks:
																																			_lowerCAmelCase     :							Any    =    torch.utils.checkpoint.checkpoint(create_custom_forward(lowerCAmelCase__),			lowerCAmelCase__)
																												# middle
																												_lowerCAmelCase     :							Dict    =    torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block),			lowerCAmelCase__)

														else:
																					# down
																					for down_block in self.down_blocks:
																												_lowerCAmelCase     :							Optional[int]    =    down_block(lowerCAmelCase__)

																					# middle
																					_lowerCAmelCase     :							str    =    self.mid_block(lowerCAmelCase__)

														# post-process
														_lowerCAmelCase     :							Union[str, Any]    =    self.conv_norm_out(lowerCAmelCase__)
														_lowerCAmelCase     :							Any    =    self.conv_act(lowerCAmelCase__)
														_lowerCAmelCase     :							Optional[int]    =    self.conv_out(lowerCAmelCase__)

														return sample



class   UpperCAmelCase_    (       nn.Module):







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

														'''simple docstring'''

														super().__init__()
														_lowerCAmelCase     :							Optional[Any]    =    layers_per_block

														_lowerCAmelCase     :							Dict    =    nn.Convad(
														    lowerCAmelCase__,			block_out_channels[-1],			kernel_size=3,			stride=1,			padding=1,			)

														_lowerCAmelCase     :							Dict    =    None
														_lowerCAmelCase     :							Any    =    nn.ModuleList([])

														_lowerCAmelCase     :							Tuple    =    in_channels if norm_type == "spatial" else None

														# mid
														_lowerCAmelCase     :							str    =    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
														_lowerCAmelCase     :							List[str]    =    list(reversed(lowerCAmelCase__))
														_lowerCAmelCase     :							List[str]    =    reversed_block_out_channels[0]
														for i, up_block_type in enumerate(lowerCAmelCase__):
																					_lowerCAmelCase     :							int    =    output_channel
																					_lowerCAmelCase     :							List[Any]    =    reversed_block_out_channels[i]

																					_lowerCAmelCase     :							Any    =    i == len(lowerCAmelCase__) - 1

																					_lowerCAmelCase     :							Union[str, Any]    =    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__)
																					_lowerCAmelCase     :							Any    =    output_channel

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

														_lowerCAmelCase     :							Union[str, Any]    =    False







							def    snake_case__      (   self,			__a,			__a=None):

														'''simple docstring'''

														_lowerCAmelCase     :							Any    =    z
														_lowerCAmelCase     :							str    =    self.conv_in(lowerCAmelCase__)

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

																					def create_custom_forward(__a):
																												def custom_forward(*__a):
																																			return module(*lowerCAmelCase__)

																												return custom_forward

																					if is_torch_version(">=",			"1.11.0"):
																												# middle
																												_lowerCAmelCase     :							Dict    =    torch.utils.checkpoint.checkpoint(
																												    create_custom_forward(self.mid_block),			lowerCAmelCase__,			lowerCAmelCase__,			use_reentrant=lowerCAmelCase__)
																												_lowerCAmelCase     :							Any    =    sample.to(lowerCAmelCase__)

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

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

																					# up
																					for up_block in self.up_blocks:
																												_lowerCAmelCase     :							List[Any]    =    up_block(lowerCAmelCase__,			lowerCAmelCase__)

        # post-process
														if latent_embeds is None:
																					_lowerCAmelCase     :							Union[str, Any]    =    self.conv_norm_out(lowerCAmelCase__)
														else:
																					_lowerCAmelCase     :							List[str]    =    self.conv_norm_out(lowerCAmelCase__,			lowerCAmelCase__)
														_lowerCAmelCase     :							Optional[int]    =    self.conv_act(lowerCAmelCase__)
														_lowerCAmelCase     :							str    =    self.conv_out(lowerCAmelCase__)

														return sample



class   UpperCAmelCase_    (       nn.Module):







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

														'''simple docstring'''

														super().__init__()
														_lowerCAmelCase     :							Any    =    n_e
														_lowerCAmelCase     :							Union[str, Any]    =    vq_embed_dim
														_lowerCAmelCase     :							Tuple    =    beta
														_lowerCAmelCase     :							Dict    =    legacy

														_lowerCAmelCase     :							List[Any]    =    nn.Embedding(self.n_e,			self.vq_embed_dim)
														self.embedding.weight.data.uniform_(-1.0 / self.n_e,			1.0 / self.n_e)

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

														_lowerCAmelCase     :							List[str]    =    sane_index_shape







							def    snake_case__      (   self,			__a):

														'''simple docstring'''

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







							def    snake_case__      (   self,			__a):

														'''simple docstring'''

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







							def    snake_case__      (   self,			__a):

														'''simple docstring'''

														_lowerCAmelCase     :							int    =    z.permute(0,			2,			3,			1).contiguous()
														_lowerCAmelCase     :							Optional[Any]    =    z.view(-1,			self.vq_embed_dim)

														# distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
														_lowerCAmelCase     :							Optional[int]    =    torch.argmin(torch.cdist(lowerCAmelCase__,			self.embedding.weight),			dim=1)

														_lowerCAmelCase     :							Tuple    =    self.embedding(lowerCAmelCase__).view(z.shape)
														_lowerCAmelCase     :							List[str]    =    None
														_lowerCAmelCase     :							Optional[Any]    =    None

														# compute loss for embedding
														if not self.legacy:
																					_lowerCAmelCase     :							Union[str, Any]    =    self.beta * torch.mean((z_q.detach() - z) ** 2) + torch.mean((z_q - z.detach()) ** 2)
														else:
																					_lowerCAmelCase     :							int    =    torch.mean((z_q.detach() - z) ** 2) + self.beta * torch.mean((z_q - z.detach()) ** 2)

														# preserve gradients
														_lowerCAmelCase     :							Tuple    =    z + (z_q - z).detach()

														# reshape back to match original input shape
														_lowerCAmelCase     :							List[str]    =    z_q.permute(0,			3,			1,			2).contiguous()

														if self.remap is not None:
																					_lowerCAmelCase     :							Union[str, Any]    =    min_encoding_indices.reshape(z.shape[0],			-1)  # add batch axis
																					_lowerCAmelCase     :							Dict    =    self.remap_to_used(lowerCAmelCase__)
																					_lowerCAmelCase     :							Tuple    =    min_encoding_indices.reshape(-1,			1)  # flatten

														if self.sane_index_shape:
																					_lowerCAmelCase     :							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    snake_case__      (   self,			__a,			__a):

														'''simple docstring'''

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

														# get quantized latent vectors
														_lowerCAmelCase     :							List[str]    =    self.embedding(lowerCAmelCase__)

														if shape is not None:
																					_lowerCAmelCase     :							Optional[int]    =    z_q.view(lowerCAmelCase__)
																					# reshape back to match original input shape
																					_lowerCAmelCase     :							Any    =    z_q.permute(0,			3,			1,			2).contiguous()

														return z_q



class   UpperCAmelCase_    (       lowerCamelCase_):







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

														'''simple docstring'''

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







							def    snake_case__      (   self,			__a = None):

														'''simple docstring'''

														_lowerCAmelCase     :							List[Any]    =    randn_tensor(
														    self.mean.shape,			generator=lowerCAmelCase__,			device=self.parameters.device,			dtype=self.parameters.dtype)
														_lowerCAmelCase     :							Tuple    =    self.mean + self.std * sample
														return x







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

														'''simple docstring'''

														if self.deterministic:
																					return torch.Tensor([0.0])
														_lowerCAmelCase     :							Tuple    =    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    snake_case__      (   self):

														'''simple docstring'''

														return self.mean




"""simple docstring"""




import unittest

import numpy as np

from diffusers import OnnxStableDiffusionInpaintPipelineLegacy
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    load_numpy,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)


if is_onnx_available():
	import onnxruntime as ort
@nightly
@require_onnxruntime
@require_torch_gpu
class 					lowerCAmelCase	(					unittest.TestCase    ):




		'''simple docstring'''







		@property
		def 							__A						(	self				)     ->							List[Any]:
								return (
								    "CUDAExecutionProvider",
								    {
								        "gpu_mem_limit": "15000000000",  # 15GB
								        "arena_extend_strategy": "kSameAsRequested",
								    },
								)







		@property
		def 							__A						(	self				)     ->							List[str]:
								SCREAMING_SNAKE_CASE													=					ort.SessionOptions()
								SCREAMING_SNAKE_CASE													=					False
								return options







		def 							__A						(	self				)     ->							Dict:
								SCREAMING_SNAKE_CASE													=					load_image(
								    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
								    '/in_paint/overture-creations-5sI6fQgYIuo.png'				)
								SCREAMING_SNAKE_CASE													=					load_image(
								    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
								    '/in_paint/overture-creations-5sI6fQgYIuo_mask.png'				)
								SCREAMING_SNAKE_CASE													=					load_numpy(
								    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
								    '/in_paint/red_cat_sitting_on_a_park_bench_onnx.npy'				)

								# using the PNDM scheduler by default
								SCREAMING_SNAKE_CASE													=					OnnxStableDiffusionInpaintPipelineLegacy.from_pretrained(
								    'CompVis/stable-diffusion-v1-4'						,       revision='onnx'						,       safety_checker=lowerCAmelCase__						,       feature_extractor=lowerCAmelCase__						,       provider=self.gpu_provider						,       sess_options=self.gpu_options						,       )
								pipe.set_progress_bar_config(disable=lowerCAmelCase__				)

								SCREAMING_SNAKE_CASE													=					'A red cat sitting on a park bench'

								SCREAMING_SNAKE_CASE													=					np.random.RandomState(0				)
								SCREAMING_SNAKE_CASE													=					pipe(
								    prompt=lowerCAmelCase__						,       image=lowerCAmelCase__						,       mask_image=lowerCAmelCase__						,       strength=0.75						,       guidance_scale=7.5						,       num_inference_steps=15						,       generator=lowerCAmelCase__						,       output_type='np'						,       )

								SCREAMING_SNAKE_CASE													=					output.images[0]

								assert image.shape == (512, 512, 3)
								assert np.abs(expected_image - image				).max() < 1e-2


