Add copy of github repo

Paella/src/modules.py

@@ -0,0 +1,283 @@
+import math
+import torch
+import numpy as np
+from torch import nn
+
+
+class Attention2D(nn.Module):
+    def __init__(self, c, nhead, dropout=0.0):
+        super().__init__()
+        self.attn = torch.nn.MultiheadAttention(c, nhead, dropout=dropout, bias=True, batch_first=True)
+
+    def forward(self, x, kv, self_attn=False):
+        orig_shape = x.shape
+        x = x.view(x.size(0), x.size(1), -1).permute(0, 2, 1)
+        if self_attn:
+            kv = torch.cat([x, kv], dim=1)
+        x = self.attn(x, kv, kv, need_weights=False)[0]
+        x = x.permute(0, 2, 1).view(*orig_shape)
+        return x
+
+
+class LayerNorm2d(nn.LayerNorm):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, x):
+        return super().forward(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+
+
+class GlobalResponseNorm(nn.Module):
+    "Taken from https://github.com/facebookresearch/ConvNeXt-V2/blob/3608f67cc1dae164790c5d0aead7bf2d73d9719b/models/utils.py#L105"
+    def __init__(self, dim):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))
+        self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))
+
+    def forward(self, x):
+        Gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True)
+        Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
+        return self.gamma * (x * Nx) + self.beta + x
+
+
+class ResBlock(nn.Module):
+    def __init__(self, c, c_skip=None, kernel_size=3, dropout=0.0):
+        super().__init__()
+        self.depthwise = nn.Conv2d(c + c_skip, c, kernel_size=kernel_size, padding=kernel_size // 2, groups=c)
+        self.norm = LayerNorm2d(c, elementwise_affine=False, eps=1e-6)
+        self.channelwise = nn.Sequential(
+            nn.Linear(c, c * 4),
+            nn.GELU(),
+            GlobalResponseNorm(c * 4),
+            nn.Dropout(dropout),
+            nn.Linear(c * 4, c)
+        )
+
+    def forward(self, x, x_skip=None):
+        x_res = x
+        if x_skip is not None:
+            x = torch.cat([x, x_skip], dim=1)
+        x = self.norm(self.depthwise(x)).permute(0, 2, 3, 1)
+        x = self.channelwise(x).permute(0, 3, 1, 2)
+        return x + x_res
+
+
+class AttnBlock(nn.Module):
+    def __init__(self, c, c_cond, nhead, self_attn=True, dropout=0.0):
+        super().__init__()
+        self.self_attn = self_attn
+        self.norm = LayerNorm2d(c, elementwise_affine=False, eps=1e-6)
+        self.attention = Attention2D(c, nhead, dropout)
+        self.kv_mapper = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(c_cond, c)
+        )
+
+    def forward(self, x, kv):
+        kv = self.kv_mapper(kv)
+        x = x + self.attention(self.norm(x), kv, self_attn=self.self_attn)
+        return x
+
+
+class FeedForwardBlock(nn.Module):
+    def __init__(self, c, dropout=0.0):
+        super().__init__()
+        self.norm = LayerNorm2d(c, elementwise_affine=False, eps=1e-6)
+        self.channelwise = nn.Sequential(
+            nn.Linear(c, c * 4),
+            nn.GELU(),
+            GlobalResponseNorm(c * 4),
+            nn.Dropout(dropout),
+            nn.Linear(c * 4, c)
+        )
+
+    def forward(self, x):
+        x = x + self.channelwise(self.norm(x).permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+        return x
+
+
+class TimestepBlock(nn.Module):
+    def __init__(self, c, c_timestep):
+        super().__init__()
+        self.mapper = nn.Linear(c_timestep, c * 2)
+
+    def forward(self, x, t):
+        a, b = self.mapper(t)[:, :, None, None].chunk(2, dim=1)
+        return x * (1 + a) + b
+
+
+class Paella(nn.Module):
+    def __init__(self, c_in=256, c_out=256, num_labels=8192, c_r=64, patch_size=2, c_cond=1024,
+                 c_hidden=[640, 1280, 1280], nhead=[-1, 16, 16], blocks=[6, 16, 6], level_config=['CT', 'CTA', 'CTA'],
+                 clip_embd=1024, byt5_embd=1536, clip_seq_len=4, kernel_size=3, dropout=0.1, self_attn=True):
+        super().__init__()
+        self.c_r = c_r
+        self.c_cond = c_cond
+        self.num_labels = num_labels
+        if not isinstance(dropout, list):
+            dropout = [dropout] * len(c_hidden)
+
+        # CONDITIONING
+        self.byt5_mapper = nn.Linear(byt5_embd, c_cond)
+        self.clip_mapper = nn.Linear(clip_embd, c_cond * clip_seq_len)
+        self.clip_image_mapper = nn.Linear(clip_embd, c_cond * clip_seq_len)
+        self.seq_norm = nn.LayerNorm(c_cond, elementwise_affine=False, eps=1e-6)
+
+        self.in_mapper = nn.Sequential(
+            nn.Embedding(num_labels, c_in),
+            nn.LayerNorm(c_in, elementwise_affine=False, eps=1e-6)
+        )
+        self.embedding = nn.Sequential(
+            nn.PixelUnshuffle(patch_size),
+            nn.Conv2d(c_in * (patch_size ** 2), c_hidden[0], kernel_size=1),
+            LayerNorm2d(c_hidden[0], elementwise_affine=False, eps=1e-6)
+        )
+
+        def get_block(block_type, c_hidden, nhead, c_skip=0, dropout=0):
+            if block_type == 'C':
+                return ResBlock(c_hidden, c_skip, kernel_size=kernel_size, dropout=dropout)
+            elif block_type == 'A':
+                return AttnBlock(c_hidden, c_cond, nhead, self_attn=self_attn, dropout=dropout)
+            elif block_type == 'F':
+                return FeedForwardBlock(c_hidden, dropout=dropout)
+            elif block_type == 'T':
+                return TimestepBlock(c_hidden, c_r)
+            else:
+                raise Exception(f'Block type {block_type} not supported')
+
+        # DOWN BLOCKS
+        self.down_blocks = nn.ModuleList()
+        for i in range(len(c_hidden)):
+            down_block = nn.ModuleList()
+            if i > 0:
+                down_block.append(nn.Sequential(
+                    LayerNorm2d(c_hidden[i - 1], elementwise_affine=False, eps=1e-6),
+                    nn.Conv2d(c_hidden[i - 1], c_hidden[i], kernel_size=2, stride=2),
+                ))
+            for _ in range(blocks[i]):
+                for block_type in level_config[i]:
+                    down_block.append(get_block(block_type, c_hidden[i], nhead[i], dropout=dropout[i]))
+            self.down_blocks.append(down_block)
+
+        # UP BLOCKS
+        self.up_blocks = nn.ModuleList()
+        for i in reversed(range(len(c_hidden))):
+            up_block = nn.ModuleList()
+            for j in range(blocks[i]):
+                for k, block_type in enumerate(level_config[i]):
+                    up_block.append(get_block(block_type, c_hidden[i], nhead[i],
+                                              c_skip=c_hidden[i] if i < len(c_hidden) - 1 and j == k == 0 else 0,
+                                              dropout=dropout[i]))
+            if i > 0:
+                up_block.append(nn.Sequential(
+                    LayerNorm2d(c_hidden[i], elementwise_affine=False, eps=1e-6),
+                    nn.ConvTranspose2d(c_hidden[i], c_hidden[i - 1], kernel_size=2, stride=2),
+                ))
+            self.up_blocks.append(up_block)
+
+        # OUTPUT
+        self.clf = nn.Sequential(
+            LayerNorm2d(c_hidden[0], elementwise_affine=False, eps=1e-6),
+            nn.Conv2d(c_hidden[0], c_out * (patch_size ** 2), kernel_size=1),
+            nn.PixelShuffle(patch_size),
+        )
+        self.out_mapper = nn.Sequential(
+            LayerNorm2d(c_out, elementwise_affine=False, eps=1e-6),
+            nn.Conv2d(c_out, num_labels, kernel_size=1, bias=False)
+        )
+
+        # --- WEIGHT INIT ---
+        self.apply(self._init_weights)  # General init
+        nn.init.normal_(self.byt5_mapper.weight, std=0.02)
+        nn.init.normal_(self.clip_mapper.weight, std=0.02)
+        nn.init.normal_(self.clip_image_mapper.weight, std=0.02)
+        torch.nn.init.xavier_uniform_(self.embedding[1].weight, 0.02)
+        nn.init.constant_(self.clf[1].weight, 0)
+        nn.init.normal_(self.in_mapper[0].weight, std=np.sqrt(1 / num_labels))
+        self.out_mapper[-1].weight.data = self.in_mapper[0].weight.data[:, :, None, None].clone()
+
+        for level_block in self.down_blocks + self.up_blocks:
+            for block in level_block:
+                if isinstance(block, ResBlock) or isinstance(block, FeedForwardBlock):
+                    block.channelwise[-1].weight.data *= np.sqrt(1 / sum(blocks))
+                elif isinstance(block, TimestepBlock):
+                    nn.init.constant_(block.mapper.weight, 0)
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Conv2d, nn.Linear)):
+            torch.nn.init.xavier_uniform_(m.weight)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def gen_r_embedding(self, r, max_positions=10000):
+        r = r * max_positions
+        half_dim = self.c_r // 2
+        emb = math.log(max_positions) / (half_dim - 1)
+        emb = torch.arange(half_dim, device=r.device).float().mul(-emb).exp()
+        emb = r[:, None] * emb[None, :]
+        emb = torch.cat([emb.sin(), emb.cos()], dim=1)
+        if self.c_r % 2 == 1:
+            emb = nn.functional.pad(emb, (0, 1), mode='constant')
+        return emb
+
+    def gen_c_embeddings(self, byt5, clip, clip_image):
+        seq = self.byt5_mapper(byt5)
+        if clip is not None:
+            clip = self.clip_mapper(clip).view(clip.size(0), -1, self.c_cond)
+            seq = torch.cat([seq, clip], dim=1)
+        if clip_image is not None:
+            clip_image = self.clip_image_mapper(clip_image).view(clip_image.size(0), -1, self.c_cond)
+            seq = torch.cat([seq, clip_image], dim=1)
+        seq = self.seq_norm(seq)
+        return seq
+
+    def _down_encode(self, x, r_embed, c_embed):
+        level_outputs = []
+        for down_block in self.down_blocks:
+            for block in down_block:
+                if isinstance(block, ResBlock):
+                    x = block(x)
+                elif isinstance(block, AttnBlock):
+                    x = block(x, c_embed)
+                elif isinstance(block, TimestepBlock):
+                    x = block(x, r_embed)
+                else:
+                    x = block(x)
+            level_outputs.insert(0, x)
+        return level_outputs
+
+    def _up_decode(self, level_outputs, r_embed, c_embed):
+        x = level_outputs[0]
+        for i, up_block in enumerate(self.up_blocks):
+            for j, block in enumerate(up_block):
+                if isinstance(block, ResBlock):
+                    x = block(x, level_outputs[i] if j == 0 and i > 0 else None)
+                elif isinstance(block, AttnBlock):
+                    x = block(x, c_embed)
+                elif isinstance(block, TimestepBlock):
+                    x = block(x, r_embed)
+                else:
+                    x = block(x)
+        return x
+
+    def forward(self, x, r, byt5, clip=None, clip_image=None, x_cat=None):
+        if x_cat is not None:
+            x = torch.cat([x, x_cat], dim=1)
+        # Process the conditioning embeddings
+        r_embed = self.gen_r_embedding(r)
+        c_embed = self.gen_c_embeddings(byt5, clip, clip_image)
+
+        # Model Blocks
+        x = self.embedding(self.in_mapper(x).permute(0, 3, 1, 2))
+        level_outputs = self._down_encode(x, r_embed, c_embed)
+        x = self._up_decode(level_outputs, r_embed, c_embed)
+        x = self.out_mapper(self.clf(x))
+        return x
+
+    def add_noise(self, x, t, mask=None, random_x=None):
+        if mask is None:
+            mask = (torch.rand_like(x.float()) <= t[:, None, None]).long()
+        if random_x is None:
+            random_x = torch.randint_like(x, 0, self.num_labels)
+        x = x * (1 - mask) + random_x * mask
+        return x, mask

