Initial commit

Models.py

@@ -0,0 +1,1159 @@
+import json
+from pathlib import Path
+from typing import Optional
+import torch
+import torch.backends.cuda
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+
+from transformers.activations import QuickGELUActivation
+import math
+from einops.layers.torch import Rearrange
+import einops
+
+
+MODEL_CONFIGS = {
+	# Custom models trained from scratch
+	# "Standard" definitions:
+	# name | layers | width | heads
+	#  B   |   12   |  768  |   12
+	#  L   |   24   | 1024  |   16
+	#  H   |   32   | 1280  |   16
+	#  G   |   48   | 1664  |   16
+	#  e   |   56   | 1792  |   16
+	#  22  |   48   | 6144  |   48
+
+	# B/16, 224, PaLM, GELU
+	'CustomTest6': {
+		'class': 'CLIPLikeModel',
+		'embedding_dim': 768,
+		'num_attention_heads': 12,
+		'activation_cls': nn.GELU,
+		'num_channels': 3,
+		'patch_size': 16,
+		'use_palm_alt': True,
+		'num_layers': 12,
+		'use_mha_alt': False,
+		'good_dropout': False,
+	},
+
+	# GAP head + Sinusoidal positional embeddings + 448 image size
+	'CustomTest18': {
+		'class': 'CLIPLikeModel',
+		'embedding_dim': 768,
+		'num_attention_heads': 12,
+		'activation_cls': nn.GELU,
+		'num_channels': 3,
+		'patch_size': 16,
+		'use_palm_alt': True,
+		'num_layers': 12,
+		'use_mha_alt': False,
+		'good_dropout': False,
+		'use_gap_head': True,
+		'sine_positional_embeddings': True,
+	},
+
+	# SW Model + B/16 + ASL + 448 image size
+	# cutout_max_pct = 0
+	# mixup_alpha = 0.8
+	# noise_level = 2
+	# random_resize_method = true
+	# total_labels = 6549
+	'SWModel1': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': False},
+
+	# Sinusoidal positional embeddings 
+	'SWModel2': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True},
+
+	# Sinusoidal positional embeddings + 224 image size + L/14
+	'SWModel3': {'class': 'ViT', 'num_blocks': 24, 'patch_size': 14, 'd_model': 1024, 'mlp_dim': 1024*4, 'num_heads': 16, 'stochdepth_rate': 0.05, 'layerscale_init': 1e-1, 'use_sine': True},
+
+	# Sinusoidal positional embeddings + 224 image size + G/14
+	'SWModel4': {'class': 'ViT', 'num_blocks': 48, 'patch_size': 14, 'd_model': 1664, 'mlp_dim': 1664*4, 'num_heads': 16, 'stochdepth_rate': 0.05, 'layerscale_init': 1e-1, 'use_sine': True},
+
+	# Sinusoidal positional embeddings + focal loss
+	'SWModel5': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True},
+
+	'SWModel6': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True},
+
+	'SWModel7': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True},
+	'SWModel8': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True},
+	'SWModel9': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True},
+	'SWModel10': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True},
+	'SWModel11': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0, 'use_sine': True},
+
+	# Trying head_mean_after
+	'SWModel12': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True, 'head_mean_after': True},
+
+	# Fat boy
+	'SWModel13': {'class': 'ViT', 'num_blocks': 6, 'patch_size': 16, 'd_model': 1536, 'mlp_dim': 1536*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True},
+
+	# L/14
+	'SWModel14': {'class': 'ViT', 'num_blocks': 24, 'patch_size': 14, 'd_model': 1024, 'mlp_dim': 1024*4, 'num_heads': 16, 'stochdepth_rate': 0.05, 'layerscale_init': 1e-1, 'use_sine': True},
+	'SWModel15': {'class': 'ViT', 'num_blocks': 24, 'patch_size': 14, 'd_model': 1024, 'mlp_dim': 1024*4, 'num_heads': 16, 'stochdepth_rate': 0.05, 'layerscale_init': 1e-5, 'use_sine': True},
+	'SWModel16': {'class': 'ViT', 'num_blocks': 24, 'patch_size': 14, 'd_model': 1024, 'mlp_dim': 1024*4, 'num_heads': 16, 'stochdepth_rate': 0.10, 'layerscale_init': 1e-1, 'use_sine': True},
+	'SWModel16f': {'class': 'ViT', 'num_blocks': 24, 'patch_size': 14, 'd_model': 1024, 'mlp_dim': 1024*4, 'num_heads': 16, 'stochdepth_rate': 0.10, 'layerscale_init': 1e-1, 'use_sine': True},
+	'SWModel22': {'class': 'ViT', 'num_blocks': 24, 'patch_size': 14, 'd_model': 1024, 'mlp_dim': 1024*4, 'num_heads': 16, 'stochdepth_rate': 0.20, 'layerscale_init': 1e-1, 'use_sine': True},
+	'SWModel25': {'class': 'ViT', 'num_blocks': 24, 'patch_size': 16, 'd_model': 1024, 'mlp_dim': 1024*4, 'num_heads': 16, 'stochdepth_rate': 0.15, 'layerscale_init': 1e-1, 'use_sine': True, 'cnn_stem': 'conv:c=128;ln;relu;conv:c=256;ln;relu;conv:c=512;ln;relu;conv:c=1024;ln;relu;conv:c=1024,s=1,k=1,p=0'},
+
+	# CNN stem
+	'SWModel18': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True, 'cnn_stem': 'conv:c=64;bn;relu;conv:c=128;bn;relu;conv:c=256;bn;relu;conv:c=512;bn;relu;conv:c=768,s=1,k=1'},
+	'SWModel19': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True, 'cnn_stem': 'conv:c=64;bn;relu;conv:c=128;bn;relu;conv:c=128,s=1;bn;relu;conv:c=256;bn;relu;conv:c=256,s=1;bn;relu;conv:c=512;bn;relu;conv:c=768,s=1,k=1,p=0'},
+	'SWModel20': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True, 'cnn_stem': 'conv:c=64;ln;relu;conv:c=128;ln;relu;conv:c=256;ln;relu;conv:c=512;ln;relu;conv:c=768,s=1,k=1,p=0'},
+	'SWModel21': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True, 'cnn_stem': 'conv:c=64;ln;gelu;conv:c=128;ln;gelu;conv:c=256;ln;gelu;conv:c=512;ln;gelu;conv:c=768,s=1,k=1,p=0'},
+	'SWModel23': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True, 'cnn_stem': 'conv:c=64;ln;relu;conv:c=128;ln;relu;conv:c=256;ln;relu;conv:c=512;ln;relu;conv:c=768,s=1,k=1,p=0'},
+	'SWModel24': {'class': 'ViT', 'num_blocks': 12, 'patch_size': 16, 'd_model': 768, 'mlp_dim': 768*4, 'num_heads': 12, 'stochdepth_rate': 0.05, 'use_sine': True, 'cnn_stem': 'conv:c=64;ln;relu;conv:c=128;ln;relu;conv:c=256;ln;relu;conv:c=512;ln;relu;conv:c=768,s=1,k=1,p=0'},
+
+	# H/14
+	'SWModel17': {'class': 'ViT', 'num_blocks': 32, 'patch_size': 14, 'd_model': 1280, 'mlp_dim': 1280*4, 'num_heads': 16, 'stochdepth_rate': 0.05, 'layerscale_init': 1e-1, 'use_sine': True},
+	'SWModel26': {'class': 'ViT', 'num_blocks': 32, 'patch_size': 14, 'd_model': 1280, 'mlp_dim': 1280*4, 'num_heads': 16, 'stochdepth_rate': 0.15, 'layerscale_init': 1e-1, 'use_sine': True},
+}
+
+
+class VisionModel(nn.Module):
+	image_size: int
+	n_tags: int
+
+	def __init__(self, image_size: int, n_tags: int):
+		super().__init__()
+
+		self.image_size = image_size
+		self.n_tags = n_tags
+	
+	@staticmethod
+	def load_model(path: Path | str, device: str | None = None) -> 'VisionModel':
+		"""
+		Load a model from a directory.
+		:param path: The directory containing the model.
+		:return: The model, the image size, and the number of tags.
+		"""
+		with open(Path(path) / 'config.json', 'r') as f:
+			config = json.load(f)
+		
+		if (Path(path) / 'model.safetensors').exists():
+			from safetensors.torch import load_file
+			resume = load_file(Path(path) / 'model.safetensors', device='cpu')
+		else:
+			resume = torch.load(Path(path) / 'model.pt', map_location=torch.device('cpu'))
+
+		model_classes = VisionModel.__subclasses__()
+		model_cls = next(cls for cls in model_classes if cls.__name__ == config['class'])
+
+		model = model_cls(**{k: v for k, v in config.items() if k != 'class'})
+		model.load(resume['model'])
+		if device is not None:
+			model = model.to(device)
+
+		return model
+	
+	@staticmethod
+	def from_config(config: dict) -> 'VisionModel':
+		model_classes = VisionModel.__subclasses__()
+		model_cls = next(cls for cls in model_classes if cls.__name__ == config['class'])
+		return model_cls(**{k: v for k, v in config.items() if k != 'class'})
+	
+	def get_optimized_parameters(self, lr: float):
+		raise NotImplementedError
+	
+	def save(self):
+		raise NotImplementedError
+	
+	def load(self, state_dict):
+		raise NotImplementedError
+
+
+def basic_calculate_loss(preds: dict[str, torch.Tensor], batch: dict, pos_weight: torch.Tensor | None, loss_type: str):
+	def asl_helper(preds, target):
+		p = F.softmax(preds, dim=1)
+		xs_pos = p.clamp(min=1e-6)
+		xs_neg = (1 - p).clamp(min=1e-6)
+
+		los_pos = torch.log(torch.gather(xs_pos, 1, target.unsqueeze(1))).sum()
+		los_neg = torch.log(xs_neg)
+		los_neg = los_neg.sum() - torch.gather(los_neg, 1, target.unsqueeze(1)).sum()
+		loss = los_pos + los_neg
+
+		return -loss
+
+	if loss_type == "ce":
+		loss = F.binary_cross_entropy_with_logits(preds['tags'], batch['tags'])
+	elif loss_type == "weighted":
+		loss = F.binary_cross_entropy_with_logits(preds['tags'], batch['tags'], pos_weight=pos_weight)
+	elif loss_type == "focal":
+		gamma = 2
+		p = torch.sigmoid(preds['tags'])
+		ce_loss = F.binary_cross_entropy_with_logits(preds['tags'], batch['tags'], reduction='none')
+		p_t = p * batch['tags'] + (1 - p) * (1 - batch['tags'])
+		loss = ce_loss * ((1 - p_t) ** gamma)
+		loss = loss.mean()
+	elif loss_type == "focal2":
+		gamma = 2
+		p = torch.sigmoid(preds['tags'])
+		ce_loss = F.binary_cross_entropy_with_logits(preds['tags'], batch['tags'], reduction='none')
+		p_t = p * batch['tags'] + (1 - p) * (1 - batch['tags'])
+		loss = ce_loss * ((1 - p_t) ** gamma) * 256
+		loss = loss.mean()
+	elif loss_type == "asl":
+		p = torch.sigmoid(preds['tags'])
+		xs_pos = p
+		xs_neg = 1 - p
+
+		los_pos = batch['tags'] * torch.log(xs_pos.clamp(min=1e-6))
+		los_neg = (1 - batch['tags']) * torch.log(xs_neg.clamp(min=1e-6))
+		loss = los_pos + los_neg
+		loss = -loss.sum()
+
+		# Rating
+		loss = loss + asl_helper(preds['rating'], batch['rating'])
+
+		# Score
+		loss = loss + asl_helper(preds['score'], batch['score'])
+	elif loss_type == "asl2":
+		p = torch.sigmoid(preds['tags'])
+		xs_pos = p
+		xs_neg = 1 - p
+
+		los_pos = batch['tags'] * torch.log(xs_pos.clamp(min=1e-6))
+		los_neg = (1 - batch['tags']) * torch.log(xs_neg.clamp(min=1e-6))
+		loss = -los_pos - los_neg
+		loss = loss.sum()
+	elif loss_type == "asl3":
+		p = torch.sigmoid(preds['tags'])
+		xs_pos = p
+		xs_neg = 1 - p
+
+		los_pos = batch['tags'] * torch.log(xs_pos.clamp(min=1e-6))
+		los_neg = (1 - batch['tags']) * torch.log(xs_neg.clamp(min=1e-6))
+		loss = -los_pos - los_neg
+		loss = loss.mean()
+	elif loss_type == "asl4":
+		p = torch.sigmoid(preds['tags'])
+		xs_pos = p
+		xs_neg = 1 - p
+
+		los_pos = batch['tags'] * torch.log(xs_pos.clamp(min=1e-6))
+		los_neg = (1 - batch['tags']) * torch.log(xs_neg.clamp(min=1e-6))
+		loss = -los_pos - los_neg
+		loss = loss.mean() * 128
+	elif loss_type == "asl5":
+		loss = F.binary_cross_entropy_with_logits(preds['tags'], batch['tags'], pos_weight=pos_weight) * 128
+	elif loss_type == "asl6":
+		loss = F.binary_cross_entropy_with_logits(preds['tags'], batch['tags'], pos_weight=pos_weight) * 256
+	elif loss_type == "asl7":
+		loss = F.binary_cross_entropy_with_logits(preds['tags'], batch['tags'], pos_weight=pos_weight) * 2
+	else:
+		raise ValueError(f"Invalid loss type: {loss_type}")
+	
+	return loss
+
+
+class CLIPMlp(nn.Module):
+	def __init__(self, hidden_size: int, intermediate_size: int, activation_cls):
+		super().__init__()
+		self.activation_fn = activation_cls()
+		self.fc1 = nn.Linear(hidden_size, intermediate_size)
+		self.fc2 = nn.Linear(intermediate_size, hidden_size)
+
+	def forward(self, hidden_states: torch.Tensor):
+		hidden_states = self.fc1(hidden_states)
+		hidden_states = self.activation_fn(hidden_states)
+		hidden_states = self.fc2(hidden_states)
+		return hidden_states
+
+
+class FastCLIPAttention2(nn.Module):
+	"""Fast Attention module for CLIP-like. This is NOT a drop-in replacement for CLIPAttention, since it adds additional flexibility.  Mainly uses xformers."""
+	def __init__(self, hidden_size: int, out_dim: int, num_attention_heads: int, out_seq_len: Optional[int] = None, norm_qk: bool = False):
+		super().__init__()
+		self.out_seq_len = out_seq_len
+		self.embed_dim = hidden_size
+		self.out_dim = out_dim
+		self.norm_qk = norm_qk
+		self.num_heads = num_attention_heads
+		self.head_dim = hidden_size // num_attention_heads
+		assert self.head_dim * num_attention_heads == self.embed_dim, "embed_dim must be divisible by num_attention_heads"
+
+		self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
+		self.kv_proj = nn.Linear(self.embed_dim, self.embed_dim * 2)
+		self.out_proj = nn.Linear(self.embed_dim, self.out_dim)
+
+		if self.norm_qk:
+			self.query_norm = nn.LayerNorm(self.embed_dim)
+			self.key_norm = nn.LayerNorm(self.embed_dim)
+	
+	#def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+	#	return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).contiguous()
+	
+	def forward(self, query_states: torch.Tensor, kv_states: torch.Tensor) -> torch.Tensor:
+		bsz, src_len, embed_dim = kv_states.size()
+		if self.out_seq_len is not None:
+			tgt_len = self.out_seq_len
+		else:
+			tgt_len = src_len
+		
+		kv_states = self.kv_proj(kv_states)  # (bsz, src_len, embed_dim * 2)
+		q_states = self.q_proj(query_states[:, :tgt_len])   # (bsz, tgt_len, embed_dim)
+
+		# NOTE: It is not clear if LayerNorm should be applied to the embed_dim, or to the head_dim
+		if self.norm_qk:
+			q_states = self.query_norm(q_states).type(q_states.dtype)
+			k_states = self.key_norm(kv_states[:, :, :embed_dim]).type(kv_states.dtype)
+			v_states = kv_states[:, :, embed_dim:]
+		else:
+			k_states = kv_states[:, :, :embed_dim]
+			v_states = kv_states[:, :, embed_dim:]
+		
+		q_states = q_states.view(bsz, tgt_len, self.num_heads, self.head_dim).transpose(1, 2)  # (bsz, num_heads, tgt_len, head_dim)
+		k_states = k_states.view(bsz, src_len, self.num_heads, self.head_dim).transpose(1, 2)  # (bsz, num_heads, src_len, head_dim)
+		v_states = v_states.view(bsz, src_len, self.num_heads, self.head_dim).transpose(1, 2)  # (bsz, num_heads, src_len, head_dim)
+
+		# Performs scale of query_states, attention, and softmax
+		with torch.backends.cuda.sdp_kernel(enable_math=False):
+			x = F.scaled_dot_product_attention(q_states, k_states, v_states)   # (bsz, num_heads, tgt_len, head_dim)
+			x = x.transpose(1, 2).contiguous().view(bsz, tgt_len, embed_dim)   # (bsz, tgt_len, embed_dim)
+		
+		# Projection
+		x = self.out_proj(x)  # (bsz, tgt_len, out_dim)
+
+		return x
+
+
+class SkipInit(nn.Module):
+	def __init__(self, hidden_size: int, channel_wise: bool, init_scale: float):
+		super().__init__()
+		self.hidden_size = hidden_size
+		self.channel_wise = channel_wise
+		self.init_scale = init_scale
+
+		if self.channel_wise:
+			self.scale = nn.Parameter(torch.ones(hidden_size) * init_scale)
+		else:
+			self.scale = nn.Parameter(torch.tensor(init_scale))
+
+	def forward(self, x: torch.Tensor) -> torch.Tensor:
+		return x * self.scale
+
+
+class FastCLIPEncoderLayer(nn.Module):
+	def __init__(
+		self,
+		hidden_size: int,
+		num_attention_heads: int,
+		out_seq_len: Optional[int],
+		activation_cls = QuickGELUActivation,
+		use_palm_alt: bool = False,
+		norm_qk: bool = False,
+		skip_init: Optional[float] = None,
+		stochastic_depth: Optional[float] = None,
+	):
+		super().__init__()
+
+		self.use_palm_alt = use_palm_alt
+		self.stochastic_depth = stochastic_depth
+		
+		self.self_attn = FastCLIPAttention2(
+			hidden_size=hidden_size,
+			out_dim=hidden_size,
+			num_attention_heads=num_attention_heads,
+			out_seq_len=out_seq_len,
+			norm_qk=norm_qk,
+		)
+		self.mlp = CLIPMlp(hidden_size, 4 * hidden_size, activation_cls)
+		self.layer_norm1 = nn.LayerNorm(hidden_size)
+		if not use_palm_alt:
+			self.layer_norm2 = nn.LayerNorm(hidden_size)
+		
+		if skip_init is not None:
+			self.attn_skip_init = SkipInit(hidden_size, channel_wise=True, init_scale=skip_init)
+			self.mlp_skip_init = SkipInit(hidden_size, channel_wise=True, init_scale=skip_init)
+		else:
+			self.attn_skip_init = nn.Identity()
+			self.mlp_skip_init = nn.Identity()
+	
+	def forward(self, hidden_states: torch.Tensor):
+		residual = hidden_states
+		hidden_states = self.layer_norm1(hidden_states)
+
+		if not self.use_palm_alt:
+			hidden_states = self.self_attn(query_states=hidden_states, kv_states=hidden_states)
+			hidden_states = self.attn_skip_init(hidden_states)
+			hidden_states = hidden_states + residual[:, :hidden_states.size(1)]
+
+			residual = hidden_states
+			hidden_states = self.layer_norm2(hidden_states)
+			hidden_states = self.mlp(hidden_states)
+			hidden_states = self.mlp_skip_init(hidden_states)
+			hidden_states = hidden_states + residual
+		else:
+			# An alternative implementation inspired by the PALM paper
+			# By performing the attention and MLP in parallel it's possible to fuse the linear projections of the attention and MLP layers
+			# We don't do that here yet, but that supposedly improves efficiency without hurting performance
+			attn = self.self_attn(query_states=hidden_states, kv_states=hidden_states)
+			attn = self.attn_skip_init(attn)
+			mlp = self.mlp(hidden_states[:, :attn.size(1)])
+			mlp = self.mlp_skip_init(mlp)
+
+			if self.stochastic_depth is not None:
+				attn = torchvision.ops.stochastic_depth(attn, self.stochastic_depth, mode='row', training=self.training)
+				mlp = torchvision.ops.stochastic_depth(mlp, self.stochastic_depth, mode='row', training=self.training)
+
+			hidden_states = residual[:, :attn.size(1)] + attn + mlp
+
+		return hidden_states
+
+
+def sinusoidal_position_embedding(width: int, height: int, depth: int, dtype, device, temperature = 10000):
+	"""
+	Sinusoidal position embedding. Returns a flat tensor of shape (h * w, d).
+	"""
+	assert depth % 4 == 0, "Embedding dimension must be divisible by 4."
+
+	y, x = torch.meshgrid(torch.arange(height, device=device), torch.arange(width, device=device), indexing="ij")
+	omega = torch.arange(depth // 4, device=device) / (depth // 4 - 1)
+	omega = 1. / (temperature ** omega)
+
+	y = y.flatten()[:, None] * omega[None, :]
+	x = x.flatten()[:, None] * omega[None, :]
+	embedding = torch.cat([x.sin(), x.cos(), y.sin(), y.cos()], dim=1)
+
+	return embedding.type(dtype)
+
+
+class CLIPEmbeddingLayer(nn.Module):
+	def __init__(self, hidden_size: int, num_channels: int, image_size: int, patch_size: int, patch_dropout: float = 0.0, good_dropout: bool = False, dpn: bool = False, sine_positional_embeddings: bool = False):
+		super().__init__()
+
+		assert image_size % patch_size == 0, "Image dimensions must be divisible by the patch size."
+
+		seq_len = (image_size // patch_size) ** 2
+		self.patch_dropout = patch_dropout
+		self.hidden_size = hidden_size
+		self.good_dropout = good_dropout
+		self.dpn = dpn
+		self.sine_positional_embeddings = sine_positional_embeddings
+		self.patch_size = patch_size
+
+		self.patch_embeddings = nn.Conv2d(
+			in_channels=num_channels,
+			out_channels=hidden_size,
+			kernel_size=patch_size,
+			stride=patch_size,
+			bias=False,
+		)
+		if not self.sine_positional_embeddings:
+			self.positional_embeddings = nn.Embedding(seq_len, hidden_size)
+		self.register_buffer("position_ids", torch.arange(seq_len))
+
+		if self.dpn:
+			self.to_patch_embeddings = nn.Sequential(
+				Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=patch_size, p2=patch_size),
+				nn.LayerNorm(3 * patch_size * patch_size),
+				nn.Linear(3 * patch_size * patch_size, hidden_size),
+				nn.LayerNorm(hidden_size),
+			)
+		else:
+			self.to_patch_embeddings = nn.Conv2d(
+				in_channels=num_channels,
+				out_channels=hidden_size,
+				kernel_size=patch_size,
+				stride=patch_size,
+				bias=False,
+			)
+	
+	def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
+		B, C, H, W = pixel_values.shape
+		assert H % self.patch_size == 0, f"Input image height ({H}) needs to be divisible by the patch size ({self.patch_size})."
+		assert W % self.patch_size == 0, f"Input image width ({W}) needs to be divisible by the patch size ({self.patch_size})."
+
+		if self.dpn:
+			patches = self.to_patch_embeddings(pixel_values)
+		else:
+			patches = self.to_patch_embeddings(pixel_values)
+			patches = patches.flatten(2).transpose(1, 2)
+		
+		seq_len = patches.shape[1]
+		patch_dropout = int(math.ceil((1.0 - self.patch_dropout) * seq_len))
+		
+		if self.sine_positional_embeddings:
+			position_embeddings = sinusoidal_position_embedding(W // self.patch_size, H // self.patch_size, self.hidden_size, pixel_values.dtype, pixel_values.device)
+		else:
+			position_embeddings = self.positional_embeddings(self.position_ids)
+
+		if patch_dropout == seq_len or not self.training:
+			embeddings = patches + position_embeddings
+		elif self.good_dropout:
+			# Pick random patches to drop out
+			# The "good_dropout" variant uses random permutations for each batch item, but is slightly slower and involves more code
+
+			# The below method is a nice trick to generate a batch of random permutations.
+			# Torch (as of 1.13) doesn't have a built-in function to do this, and a for loop of torch.randperm is slow.
+			# Based on some benchmarks I measured the generation of the mask and the fetching to be only 50% slower than the non-"good_dropout" variant.
+			# And the time taken here is only a fraction of the time spent performing the embedding convolution.
+			# Generate a matrix of random numbers between 0 and 1 of shape (B, seq_len)
+			patch_mask = torch.rand(B, seq_len, device=patches.device)
+			# For each batch tensor, use argsort to convert the random numbers into a permutation of the patch indices
+			patch_mask = torch.argsort(patch_mask, dim=1)
+			# Truncate
+			patch_mask = patch_mask[:, :patch_dropout]
+
+			embeddings = patches.gather(1, patch_mask.unsqueeze(-1).expand(-1, -1, self.hidden_size)) + position_embeddings[patch_mask]
+		else:
+			# The non-"good_dropout" variant uses a single random permutation for all batch items, but is faster and uses less code
+			indices = torch.randperm(seq_len, device=pixel_values.device)[:patch_dropout]
+			embeddings = patches[:, indices, :] + position_embeddings[indices.expand(1, -1)]
+		
+		return embeddings
+
+
+class MHAPoolingHead(nn.Module):
+	def __init__(self, hidden_size: int, num_attention_heads: int, activation_cls, out_dim: int, alt_style: bool, norm_qk: bool):
+		super().__init__()
+
+		self.out_dim = out_dim if not alt_style else hidden_size
+
+		self.probe = nn.Parameter(torch.randn(hidden_size))
+
+		self.mlp = CLIPMlp(hidden_size, 4 * hidden_size, activation_cls)
+		self.layer_norm = nn.LayerNorm(hidden_size)
+		self.pooling_head = nn.Linear(hidden_size, 1)
+
+		self.self_attn = FastCLIPAttention2(
+			hidden_size=hidden_size,
+			out_dim=self.out_dim,
+			num_attention_heads=num_attention_heads,
+			out_seq_len=1,
+			norm_qk=norm_qk,
+		)
+		self.mlp = CLIPMlp(self.out_dim, 4 * self.out_dim, activation_cls)
+		self.layer_norm1 = nn.LayerNorm(hidden_size)
+		self.layer_norm2 = nn.LayerNorm(self.out_dim)
+
+		if alt_style:
+			self.final_proj = nn.Linear(hidden_size, out_dim)
+		else:
+			self.final_proj = nn.Identity()
+	
+	def forward(self, hidden_states: torch.Tensor):
+		hidden_states = self.layer_norm1(hidden_states)
+		query_states = self.probe.unsqueeze(0).unsqueeze(0).expand(hidden_states.size(0), 1, -1)
+
+		hidden_states = self.self_attn(query_states=query_states, kv_states=hidden_states)
+		# We don't use a residual connection here because the out_dim is different from the hidden_size
+
+		residual = hidden_states
+		hidden_states = self.layer_norm2(hidden_states)
+		hidden_states = self.mlp(hidden_states)
+		hidden_states = hidden_states + residual
+		hidden_states = self.final_proj(hidden_states)
+
+		return hidden_states.squeeze(1)
+
+
+class GAPHead(nn.Module):
+	def __init__(self, hidden_size: int, out_dim: int):
+		super().__init__()
+
+		self.norm = nn.LayerNorm(hidden_size)
+		self.proj = nn.Linear(hidden_size, out_dim)
+
+	def forward(self, x: torch.Tensor) -> torch.Tensor:
+		x = x.mean(dim=1)
+		x = self.norm(x)
+		x = self.proj(x)
+		return x
+
+
+class CLIPLikeModel(VisionModel):
+	def __init__(
+		self,
+		n_tags: int,
+		embedding_dim: int,
+		num_attention_heads: int,
+		activation_cls,
+		num_channels: int,
+		image_size: int,
+		patch_size: int,
+		patch_dropout: float,
+		use_palm_alt: bool,
+		num_layers: int,
+		use_mha_alt: bool,
+		loss_type: str,
+		good_dropout: bool=False,
+		dpn: bool=False,
+		sine_positional_embeddings: bool=False,
+		norm_qk: bool = False,
+		no_wd_bias: bool = False,
+		use_gap_head: bool = False,
+		skip_init: Optional[float] = None,
+		stochastic_depth: Optional[float] = None,
+	):
+		super().__init__(image_size, n_tags)
+
+		out_dim = n_tags
+		self.n_tags = n_tags
+		self.loss_type = loss_type
+		self.no_wd_bias = no_wd_bias
+		
+		stochastic_depth_space = torch.linspace(0, stochastic_depth, num_layers) if stochastic_depth is not None else None
+
+		self.embedding_layer = CLIPEmbeddingLayer(embedding_dim, num_channels, image_size, patch_size, patch_dropout, good_dropout, dpn, sine_positional_embeddings)
+		self.pre_layer_norm = nn.LayerNorm(embedding_dim)
+		self.encoder_layers = nn.ModuleList([FastCLIPEncoderLayer(
+			hidden_size=embedding_dim,
+			num_attention_heads=num_attention_heads,
+			out_seq_len=None,
+			activation_cls=activation_cls,
+			use_palm_alt=use_palm_alt,
+			norm_qk=norm_qk,
+			skip_init=skip_init,
+			stochastic_depth=stochastic_depth_space[i].item() if stochastic_depth_space is not None else None,
+		) for i in range(num_layers)])
+
+		if use_gap_head:
+			self.pooling_head = GAPHead(embedding_dim, out_dim)
+		else:
+			self.pooling_head = MHAPoolingHead(embedding_dim, num_attention_heads, activation_cls, out_dim, use_mha_alt, norm_qk=norm_qk)
+	
+	def forward(self, batch):
+		hidden_states = self.embedding_layer(batch['image'])
+		hidden_states = self.pre_layer_norm(hidden_states)
+
+		for layer in self.encoder_layers:
+			hidden_states = layer(hidden_states)
+		
+		preds = self.pooling_head(hidden_states)
+
+		result = {
+			'tags': preds,
+		}
+
+		return result
+	
+	def calculate_loss(self, preds, batch, pos_weight):
+		return basic_calculate_loss(preds, batch, pos_weight, self.loss_type)
+	
+	def get_optimized_parameters(self, lr: float):
+		if self.no_wd_bias:
+			return self.get_optimized_parameters_no_wd_bias()
+		else:
+			return self.parameters()
+	
+	def get_optimized_parameters_no_wd_bias(self):
+		decay = []
+		no_decay = []
+
+		for name, param in self.named_parameters():
+			if not param.requires_grad:
+				continue
+
+			if len(param.shape) == 1 or name.endswith(".bias"):
+				no_decay.append(param)
+				print(f'No decay: {name}')
+			else:
+				decay.append(param)
+		
+		return [
+			{'params': decay},
+			{'params': no_decay, 'weight_decay': 0.},
+		]
+	
+	def save(self):
+		return self.state_dict()
+	
+	def load(self, state_dict):
+		self.load_state_dict(state_dict)
+
+
+class MaskedAutoEncoderViT(nn.Module):
+	def __init__(
+		self,
+		n_tags: int,
+
+		embedding_dim: int,
+		num_attention_heads: int,
+		activation_cls,
+		num_channels: int,
+		image_size: int,
+		patch_size: int,
+		num_layers: int,
+		loss_type: str,
+		sine_positional_embeddings: bool=False,
+
+		decoder_embedding_dim: int = 512,
+		decoder_num_attention_heads: int = 8,
+		decoder_num_layers: int = 6,
+		decoder_force_projection: bool = False,
+
+		masking_ratio: float = 0.75,
+		mae_loss_weight: float = 1.0,
+		mae_normalize_targets: bool = False,
+		mae_post_norm: bool = False,
+	):
+		super().__init__()
+
+		self.n_tags = n_tags
+		self.seq_len = (image_size // patch_size) ** 2
+		self.embedding_dim = embedding_dim
+		self.decoder_embedding_dim = decoder_embedding_dim
+		self.sine_positional_embeddings = sine_positional_embeddings
+		self.image_size = image_size
+		self.patch_size = patch_size
+		self.masking_ratio = masking_ratio
+		self.loss_type = loss_type
+		self.mae_loss_weight = mae_loss_weight
+		self.mae_normalize_targets = mae_normalize_targets
+
+		if not self.sine_positional_embeddings:
+			self.positional_embeddings = nn.Embedding(self.seq_len, embedding_dim)
+			self.decoder_positional_embeddings = nn.Embedding(self.seq_len, decoder_embedding_dim)
+		self.register_buffer("position_ids", torch.arange(self.seq_len))
+
+		self.to_patches = Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=patch_size, p2=patch_size)
+		self.patch_embedder = nn.Linear(num_channels * patch_size * patch_size, embedding_dim)
+
+		# Encoder
+		self.pre_layer_norm = nn.LayerNorm(embedding_dim)
+		self.encoder_layers = nn.ModuleList([FastCLIPEncoderLayer(
+			hidden_size=embedding_dim,
+			num_attention_heads=num_attention_heads,
+			out_seq_len=None,
+			activation_cls=activation_cls,
+			use_palm_alt=True,
+			norm_qk=False,
+			skip_init=None,
+		) for _ in range(num_layers)])
+
+		# Head for classification
+		self.pooling_head = GAPHead(embedding_dim, n_tags)
+
+		# Decoder
+		if embedding_dim != decoder_embedding_dim or decoder_force_projection:
+			self.encoder_to_decoder_proj = nn.Linear(embedding_dim, decoder_embedding_dim)
+		else:
+			self.encoder_to_decoder_proj = nn.Identity()
+		self.decoder_pre_layer_norm = nn.LayerNorm(decoder_embedding_dim)
+		self.decoder_layers = nn.ModuleList([FastCLIPEncoderLayer(
+			hidden_size=decoder_embedding_dim,
+			num_attention_heads=decoder_num_attention_heads,
+			out_seq_len=None,
+			activation_cls=activation_cls,
+			use_palm_alt=True,
+			norm_qk=False,
+			skip_init=None,
+		) for _ in range(decoder_num_layers)])
+
+		if mae_post_norm:
+			self.decoder_to_pixel_values = nn.Sequential(
+				nn.LayerNorm(decoder_embedding_dim),
+				nn.Linear(decoder_embedding_dim, num_channels * patch_size * patch_size)
+			)
+		else:
+			self.decoder_to_pixel_values = nn.Linear(decoder_embedding_dim, num_channels * patch_size * patch_size)
+		self.mask_token = nn.Parameter(torch.zeros(decoder_embedding_dim))
+		torch.nn.init.normal_(self.mask_token, std=0.02)
+
+	def forward(self, batch):
+		pixel_values = batch['image']
+		device = pixel_values.device
+		B, C, H, W = pixel_values.shape
+		assert H % self.patch_size == 0, f"Input image height ({H}) needs to be divisible by the patch size ({self.patch_size})."
+		assert W % self.patch_size == 0, f"Input image width ({W}) needs to be divisible by the patch size ({self.patch_size})."
+
+		# Convert image to patches (B, seq_len, C * patch_size * patch_size)
+		patches = self.to_patches(pixel_values)
+		seq_len = patches.shape[1]
+		num_masked = int(self.masking_ratio * seq_len)
+
+		# For each batch tensor, use argsort to convert the random numbers into a permutation of the patch indices
+		# From this we can get the masked and unmasked indices
+		patch_mask = torch.rand(B, seq_len, device=device)
+		patch_mask = torch.argsort(patch_mask, dim=1)
+		masked_indices, unmasked_indices = patch_mask[:, :num_masked], patch_mask[:, num_masked:]
+		batch_range = torch.arange(B, device=device)[:, None]
+
+		# Masked and unmasked patches
+		unmasked_patches = patches[batch_range, unmasked_indices]
+		masked_patches = patches[batch_range, masked_indices]
+
+		# Embed unmasked patches for the encoder (B, seq_len, embedding_dim)
+		tokens = self.patch_embedder(unmasked_patches)
+
+		if self.sine_positional_embeddings:
+			position_embeddings = sinusoidal_position_embedding(W // self.patch_size, H // self.patch_size, self.embedding_dim, pixel_values.dtype, device)
+			decoder_position_embeddings = sinusoidal_position_embedding(W // self.patch_size, H // self.patch_size, self.decoder_embedding_dim, pixel_values.dtype, device)
+		else:
+			position_embeddings = self.positional_embeddings(self.position_ids)
+			decoder_position_embeddings = self.decoder_positional_embeddings(self.position_ids)
+		
+		# Add position embeddings
+		tokens = tokens + position_embeddings[unmasked_indices]
+
+		# Run the encoder
+		encoded_tokens = self.pre_layer_norm(tokens)
+
+		for layer in self.encoder_layers:
+			encoded_tokens = layer(encoded_tokens)
+		
+		# Label predictions
+		if self.training:
+			preds = self.pooling_head(encoded_tokens)
+		else:
+			# During inference, classify using the entire image
+			# But we'll do the usual for the MAE part, just so we can see how MAE is performing during validation
+			tokens = self.patch_embedder(patches)
+			tokens = tokens + position_embeddings
+			tokens = self.pre_layer_norm(tokens)
+			for layer in self.encoder_layers:
+				tokens = layer(tokens)
+			preds = self.pooling_head(tokens)
+		
+		# Projection for the decoder and position embeddings
+		decoder_tokens = self.encoder_to_decoder_proj(encoded_tokens)
+		decoder_tokens = decoder_tokens + decoder_position_embeddings[unmasked_indices]
+
+		# Fill in the masked patches
+		mask_tokens = einops.repeat(self.mask_token, 'd -> b n d', b = B, n = num_masked)
+		mask_tokens = mask_tokens + decoder_position_embeddings[masked_indices]
+		decoder_tokens = torch.cat([decoder_tokens, mask_tokens], dim=1)
+
+		# Run the decoder
+		decoded_tokens = self.decoder_pre_layer_norm(decoder_tokens)
+		
+		for layer in self.decoder_layers:
+			decoded_tokens = layer(decoded_tokens)
+
+		# Only predict the masked patches
+		# All the masked patches are at the end of the sequence
+		decoded_tokens = decoded_tokens[:, -num_masked:]
+		pred_pixel_values = self.decoder_to_pixel_values(decoded_tokens)
+
+		# Calculate the mae loss
+		if self.mae_normalize_targets:
+			# Normalize each patch by its mean and variance. The ViCHA paper says this provides better results
+			means = masked_patches.mean(dim=-1, keepdim=True)
+			vars = masked_patches.var(dim=-1, keepdim=True)
+			target = (masked_patches - means) / (vars + 1e-6)**0.5
+			mae_loss = F.mse_loss(pred_pixel_values, target)
+		else:
+			mae_loss = F.mse_loss(pred_pixel_values, masked_patches)
+		mae_loss = mae_loss * self.mae_loss_weight
+
+		return {
+			'tags': preds,
+			'mae_loss': mae_loss,
+		}
+
+	def calculate_loss(self, preds, batch, pos_weight):
+		return basic_calculate_loss(preds, batch, pos_weight, self.loss_type) + preds['mae_loss']
+	
+	def get_optimized_parameters(self, lr: float):
+		return self.parameters()
+	
+	def save(self):
+		return self.state_dict()
+	
+	def load(self, state_dict):
+		self.load_state_dict(state_dict)
+
+
+class StochDepth(nn.Module):
+	def __init__(self, drop_rate: float, scale_by_keep: bool = False):
+		super().__init__()
+		self.drop_rate = drop_rate
+		self.scale_by_keep = scale_by_keep
+	
+	def forward(self, x):
+		if not self.training:
+			return x
+		
+		batch_size = x.shape[0]
+		r = torch.rand((batch_size, 1, 1), device=x.device)
+		keep_prob = 1 - self.drop_rate
+		binary_tensor = torch.floor(keep_prob + r)
+		if self.scale_by_keep:
+			x = x / keep_prob
+		
+		return x * binary_tensor
+
+
+class SkipInitChannelwise(nn.Module):
+	def __init__(self, channels, init_val=1e-6):
+		super().__init__()
+		self.channels = channels
+		self.init_val = init_val
+		self.skip = nn.Parameter(torch.ones(channels) * init_val)
+	
+	def forward(self, x):
+		return x * self.skip
+
+
+class PosEmbedding(nn.Module):
+	def __init__(self, d_model: int, max_len: int, use_sine: bool, patch_size: int):
+		super().__init__()
+		self.d_model = d_model
+		self.max_len = max_len
+		self.use_sine = use_sine
+		self.patch_size = patch_size
+
+		if not self.use_sine:
+			self.embedding = nn.Embedding(max_len, d_model)
+			nn.init.trunc_normal_(self.embedding.weight, std=0.02)
+			self.register_buffer("position_ids", torch.arange(max_len))
+	
+	def forward(self, x, width: int, height: int):
+		if self.use_sine:
+			position_embeddings = sinusoidal_position_embedding(width // self.patch_size, height // self.patch_size, self.d_model, x.dtype, x.device)
+		else:
+			position_embeddings = self.embedding(self.position_ids)
+		
+		return x + position_embeddings
+
+
+class MLPBlock(nn.Module):
+	def __init__(self, d_model: int, d_ff: int, stochdepth_rate: float):
+		super().__init__()
+		self.linear1 = nn.Linear(d_model, d_ff)
+		self.linear2 = nn.Linear(d_ff, d_model)
+		self.activation = nn.GELU()
+		if stochdepth_rate > 0:
+			self.stochdepth = StochDepth(stochdepth_rate, scale_by_keep=True)
+		else:
+			self.stochdepth = None
+	
+	def forward(self, x):
+		x = self.linear1(x)
+		x = self.activation(x)
+		if self.stochdepth is not None:
+			x = self.stochdepth(x)
+		x = self.linear2(x)
+		return x
+
+
+class ViTBlock(nn.Module):
+	def __init__(self, num_heads: int, d_model: int, d_ff: int, layerscale_init: float, stochdepth_rate: float):
+		super().__init__()
+		self.num_heads = num_heads
+		self.d_model = d_model
+
+		assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
+
+		# MHA
+		self.norm1 = nn.LayerNorm(d_model)
+		self.qkv_proj = nn.Linear(d_model, d_model * 3)
+		self.out_proj = nn.Linear(d_model, d_model)
+		self.skip_init1 = SkipInitChannelwise(channels=d_model, init_val=layerscale_init)
+		self.stochdepth1 = StochDepth(stochdepth_rate, scale_by_keep=True) if stochdepth_rate > 0 else None
+		
+		# MLP
+		self.norm2 = nn.LayerNorm(d_model)
+		self.mlp = MLPBlock(d_model, d_ff, stochdepth_rate)
+		self.skip_init2 = SkipInitChannelwise(channels=d_model, init_val=layerscale_init)
+		self.stochdepth2 = StochDepth(stochdepth_rate, scale_by_keep=True) if stochdepth_rate > 0 else None
+	
+	def forward(self, x):
+		bsz, src_len, embed_dim = x.shape
+
+		out = x
+		out = self.norm1(out)
+
+		# MHA
+		qkv_states = self.qkv_proj(out).split(self.d_model, dim=-1)
+		q_states = qkv_states[0].view(bsz, src_len, self.num_heads, embed_dim // self.num_heads).transpose(1, 2)  # (bsz, num_heads, src_len, embed_dim // num_heads)
+		k_states = qkv_states[1].view(bsz, src_len, self.num_heads, embed_dim // self.num_heads).transpose(1, 2)  # (bsz, num_heads, src_len, embed_dim // num_heads)
+		v_states = qkv_states[2].view(bsz, src_len, self.num_heads, embed_dim // self.num_heads).transpose(1, 2)  # (bsz, num_heads, src_len, embed_dim // num_heads)
+
+		with torch.backends.cuda.sdp_kernel(enable_math=False):
+			out = F.scaled_dot_product_attention(q_states, k_states, v_states)   # (bsz, num_heads, tgt_len, head_dim)
+			out = out.transpose(1, 2).contiguous().view(bsz, src_len, embed_dim)   # (bsz, tgt_len, embed_dim)
+		
+		out = self.out_proj(out)
+
+		out = self.skip_init1(out)
+		if self.stochdepth1 is not None:
+			out = self.stochdepth1(out)
+		x = out + x
+
+		out = self.norm2(x)
+		out = self.mlp(out)
+		out = self.skip_init2(out)
+		if self.stochdepth2 is not None:
+			out = self.stochdepth2(out)
+		
+		out = out + x
+
+		return out
+
+
+def CaiT_LayerScale_init(network_depth):
+	if network_depth <= 18:
+		return 1e-1
+	elif network_depth <= 24:
+		return 1e-5
+	else:
+		return 1e-6
+
+
+class CNNLayerNorm(nn.Module):
+	def __init__(self, d_model: int):
+		super().__init__()
+		self.norm = nn.LayerNorm(d_model)
+	
+	def forward(self, x: torch.Tensor) -> torch.Tensor:
+		x = x.transpose(1, 3)
+		x = self.norm(x)
+		x = x.transpose(1, 3)
+		return x
+
+
+class CNNStem(nn.Module):
+	def __init__(self, config: str):
+		super().__init__()
+		self.config = config
+
+		layers = []
+		channels = 3
+
+		for line in config.split(";"):
+			ty, line = line.split(":") if ":" in line else (line, "")
+			options = line.split(",")
+			options = [o.split("=") for o in options] if line else []
+			options = {k: v for k, v in options}
+
+			if ty == 'conv':
+				layers.append(nn.Conv2d(
+					in_channels=channels,
+					out_channels=int(options['c']),
+					kernel_size=int(options['k'] if 'k' in options else 3),
+					stride=int(options['s'] if 's' in options else 2),
+					bias=True,
+					padding=int(options['p'] if 'p' in options else 1),
+				))
+				channels = int(options['c'])
+			elif ty == 'bn':
+				layers.append(nn.BatchNorm2d(channels))
+			elif ty == 'ln':
+				layers.append(CNNLayerNorm(channels))
+			elif ty == 'relu':
+				layers.append(nn.ReLU())
+			elif ty == 'gelu':
+				layers.append(nn.GELU())
+
+		self.conv = nn.Sequential(*layers)
+	
+	def forward(self, x: torch.Tensor) -> torch.Tensor:
+		return self.conv(x)
+
+
+class ViT(VisionModel):
+	def __init__(self,
+		n_tags: int,
+		image_size: int,
+		num_blocks: int,
+		patch_size: int,
+		d_model: int,
+		mlp_dim: int,
+		num_heads: int,
+		stochdepth_rate: float,
+		use_sine: bool,
+		loss_type: str,
+		layerscale_init: Optional[float] = None,
+		head_mean_after: bool = False,
+		cnn_stem: str | None = None,
+		patch_dropout: float = 0.0,
+	):
+		super().__init__(image_size, n_tags)
+
+		#assert image_size % patch_size == 0, "image_size must be divisible by patch_size"
+		assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
+
+		out_dim = n_tags
+		self.n_tags = n_tags
+		self.loss_type = loss_type
+		self.patch_size = patch_size
+		self.head_mean_after = head_mean_after
+		self.patch_dropout = patch_dropout
+
+		layerscale_init = CaiT_LayerScale_init(num_blocks) if layerscale_init is None else layerscale_init
+		self.patch_embeddings = nn.Conv2d(
+			in_channels=3,
+			out_channels=d_model,
+			kernel_size=patch_size,
+			stride=patch_size,
+			bias=True,
+		) if cnn_stem is None else CNNStem(cnn_stem)
+		self.pos_embedding = PosEmbedding(d_model, (image_size // patch_size) ** 2, use_sine=use_sine, patch_size=patch_size)
+
+		self.blocks = nn.ModuleList([
+			ViTBlock(num_heads, d_model, mlp_dim, layerscale_init, stochdepth_rate)
+			for _ in range(num_blocks)
+		])
+
+		self.norm = nn.LayerNorm(d_model)
+		self.head = nn.Linear(d_model, out_dim)
+
+	def forward(self, batch, return_embeddings=False, return_loss: bool = False, pos_weight = None):
+		B, C, H, W = batch['image'].shape
+		assert H % self.patch_size == 0, f"Input image height ({H}) needs to be divisible by the patch size ({self.patch_size})."
+		assert W % self.patch_size == 0, f"Input image width ({W}) needs to be divisible by the patch size ({self.patch_size})."
+
+		x = self.patch_embeddings(batch['image'])  # (bsz, d_model, patch_num, patch_num)
+		x = x.flatten(2).transpose(1, 2)  # (bsz, patch_num ** 2, d_model)
+		x = self.pos_embedding(x, W, H)   # (bsz, patch_num ** 2, d_model)
+
+		# Patch dropout
+		seq_len = x.shape[1]
+		patch_dropout = int(math.ceil((1.0 - self.patch_dropout) * seq_len))
+
+		if patch_dropout != seq_len:
+			# Generate a matrix of random numbers between 0 and 1 of shape (B, seq_len)
+			patch_mask = torch.rand(B, seq_len, device=x.device)
+			# For each batch tensor, use argsort to convert the random numbers into a permutation of the patch indices
+			patch_mask = torch.argsort(patch_mask, dim=1)
+			# Truncate
+			patch_mask = patch_mask[:, :patch_dropout]
+
+			x = x.gather(1, patch_mask.unsqueeze(-1).expand(-1, -1, x.shape[-1]))
+
+			#indices = torch.randperm(seq_len, device=x.device)[:patch_dropout]
+			#x = x[:, indices, :]
+
+		# Transformer
+		for block in self.blocks:
+			x = block(x)
+		
+		# Head
+		result = {}
+
+		x = self.norm(x)
+		if self.head_mean_after:
+			x = self.head(x)
+			x = x.mean(dim=1)
+		else:
+			x = x.mean(dim=1)
+			if return_embeddings:
+				result['embeddings'] = x
+			x = self.head(x)
+
+		result['tags'] = x
+
+		if return_loss:
+			result['loss'] = self.calculate_loss(result, batch, pos_weight)
+
+		return result
+	
+	def calculate_loss(self, preds, batch, pos_weight):
+		return basic_calculate_loss(preds, batch, pos_weight, self.loss_type)
+	
+	def get_optimized_parameters(self, lr: float):
+		return self.parameters()
+	
+	def save(self):
+		return self.state_dict()
+	
+	def load(self, state_dict):
+		if 'head.weight' in state_dict and 'head.bias' in state_dict and state_dict['head.weight'].shape[0] == (self.n_tags + 9):
+			# Support old models which included 3 rating and 6 score dimensions
+			state_dict['head.weight'] = state_dict['head.weight'][:self.n_tags]
+			state_dict['head.bias'] = state_dict['head.bias'][:self.n_tags]
+
+		self.load_state_dict(state_dict)

app.py

@@ -0,0 +1,41 @@
+import gradio as gr
+from Models import VisionModel
+import huggingface_hub
+from PIL import Image
+import torch.amp.autocast_mode
+from pathlib import Path
+
+
+MODEL_REPO = "fancyfeast/joytag"
+
+
+@torch.no_grad()
+def predict(image: Image.Image):
+	with torch.amp.autocast_mode.autocast('cuda', enabled=True):
+		preds = model(image)
+		tag_preds = preds['tags'].sigmoid().cpu()
+	
+	return {top_tags[i]: tag_preds[i] for i in range(len(top_tags))}
+
+
+print("Downloading model...")
+path = huggingface_hub.snapshot_download(MODEL_REPO)
+print("Loading model...")
+model = VisionModel.load_model(path)
+model.eval()
+
+with open(Path(path) / 'top_tags.txt', 'r') as f:
+	top_tags = [line.strip() for line in f.readlines() if line.strip()]
+
+print("Starting server...")
+
+gradio_app = gr.Interface(
+	predict,
+	inputs=gr.Image(label="Source", sources=['upload', 'webcam'], type='pil'),
+	outputs=[gr.Label(label="Result", num_top_classes=5)],
+	title="JoyTag",
+)
+
+
+if __name__ == '__main__':
+	gradio_app.launch()

requirements.txt

@@ -0,0 +1,5 @@
+torch==2.1.2
+transformers==4.36.2
+torchvision==0.16.2
+einops==0.7.0
+safetensors==0.4.1