



def   UpperCAmelCase							(		)							->					Optional[int]:
	_snake_case					:							Union[str, Any]					     =   1

	while True:
		if (
		    sorted(str(A__     )     )
		    == sorted(str(2 * i     )     )
		    == sorted(str(3 * i     )     )
		    == sorted(str(4 * i     )     )
		    == sorted(str(5 * i     )     )
		    == sorted(str(6 * i     )     )
		):
			return i

		i += 1


if __name__ == "__main__":
						print(solution())










import math
from typing import Optional

import numpy as np
import torch
from torch import nn

from .activations import get_activation
def   UpperCAmelCase							(		A__				,  A__				,  A__ = False				,  A__ = 1				,  A__ = 1				,  A__ = 1_00_00				,  )							->					Optional[int]:
	assert len(timesteps.shape     ) == 1, "Timesteps should be a 1d-array"

	_snake_case					:							Tuple					     =   embedding_dim // 2
	_snake_case					:							Optional[Any]					     =   -math.log(A__     ) * torch.arange(
	    start=0				,  end=A__				,  dtype=torch.floataa				,  device=timesteps.device     )
	_snake_case					:							Dict					     =   exponent / (half_dim - downscale_freq_shift)

	_snake_case					:							Optional[int]					     =   torch.exp(A__     )
	_snake_case					:							List[str]					     =   timesteps[:, None].float() * emb[None, :]

	# scale embeddings
	_snake_case					:							Optional[int]					     =   scale * emb

	# concat sine and cosine embeddings
	_snake_case					:							Union[str, Any]					     =   torch.cat([torch.sin(A__     ), torch.cos(A__     )]				,  dim=-1     )

	# flip sine and cosine embeddings
	if flip_sin_to_cos:
		_snake_case					:							Optional[Any]					     =   torch.cat([emb[:, half_dim:], emb[:, :half_dim]]				,  dim=-1     )

	# zero pad
	if embedding_dim % 2 == 1:
		_snake_case					:							Optional[Any]					     =   torch.nn.functional.pad(A__				,  (0, 1, 0, 0)     )
	return emb
def   UpperCAmelCase							(		A__				,  A__				,  A__=False				,  A__=0     )							->					Optional[Any]:
	_snake_case					:							str					     =   np.arange(A__				,  dtype=np.floataa     )
	_snake_case					:							int					     =   np.arange(A__				,  dtype=np.floataa     )
	_snake_case					:							Union[str, Any]					     =   np.meshgrid(A__				,  A__     )  # here w goes first
	_snake_case					:							Optional[int]					     =   np.stack(A__				,  axis=0     )

	_snake_case					:							List[Any]					     =   grid.reshape([2, 1, grid_size, grid_size]     )
	_snake_case					:							Optional[int]					     =   get_ad_sincos_pos_embed_from_grid(A__				,  A__     )
	if cls_token and extra_tokens > 0:
		_snake_case					:							Any					     =   np.concatenate([np.zeros([extra_tokens, embed_dim]     ), pos_embed]				,  axis=0     )
	return pos_embed
