Models.py

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


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']

		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)
		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)