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




			"""simple docstring"""







			snake_case	     =				42





class __lowerCAmelCase				(					nn.Module      ):




			"""simple docstring"""







			def __init__(				self						:					Any       ,   _snake_case						:					Union[str, Any]=3       ,   _snake_case						:					Dict=3       ,   _snake_case						:					Union[str, Any]=("DownEncoderBlock2D",)       ,   _snake_case						:					int=(64,)       ,   _snake_case						:					int=2       ,   _snake_case						:					Any=32       ,   _snake_case						:					Any="silu"       ,   _snake_case						:					Any=True       ,   )      -> Any:







								"""simple docstring"""
								super().__init__()
								A_				 =						layers_per_block

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

													A_				 =						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
								A_				 =						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
								A_				 =						nn.GroupNorm(num_channels=block_out_channels[-1]       ,   num_groups=UpperCAmelCase_       ,   eps=1e-6	)
								A_				 =						nn.SiLU()

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

								A_				 =						False






			def      lowerCamelCase__     (				self						:					Any       ,   _snake_case						:					List[str]	)      -> Union[str, Any]:







								"""simple docstring"""
								A_				 =						x
								A_				 =						self.conv_in(UpperCAmelCase_	)

								if self.training and self.gradient_checkpointing:

													def create_custom_forward(_snake_case						:					Tuple	):
																		def custom_forward(*_snake_case						:					Optional[Any]	):
																							return module(*UpperCAmelCase_	)

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

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

													# middle
													A_				 =						self.mid_block(UpperCAmelCase_	)

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

								return sample





class __lowerCAmelCase				(					nn.Module      ):




			"""simple docstring"""







			def __init__(				self						:					List[Any]       ,   _snake_case						:					str=3       ,   _snake_case						:					Tuple=3       ,   _snake_case						:					Union[str, Any]=("UpDecoderBlock2D",)       ,   _snake_case						:					str=(64,)       ,   _snake_case						:					Union[str, Any]=2       ,   _snake_case						:					Union[str, Any]=32       ,   _snake_case						:					Optional[Any]="silu"       ,   _snake_case						:					Any="group"       ,   )      -> Union[str, Any]:







								"""simple docstring"""
								super().__init__()
								A_				 =						layers_per_block

								A_				 =						nn.Convad(
								    UpperCAmelCase_       ,   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=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
								A_				 =						list(reversed(UpperCAmelCase_	)	)
								A_				 =						reversed_block_out_channels[0]
								for i, up_block_type in enumerate(UpperCAmelCase_	):
													A_				 =						output_channel
													A_				 =						reversed_block_out_channels[i]

													A_				 =						i == len(UpperCAmelCase_	) - 1

													A_				 =						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_	)
													A_				 =						output_channel

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

								A_				 =						False






			def      lowerCamelCase__     (				self						:					Tuple       ,   _snake_case						:					str       ,   _snake_case						:					str=None	)      -> int:







								"""simple docstring"""
								A_				 =						z
								A_				 =						self.conv_in(UpperCAmelCase_	)

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

													def create_custom_forward(_snake_case						:					Optional[int]	):
																		def custom_forward(*_snake_case						:					Tuple	):
																							return module(*UpperCAmelCase_	)

																		return custom_forward

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

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

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

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

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

								return sample





class __lowerCAmelCase				(					nn.Module      ):




			"""simple docstring"""







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







								"""simple docstring"""
								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      lowerCamelCase__     (				self						:					int       ,   _snake_case						:					Optional[int]	)      -> Dict:







								"""simple docstring"""
								A_				 =						inds.shape
								assert len(UpperCAmelCase_	) > 1
								A_				 =						inds.reshape(ishape[0]       ,   -1	)
								A_				 =						self.used.to(UpperCAmelCase_	)
								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(UpperCAmelCase_	)




			def      lowerCamelCase__     (				self						:					int       ,   _snake_case						:					Dict	)      -> Any:







								"""simple docstring"""
								A_				 =						inds.shape
								assert len(UpperCAmelCase_	) > 1
								A_				 =						inds.reshape(ishape[0]       ,   -1	)
								A_				 =						self.used.to(UpperCAmelCase_	)
								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       ,   UpperCAmelCase_	)
								return back.reshape(UpperCAmelCase_	)




			def      lowerCamelCase__     (				self						:					Dict       ,   _snake_case						:					Optional[int]	)      -> Tuple:







								"""simple docstring"""
								# 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(UpperCAmelCase_       ,   self.embedding.weight	)       ,   dim=1	)

								A_				 =						self.embedding(UpperCAmelCase_	).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(UpperCAmelCase_	)
													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      lowerCamelCase__     (				self						:					int       ,   _snake_case						:					List[Any]       ,   _snake_case						:					Any	)      -> Tuple:







								"""simple docstring"""
								# 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(UpperCAmelCase_	)
													A_				 =						indices.reshape(-1	)  # flatten again

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

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

								return z_q





class __lowerCAmelCase				(					__lowerCamelCase      ):




			"""simple docstring"""







			def __init__(				self						:					List[str]       ,   _snake_case						:					List[Any]       ,   _snake_case						:					Tuple=False	)      -> Optional[Any]:







								"""simple docstring"""
								A_				 =						parameters
								A_				 =						torch.chunk(UpperCAmelCase_       ,   2       ,   dim=1	)
								A_				 =						torch.clamp(self.logvar       ,   -3_0.0       ,   2_0.0	)
								A_				 =						deterministic
								A_				 =						torch.exp(0.5 * self.logvar	)
								A_				 =						torch.exp(self.logvar	)
								if self.deterministic:
													A_				 =						torch.zeros_like(
													    self.mean       ,   device=self.parameters.device       ,   dtype=self.parameters.dtype	)




			def      lowerCamelCase__     (				self						:					Union[str, Any]       ,   _snake_case						:					Optional[torch.Generator] = None	)      -> Any:







								"""simple docstring"""
								# make sure sample is on the same device as the parameters and has same dtype
								A_				 =						randn_tensor(
								    self.mean.shape       ,   generator=UpperCAmelCase_       ,   device=self.parameters.device       ,   dtype=self.parameters.dtype	)
								A_				 =						self.mean + self.std * sample
								return x




			def      lowerCamelCase__     (				self						:					Tuple       ,   _snake_case						:					Dict=None	)      -> int:







								"""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						:					Dict       ,   _snake_case						:					Union[str, Any]       ,   _snake_case						:					Optional[int]=[1, 2, 3]	)      -> Optional[int]:







								"""simple docstring"""
								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=UpperCAmelCase_	)






			def      lowerCamelCase__     (				self						:					Optional[int]	)      -> Tuple:







								"""simple docstring"""
								return self.mean






"""simple docstring"""





# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


UpperCamelCase_					:						str				= {'''configuration_timm_backbone''': ['''TimmBackboneConfig''']}


try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		UpperCamelCase_					:						Union[str, Any]				= ['''TimmBackbone''']


if TYPE_CHECKING:
		from .configuration_timm_backbone import TimmBackboneConfig

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_timm_backbone import TimmBackbone

else:
		import sys

		UpperCamelCase_					:						List[Any]				= _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)




