



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     (	lowerCAmelCase__   ):
	lowerCAmelCase__       					=  4_2



class  __lowerCAmelCase     (	nn.Module   ):




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



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

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

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

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

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

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

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

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

				__lowerCamelCase													=    False




	def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



				'''simple docstring'''
				__lowerCamelCase													=    x
				__lowerCamelCase													=    self.conv_in(__UpperCAmelCase				)

				if self.training and self.gradient_checkpointing:

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

										return custom_forward

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

				else:
							# down
							for down_block in self.down_blocks:
										__lowerCamelCase													=    down_block(__UpperCAmelCase				)

							# middle
							__lowerCamelCase													=    self.mid_block(__UpperCAmelCase				)

				# post-process
				__lowerCamelCase													=    self.conv_norm_out(__UpperCAmelCase				)
				__lowerCamelCase													=    self.conv_act(__UpperCAmelCase				)
				__lowerCamelCase													=    self.conv_out(__UpperCAmelCase				)

				return sample



class  __lowerCAmelCase     (	nn.Module   ):




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



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

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

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

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

				# mid
				__lowerCamelCase													=    UNetMidBlockaD(
				    in_channels=block_out_channels[-1]	,				resnet_eps=1E-6	,				resnet_act_fn=__UpperCAmelCase	,				output_scale_factor=1	,				resnet_time_scale_shift='''default''' if norm_type == '''group''' else norm_type	,				attention_head_dim=block_out_channels[-1]	,				resnet_groups=__UpperCAmelCase	,				temb_channels=__UpperCAmelCase	,				)

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

							__lowerCamelCase													=    i == len(__UpperCAmelCase				) - 1

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

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

				__lowerCamelCase													=    False




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



				'''simple docstring'''
				__lowerCamelCase													=    z
				__lowerCamelCase													=    self.conv_in(__UpperCAmelCase				)

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

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

										return custom_forward

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

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

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

							# up
							for up_block in self.up_blocks:
										__lowerCamelCase													=    up_block(__UpperCAmelCase	,				__UpperCAmelCase				)

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

				return sample



class  __lowerCAmelCase     (	nn.Module   ):




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



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

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

				__lowerCamelCase													=    sane_index_shape




	def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



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




	def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



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




	def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



				'''simple docstring'''
				__lowerCamelCase													=    z.permute(0	,				2	,				3	,				1				).contiguous()
				__lowerCamelCase													=    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													=    torch.argmin(torch.cdist(__UpperCAmelCase	,				self.embedding.weight				)	,				dim=1				)

				__lowerCamelCase													=    self.embedding(__UpperCAmelCase				).view(z.shape				)
				__lowerCamelCase													=    None
				__lowerCamelCase													=    None

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

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

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

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

				if self.sane_index_shape:
							__lowerCamelCase													=    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	,				__UpperCAmelCase	,				__UpperCAmelCase				):



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

				# get quantized latent vectors
				__lowerCamelCase													=    self.embedding(__UpperCAmelCase				)

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

				return z_q



class  __lowerCAmelCase     (	lowerCAmelCase__   ):




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



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




	def        lowerCamelCase      (						self	,				__UpperCAmelCase = None				):



				'''simple docstring'''
				__lowerCamelCase													=    randn_tensor(
				    self.mean.shape	,				generator=__UpperCAmelCase	,				device=self.parameters.device	,				dtype=self.parameters.dtype				)
				__lowerCamelCase													=    self.mean + self.std * sample
				return x




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



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




	def        lowerCamelCase      (						self				):



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











from ...configuration_utils import PretrainedConfig
from ...utils import logging


a_			=					logging.get_logger(__name__)

a_			=					{
    """uclanlp/visualbert-vqa""": """https://huggingface.co/uclanlp/visualbert-vqa/resolve/main/config.json""",
    """uclanlp/visualbert-vqa-pre""": """https://huggingface.co/uclanlp/visualbert-vqa-pre/resolve/main/config.json""",
    """uclanlp/visualbert-vqa-coco-pre""": (
        """https://huggingface.co/uclanlp/visualbert-vqa-coco-pre/resolve/main/config.json"""
    ),
    """uclanlp/visualbert-vcr""": """https://huggingface.co/uclanlp/visualbert-vcr/resolve/main/config.json""",
    """uclanlp/visualbert-vcr-pre""": """https://huggingface.co/uclanlp/visualbert-vcr-pre/resolve/main/config.json""",
    """uclanlp/visualbert-vcr-coco-pre""": (
        """https://huggingface.co/uclanlp/visualbert-vcr-coco-pre/resolve/main/config.json"""
    ),
    """uclanlp/visualbert-nlvr2""": """https://huggingface.co/uclanlp/visualbert-nlvr2/resolve/main/config.json""",
    """uclanlp/visualbert-nlvr2-pre""": """https://huggingface.co/uclanlp/visualbert-nlvr2-pre/resolve/main/config.json""",
    """uclanlp/visualbert-nlvr2-coco-pre""": (
        """https://huggingface.co/uclanlp/visualbert-nlvr2-coco-pre/resolve/main/config.json"""
    )
    # See all VisualBERT models at https://huggingface.co/models?filter=visual_bert
}



class  __lowerCAmelCase     (	lowerCAmelCase__   ):
 lowerCAmelCase__       					=  """visual_bert"""




 def __init__(						self	,				__UpperCAmelCase=30522	,				__UpperCAmelCase=768	,				__UpperCAmelCase=512	,				__UpperCAmelCase=12	,				__UpperCAmelCase=12	,				__UpperCAmelCase=3072	,				__UpperCAmelCase="gelu"	,				__UpperCAmelCase=0.1	,				__UpperCAmelCase=0.1	,				__UpperCAmelCase=512	,				__UpperCAmelCase=2	,				__UpperCAmelCase=0.02	,				__UpperCAmelCase=1E-1_2	,				__UpperCAmelCase=False	,				__UpperCAmelCase=True	,				__UpperCAmelCase=1	,				__UpperCAmelCase=0	,				__UpperCAmelCase=2	,				**__UpperCAmelCase	,				):



    '''simple docstring'''
    super().__init__(pad_token_id=__UpperCAmelCase	,				bos_token_id=__UpperCAmelCase	,				eos_token_id=__UpperCAmelCase	,				**__UpperCAmelCase				)

    __lowerCamelCase													=    vocab_size
    __lowerCamelCase													=    max_position_embeddings
    __lowerCamelCase													=    hidden_size
    __lowerCamelCase													=    visual_embedding_dim
    __lowerCamelCase													=    num_hidden_layers
    __lowerCamelCase													=    num_attention_heads
    __lowerCamelCase													=    intermediate_size
    __lowerCamelCase													=    hidden_act
    __lowerCamelCase													=    hidden_dropout_prob
    __lowerCamelCase													=    attention_probs_dropout_prob
    __lowerCamelCase													=    initializer_range
    __lowerCamelCase													=    type_vocab_size
    __lowerCamelCase													=    layer_norm_eps
    __lowerCamelCase													=    bypass_transformer
    __lowerCamelCase													=    special_visual_initialize