def   UpperCAmelCase							(		A__				,  A__     )							->					List[str]:
	if embed_dim % 2 != 0:
		raise ValueError("""embed_dim must be divisible by 2"""     )

	# use half of dimensions to encode grid_h
	_snake_case					:							int					     =   get_ad_sincos_pos_embed_from_grid(embed_dim // 2				,  grid[0]     )  # (H*W, D/2)
	_snake_case					:							List[str]					     =   get_ad_sincos_pos_embed_from_grid(embed_dim // 2				,  grid[1]     )  # (H*W, D/2)

	_snake_case					:							List[Any]					     =   np.concatenate([emb_h, emb_w]				,  axis=1     )  # (H*W, D)
	return emb
def   UpperCAmelCase							(		A__				,  A__     )							->					str:
	if embed_dim % 2 != 0:
		raise ValueError("""embed_dim must be divisible by 2"""     )

	_snake_case					:							Union[str, Any]					     =   np.arange(embed_dim // 2				,  dtype=np.floataa     )
	omega /= embed_dim / 2.0
	_snake_case					:							Optional[Any]					     =   1.0 / 1_00_00**omega  # (D/2,)

	_snake_case					:							Union[str, Any]					     =   pos.reshape(-1     )  # (M,)
	_snake_case					:							Dict					     =   np.einsum("""m,d->md"""				,  A__				,  A__     )  # (M, D/2), outer product

	_snake_case					:							Union[str, Any]					     =   np.sin(A__     )  # (M, D/2)
	_snake_case					:							Tuple					     =   np.cos(A__     )  # (M, D/2)

	_snake_case					:							Union[str, Any]					     =   np.concatenate([emb_sin, emb_cos]				,  axis=1     )  # (M, D)
	return emb





class 		__SCREAMING_SNAKE_CASE      (				nn.Module      ):






				"""simple docstring"""

				def __init__(				self						,	SCREAMING_SNAKE_CASE__=2_24						,	SCREAMING_SNAKE_CASE__=2_24						,	SCREAMING_SNAKE_CASE__=16						,	SCREAMING_SNAKE_CASE__=3						,	SCREAMING_SNAKE_CASE__=7_68						,	SCREAMING_SNAKE_CASE__=False						,	SCREAMING_SNAKE_CASE__=True						,	SCREAMING_SNAKE_CASE__=True						,	):






					"""simple docstring"""

					super().__init__()

					_snake_case					:							Tuple					     =   (height // patch_size) * (width // patch_size)
					_snake_case					:							Union[str, Any]					     =   flatten
					_snake_case					:							str					     =   layer_norm

					_snake_case					:							List[Any]					     =   nn.Convad(
					    SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	kernel_size=(patch_size, patch_size)						,	stride=SCREAMING_SNAKE_CASE__						,	bias=SCREAMING_SNAKE_CASE__      )
					if layer_norm:
						_snake_case					:							Dict					     =   nn.LayerNorm(SCREAMING_SNAKE_CASE__						,	elementwise_affine=SCREAMING_SNAKE_CASE__						,	eps=1e-6      )
					else:
						_snake_case					:							Union[str, Any]					     =   None

					_snake_case					:							Tuple					     =   get_ad_sincos_pos_embed(SCREAMING_SNAKE_CASE__						,	int(num_patches**0.5      )      )
					self.register_buffer("""pos_embed"""						,	torch.from_numpy(SCREAMING_SNAKE_CASE__      ).float().unsqueeze(0      )						,	persistent=SCREAMING_SNAKE_CASE__      )







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__      ):






					"""simple docstring"""

					_snake_case					:							List[str]					     =   self.proj(SCREAMING_SNAKE_CASE__      )
					if self.flatten:
						_snake_case					:							Any					     =   latent.flatten(2      ).transpose(1						,	2      )  # BCHW -> BNC
					if self.layer_norm:
						_snake_case					:							Any					     =   self.norm(SCREAMING_SNAKE_CASE__      )
					return latent + self.pos_embed





class 		__SCREAMING_SNAKE_CASE      (				nn.Module      ):






				"""simple docstring"""

				def __init__(				self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__ = "silu"						,	SCREAMING_SNAKE_CASE__ = None						,	SCREAMING_SNAKE_CASE__ = None						,	SCREAMING_SNAKE_CASE__=None						,	):






					"""simple docstring"""

					super().__init__()

					_snake_case					:							List[Any]					     =   nn.Linear(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      )

					if cond_proj_dim is not None:
						_snake_case					:							Dict					     =   nn.Linear(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	bias=SCREAMING_SNAKE_CASE__      )
					else:
						_snake_case					:							Tuple					     =   None

					_snake_case					:							int					     =   get_activation(SCREAMING_SNAKE_CASE__      )

					if out_dim is not None:
						_snake_case					:							Any					     =   out_dim
					else:
						_snake_case					:							int					     =   time_embed_dim
					_snake_case					:							str					     =   nn.Linear(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      )

					if post_act_fn is None:
						_snake_case					:							Tuple					     =   None
					else:
						_snake_case					:							Optional[Any]					     =   get_activation(SCREAMING_SNAKE_CASE__      )







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__=None      ):






					"""simple docstring"""

					if condition is not None:
						_snake_case					:							int					     =   sample + self.cond_proj(SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Dict					     =   self.linear_a(SCREAMING_SNAKE_CASE__      )

					if self.act is not None:
						_snake_case					:							Optional[int]					     =   self.act(SCREAMING_SNAKE_CASE__      )

					_snake_case					:							Any					     =   self.linear_a(SCREAMING_SNAKE_CASE__      )

					if self.post_act is not None:
						_snake_case					:							str					     =   self.post_act(SCREAMING_SNAKE_CASE__      )
					return sample





class 		__SCREAMING_SNAKE_CASE      (				nn.Module      ):






				"""simple docstring"""

				def __init__(				self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      ):






					"""simple docstring"""

					super().__init__()
					_snake_case					:							List[str]					     =   num_channels
					_snake_case					:							List[Any]					     =   flip_sin_to_cos
					_snake_case					:							List[str]					     =   downscale_freq_shift







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__      ):






					"""simple docstring"""

					_snake_case					:							Any					     =   get_timestep_embedding(
					    SCREAMING_SNAKE_CASE__						,	self.num_channels						,	flip_sin_to_cos=self.flip_sin_to_cos						,	downscale_freq_shift=self.downscale_freq_shift						,	)
					return t_emb





class 		__SCREAMING_SNAKE_CASE      (				nn.Module      ):






				"""simple docstring"""

				def __init__(				self						,	SCREAMING_SNAKE_CASE__ = 2_56						,	SCREAMING_SNAKE_CASE__ = 1.0						,	SCREAMING_SNAKE_CASE__=True						,	SCREAMING_SNAKE_CASE__=True						,	SCREAMING_SNAKE_CASE__=False      ):






					"""simple docstring"""

					super().__init__()
					_snake_case					:							Any					     =   nn.Parameter(torch.randn(SCREAMING_SNAKE_CASE__      ) * scale						,	requires_grad=SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Union[str, Any]					     =   log
					_snake_case					:							Any					     =   flip_sin_to_cos

					if set_W_to_weight:
						# to delete later
						_snake_case					:							int					     =   nn.Parameter(torch.randn(SCREAMING_SNAKE_CASE__      ) * scale						,	requires_grad=SCREAMING_SNAKE_CASE__      )

						_snake_case					:							List[Any]					     =   self.W







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__      ):






					"""simple docstring"""

					if self.log:
						_snake_case					:							Any					     =   torch.log(SCREAMING_SNAKE_CASE__      )

					_snake_case					:							Any					     =   x[:, None] * self.weight[None, :] * 2 * np.pi

					if self.flip_sin_to_cos:
						_snake_case					:							Dict					     =   torch.cat([torch.cos(SCREAMING_SNAKE_CASE__      ), torch.sin(SCREAMING_SNAKE_CASE__      )]						,	dim=-1      )
					else:
						_snake_case					:							int					     =   torch.cat([torch.sin(SCREAMING_SNAKE_CASE__      ), torch.cos(SCREAMING_SNAKE_CASE__      )]						,	dim=-1      )
					return out





class 		__SCREAMING_SNAKE_CASE      (				nn.Module      ):






				"""simple docstring"""

				def __init__(				self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	):






					"""simple docstring"""

					super().__init__()

					_snake_case					:							List[Any]					     =   height
					_snake_case					:							Optional[int]					     =   width
					_snake_case					:							Any					     =   num_embed
					_snake_case					:							int					     =   embed_dim

					_snake_case					:							Optional[Any]					     =   nn.Embedding(self.num_embed						,	SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Any					     =   nn.Embedding(self.height						,	SCREAMING_SNAKE_CASE__      )
					_snake_case					:							List[Any]					     =   nn.Embedding(self.width						,	SCREAMING_SNAKE_CASE__      )







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__      ):






					"""simple docstring"""

					_snake_case					:							Any					     =   self.emb(SCREAMING_SNAKE_CASE__      )

					_snake_case					:							List[str]					     =   self.height_emb(torch.arange(self.height						,	device=index.device      ).view(1						,	self.height      )      )

					# 1 x H x D -> 1 x H x 1 x D
					_snake_case					:							List[str]					     =   height_emb.unsqueeze(2      )

					_snake_case					:							List[Any]					     =   self.width_emb(torch.arange(self.width						,	device=index.device      ).view(1						,	self.width      )      )

					# 1 x W x D -> 1 x 1 x W x D
					_snake_case					:							List[str]					     =   width_emb.unsqueeze(1      )

					_snake_case					:							int					     =   height_emb + width_emb

					# 1 x H x W x D -> 1 x L xD
					_snake_case					:							int					     =   pos_emb.view(1						,	self.height * self.width						,	-1      )

					_snake_case					:							str					     =   emb + pos_emb[:, : emb.shape[1], :]

					return emb





class 		__SCREAMING_SNAKE_CASE      (				nn.Module      ):






				"""simple docstring"""

				def __init__(				self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      ):






					"""simple docstring"""

					super().__init__()
					_snake_case					:							Any					     =   dropout_prob > 0
					_snake_case					:							Optional[int]					     =   nn.Embedding(num_classes + use_cfg_embedding						,	SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Union[str, Any]					     =   num_classes
					_snake_case					:							Tuple					     =   dropout_prob







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__=None      ):






					"""simple docstring"""

					if force_drop_ids is None:
						_snake_case					:							Dict					     =   torch.rand(labels.shape[0]						,	device=labels.device      ) < self.dropout_prob
					else:
						_snake_case					:							List[str]					     =   torch.tensor(force_drop_ids == 1      )
					_snake_case					:							List[Any]					     =   torch.where(SCREAMING_SNAKE_CASE__						,	self.num_classes						,	SCREAMING_SNAKE_CASE__      )
					return labels







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__=None      ):






					"""simple docstring"""

					_snake_case					:							Union[str, Any]					     =   self.dropout_prob > 0
					if (self.training and use_dropout) or (force_drop_ids is not None):
						_snake_case					:							Dict					     =   self.token_drop(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Optional[Any]					     =   self.embedding_table(SCREAMING_SNAKE_CASE__      )
					return embeddings





class 		__SCREAMING_SNAKE_CASE      (				nn.Module      ):






				"""simple docstring"""

				def __init__(				self						,	SCREAMING_SNAKE_CASE__ = 10_24						,	SCREAMING_SNAKE_CASE__ = 7_68						,	SCREAMING_SNAKE_CASE__ = 7_68						,	SCREAMING_SNAKE_CASE__ = 10						,	):






					"""simple docstring"""

					super().__init__()

					_snake_case					:							str					     =   num_image_text_embeds
					_snake_case					:							Optional[Any]					     =   nn.Linear(SCREAMING_SNAKE_CASE__						,	self.num_image_text_embeds * cross_attention_dim      )
					_snake_case					:							Optional[Any]					     =   nn.Linear(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      )







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      ):






					"""simple docstring"""

					_snake_case					:							Optional[int]					     =   text_embeds.shape[0]

					# image
					_snake_case					:							Optional[Any]					     =   self.image_embeds(SCREAMING_SNAKE_CASE__      )
					_snake_case					:							List[str]					     =   image_text_embeds.reshape(SCREAMING_SNAKE_CASE__						,	self.num_image_text_embeds						,	-1      )

					# text
					_snake_case					:							Dict					     =   self.text_proj(SCREAMING_SNAKE_CASE__      )

					return torch.cat([image_text_embeds, text_embeds]						,	dim=1      )





class 		__SCREAMING_SNAKE_CASE      (				nn.Module      ):






				"""simple docstring"""

				def __init__(				self						,	SCREAMING_SNAKE_CASE__ = 7_68						,	SCREAMING_SNAKE_CASE__ = 7_68						,	SCREAMING_SNAKE_CASE__ = 32						,	):






					"""simple docstring"""

					super().__init__()

					_snake_case					:							Dict					     =   num_image_text_embeds
					_snake_case					:							int					     =   nn.Linear(SCREAMING_SNAKE_CASE__						,	self.num_image_text_embeds * cross_attention_dim      )
					_snake_case					:							Optional[int]					     =   nn.LayerNorm(SCREAMING_SNAKE_CASE__      )







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__      ):






					"""simple docstring"""

					_snake_case					:							Tuple					     =   image_embeds.shape[0]

					# image
					_snake_case					:							Tuple					     =   self.image_embeds(SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Tuple					     =   image_embeds.reshape(SCREAMING_SNAKE_CASE__						,	self.num_image_text_embeds						,	-1      )
					_snake_case					:							Union[str, Any]					     =   self.norm(SCREAMING_SNAKE_CASE__      )
					return image_embeds





class 		__SCREAMING_SNAKE_CASE      (				nn.Module      ):






				"""simple docstring"""

				def __init__(				self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__=0.1      ):






					"""simple docstring"""

					super().__init__()

					_snake_case					:							Optional[Any]					     =   Timesteps(num_channels=2_56						,	flip_sin_to_cos=SCREAMING_SNAKE_CASE__						,	downscale_freq_shift=1      )
					_snake_case					:							Optional[Any]					     =   TimestepEmbedding(in_channels=2_56						,	time_embed_dim=SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Dict					     =   LabelEmbedding(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      )







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__=None      ):






					"""simple docstring"""

					_snake_case					:							Optional[int]					     =   self.time_proj(SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Optional[int]					     =   self.timestep_embedder(timesteps_proj.to(dtype=SCREAMING_SNAKE_CASE__      )      )  # (N, D)

					_snake_case					:							Tuple					     =   self.class_embedder(SCREAMING_SNAKE_CASE__      )  # (N, D)

					_snake_case					:							Tuple					     =   timesteps_emb + class_labels  # (N, D)

					return conditioning





class 		__SCREAMING_SNAKE_CASE      (				nn.Module      ):






				"""simple docstring"""

				def __init__(				self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__ = 64      ):






					"""simple docstring"""

					super().__init__()
					_snake_case					:							List[Any]					     =   nn.LayerNorm(SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Any					     =   AttentionPooling(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      )
					_snake_case					:							List[Any]					     =   nn.Linear(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Optional[Any]					     =   nn.LayerNorm(SCREAMING_SNAKE_CASE__      )







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__      ):






					"""simple docstring"""

					_snake_case					:							str					     =   self.norma(SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Any					     =   self.pool(SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Optional[Any]					     =   self.proj(SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Tuple					     =   self.norma(SCREAMING_SNAKE_CASE__      )
					return hidden_states





class 		__SCREAMING_SNAKE_CASE      (				nn.Module      ):






				"""simple docstring"""

				def __init__(				self						,	SCREAMING_SNAKE_CASE__ = 7_68						,	SCREAMING_SNAKE_CASE__ = 7_68						,	SCREAMING_SNAKE_CASE__ = 15_36      ):






					"""simple docstring"""

					super().__init__()
					_snake_case					:							Optional[Any]					     =   nn.Linear(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      )
					_snake_case					:							str					     =   nn.LayerNorm(SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Optional[int]					     =   nn.Linear(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      )







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      ):






					"""simple docstring"""

					_snake_case					:							int					     =   self.text_proj(SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Union[str, Any]					     =   self.text_norm(SCREAMING_SNAKE_CASE__      )

					# image
					_snake_case					:							Any					     =   self.image_proj(SCREAMING_SNAKE_CASE__      )

					return time_image_embeds + time_text_embeds





class 		__SCREAMING_SNAKE_CASE      (				nn.Module      ):






				"""simple docstring"""

				def __init__(				self						,	SCREAMING_SNAKE_CASE__ = 7_68						,	SCREAMING_SNAKE_CASE__ = 15_36      ):






					"""simple docstring"""

					super().__init__()
					_snake_case					:							Union[str, Any]					     =   nn.Linear(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Tuple					     =   nn.LayerNorm(SCREAMING_SNAKE_CASE__      )







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__      ):






					"""simple docstring"""

					_snake_case					:							Optional[int]					     =   self.image_proj(SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Any					     =   self.image_norm(SCREAMING_SNAKE_CASE__      )
					return time_image_embeds





class 		__SCREAMING_SNAKE_CASE      (				nn.Module      ):






				"""simple docstring"""

				def __init__(				self						,	SCREAMING_SNAKE_CASE__ = 7_68						,	SCREAMING_SNAKE_CASE__ = 15_36      ):






					"""simple docstring"""

					super().__init__()
					_snake_case					:							Optional[Any]					     =   nn.Linear(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Optional[int]					     =   nn.LayerNorm(SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Any					     =   nn.Sequential(
					    nn.Convad(3						,	16						,	3						,	padding=1      )						,	nn.SiLU()						,	nn.Convad(16						,	16						,	3						,	padding=1      )						,	nn.SiLU()						,	nn.Convad(16						,	32						,	3						,	padding=1						,	stride=2      )						,	nn.SiLU()						,	nn.Convad(32						,	32						,	3						,	padding=1      )						,	nn.SiLU()						,	nn.Convad(32						,	96						,	3						,	padding=1						,	stride=2      )						,	nn.SiLU()						,	nn.Convad(96						,	96						,	3						,	padding=1      )						,	nn.SiLU()						,	nn.Convad(96						,	2_56						,	3						,	padding=1						,	stride=2      )						,	nn.SiLU()						,	nn.Convad(2_56						,	4						,	3						,	padding=1      )						,	)







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      ):






					"""simple docstring"""

					_snake_case					:							Dict					     =   self.image_proj(SCREAMING_SNAKE_CASE__      )
					_snake_case					:							int					     =   self.image_norm(SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Optional[Any]					     =   self.input_hint_block(SCREAMING_SNAKE_CASE__      )
					return time_image_embeds, hint





class 		__SCREAMING_SNAKE_CASE      (				nn.Module      ):






				"""simple docstring"""

				def __init__(				self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__=None      ):






					"""simple docstring"""

					super().__init__()
					_snake_case					:							Union[str, Any]					     =   dtype
					_snake_case					:							Optional[int]					     =   nn.Parameter(torch.randn(1						,	SCREAMING_SNAKE_CASE__      ) / embed_dim**0.5      )
					_snake_case					:							Optional[int]					     =   nn.Linear(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	dtype=self.dtype      )
					_snake_case					:							Dict					     =   nn.Linear(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	dtype=self.dtype      )
					_snake_case					:							Optional[int]					     =   nn.Linear(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	dtype=self.dtype      )
					_snake_case					:							Dict					     =   num_heads
					_snake_case					:							Tuple					     =   embed_dim // self.num_heads







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__      ):






					"""simple docstring"""

					_snake_case     , _snake_case     , _snake_case					:							Dict					     =   x.size()

					def shape(SCREAMING_SNAKE_CASE__      ):
						# (bs, length, width) --> (bs, length, n_heads, dim_per_head)
						_snake_case					:							Tuple					     =   x.view(SCREAMING_SNAKE_CASE__						,	-1						,	self.num_heads						,	self.dim_per_head      )
						# (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head)
						_snake_case					:							Optional[int]					     =   x.transpose(1						,	2      )
						# (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head)
						_snake_case					:							Optional[int]					     =   x.reshape(bs * self.num_heads						,	-1						,	self.dim_per_head      )
						# (bs*n_heads, length, dim_per_head) --> (bs*n_heads, dim_per_head, length)
						_snake_case					:							Optional[int]					     =   x.transpose(1						,	2      )
						return x

					_snake_case					:							int					     =   x.mean(dim=1						,	keepdim=SCREAMING_SNAKE_CASE__      ) + self.positional_embedding.to(x.dtype      )
					_snake_case					:							Any					     =   torch.cat([class_token, x]						,	dim=1      )  # (bs, length+1, width)

					# (bs*n_heads, class_token_length, dim_per_head)
					_snake_case					:							Optional[int]					     =   shape(self.q_proj(SCREAMING_SNAKE_CASE__      )      )
					# (bs*n_heads, length+class_token_length, dim_per_head)
					_snake_case					:							str					     =   shape(self.k_proj(SCREAMING_SNAKE_CASE__      )      )
					_snake_case					:							List[str]					     =   shape(self.v_proj(SCREAMING_SNAKE_CASE__      )      )

					# (bs*n_heads, class_token_length, length+class_token_length):
					_snake_case					:							Dict					     =   1 / math.sqrt(math.sqrt(self.dim_per_head      )      )
					_snake_case					:							Any					     =   torch.einsum("""bct,bcs->bts"""						,	q * scale						,	k * scale      )  # More stable with f16 than dividing afterwards
					_snake_case					:							Tuple					     =   torch.softmax(weight.float()						,	dim=-1      ).type(weight.dtype      )

					# (bs*n_heads, dim_per_head, class_token_length)
					_snake_case					:							Optional[int]					     =   torch.einsum("""bts,bcs->bct"""						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      )

					# (bs, length+1, width)
					_snake_case					:							Tuple					     =   a.reshape(SCREAMING_SNAKE_CASE__						,	-1						,	1      ).transpose(1						,	2      )

					return a[:, 0, :]  # cls_token