Paella/src/train.py

@@ -0,0 +1,80 @@
+import os
+import torch
+import numpy as np
+from tqdm import tqdm
+from modules import Paella
+from torch import nn, optim
+from warmup_scheduler import GradualWarmupScheduler
+from utils import get_dataloader, load_conditional_models
+
+steps = 100_000
+warmup_updates = 10000
+batch_size = 16
+checkpoint_frequency = 2000
+lr = 1e-4
+
+train_device = "cuda"
+dataset_path = ""
+byt5_model_name = "google/byt5-xl"
+vqmodel_path = ""
+run_name = "Paella-ByT5-XL-v1"
+output_path = "output"
+checkpoint_path = f"{run_name}.pt"
+
+
+def train():
+    os.makedirs(output_path, exist_ok=True)
+    device = torch.device(train_device)
+
+    dataloader = get_dataloader(dataset_path, batch_size=batch_size)
+    checkpoint = torch.load(checkpoint_path, map_location=device) if os.path.exists(checkpoint_path) else None
+
+    model = Paella(byt5_embd=2560).to(device)
+    vqgan, (byt5_tokenizer, byt5) = load_conditional_models(byt5_model_name, vqmodel_path, device)
+    optimizer = optim.AdamW(model.parameters(), lr=lr)
+    scheduler = GradualWarmupScheduler(optimizer, multiplier=1, total_epoch=warmup_updates)
+    criterion = nn.CrossEntropyLoss(label_smoothing=0.1, reduction='none')
+
+    start_iter = 1
+    if checkpoint is not None:
+        model.load_state_dict(checkpoint['state_dict'])
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        scheduler.last_epoch = checkpoint['scheduler_last_step']
+        start_iter = checkpoint['scheduler_last_step'] + 1
+        del checkpoint
+
+    pbar = tqdm(range(start_iter, steps+1))
+    model.train()
+    for i, (images, captions) in enumerate(dataloader):
+        images = images.to(device)
+
+        with torch.no_grad():
+            if np.random.rand() < 0.05:
+                byt5_captions = [''] * len(captions)
+            else:
+                byt5_captions = captions
+            byt5_tokens = byt5_tokenizer(byt5_captions, padding="longest", return_tensors="pt", max_length=768, truncation=True).input_ids.to(device)
+            byt_embeddings = byt5(input_ids=byt5_tokens).last_hidden_state
+
+            t = (1-torch.rand(images.size(0), device=device))
+            latents = vqgan.encode(images)[2]
+            noised_latents, _ = model.add_noise(latents, t)
+
+        pred = model(noised_latents, t, byt_embeddings)
+        loss = criterion(pred, latents)
+
+        loss.backward()
+        grad_norm = nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        scheduler.step()
+        optimizer.zero_grad()
+
+        acc = (pred.argmax(1) == latents).float().mean()
+
+        pbar.set_postfix({'bs': images.size(0), 'loss': loss.item(), 'acc': acc.item(), 'grad_norm': grad_norm.item(), 'lr': optimizer.param_groups[0]['lr'], 'total_steps': scheduler.last_epoch})
+
+        if i % checkpoint_frequency == 0:
+            torch.save({'state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'scheduler_last_step': scheduler.last_epoch, 'iter' : i}, checkpoint_path)
+
+
+if __name__ == '__main__':
+    train()

Paella/src/utils.py

@@ -0,0 +1,55 @@
+import torch
+import torchvision
+from vqgan import VQModel
+from torch.utils.data import Dataset, DataLoader
+from transformers import T5EncoderModel, AutoTokenizer
+
+transforms = torchvision.transforms.Compose([
+    torchvision.transforms.ToTensor(),
+    torchvision.transforms.Resize(256),
+    torchvision.transforms.RandomCrop(256),
+])
+
+
+class YOUR_DATASET(Dataset):
+    def __init__(self, dataset_path):
+        pass
+
+
+def get_dataloader(dataset_path, batch_size):
+    dataset = YOUR_DATASET(dataset_path)
+    return DataLoader(dataset, batch_size=batch_size, num_workers=8, pin_memory=True)
+
+
+def load_conditional_models(byt5_model_name, vqgan_path, device):
+    vqgan = VQModel().to(device)
+    vqgan.load_state_dict(torch.load(vqgan_path, map_location=device)['state_dict'])
+    vqgan.eval().requires_grad_(False)
+
+    byt5 = T5EncoderModel.from_pretrained(byt5_model_name).to(device).eval().requires_grad_(False)
+    byt5_tokenizer = AutoTokenizer.from_pretrained(byt5_model_name)
+
+    return vqgan, (byt5_tokenizer, byt5)
+
+
+def sample(model, model_inputs, latent_shape, unconditional_inputs=None, steps=12, renoise_steps=11, temperature=(1.0, 0.2), cfg=8.0, t_start=1.0, t_end=0.0, device="cuda"):
+    with torch.inference_mode():
+        sampled = torch.randint(0, model.num_labels, size=latent_shape, device=device)
+        init_noise = sampled.clone()
+        t_list = torch.linspace(t_start, t_end, steps+1)
+        temperatures = torch.linspace(temperature[0], temperature[1], steps)
+        for i, t in enumerate(t_list[:steps]):
+            t = torch.ones(latent_shape[0], device=device) * t
+
+            logits = model(sampled, t, **model_inputs)
+            if cfg:
+                logits = logits * cfg + model(sampled, t, **unconditional_inputs) * (1-cfg)
+            scores = logits.div(temperatures[i]).softmax(dim=1)
+
+            sampled = scores.permute(0, 2, 3, 1).reshape(-1, logits.size(1))
+            sampled = torch.multinomial(sampled, 1)[:, 0].view(logits.size(0), *logits.shape[2:])
+
+            if i < renoise_steps:
+                t_next = torch.ones(latent_shape[0], device=device) * t_list[i+1]
+                sampled = model.add_noise(sampled, t_next, random_x=init_noise)[0]
+    return sampled

Paella/src/vqgan.py

@@ -0,0 +1,140 @@
+import torch
+from torch import nn
+from torchtools.nn import VectorQuantize
+
+
+class ResBlock(nn.Module):
+    def __init__(self, c, c_hidden):
+        super().__init__()
+        # depthwise/attention
+        self.norm1 = nn.LayerNorm(c, elementwise_affine=False, eps=1e-6)
+        self.depthwise = nn.Sequential(
+            nn.ReplicationPad2d(1),
+            nn.Conv2d(c, c, kernel_size=3, groups=c)
+        )
+
+        self.norm2 = nn.LayerNorm(c, elementwise_affine=False, eps=1e-6)
+        self.channelwise = nn.Sequential(
+            nn.Linear(c, c_hidden),
+            nn.GELU(),
+            nn.Linear(c_hidden, c),
+        )
+
+        self.gammas = nn.Parameter(torch.zeros(6), requires_grad=True)
+
+        def _basic_init(module):
+            if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+
+        self.apply(_basic_init)
+
+    def _norm(self, x, norm):
+        return norm(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+
+    def forward(self, x):
+        mods = self.gammas
+        x_temp = self._norm(x, self.norm1) * (1 + mods[0]) + mods[1]
+        x = x + self.depthwise(x_temp) * mods[2]
+        x_temp = self._norm(x, self.norm2) * (1 + mods[3]) + mods[4]
+        x = x + self.channelwise(x_temp.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) * mods[5]
+        return x
+
+
+class VQModel(nn.Module):
+    def __init__(self, levels=2, bottleneck_blocks=12, c_hidden=384, c_latent=4, codebook_size=8192,
+                 scale_factor=0.3764):  # 1.0
+        super().__init__()
+        self.c_latent = c_latent
+        self.scale_factor = scale_factor
+        c_levels = [c_hidden // (2 ** i) for i in reversed(range(levels))]
+
+        # Encoder blocks
+        self.in_block = nn.Sequential(
+            nn.PixelUnshuffle(2),
+            nn.Conv2d(3 * 4, c_levels[0], kernel_size=1)
+        )
+        down_blocks = []
+        for i in range(levels):
+            if i > 0:
+                down_blocks.append(nn.Conv2d(c_levels[i - 1], c_levels[i], kernel_size=4, stride=2, padding=1))
+            block = ResBlock(c_levels[i], c_levels[i] * 4)
+            down_blocks.append(block)
+        down_blocks.append(nn.Sequential(
+            nn.Conv2d(c_levels[-1], c_latent, kernel_size=1, bias=False),
+            nn.BatchNorm2d(c_latent),  # then normalize them to have mean 0 and std 1
+        ))
+        self.down_blocks = nn.Sequential(*down_blocks)
+
+        self.codebook_size = codebook_size
+        self.vquantizer = VectorQuantize(c_latent, k=codebook_size)
+
+        # Decoder blocks
+        up_blocks = [nn.Sequential(
+            nn.Conv2d(c_latent, c_levels[-1], kernel_size=1)
+        )]
+        for i in range(levels):
+            for j in range(bottleneck_blocks if i == 0 else 1):
+                block = ResBlock(c_levels[levels - 1 - i], c_levels[levels - 1 - i] * 4)
+                up_blocks.append(block)
+            if i < levels - 1:
+                up_blocks.append(
+                    nn.ConvTranspose2d(c_levels[levels - 1 - i], c_levels[levels - 2 - i], kernel_size=4, stride=2,
+                                       padding=1))
+        self.up_blocks = nn.Sequential(*up_blocks)
+        self.out_block = nn.Sequential(
+            nn.Conv2d(c_levels[0], 3 * 4, kernel_size=1),
+            nn.PixelShuffle(2),
+        )
+
+    def encode(self, x):
+        x = self.in_block(x)
+        x = self.down_blocks(x)
+        qe, (vq_loss, commit_loss), indices = self.vquantizer.forward(x, dim=1)
+        return qe / self.scale_factor, x / self.scale_factor, indices, vq_loss + commit_loss * 0.25
+
+    def decode(self, x):
+        x = x * self.scale_factor
+        x = self.up_blocks(x)
+        x = self.out_block(x)
+        return x
+
+    def decode_indices(self, x):
+        x = self.vquantizer.idx2vq(x, dim=1)
+        x = self.up_blocks(x)
+        x = self.out_block(x)
+        return x
+
+    def forward(self, x, quantize=False):
+        qe, x, _, vq_loss = self.encode(x, quantize)
+        x = self.decode(qe)
+        return x, vq_loss
+
+
+class Discriminator(nn.Module):
+    def __init__(self, c_in=3, c_cond=0, c_hidden=512, depth=6):
+        super().__init__()
+        d = max(depth - 3, 3)
+        layers = [
+            nn.utils.spectral_norm(nn.Conv2d(c_in, c_hidden // (2 ** d), kernel_size=3, stride=2, padding=1)),
+            nn.LeakyReLU(0.2),
+        ]
+        for i in range(depth - 1):
+            c_in = c_hidden // (2 ** max((d - i), 0))
+            c_out = c_hidden // (2 ** max((d - 1 - i), 0))
+            layers.append(nn.utils.spectral_norm(nn.Conv2d(c_in, c_out, kernel_size=3, stride=2, padding=1)))
+            layers.append(nn.InstanceNorm2d(c_out))
+            layers.append(nn.LeakyReLU(0.2))
+        self.encoder = nn.Sequential(*layers)
+        self.shuffle = nn.Conv2d((c_hidden + c_cond) if c_cond > 0 else c_hidden, 1, kernel_size=1)
+        self.logits = nn.Sigmoid()
+
+    def forward(self, x, cond=None):
+        x = self.encoder(x)
+        if cond is not None:
+            cond = cond.view(cond.size(0), cond.size(1), 1, 1, ).expand(-1, -1, x.size(-2), x.size(-1))
+            x = torch.cat([x, cond], dim=1)
+        x = self.shuffle(x)
+        x = self.logits(x)
+        return x

Paella/utils/alter_attention.py

@@ -0,0 +1,53 @@
+import torch
+from torch import nn
+
+class CustomMultiheadAttention(nn.MultiheadAttention):
+    def forward(self, *args, attn_weights=None, **kwargs):
+        q, k, v = args[:3]
+        need_weights = kwargs.get('need_weights', False)
+        
+        w = self.in_proj_weight.chunk(3, dim=0)
+        b = self.in_proj_bias.chunk(3, dim=0)
+        
+        if not self.batch_first:
+            q, k, v = q.permute(0, 1), k.permute(0, 1), v.permute(0, 1)
+        
+        q = nn.functional.linear(q, w[0], bias=b[0]).view(q.size(0), q.size(1), self.num_heads, -1).permute(0, 2, 1, 3)
+        k = nn.functional.linear(k, w[1], bias=b[1]).view(k.size(0), k.size(1), self.num_heads, -1).permute(0, 2, 1, 3)
+        v = nn.functional.linear(v, w[2], bias=b[2]).view(v.size(0), v.size(1), self.num_heads, -1).permute(0, 2, 1, 3)
+
+        scores = (q @ k.transpose(-2, -1)) / (q.size(-1) ** 0.5)
+        attention = scores.softmax(dim=-1)
+        # print(attention.shape)
+        
+        if attn_weights is not None:
+            # print("q ", q.shape)
+            # print("k ", k.shape)
+            weights = torch.ones((attention.shape[2], attention.shape[3])).to(q.device)
+            # print("Weights: ", weights.shape)
+            attn_weights = attn_weights.expand(attention.shape[2], attn_weights.shape[0])
+            weights[-attn_weights.shape[0]:, -attn_weights.shape[1]:] = attn_weights
+            # print(f"{-attn_weights.shape[0]}, {-attn_weights.shape[1]}")
+            attn_weights = weights.clone()
+            # print("Attn Weights: ", weights.shape)
+            # print("weight", attn_weights.shape)
+            attention = attention * attn_weights
+        
+        x = attention @ v
+        x = x.permute(0, 2, 1, 3).reshape(x.size(0), x.size(2), -1)
+        x = self.out_proj(x)
+        
+        if not self.batch_first:
+            x = x.permute(0, 1)
+        
+        return (x, attention if need_weights else None) 
+    
+def replace_attention_layers(model):
+    for n, module in model.named_children():
+        if len(list(module.children())) > 0:
+            replace_attention_layers(module)
+            
+        if isinstance(module, nn.MultiheadAttention):
+            new_module = CustomMultiheadAttention(module.embed_dim, module.num_heads, dropout=module.dropout, bias=True, batch_first=module.batch_first)
+            new_module.load_state_dict(module.state_dict())
+            setattr(model, n, new_module)

Paella/utils/modules.py

@@ -0,0 +1,291 @@
+import torch
+from torch import nn
+import numpy as np
+import math
+
+
+class Attention2D(nn.Module):
+    def __init__(self, c, nhead, dropout=0.0):
+        super().__init__()
+        self.attn = torch.nn.MultiheadAttention(c, nhead, dropout=dropout, bias=True, batch_first=True)
+
+    def forward(self, x, kv, self_attn=False, **kwargs):
+        orig_shape = x.shape
+        x = x.view(x.size(0), x.size(1), -1).permute(0, 2, 1)  # Bx4xHxW -> Bx(HxW)x4
+        if self_attn:
+            kv = torch.cat([x, kv], dim=1)
+        x = self.attn(x, kv, kv, need_weights=False, **kwargs)[0]
+        x = x.permute(0, 2, 1).view(*orig_shape)
+        return x
+
+
+class LayerNorm2d(nn.LayerNorm):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, x):
+        return super().forward(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+
+
+class GlobalResponseNorm(nn.Module):
+    "Taken from https://github.com/facebookresearch/ConvNeXt-V2/blob/3608f67cc1dae164790c5d0aead7bf2d73d9719b/models/utils.py#L105"
+    def __init__(self, dim):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))
+        self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))
+
+    def forward(self, x):
+        Gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True)
+        Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
+        return self.gamma * (x * Nx) + self.beta + x
+
+
+class ResBlock(nn.Module):
+    def __init__(self, c, c_skip=None, kernel_size=3, dropout=0.0):
+        super().__init__()
+        self.depthwise = nn.Conv2d(c + c_skip, c, kernel_size=kernel_size, padding=kernel_size // 2, groups=c)
+        self.norm = LayerNorm2d(c, elementwise_affine=False, eps=1e-6)
+        self.channelwise = nn.Sequential(
+            nn.Linear(c, c * 4),
+            nn.GELU(),
+            GlobalResponseNorm(c * 4),
+            nn.Dropout(dropout),
+            nn.Linear(c * 4, c)
+        )
+
+    def forward(self, x, x_skip=None):
+        x_res = x
+        if x_skip is not None:
+            x = torch.cat([x, x_skip], dim=1)
+        x = self.norm(self.depthwise(x)).permute(0, 2, 3, 1)
+        x = self.channelwise(x).permute(0, 3, 1, 2)
+        return x + x_res
+
+
+class AttnBlock(nn.Module):
+    def __init__(self, c, c_cond, nhead, self_attn=True, dropout=0.0):
+        super().__init__()
+        self.self_attn = self_attn
+        self.norm = LayerNorm2d(c, elementwise_affine=False, eps=1e-6)
+        self.attention = Attention2D(c, nhead, dropout)
+        self.kv_mapper = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(c_cond, c)
+        )
+
+    def forward(self, x, kv, **kwargs):
+        kv = self.kv_mapper(kv)
+        x = x + self.attention(self.norm(x), kv, self_attn=self.self_attn, **kwargs)
+        return x
+
+
+class FeedForwardBlock(nn.Module):
+    def __init__(self, c, dropout=0.0):
+        super().__init__()
+        self.norm = LayerNorm2d(c, elementwise_affine=False, eps=1e-6)
+        self.channelwise = nn.Sequential(
+            nn.Linear(c, c * 4),
+            nn.GELU(),
+            GlobalResponseNorm(c * 4),
+            nn.Dropout(dropout),
+            nn.Linear(c * 4, c)
+        )
+
+    def forward(self, x):
+        x = x + self.channelwise(self.norm(x).permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+        return x
+
+
+class TimestepBlock(nn.Module):
+    def __init__(self, c, c_timestep):
+        super().__init__()
+        self.mapper = nn.Linear(c_timestep, c * 2)
+
+    def forward(self, x, t):
+        a, b = self.mapper(t)[:, :, None, None].chunk(2, dim=1)
+        return x * (1 + a) + b
+
+
+class Paella(nn.Module):
+    def __init__(self, c_in=256, c_out=256, num_labels=8192, c_r=64, patch_size=2, c_cond=1024,
+                 c_hidden=[640, 1280, 1280], nhead=[-1, 16, 16], blocks=[6, 16, 6], level_config=['CT', 'CTA', 'CTA'],
+                 clip_embd=1024, byt5_embd=1536, clip_seq_len=4, kernel_size=3, dropout=0.1, self_attn=True):
+        super().__init__()
+        self.c_r = c_r
+        self.c_cond = c_cond
+        self.num_labels = num_labels
+        if not isinstance(dropout, list):
+            dropout = [dropout] * len(c_hidden)
+
+        # CONDITIONING
+        self.byt5_mapper = nn.Linear(byt5_embd, c_cond)
+        self.clip_mapper = nn.Linear(clip_embd, c_cond * clip_seq_len)
+        self.clip_image_mapper = nn.Linear(clip_embd, c_cond * clip_seq_len)
+        self.seq_norm = nn.LayerNorm(c_cond, elementwise_affine=False, eps=1e-6)
+
+        self.in_mapper = nn.Sequential(
+            nn.Embedding(num_labels, c_in),
+            nn.LayerNorm(c_in, elementwise_affine=False, eps=1e-6)
+        )
+        self.embedding = nn.Sequential(
+            nn.PixelUnshuffle(patch_size),
+            nn.Conv2d(c_in * (patch_size ** 2), c_hidden[0], kernel_size=1),
+            LayerNorm2d(c_hidden[0], elementwise_affine=False, eps=1e-6)
+        )
+
+        def get_block(block_type, c_hidden, nhead, c_skip=0, dropout=0):
+            if block_type == 'C':
+                return ResBlock(c_hidden, c_skip, kernel_size=kernel_size, dropout=dropout)
+            elif block_type == 'A':
+                return AttnBlock(c_hidden, c_cond, nhead, self_attn=self_attn, dropout=dropout)
+            elif block_type == 'F':
+                return FeedForwardBlock(c_hidden, dropout=dropout)
+            elif block_type == 'T':
+                return TimestepBlock(c_hidden, c_r)
+            else:
+                raise Exception(f'Block type {block_type} not supported')
+
+        # DOWN BLOCK
+        self.down_blocks = nn.ModuleList()
+        for i in range(len(c_hidden)):
+            down_block = nn.ModuleList()
+            if i > 0:
+                down_block.append(nn.Sequential(
+                    LayerNorm2d(c_hidden[i - 1], elementwise_affine=False, eps=1e-6),
+                    nn.Conv2d(c_hidden[i - 1], c_hidden[i], kernel_size=2, stride=2),
+                ))
+            for _ in range(blocks[i]):
+                for block_type in level_config[i]:
+                    down_block.append(get_block(block_type, c_hidden[i], nhead[i], dropout=dropout[i]))
+            self.down_blocks.append(down_block)
+
+        # UP BLOCKS
+        self.up_blocks = nn.ModuleList()
+        for i in reversed(range(len(c_hidden))):
+            up_block = nn.ModuleList()
+            for j in range(blocks[i]):
+                for k, block_type in enumerate(level_config[i]):
+                    up_block.append(get_block(block_type, c_hidden[i], nhead[i],
+                                              c_skip=c_hidden[i] if i < len(c_hidden) - 1 and j == k == 0 else 0,
+                                              dropout=dropout[i]))
+            if i > 0:
+                up_block.append(nn.Sequential(
+                    LayerNorm2d(c_hidden[i], elementwise_affine=False, eps=1e-6),
+                    nn.ConvTranspose2d(c_hidden[i], c_hidden[i - 1], kernel_size=2, stride=2),
+                ))
+            self.up_blocks.append(up_block)
+
+        # OUTPUT
+        self.clf = nn.Sequential(
+            LayerNorm2d(c_hidden[0], elementwise_affine=False, eps=1e-6),
+            nn.Conv2d(c_hidden[0], c_out * (patch_size ** 2), kernel_size=1),
+            nn.PixelShuffle(patch_size),
+        )
+        self.out_mapper = nn.Sequential(
+            LayerNorm2d(c_out, elementwise_affine=False, eps=1e-6),
+            nn.Conv2d(c_out, num_labels, kernel_size=1, bias=False)
+        )
+
+        # --- WEIGHT INIT ---
+        self.apply(self._init_weights)
+        nn.init.normal_(self.byt5_mapper.weight, std=0.02)
+        nn.init.normal_(self.clip_mapper.weight, std=0.02)
+        nn.init.normal_(self.clip_image_mapper.weight, std=0.02)
+        torch.nn.init.xavier_uniform_(self.embedding[1].weight, 0.02)  # inputs
+        nn.init.constant_(self.clf[1].weight, 0)  # outputs
+        nn.init.normal_(self.in_mapper[0].weight, std=np.sqrt(1 / num_labels))  # out mapper
+        self.out_mapper[-1].weight.data = self.in_mapper[0].weight.data[:, :, None, None].clone()
+
+        for level_block in self.down_blocks + self.up_blocks:
+            for block in level_block:
+                if isinstance(block, ResBlock) or isinstance(block, FeedForwardBlock):
+                    block.channelwise[-1].weight.data *= np.sqrt(1 / sum(blocks))
+                elif isinstance(block, TimestepBlock):
+                    nn.init.constant_(block.mapper.weight, 0)
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Conv2d, nn.Linear)):
+            torch.nn.init.xavier_uniform_(m.weight)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def gen_r_embedding(self, r, max_positions=10000):
+        r = r * max_positions
+        half_dim = self.c_r // 2
+        emb = math.log(max_positions) / (half_dim - 1)
+        emb = torch.arange(half_dim, device=r.device).float().mul(-emb).exp()
+        emb = r[:, None] * emb[None, :]
+        emb = torch.cat([emb.sin(), emb.cos()], dim=1)
+        if self.c_r % 2 == 1:  # zero pad
+            emb = nn.functional.pad(emb, (0, 1), mode='constant')
+        return emb
+
+    def gen_c_embeddings(self, byt5, clip, clip_image):
+        seq = self.byt5_mapper(byt5)
+        if clip is not None:
+            clip = self.clip_mapper(clip).view(clip.size(0), -1, self.c_cond)
+            seq = torch.cat([seq, clip], dim=1)
+        if clip_image is not None:
+            if isinstance(clip_image, list):
+                for ci in clip_image:
+                    ci = self.clip_image_mapper(ci).view(ci.size(0), -1, self.c_cond)
+                    seq = torch.cat([seq, ci], dim=1)
+            else:
+                clip_image = self.clip_image_mapper(clip_image).view(clip_image.size(0), -1, self.c_cond)
+                seq = torch.cat([seq, clip_image], dim=1)
+        seq = self.seq_norm(seq)
+        return seq
+
+    def _down_encode(self, x, r_embed, c_embed, **kwargs):
+        level_outputs = []
+        for down_block in self.down_blocks:
+            for block in down_block:
+                if isinstance(block, ResBlock):
+                    x = block(x)
+                elif isinstance(block, AttnBlock):
+                    x = block(x, c_embed, **kwargs)
+                elif isinstance(block, TimestepBlock):
+                    x = block(x, r_embed)
+                else:
+                    x = block(x)
+            level_outputs.insert(0, x)
+        return level_outputs
+
+    def _up_decode(self, level_outputs, r_embed, c_embed, **kwargs):
+        x = level_outputs[0]
+        for i, up_block in enumerate(self.up_blocks):
+            for j, block in enumerate(up_block):
+                if isinstance(block, ResBlock):
+                    x = block(x, level_outputs[i] if j == 0 and i > 0 else None)
+                elif isinstance(block, AttnBlock):
+                    x = block(x, c_embed, **kwargs)
+                elif isinstance(block, TimestepBlock):
+                    x = block(x, r_embed)
+                else:
+                    x = block(x)
+        return x
+
+    def forward(self, x, r, byt5, clip=None, clip_image=None, x_cat=None, **kwargs):
+        if x_cat is not None:
+            x = torch.cat([x, x_cat], dim=1)
+        # Process the conditioning embeddings
+        r_embed = self.gen_r_embedding(r)
+        c_embed = self.gen_c_embeddings(byt5, clip, clip_image)
+
+        # Model Blocks
+        x = self.embedding(self.in_mapper(x).permute(0, 3, 1, 2))
+        level_outputs = self._down_encode(x, r_embed, c_embed, **kwargs)
+        x = self._up_decode(level_outputs, r_embed, c_embed, **kwargs)
+        x = self.out_mapper(self.clf(x))
+        return x
+
+    def add_noise(self, x, t, mask=None, random_x=None):
+        if mask is None:
+            mask = (torch.rand_like(x.float()) <= t[:, None, None]).long()
+        if random_x is None:
+            random_x = torch.randint_like(x, 0, self.num_labels)
+        x = x * (1 - mask) + random_x * mask
+        return x, mask
+
+    def get_loss_weight(self, t, mask, min_val=0.3):
+        return 1 - (1 - mask) * ((1 - t) * (1 - min_val))[:, None, None]