mirror of https://github.com/vladmandic/automatic
232 lines
8.4 KiB
Python
232 lines
8.4 KiB
Python
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from torch import einsum
|
|
from torchvision import transforms
|
|
|
|
from PIL import Image
|
|
from einops import rearrange
|
|
|
|
from .modeling_vit import create_siglip_vit
|
|
|
|
|
|
def create_anyres_preprocess(
|
|
short_size=384,
|
|
long_size=1152,
|
|
patch_size=16,
|
|
random_ratio=None,
|
|
min_short_size=128,
|
|
max_aspect_ratio=3.,
|
|
filtering=True
|
|
):
|
|
|
|
def resize_and_filtering(pil_image):
|
|
pil_image = pil_image.convert('RGB')
|
|
width, height = pil_image.size
|
|
ss, ls = min(width, height), max(width, height)
|
|
aspect_ratio = ls / ss
|
|
if filtering and (ss < min_short_size or aspect_ratio > max_aspect_ratio):
|
|
return None
|
|
target_width, target_height = width, height
|
|
if random_ratio is not None:
|
|
log_ratio = torch.log(torch.tensor(random_ratio))
|
|
sqrt_ratio = torch.exp(0.5 * torch.empty(1).uniform_(log_ratio[0], log_ratio[1])).item()
|
|
target_width = int(round(target_width * sqrt_ratio))
|
|
target_height = int(round(target_height / sqrt_ratio))
|
|
|
|
ss = min(target_width, target_height)
|
|
if ss < short_size:
|
|
target_width = target_width * (short_size / ss)
|
|
target_height = target_height * (short_size / ss)
|
|
|
|
ls = max(target_width, target_height)
|
|
if ls > long_size:
|
|
target_width = target_width * (long_size / ls)
|
|
target_height = target_height * (long_size / ls)
|
|
|
|
target_width = int(round(target_width / patch_size)) * patch_size
|
|
target_height = int(round(target_height / patch_size)) * patch_size
|
|
pil_image = pil_image.resize((target_width, target_height), resample=Image.BICUBIC)
|
|
|
|
to_tensor = transforms.Compose([
|
|
transforms.ToTensor(),
|
|
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
|
|
])
|
|
return to_tensor(pil_image)
|
|
|
|
transform = transforms.Lambda(resize_and_filtering)
|
|
return transform
|
|
|
|
|
|
class IBQ(nn.Module):
|
|
def __init__(self, n_e, e_dim, skip_quantization_prob=0.0, quantization_temp=2.0, beta=0.25, sane_index_shape=False, l2_norm=True):
|
|
super().__init__()
|
|
self.n_e = n_e
|
|
self.e_dim = e_dim
|
|
self.quantization_temp = quantization_temp
|
|
self.skip_quantization_prob = skip_quantization_prob
|
|
self.beta = beta
|
|
self.sane_index_shape = sane_index_shape
|
|
self.l2_norm = l2_norm
|
|
|
|
self.embedding = nn.Embedding(self.n_e, self.e_dim)
|
|
self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e)
|
|
if self.l2_norm:
|
|
self.embedding.weight.data = F.normalize(self.embedding.weight.data, p=2, dim=-1)
|
|
|
|
def forward(self, z, temp=None, rescale_logits=False, return_logits=False, **kwargs):
|
|
assert temp is None or temp == 1.0, "Only for interface compatible with Gumbel"
|
|
assert rescale_logits == False, "Only for interface compatible with Gumbel"
|
|
assert return_logits == False, "Only for interface compatible with Gumbel"
|
|
# reshape z -> (batch, height, width, channel) and flatten
|
|
z = rearrange(z, 'b c h w -> b h w c').contiguous()
|
|
assert z.shape[-1] == self.e_dim
|
|
z_flattened = z.view(-1, self.e_dim)
|
|
# distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
|
|
|
|
if self.l2_norm:
|
|
z = F.normalize(z, p=2, dim=-1)
|
|
z_flattened = F.normalize(z_flattened, p=2, dim=-1)
|
|
embedding = F.normalize(self.embedding.weight, p=2, dim=-1)
|
|
else:
|
|
embedding = self.embedding.weight
|
|
|
|
d = torch.sum(z_flattened ** 2, dim=1, keepdim=True) + \
|
|
torch.sum(embedding**2, dim=1) - 2 * \
|
|
torch.einsum('bd,dn->bn', z_flattened, torch.einsum('n d -> d n', embedding))
|
|
|
|
if self.training:
|
|
logits = -d / self.quantization_temp
|
|
soft_one_hot = F.softmax(logits, dim=1)
|
|
min_encoding_indices = soft_one_hot.max(1, keepdim=True)[1]
|
|
hard_one_hot = torch.zeros_like(logits, memory_format=torch.legacy_contiguous_format).scatter_(1, min_encoding_indices, 1.0)
|
|
one_hot = hard_one_hot - soft_one_hot.detach() + soft_one_hot
|
|
|
|
z_q = einsum('b n, n d -> b d', one_hot, self.embedding.weight).view(z.shape)
|
|
z_q_2 = einsum('b n, n d -> b d', hard_one_hot, self.embedding.weight).view(z.shape)
|
|
|
|
# compute loss for embedding
|
|
commit_loss = torch.mean((z_q - z) ** 2) + torch.mean((z_q_2.detach() - z) ** 2) + self.beta * \
|
|
torch.mean((z_q_2 - z.detach()) ** 2)
|
|
else:
|
|
min_encoding_indices = torch.argmin(d, dim=1)
|
|
z_q = embedding[min_encoding_indices].view(z.shape)
|
|
commit_loss = None
|
|
|
|
if self.training and self.skip_quantization_prob > 0.0:
|
|
z_q = torch.where(
|
|
torch.rand_like(z_q[:, 0:1, 0:1, 0:1]).expand_as(z_q) <= self.skip_quantization_prob,
|
|
z, z_q,
|
|
)
|
|
|
|
# reshape back to match original input shape
|
|
z_q = rearrange(z_q, 'b h w c -> b c h w').contiguous()
|
|
|
|
if self.sane_index_shape:
|
|
min_encoding_indices = min_encoding_indices.reshape(z_q.shape[0], z_q.shape[2], z_q.shape[3])
|
|
|
|
return (z_q, None, min_encoding_indices), commit_loss
|
|
|
|
def get_codebook_entry(self, indices, bhwc):
|
|
# shape specifying (batch, height, width, channel)
|
|
# get quantized latent vectors
|
|
z_q = self.embedding(indices)
|
|
|
|
if bhwc is not None:
|
|
z_q = z_q.view(bhwc)
|
|
# reshape back to match original input shape
|
|
z_q = z_q.permute(0, 3, 1, 2).contiguous()
|
|
|
|
return z_q
|
|
|
|
|
|
class ResidualBlock(nn.Module):
|
|
def __init__(self, channels, num_groups=32):
|
|
super().__init__()
|
|
self.conv1 = nn.Conv2d(channels, channels, 3, padding='same')
|
|
self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=channels)
|
|
self.activate = nn.GELU()
|
|
self.conv2 = nn.Conv2d(channels, channels, 3, padding='same')
|
|
self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=channels)
|
|
|
|
def forward(self, x):
|
|
res = x
|
|
x = self.norm1(x)
|
|
x = self.activate(x)
|
|
x = self.conv1(x)
|
|
x = self.norm2(x)
|
|
x = self.activate(x)
|
|
x = self.conv2(x)
|
|
return x + res
|
|
|
|
|
|
class VQConvProjector(nn.Module):
|
|
def __init__(
|
|
self,
|
|
z_channels=1536,
|
|
codebook_size=16384,
|
|
codebook_dim=2048,
|
|
conv_layers=2,
|
|
with_norm=True,
|
|
skip_quant_prob=0.1,
|
|
):
|
|
super().__init__()
|
|
self.quant_conv = nn.Conv2d(z_channels, codebook_dim, 1)
|
|
self.quantize = IBQ(codebook_size, codebook_dim, skip_quant_prob, sane_index_shape=True)
|
|
self.post_quant_conv = nn.Conv2d(codebook_dim, z_channels, 1)
|
|
block = ResidualBlock
|
|
self.post_conv = nn.Sequential(*[block(z_channels) for _ in range(conv_layers)])
|
|
|
|
def forward(self, x, h, w):
|
|
x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w)
|
|
z = self.quant_conv(x)
|
|
(z_q, _, _), codebook_loss = self.quantize(z)
|
|
z = self.post_quant_conv(z_q)
|
|
z = self.post_conv(z)
|
|
z = rearrange(z, 'b c h w -> b (h w) c')
|
|
return z, codebook_loss
|
|
|
|
def encode(self, x, h, w):
|
|
x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w)
|
|
z = self.quant_conv(x)
|
|
(_, _, tokens), _ = self.quantize(z)
|
|
return tokens
|
|
|
|
def decode(self, tokens, bhwc):
|
|
z_q = self.quantize.get_codebook_entry(tokens, bhwc)
|
|
z = self.post_quant_conv(z_q)
|
|
z = self.post_conv(z)
|
|
return z
|
|
|
|
|
|
class SiglipTokenizer(nn.Module):
|
|
def __init__(
|
|
self,
|
|
siglip_name,
|
|
siglip_path,
|
|
projector_path,
|
|
z_channels=1536,
|
|
codebook_size=16384,
|
|
codebook_dim=2048,
|
|
with_norm=True
|
|
):
|
|
super().__init__()
|
|
self.vit = create_siglip_vit(model_name=siglip_name, path=siglip_path)
|
|
self.vqproj = VQConvProjector(
|
|
z_channels=z_channels,
|
|
codebook_size=codebook_size,
|
|
codebook_dim=codebook_dim,
|
|
with_norm=with_norm
|
|
)
|
|
self.vqproj.load_state_dict(torch.load(projector_path, map_location='cpu'), strict=True)
|
|
|
|
def encode(self, x):
|
|
features, (h, w), _ = self.vit(x)
|
|
tokens = self.vqproj.encode(features, h, w)
|
|
return tokens
|
|
|
|
def decode(self, tokens, bhwc):
|
|
return self.vqproj.decode(tokens, bhwc)
|