file upload

app.py

@@ -1,7 +1,3 @@
 import gradio as gr
 
-def greet(name):
-    return "Hello " + name + "!!"
-
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+demo = gr.load("HuangLab/CELL-E_2_HPA_Finetuned_480", src="models")

celle/__init__.py

@@ -0,0 +1,4 @@
+from celle.celle import CELLE
+from celle.vae import VQGanVAE
+
+__version__ = "2.0.0"

celle/attention.py

@@ -0,0 +1,253 @@
+import torch
+from torch import nn, einsum
+import torch.nn.functional as F
+from einops import rearrange, repeat
+
+from rotary_embedding_torch import apply_rotary_emb
+from celle.utils import exists, default, max_neg_value
+
+
+# helpers
+def stable_softmax(t, dim=-1, alpha=32**2):
+    t = t / alpha
+    t = t - torch.amax(t, dim=dim, keepdim=True).detach()
+    return (t * alpha).softmax(dim=dim)
+
+
+def apply_pos_emb(pos_emb, qkv):
+    n = qkv[0].shape[-2]
+    pos_emb = pos_emb[..., :n, :]
+    return tuple(map(lambda t: apply_rotary_emb(pos_emb, t), qkv))
+
+
+# classes
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        seq_len,
+        causal=False,
+        heads=8,
+        dim_head=64,
+        dropout=0.0,
+        stable=False,
+        static_mask=None,
+    ):
+        super().__init__()
+        inner_dim = dim_head * heads
+        self.heads = heads
+        self.seq_len = seq_len
+        self.scale = dim_head**-0.5
+        self.stable = stable
+        self.causal = causal
+        self.register_buffer("static_mask", static_mask, persistent=False)
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
+        self.to_out = nn.Sequential(nn.Linear(inner_dim, dim), nn.Dropout(dropout))
+        self.save_attn = nn.Identity()
+
+    def forward(self, x, context_mask=None, rotary_pos_emb=None):
+        # x: [batch_size, seq_len, dim]
+        b, n, _, h = *x.shape, self.heads
+        device = x.device
+
+        softmax = torch.softmax if not self.stable else stable_softmax
+
+        # qkv: 3 tensors of shape [batch_size, seq_len, inner_dim]
+        qkv = self.to_qkv(x).chunk(3, dim=-1)
+
+        # q,k,v: [batch_size, heads, seq_len, dim_head]
+        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), qkv)
+
+        if exists(rotary_pos_emb):
+            q, k, v = apply_pos_emb(rotary_pos_emb[..., :, :], (q, k, v))
+
+        q *= self.scale
+
+        # dots: [batch_size, heads, seq_len_i ,seq_len_j]
+        dots = torch.einsum("b h i d, b h j d -> b h i j", q, k)
+        mask_value = max_neg_value(dots)
+
+        if exists(context_mask):
+            # context_mask: [batch_size ,1 ,1 ,seq_len_j]
+            context_mask = rearrange(context_mask, "b j -> b 1 1 j")
+            context_mask = F.pad(context_mask, (1, 0), value=True)
+
+            mask_value = -torch.finfo(dots.dtype).max
+            dots = dots.masked_fill(~context_mask, mask_value)
+
+        if self.causal:
+            i, j = dots.shape[-2:]
+            context_mask = torch.ones(i, j, device=device).triu_(j - i + 1).bool()
+            dots.masked_fill_(context_mask, mask_value)
+
+        if exists(self.static_mask):
+            dots.masked_fill_(~self.static_mask[:n, :n], mask_value)
+
+        # attn: [batch_size ,heads ,seq_len_i ,seq_len_j]
+        attn = softmax(dots, dim=-1)
+        attn = self.save_attn(attn)
+
+        # out: [batch_size ,heads ,seq_len_i ,dim_head]
+        out = torch.einsum("b h n j, b h j d -> b h n d", attn, v)
+
+        # out: [batch_size ,seq_len_i ,(heads*dim_head)]
+        out = rearrange(out, "b h n d -> b n (h d)")
+
+        # out: [batch_size ,seq_len_i ,dim]
+        out = self.to_out(out)
+
+        return out
+
+
+# sparse attention with convolutional pattern, as mentioned in the blog post. customizable kernel size and dilation
+
+
+class SparseConvCausalAttention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        seq_len,
+        image_size=32,
+        kernel_size=5,
+        dilation=1,
+        heads=8,
+        dim_head=64,
+        dropout=0.0,
+        stable=False,
+        **kwargs,
+    ):
+        super().__init__()
+        assert kernel_size % 2 == 1, "kernel size must be odd"
+
+        inner_dim = dim_head * heads
+        self.seq_len = seq_len
+        self.heads = heads
+        self.scale = dim_head**-0.5
+        self.image_size = image_size
+        self.kernel_size = kernel_size
+        self.dilation = dilation
+
+        self.stable = stable
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
+
+        self.to_out = nn.Sequential(nn.Linear(inner_dim, dim), nn.Dropout(dropout))
+
+    def forward(self, x, mask=None, rotary_pos_emb=None):
+        b, n, _, h, img_size, kernel_size, dilation, seq_len, device = (
+            *x.shape,
+            self.heads,
+            self.image_size,
+            self.kernel_size,
+            self.dilation,
+            self.seq_len,
+            x.device,
+        )
+        softmax = torch.softmax if not self.stable else stable_softmax
+
+        img_seq_len = img_size**2
+        text_len = seq_len + 1 - img_seq_len
+
+        # padding
+
+        padding = seq_len - n + 1
+        mask = default(mask, lambda: torch.ones(b, text_len, device=device).bool())
+
+        x = F.pad(x, (0, 0, 0, padding), value=0)
+        mask = mask[:, :text_len]
+
+        # derive query / keys / values
+
+        qkv = self.to_qkv(x).chunk(3, dim=-1)
+        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), qkv)
+
+        if exists(rotary_pos_emb):
+            q, k, v = apply_pos_emb(rotary_pos_emb, (q, k, v))
+
+        q *= self.scale
+
+        ((q_text, q_img), (k_text, k_img), (v_text, v_img)) = map(
+            lambda t: (t[:, :-img_seq_len], t[:, -img_seq_len:]), (q, k, v)
+        )
+
+        # text attention
+
+        dots_text = einsum("b i d, b j d -> b i j", q_text, k_text)
+        mask_value = max_neg_value(dots_text)
+
+        i, j = dots_text.shape[-2:]
+        text_causal_mask = torch.ones(i, j, device=device).triu_(j - i + 1).bool()
+        dots_text.masked_fill_(text_causal_mask, mask_value)
+
+        attn_text = softmax(dots_text, dim=-1)
+        out_text = einsum("b i j, b j d -> b i d", attn_text, v_text)
+
+        # image attention
+
+        effective_kernel_size = (kernel_size - 1) * dilation + 1
+        padding = effective_kernel_size // 2
+
+        k_img, v_img = map(
+            lambda t: rearrange(t, "b (h w) c -> b c h w", h=img_size), (k_img, v_img)
+        )
+        k_img, v_img = map(
+            lambda t: F.unfold(t, kernel_size, padding=padding, dilation=dilation),
+            (k_img, v_img),
+        )
+        k_img, v_img = map(
+            lambda t: rearrange(t, "b (d j) i -> b i j d", j=kernel_size**2),
+            (k_img, v_img),
+        )
+
+        # let image attend to all of text
+
+        dots_image = einsum("b i d, b i j d -> b i j", q_img, k_img)
+        dots_image_to_text = einsum("b i d, b j d -> b i j", q_img, k_text)
+
+        # calculate causal attention for local convolution
+
+        i, j = dots_image.shape[-2:]
+        img_seq = torch.arange(img_seq_len, device=device)
+        k_img_indices = rearrange(img_seq.float(), "(h w) -> () () h w", h=img_size)
+        k_img_indices = F.pad(
+            k_img_indices, (padding,) * 4, value=img_seq_len
+        )  # padding set to be max, so it is never attended to
+        k_img_indices = F.unfold(k_img_indices, kernel_size, dilation=dilation)
+        k_img_indices = rearrange(k_img_indices, "b j i -> b i j")
+
+        # mask image attention
+
+        q_img_indices = rearrange(img_seq, "i -> () i ()")
+        causal_mask = q_img_indices < k_img_indices
+
+        # concat text mask with image causal mask
+
+        causal_mask = repeat(causal_mask, "() i j -> b i j", b=b * h)
+        mask = repeat(mask, "b j -> (b h) i j", i=i, h=h)
+        mask = torch.cat((~mask, causal_mask), dim=-1)
+
+        # image can attend to all of text
+
+        dots = torch.cat((dots_image_to_text, dots_image), dim=-1)
+        dots.masked_fill_(mask, mask_value)
+
+        attn = softmax(dots, dim=-1)
+
+        # aggregate
+
+        attn_image_to_text, attn_image = attn[..., :text_len], attn[..., text_len:]
+
+        out_image_to_image = einsum("b i j, b i j d -> b i d", attn_image, v_img)
+        out_image_to_text = einsum("b i j, b j d -> b i d", attn_image_to_text, v_text)
+
+        out_image = out_image_to_image + out_image_to_text
+
+        # combine attended values for both text and image
+
+        out = torch.cat((out_text, out_image), dim=1)
+
+        out = rearrange(out, "(b h) n d -> b n (h d)", h=h)
+
+        out = self.to_out(out)
+
+        return out[:, :n]

celle/celle.py

@@ -0,0 +1,1060 @@
+# Import necessary packages and modules
+from math import floor, ceil
+import torch
+from torch import nn
+import torch.nn.functional as F
+from axial_positional_embedding import AxialPositionalEmbedding
+from einops import rearrange
+from celle.utils import (
+    exists,
+    always,
+    eval_decorator,
+    gumbel_sample,
+    top_k,
+    gamma_func,
+    DivideMax,
+)
+
+# Import additional modules from within the codebase
+from celle.transformer import Transformer
+
+
+def generate_mask(gamma_func, batch_size, length, device):
+    # Get the number of `True` values in the mask for each batch element
+    num_true_values = floor(gamma_func(torch.rand(1)) * length)
+
+    # Generate a random sample of indices to set to `True` in the mask
+    # The number of indices in the sample is determined by `num_true_values`
+    indices = (
+        torch.rand((batch_size, length), device=device)
+        .topk(num_true_values, dim=1)
+        .indices
+    )
+
+    # Create a binary mask tensor with `True` values at the sampled indices
+    mask = torch.zeros((batch_size, length), dtype=torch.bool, device=device)
+    mask.scatter_(dim=1, index=indices, value=True)
+
+    return mask
+
+
+def match_batch_size(text, condition, image, batch_size):
+    """
+    This function ensures all inputs to the sample function have the same batch size.
+    """
+    if text.shape[0] != batch_size:
+        text = text.repeat(batch_size, 1)
+
+    if condition.shape[0] != batch_size:
+        condition = condition.repeat(batch_size, 1)
+
+    if image.shape[0] != batch_size:
+        image = image.repeat(batch_size, 1)
+
+    return text, condition, image
+
+
+def calc_unmask_probs(timestep, timesteps, gamma_func):
+    if timestep == 1 or timesteps == 1:
+        unmask_prob = 1
+    else:
+        unmask_prob = 1 - gamma_func(timestep)
+    return unmask_prob
+
+
+def calculate_logits(
+    input_tokens, input_mask, logits_function, filter_thres, temperature
+):
+    logits, _, _ = logits_function(input_tokens, input_mask, return_encoding=False)
+    filtered_logits = top_k(logits, thres=filter_thres)
+    sample = gumbel_sample(filtered_logits, temperature=temperature, dim=-1)
+
+    return logits, sample
+
+
+def unmask_tokens(
+    input_tokens,
+    input_mask,
+    num_masked_tokens,
+    logits,
+    sample,
+    timestep,
+    timesteps,
+    gamma,
+    filter_func=None,
+    pad_token=None,
+    mask_token=None,
+    force_aas=True,
+):
+    sample = sample.masked_fill(~input_mask.unsqueeze(-1), -torch.inf)
+    if filter_func:
+        sample = filter_func(
+            input_tokens, sample, force_aas, pad_token=pad_token, mask_token=mask_token
+        )
+    selected_token_probs, selected_tokens = torch.max(sample, dim=-1)
+
+    unmask_prob = calc_unmask_probs(timestep, timesteps, gamma)
+    num_tokens_to_unmask = max(1, ceil(unmask_prob * num_masked_tokens))
+
+    _, top_k_indices = torch.topk(selected_token_probs, num_tokens_to_unmask, dim=-1)
+
+    sample_mask = torch.zeros(
+        input_tokens.shape, dtype=torch.bool, device=input_tokens.device
+    )
+    sample_mask.scatter_(dim=1, index=top_k_indices, value=True)
+
+    unmasked_tokens = torch.where(sample_mask, selected_tokens, input_tokens)
+    full_logits = torch.where(
+        sample_mask.unsqueeze(-1), logits, torch.zeros_like(logits)
+    )
+    return unmasked_tokens, full_logits
+
+
+def suppress_invalid_text_tokens(
+    text,
+    logits,
+    start_token=None,
+    end_token=None,
+    pad_token=None,
+    mask_token=None,
+    force_aas=False,
+):
+    # Find the indices of start_token and end_token in tensor text along axis=1
+    idx_start = (text == start_token).nonzero(as_tuple=True)[1]
+    idx_end = (text == end_token).nonzero(as_tuple=True)[1]
+
+    # For every position other than the index corresponding to the start index, set the values on the start index of dimension=2 to -torch.inf
+    if idx_start.nelement() != start_token:
+        try:
+            mask = idx_start.unsqueeze(1) != torch.arange(
+                logits.size(1), device=text.device
+            )
+            indices = torch.where(mask)
+            logits[indices[0], indices[1], start_token] = -torch.inf
+        except:
+            pass
+
+    # else:
+    #     idx_start = torch.zeros(text.size(0), dtype=torch.long)
+
+    # Similarly, for every position other than the index corresponding to the end index, set the values on the end index of dimension=2 to -torch.inf
+    if idx_end.nelement() != 0:
+        try:
+            mask = idx_end.unsqueeze(1) != torch.arange(
+                logits.size(1), device=text.device
+            )
+            indices = torch.where(mask)
+            logits[indices[0], indices[1], end_token] = -torch.inf
+        except:
+            pass
+
+    # else:
+    #     idx_end = torch.full((text.size(0),), text.size(1) - 1, dtype=torch.long)
+
+    if pad_token:
+        if idx_start.nelement() != 0 and idx_end.nelement() != 0:
+            try:
+                # For every position between the indices of start_token and end_token, set the values for 1st index of dimension=2 equal to -torch.inf. Any value outside of that range should be set to torch.inf.
+                mask = (
+                    torch.arange(logits.size(1), device=text.device)
+                    >= idx_start.unsqueeze(1)
+                ) & (
+                    torch.arange(logits.size(1), device=text.device)
+                    <= idx_end.unsqueeze(1)
+                )
+
+                indices = torch.where(mask)
+                logits[indices[0], indices[1], pad_token] = -torch.inf
+
+                indices = torch.where(~mask)
+                logits[indices[0], indices[1], pad_token] = torch.inf
+
+            except:
+                pass
+
+        elif idx_start.nelement() != 0:
+            try:
+                mask = torch.arange(
+                    logits.size(1), device=text.device
+                ) < idx_start.unsqueeze(1)
+                logits[indices[0], indices[1], pad_token] = torch.inf
+            except:
+                pass
+
+        elif idx_end.nelement() != 0:
+            try:
+                mask = torch.arange(
+                    logits.size(1), device=text.device
+                ) > idx_end.unsqueeze(1)
+                logits[indices[0], indices[1], pad_token] = torch.inf
+            except:
+                pass
+
+    if force_aas:
+        if pad_token:
+            logits[:, :, pad_token] = -torch.inf
+        logits[:, :, 3] = -torch.inf
+        logits[:, :, 29:] = -torch.inf
+
+    if mask_token:
+        logits[:, :, mask_token] = -torch.inf
+
+    return logits
+
+
+def detokenize_text(text_embedding, sequence):
+    if text_embedding == "esm1b" or text_embedding == "esm2":
+        from esm import Alphabet
+
+        alphabet = (
+            Alphabet.from_architecture("ESM-1b").get_batch_converter().alphabet.all_toks
+        )
+    else:
+        assert NameError("Detokenization only available for ESM mdodels")
+
+    output_seqs = []
+
+    for batch in sequence:
+        converted_seq = [alphabet[idx] for idx in batch]
+        converted_seq = "".join(converted_seq)
+        output_seqs.append(converted_seq)
+
+    return output_seqs
+
+class ImageEmbedding(nn.Module):
+    def __init__(self, num_tokens, dim):
+        super(ImageEmbedding, self).__init__()
+        self.image_embedding = nn.Embedding(num_tokens, dim)
+        
+    def forward(self, image):
+        return self.image_embedding(image)
+
+
+class ModelExtender(nn.Module):
+    def __init__(self, vocab, out_features, fixed_embedding=False):
+        super(ModelExtender, self).__init__()
+
+        # Initialize the model according to the given vocabulary
+        self.vocab = vocab
+        
+        if vocab == "esm1b":
+            from esm import pretrained
+
+            self.model, _ = pretrained.esm1b_t33_650M_UR50S()
+            self.in_features = 1280
+        elif vocab == "esm2":
+            from esm import pretrained
+
+            if out_features == 320:
+                self.model, _ = pretrained.esm2_t6_8M_UR50D()
+            elif out_features == 480:
+                self.model, _ = pretrained.esm2_t12_35M_UR50D()
+            elif out_features == 640:
+                self.model, _ = pretrained.esm2_t30_150M_UR50D()
+            elif out_features == 1280:
+                self.model, _ = pretrained.esm2_t33_650M_UR50D()
+            elif out_features == 2560:
+                self.model, _ = pretrained.esm2_t36_3B_UR50D()
+            else:
+                self.model, _ = pretrained.esm2_t33_650M_UR50D()
+            self.in_features = self.model.embed_dim
+
+        # Set the number of output features and initialize the scaling layer
+        self.out_features = out_features
+        self.scale_layer = nn.Linear(self.in_features, self.out_features)
+
+        # Determine whether to freeze the model's parameters
+        self.fixed_embedding = fixed_embedding
+        if self.fixed_embedding:
+            self.model = self.model.eval()
+
+    def forward(self, x, **kwargs):
+        # If the model's parameters are fixed, use torch.no_grad()
+        if self.fixed_embedding:
+            with torch.no_grad():
+                if self.vocab == "esm1b" or self.vocab == "esm2":
+                    # Reduce sequence length dimension, get top layer representation tensor
+                    x = self.model(x.squeeze(1), repr_layers=[self.model.num_layers])[
+                        "representations"
+                    ][self.model.num_layers]
+                    # Tensor shape: (batch_size, hidden_size)
+                else:
+                    # Get top layer representation tensor
+                    x = self.model(x, **kwargs)[0]
+                    # Tensor shape: (batch_size, sequence_length, hidden_size)
+        else:
+            if self.vocab == "esm1b" or self.vocab == "esm2":
+                # Reduce sequence length dimension, get top layer representation tensor
+                x = self.model(x.squeeze(1), repr_layers=[self.model.num_layers])[
+                    "representations"
+                ][self.model.num_layers]
+                # Tensor shape: (batch_size, hidden_size)
+            else:
+                # Get top layer representation tensor
+                x = self.model(x, **kwargs)[0]
+                # Tensor shape: (batch_size, sequence_length, hidden_size)
+
+        # Scale the representation tensor if necessary
+        if self.out_features != self.in_features:
+            x = self.scale_layer(x)
+            # Tensor shape: (batch_size, out_features)
+
+        return x
+
+class CELLE(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        vae,  # The VAE model used to encode/decode images
+        condition_vae=None,  # An optional VAE model used to condition the image generation
+        num_images=2,  # Number of images to generate
+        num_text_tokens=30,  # Number of tokens in the text vocabulary
+        text_seq_len=1000,  # Maximum length of input text sequence
+        depth=16,  # Number of layers in the transformer model
+        heads=16,  # Number of attention heads
+        dim_head=64,  # Dimensionality of each attention head
+        attn_dropout=0.1,  # Dropout rate for attention weights
+        ff_dropout=0.1,  # Dropout rate for feedforward layers
+        attn_types=None,  # Types of attention to use in the transformer
+        causal=False,  # Whether to use causal attention
+        loss_cond_weight=1,  # Weight of conditioning loss
+        loss_img_weight=1,  # Weight of image generation loss
+        stable=False,  # Whether to use divide-by-max normalization in the transformer
+        rotary_emb=True,  # Whether to use rotary positional embeddings
+        text_embedding="esm2",  # Text embedding to use (esm1b, esm2)
+        fixed_embedding=True,  # Whether to fix the text embedding or learn it
+        sampling_mode="cosine",  # Sampling mode for the VAE
+        linear_project=False,  # Whether to project embeddings linearly
+        **kwargs,
+    ):
+        super().__init__()
+
+        # Set the stable flag
+        self.stable = stable
+
+        # If the stable flag is set, initialize the DivideMax layer for normalization
+        if stable:
+            self.norm_by_max = DivideMax(dim=-1)
+
+        ### Initializing text parameters ###
+
+        # Initialize the text and fixed embeddings
+        self.text_embedding = text_embedding
+        self.fixed_embedding = fixed_embedding
+
+        # Offset logits index and calculate cross entropy loss
+        self.num_text_tokens = num_text_tokens
+        self.linear_project = linear_project
+
+        # Add <BOS> and <EOS> tokens to the beginning and end of text sequences
+        if text_embedding.lower() in ("esm1b", "esm2"):
+            self.text_seq_len = text_seq_len + 2
+        else:
+            self.text_seq_len = text_seq_len
+
+        # Initialize embeddings for <SEP> token
+        self.sep_emb = nn.Embedding(1, dim)
+
+        # Initialize positional embeddings for text sequences and <SEP> token
+        self.text_pos_emb = (
+            nn.Embedding(self.text_seq_len + 1, dim) if not rotary_emb else always(0)
+        )  # +1 for <SEP>
+
+        ### ###
+
+        self.num_images = num_images
+
+        ### Initializing condition parameters ###
+
+        # Initialize the number of condition tokens, condition sequence length, and condition embedding
+        if exists(condition_vae):
+            condition_size = condition_vae.image_size
+            num_condition_tokens = condition_vae.num_tokens
+            self.num_condition_tokens = num_condition_tokens
+            condition_fmap_size = condition_vae.image_size // (
+                2**condition_vae.num_layers
+            )
+            condition_seq_len = condition_fmap_size**2
+
+            # Initialize ImageEmbedding for condition embedding
+            self.condition_emb = ImageEmbedding(num_condition_tokens + 1, dim)
+
+            # Initialize positional embeddings for condition embedding
+            self.condition_pos_emb = (
+                AxialPositionalEmbedding(
+                    dim, axial_shape=(condition_fmap_size, condition_fmap_size)
+                )
+                if not rotary_emb
+                else always(0)
+            )
+
+        else:
+            condition_fmap_size = 0
+            condition_seq_len = 0
+            num_condition_tokens = 0
+
+        ### ####
+
+        ### Initializing image parameters ###
+
+        # Initialize the image size, image token size, and sequence length
+        self.image_size = vae.image_size
+        num_image_tokens = vae.num_tokens
+        image_fmap_size = vae.image_size // (2**vae.num_layers)
+        image_seq_len = image_fmap_size**2
+        self.image_seq_len = image_seq_len
+        self.num_image_tokens = num_image_tokens
+
+        # Initialize ImageEmbedding and positional embeddings for image embedding
+        self.image_emb = ImageEmbedding(num_image_tokens + 1, dim) # +1 for <IM_MASK>
+
+        self.image_pos_emb = (
+            AxialPositionalEmbedding(
+                dim, axial_shape=(image_fmap_size, image_fmap_size)
+            )
+            if not rotary_emb
+            else always(0)
+        )
+
+        # Set total sequence length and total tokens
+        self.num_condition_tokens = num_condition_tokens
+        self.condition_seq_len = condition_seq_len
+        # Text Length + <SEP> + Condition Tokens + Image Tokens
+        seq_len = self.text_seq_len + 1 + self.condition_seq_len + self.image_seq_len
+        total_tokens = (
+            num_text_tokens + 1 + num_condition_tokens + 1 + num_image_tokens + 1
+        )
+        self.total_tokens = total_tokens
+        self.total_seq_len = seq_len
+
+        # Set the VAE and condition VAE for the model
+        self.vae = vae.eval()
+        self.condition_vae = condition_vae.eval()
+
+        ### ###
+
+        ### Setting discrete ids ###
+        # Initialize text embedding based on the given text_embedding parameter
+        if text_embedding == "esm1b" or text_embedding == "esm2":
+            self.text_mask_token = 32
+            self.pad_token = 1
+            self.text_emb = ModelExtender(text_embedding, dim, fixed_embedding)
+        else:
+            raise ValueError("Only ESM models are supported.")
+
+        # Set token indices for text, condition, and image sequences
+        self.sep_token = num_text_tokens
+        self.cond_mask_token = num_condition_tokens
+        self.image_mask_token = num_image_tokens
+
+        # Create indices for sequence and logits dimensions
+        self.seq_range = torch.arange(seq_len)
+        self.logits_range = torch.arange(total_tokens)
+
+        # Reshape sequence and logits indices
+        self.seq_range = rearrange(self.seq_range, "n -> () n ()")
+        self.logits_range = rearrange(self.logits_range, "d -> () () d")
+
+        # Create a mask to exclude invalid token positions from the model output
+        # e.g. no image tokens where sequence tokens should be
+        logits_mask = (
+            # Mask text tokens beyond text_seq_len and invalid logits_range
+            (
+                (self.seq_range < self.text_seq_len)
+                & (self.logits_range < num_text_tokens)
+                & (self.logits_range != self.text_mask_token)
+            )
+            |
+            # Mask [SEP] token after text
+            (
+                (self.seq_range == self.text_seq_len)
+                & (self.logits_range == num_text_tokens)
+            )
+            |
+            # Mask condition tokens beyond text_seq_len+1 ([SEP]) and invalid logits_range
+            (
+                (self.seq_range >= self.text_seq_len + 1)
+                & (self.seq_range < self.text_seq_len + 1 + condition_seq_len)
+                & (self.logits_range >= num_text_tokens + 1)
+                & (self.logits_range < num_text_tokens + 1 + num_condition_tokens)
+            )
+            |
+            # Mask image tokens beyond num_text_tokens+num_condition_tokens+1
+            (
+                (self.seq_range >= self.text_seq_len + 1 + condition_seq_len)
+                & (self.logits_range >= num_text_tokens + 1 + num_condition_tokens + 1)
+                & (
+                    self.logits_range
+                    < num_text_tokens + 1 + num_condition_tokens + 1 + num_image_tokens
+                )
+            )
+        )
+
+        # Invert the mask
+        logits_mask = ~logits_mask
+
+        # Register the buffer with the logits_mask
+        self.register_buffer("logits_mask", logits_mask, persistent=False)
+
+        ### ###
+
+        # Initialize the Transformer model with given parameters
+        self.transformer = Transformer(
+            dim=dim,
+            causal=causal,
+            seq_len=seq_len,
+            depth=depth,
+            heads=heads,
+            dim_head=dim_head,
+            attn_dropout=attn_dropout,
+            ff_dropout=ff_dropout,
+            image_fmap_size=image_fmap_size + condition_fmap_size,
+            num_images=num_images,
+            stable=stable,
+            rotary_emb=rotary_emb,
+        )
+
+        # Initialize the linear layers for converting transformer output to logits
+        self.to_logits = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, self.total_tokens),
+        )
+
+        # Set instance variables for weights and critic
+        self.loss_img_weight = loss_img_weight
+        self.loss_cond_weight = loss_cond_weight
+        self.gamma = gamma_func(sampling_mode)
+
+    def embed_and_transform(self, inputs, masks, return_encoding=False):
+        text, condition, image = inputs
+        device = text.device
+        text_mask, _, image_mask = masks
+
+        text_labels = text.clone()
+        text = torch.where(
+            text_mask, self.text_mask_token * torch.ones_like(text, device=device), text
+        )
+
+        tokens = self.text_emb(text)
+
+        # Add SEP token
+
+        sep_token_emb = self.sep_emb(
+            torch.zeros((tokens.shape[0], 1), dtype=torch.long, device=device)
+        )
+        tokens = torch.cat((tokens, sep_token_emb), dim=1)
+        tokens += self.text_pos_emb(torch.arange(text.shape[1] + 1, device=device))
+
+        with torch.no_grad():
+            if self.linear_project:
+                b = condition.shape[0]
+                condition, _, [_, _, condition_labels] = self.condition_vae.encode(
+                    condition
+                )
+                condition_labels = rearrange(condition_labels, "(b n) -> b n", b=b)
+
+            else:
+                condition_labels = condition
+                if condition.dtype == torch.float:
+                    condition_labels = self.condition_vae.get_codebook_indices(
+                        condition
+                    )
+                condition = condition_labels.clone()
+
+        condition_emb = self.condition_emb(condition)
+        condition_emb += self.condition_pos_emb(condition_emb)
+        tokens = torch.cat((tokens, condition_emb), dim=1)
+
+        with torch.no_grad():
+            if self.linear_project:
+                b = image.shape[0]
+                image, _, [_, _, image_labels] = self.vae.encode(image)
+                image_labels = rearrange(image_labels, "(b n) -> b n", b=b)
+
+            else:
+                image_labels = image
+                if image.dtype == torch.float:
+                    image_labels = self.vae.get_codebook_indices(image)
+                image = torch.where(
+                    image_mask,
+                    self.image_mask_token
+                    * torch.ones_like(image_labels, device=device),
+                    image_labels,
+                )
+
+        image_emb = self.image_emb(image)
+
+        image_emb += self.image_pos_emb(image_emb)
+        tokens = torch.cat((tokens, image_emb), dim=1)
+
+        if self.stable:
+            alpha = 0.1
+            tokens = tokens * alpha + tokens.detach() * (1 - alpha)
+
+        out = self.transformer(tokens)
+
+        if self.stable:
+            out = self.norm_by_max(out)
+
+        logits = self.to_logits(out)
+
+        max_neg_value = -torch.finfo(logits.dtype).max
+        logits.masked_fill_(self.logits_mask, max_neg_value)
+
+        if return_encoding:
+            return logits, out, [text_labels, condition_labels, image_labels]
+        else:
+            return logits, None, [text_labels, condition_labels, image_labels]
+
+    def forward(
+        self,
+        text,
+        condition=None,
+        image=None,
+        return_loss=False,
+        return_encoding=False,
+    ):
+        batch_size, device = text.shape[0], text.device
+
+        # Check that image is supplied when training
+        assert exists(image), "when training, image must be supplied"
+
+        # Check that image dimensions match the expected dimensions
+        assert tuple(image.shape[1:]) == (
+            self.vae.channels,
+            self.image_size,
+            self.image_size,
+        ), f"invalid image of dimensions {image.shape} passed in during training"
+
+        # Generate masks for text, condition, and image
+
+        # text_mask = generate_mask(self.gamma, batch_size, self.text_seq_len, device)
+
+        text_mask = generate_mask(
+            gamma_func("scaled-cosine"), batch_size, self.text_seq_len, device
+        )
+
+        image_mask = generate_mask(self.gamma, batch_size, self.image_seq_len, device)
+
+        # Embed and transform inputs
+        logits, _, labels = self.embed_and_transform(
+            [text, condition, image],
+            [text_mask, None, image_mask],
+            return_encoding,
+            device,
+        )
+
+        # If not returning loss, return the logits
+        if not return_loss:
+            return logits
+
+        # Separate labels
+        text, condition, image = labels
+
+        # Add SEP token to end of text label
+        sep_token = torch.tensor(self.sep_token, device=device).repeat(
+            labels.shape[0], 1
+        )
+        labels = torch.cat([labels, sep_token], dim=1)
+
+        # If condition exists and condition vae is defined, add the condition to the labels
+        if exists(condition) and exists(self.condition_vae):
+            offsetted_condition = condition + self.num_text_tokens + 1
+            labels = torch.cat((labels, offsetted_condition), dim=1)
+
+        # Add image to the labels
+        offsetted_image = (
+            image + self.num_text_tokens + 1 + self.num_condition_tokens + 1
+        )
+        labels = torch.cat((labels, offsetted_image), dim=1)
+
+        # Rearrange logits for cross-entropy loss calculation
+        # Logits size: (batch_size, vocab_size, total_seq_len)
+        # Labels size: (batch_size, total_seq_len)
+        logits = rearrange(logits, "b n c -> b c n")
+
+        # Calculate cross-entropy loss for text and image
+        loss_text = F.cross_entropy(
+            logits[:, :, : self.text_seq_len],
+            labels[:, : self.text_seq_len],
+            reduction="none",
+        )[text_mask].mean()
+
+        loss_img = F.cross_entropy(
+            logits[:, :, self.text_seq_len + 1 + self.condition_seq_len :],
+            labels[:, self.text_seq_len + 1 + self.condition_seq_len :],
+            reduction="none",
+        )[image_mask].mean()
+
+        # Calculate total loss
+        loss = (loss_text + self.loss_img_weight * loss_img) / (
+            self.loss_img_weight + 1
+        )
+
+        loss_dict = {
+            "loss_text": loss_text,
+            # "loss_cond": loss_cond,
+            "loss_img": loss_img,
+            "loss": torch.nan_to_num(loss, 0.0, 0.0, 0.0),
+        }
+
+        return loss, loss_dict, None
+
+    def create_tensors(self, text, condition, image):
+        """
+        This function creates tensors for text, condition, and image when they are not provided as inputs to the sample function.
+        """
+        device = next(
+            filter(lambda x: isinstance(x, torch.Tensor), [text, condition, image]),
+            None,
+        ).device
+
+        if not isinstance(text, torch.Tensor):
+            text = (
+                torch.ones(1, self.text_seq_len, device=device, dtype=torch.long)
+                * self.text_mask_token
+            )
+
+        if not isinstance(condition, torch.Tensor):
+            condition = (
+                torch.ones(1, self.condition_seq_len, device=device, dtype=torch.long)
+                * self.cond_mask_token
+            )
+        else:
+            with torch.no_grad():
+                condition = self.condition_vae.get_codebook_indices(condition)
+
+        if not isinstance(image, torch.Tensor):
+            image = (
+                torch.ones(1, self.image_seq_len, device=device, dtype=torch.long)
+                * self.image_mask_token
+            )
+        else:
+            with torch.no_grad():
+                image = self.vae.get_codebook_indices(image)
+
+        return text, condition, image
+
+    @torch.no_grad()
+    @eval_decorator
+    def sample(
+        self,
+        text=None,
+        condition=None,
+        image=None,
+        temperature=1.0,
+        filter_thres=0.9,
+        progress=False,
+        timesteps=1,
+        force_aas=True,
+    ):
+        # ensure timesteps is a positive integer
+        assert int(timesteps) > 0
+        # set model and VAEs to evaluation mode
+        self.eval()
+        vae = self.vae.eval()
+        if progress == True:
+            progress = tqdm
+        else:
+            progress = lambda x: x
+
+
+        # ensure that at least one of text, condition, or image is supplied
+        assert (
+            isinstance(text, torch.Tensor)
+            or isinstance(condition, torch.Tensor)
+            or isinstance(image, torch.Tensor)
+        ), "some data must be supplied"
+
+        # convert text, condition, and image to tensors if they aren't already
+        text, condition, image = self.create_tensors(text, condition, image)
+
+        # determine the maximum batch size of the input tensors
+        batch_size = max(text.shape[0], condition.shape[0], image.shape[0])
+
+        # match the batch sizes of text, condition, and image
+        text, condition, image = match_batch_size(text, condition, image, batch_size)
+
+        # determine the device of the tensors
+        device = next(
+            filter(lambda x: isinstance(x, torch.Tensor), [text, condition, image]),
+            None,
+        ).device
+
+        assert text.shape[0] == condition.shape[0] == image.shape[0]
+
+        # Create a tensor of zeros of size (batch_size, image_seq_len, num_image_tokens + 1) and set it to device
+
+        # full_text_logits = torch.zeros(batch_size, self.text_seq_len, self.num_text_tokens+3).to(device)
+        full_text_logits = torch.zeros(
+            batch_size, self.text_seq_len, self.num_text_tokens
+        ).to(device)
+
+        # Use scatter_ to fill the tensor with 1 values at the indices given by the image tensor
+        full_text_logits = full_text_logits.scatter_(
+            dim=-1, index=text.unsqueeze(-1), value=1
+        )
+        # Use scatter_ to fill the tensor with 1 values at the indices given by the image tensor
+        full_image_logits = torch.zeros(
+            batch_size, self.image_seq_len, self.num_image_tokens + 1
+        ).to(device)
+
+        # Remove the last token from each image sequence by setting full_image_logits to its first num_image_tokens elements
+        full_image_logits = full_image_logits.scatter_(
+            dim=-1, index=image.unsqueeze(-1), value=1
+        )
+
+        # cut off mask token
+        full_image_logits = full_image_logits[:, :, : self.num_image_tokens]
+
+        count = 0
+
+        for timestep in progress(torch.linspace(0, 1, timesteps)):
+            # Create masks for the text, condition, and image tensors
+            text_mask = text == self.text_mask_token
+            cond_mask = condition == self.cond_mask_token
+            image_mask = image == self.image_mask_token
+
+            # Calculate logits and samples using the calculate_logits function
+            logits, sample = calculate_logits(
+                [text, condition, image],
+                [text_mask, cond_mask, image_mask],
+                self.embed_and_transform,
+                filter_thres,
+                temperature,
+            )
+
+            # Calculate the number of masked tokens in the text and image tensors
+            num_masked_text_tokens = torch.sum(text_mask, dim=1)[0]
+            num_masked_image_tokens = torch.sum(image_mask, dim=1)[0]
+
+            # If there are masked text tokens, unmask them using unmask_tokens and fill the full text logits tensor with -inf for unmasked tokens
+            if num_masked_text_tokens.any() > 0:
+                text, full_text_logits = unmask_tokens(
+                    text,
+                    text_mask,
+                    num_masked_text_tokens,
+                    logits[:, : self.text_seq_len, : self.num_text_tokens],
+                    sample[:, : self.text_seq_len, : self.num_text_tokens],
+                    timestep,
+                    timesteps,
+                    self.gamma,
+                    suppress_invalid_text_tokens,
+                    self.pad_token,
+                    self.text_mask_token,
+                    force_aas=force_aas,
+                )
+                full_text_logits = full_text_logits.masked_fill(
+                    ~text_mask.unsqueeze(-1), -torch.inf
+                )
+
+            # If there are masked image tokens, unmask them using unmask_tokens and fill the full image logits tensor with -inf for unmasked tokens
+            if num_masked_image_tokens > 0:
+                image, full_image_logits = unmask_tokens(
+                    image,
+                    image_mask,
+                    num_masked_image_tokens,
+                    logits[:, -self.image_seq_len :, -(self.num_image_tokens + 1) : -1],
+                    sample[:, -self.image_seq_len :, -(self.num_image_tokens + 1) : -1],
+                    timestep,
+                    timesteps,
+                    self.gamma,
+                )
+                full_text_logits = full_text_logits.masked_fill(
+                    ~text_mask.unsqueeze(-1), -torch.inf
+                )
+                
+        # Generate heatmap
+        with torch.no_grad():
+            # Normalize full image logits tensor
+            full_image_logits /= torch.max(
+                torch.abs(full_image_logits), dim=-1, keepdim=True
+            ).values
+
+            # Apply quantize embedding to full image logits tensor
+            full_image_logits = torch.matmul(
+                full_image_logits, self.vae.model.quantize.embedding.weight
+            )
+
+            # Rearrange full image logits tensor
+            h = int(self.image_seq_len**0.5)
+            full_image_logits = rearrange(
+                full_image_logits, "b (h w) c -> b c h w", h=h
+            )
+
+            # Decode full image logits tensor
+            full_image_logits = self.vae.model.decode(full_image_logits)
+
+            # Add clipping to full image logits tensor
+            max_val = torch.max(full_image_logits.view(batch_size, -1), dim=-1)[0]
+            min_val = torch.min(full_image_logits.view(batch_size, -1), dim=-1)[0]
+            full_image_logits += torch.clip(1 - max_val, 0, float("inf")).view(
+                batch_size, 1, 1, 1
+            )
+            full_image_logits += torch.clip(0 - min_val, float("-inf"), 0).view(
+                batch_size, 1, 1, 1
+            )
+
+            # Clip full image logits tensor values to the range [0, 1]
+            full_image_logits = torch.clip(full_image_logits, 0, 1)
+
+        # Return text tensor, detokenized text tensor, full text logits tensor,
+        # binary image tensor, and full image logits tensor
+        return (
+            text,
+            detokenize_text(self.text_embedding, text),
+            full_text_logits,
+            1.0 * (vae.decode(image) > 0.5),
+            full_image_logits,
+        )
+
+    @torch.no_grad()
+    @eval_decorator
+    def sample_text(
+        self,
+        text=False,
+        condition=False,
+        image=False,
+        temperature=1.0,
+        filter_thres=0.9,
+        progress=False,
+        n_unmask=1,
+        place_amino=True,
+        force_aas=False,
+    ):
+        # set model and VAEs to evaluation mode
+        self.eval()
+
+        # ensure that at least one of text, condition, or image is supplied
+        assert (
+            isinstance(text, torch.Tensor)
+            or isinstance(condition, torch.Tensor)
+            or isinstance(image, torch.Tensor)
+        ), "some data must be supplied"
+
+        # convert text, condition, and image to tensors if they aren't already
+        text, condition, image = self.create_tensors(text, condition, image)
+
+        # determine the maximum batch size of the input tensors
+        batch_size = max(text.shape[0], condition.shape[0], image.shape[0])
+
+        # match the batch sizes of text, condition, and image
+        text, condition, image = match_batch_size(text, condition, image, batch_size)
+
+        # determine the device of the tensors
+        device = next(
+            filter(lambda x: isinstance(x, torch.Tensor), [text, condition, image]),
+            None,
+        ).device
+
+        assert text.shape[0] == condition.shape[0] == image.shape[0]
+
+        # Create a tensor of zeros of size (batch_size, image_seq_len, num_image_tokens + 1) and set it to device
+
+        # full_text_logits = torch.zeros(batch_size, self.text_seq_len, self.num_text_tokens+3).to(device)
+        full_text_logits = torch.zeros(
+            batch_size, self.text_seq_len, self.num_text_tokens
+        ).to(device)
+
+        # Use scatter_ to fill the tensor with 1 values at the indices given by the image tensor
+        full_text_logits = full_text_logits.scatter_(
+            dim=-1, index=text.unsqueeze(-1), value=1
+        )
+
+        text_mask = text == self.text_mask_token
+        cond_mask = condition == self.cond_mask_token
+        image_mask = image == self.image_mask_token
+
+        mask_indices = text_mask.nonzero()
+        non_mask_indices = (~text_mask).nonzero()
+
+        # figure out the center of the amino acids to determine generation direction
+        central_protein_index = torch.tensor(
+            [
+                torch.median(
+                    non_mask_indices[torch.where(non_mask_indices[:, 0] == idx)][:, -1]
+                )
+                for idx in range(batch_size)
+            ]
+        )
+
+        count = 1
+
+        run_mask = text_mask
+        if progress:
+            pbar = progress(total=torch.sum(run_mask).item())
+        while torch.sum(run_mask) > 0:
+            logits, sample = calculate_logits(
+                [text, condition, image],
+                [text_mask, cond_mask, image_mask],
+                self.embed_and_transform,
+                filter_thres,
+                temperature,
+            )
+
+            # sub_sample: [batch_size ,text_seq_len ,num_text_tokens]
+            sub_sample = sample[:, : self.text_seq_len, : self.num_text_tokens]
+            sub_sample = sub_sample.masked_fill(~text_mask.unsqueeze(-1), -torch.inf)
+            sub_sample = suppress_invalid_text_tokens(
+                text, sub_sample, 0, 2, self.pad_token, self.text_mask_token, force_aas
+            )
+            # calculate % to  unmasked
+            # get most likely token and probability for each position
+
+            for idx in range(batch_size):
+                selected_mask_indices = mask_indices[
+                    torch.where(mask_indices[:, 0] == idx)
+                ][:, -1]
+
+                # Generate to the left
+                if selected_mask_indices[-count] < central_protein_index[idx]:
+                    unmask_index = selected_mask_indices[-count]
+                    left_sample = max(0, (unmask_index + 1) - n_unmask)
+                    right_sample = min(unmask_index + 1, self.text_seq_len - 1)
+                    central_protein_index[idx] = max(
+                        0, central_protein_index[idx] - 0.5 * n_unmask
+                    )
+
+                # Generate to the right
+                elif selected_mask_indices[count - 1] > central_protein_index[idx]:
+                    unmask_index = selected_mask_indices[count - 1]
+                    left_sample = max(0, unmask_index)
+                    right_sample = min(unmask_index + n_unmask, self.text_seq_len - 1)
+                    central_protein_index[idx] = min(
+                        central_protein_index[idx] + 0.5 * n_unmask,
+                        self.text_seq_len - 1,
+                    )
+
+                # save logits for relevant position
+                full_text_logits[
+                    idx, left_sample:right_sample, : self.text_seq_len - 1
+                ] = logits[idx, left_sample:right_sample, : self.num_text_tokens]
+
+                run_mask[idx, left_sample:right_sample] = False
+
+                # you may want to resample the amion acids or calculate marginal probs
+                # if so, set place_amino to false
+                if place_amino:
+                    text[idx, left_sample:right_sample] = torch.where(
+                        text[idx, left_sample:right_sample] == self.text_mask_token,
+                        sub_sample[
+                            idx, left_sample:right_sample, : self.num_text_tokens
+                        ].argmax(dim=-1),
+                        text[idx, left_sample:right_sample],
+                    )
+
+                    text_mask = run_mask
+
+            count += n_unmask
+            
+            if progress:
+                pbar.update(n_unmask)
+        if progress:
+            pbar.close()
+
+        return (
+            text,
+            detokenize_text(self.text_embedding, text),
+            full_text_logits,
+        )

celle/reversible.py

@@ -0,0 +1,36 @@
+import torch.nn as nn
+
+# for routing arguments into the functions of the reversible layer
+def route_args(router, args, depth):
+    routed_args = [(dict(), dict()) for _ in range(depth)]
+    matched_keys = [key for key in args.keys() if key in router]
+
+    for key in matched_keys:
+        val = args[key]
+        for depth, ((f_args, g_args), routes) in enumerate(
+            zip(routed_args, router[key])
+        ):
+            new_f_args, new_g_args = map(
+                lambda route: ({key: val} if route else {}), routes
+            )
+            routed_args[depth] = ({**f_args, **new_f_args}, {**g_args, **new_g_args})
+    return routed_args
+
+class SequentialSequence(nn.Module):
+    def __init__(self, layers, args_route={}, layer_dropout=0.0):
+        super().__init__()
+        assert all(
+            len(route) == len(layers) for route in args_route.values()
+        ), "each argument route map must have the same depth as the number of sequential layers"
+        self.layers = layers
+        self.args_route = args_route
+        self.layer_dropout = layer_dropout
+
+    def forward(self, x, **kwargs):
+        args = route_args(self.args_route, kwargs, len(self.layers))
+        layers_and_args = list(zip(self.layers, args))
+
+        for (f, g), (f_args, g_args) in layers_and_args:
+            x = x + f(x, **f_args)
+            x = x + g(x, **g_args)
+        return x

celle/transformer.py

@@ -0,0 +1,213 @@
+from functools import partial
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from celle.reversible import SequentialSequence
+from celle.attention import Attention
+
+from rotary_embedding_torch import RotaryEmbedding, broadcat
+from celle.utils import exists, default, cast_tuple
+
+# https://arxiv.org/abs/2103.17239
+class LayerScale(nn.Module):
+    def __init__(self, dim, depth, fn):
+        super().__init__()
+        if depth <= 18:
+            init_eps = 0.1
+        elif depth > 18 and depth <= 24:
+            init_eps = 1e-5
+        else:
+            init_eps = 1e-6
+
+        scale = torch.zeros(1, 1, dim).fill_(init_eps)
+        self.scale = nn.Parameter(scale)
+        self.fn = fn
+
+    def forward(self, x, **kwargs):
+        return self.fn(x, **kwargs) * self.scale
+
+
+# layer norm
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.norm_out = nn.Identity()
+        self.fn = fn
+
+    def forward(self, x, **kwargs):
+        x = self.norm(x)
+        x = self.fn(x, **kwargs)
+        return self.norm_out(x)
+
+
+# feed forward
+
+
+class GEGLU(nn.Module):
+    def forward(self, x):
+        x, gates = x.chunk(2, dim=-1)
+        return x * F.gelu(gates)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, dropout=0.0, mult=4.0):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, dim * mult * 2),
+            GEGLU(),
+            nn.Dropout(dropout),
+            nn.Linear(dim * mult, dim),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+# main transformer class
+class Transformer(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        depth,
+        seq_len,
+        causal=True,
+        heads=8,
+        dim_head=64,
+        ff_mult=4,
+        attn_dropout=0.0,
+        ff_dropout=0.0,
+        image_fmap_size=None,
+        num_images=None,
+        stable=False,
+        rotary_emb=True,
+    ):
+        super().__init__()
+        layers = nn.ModuleList([])
+
+        self.seq_len = seq_len
+        self.image_fmap_size = image_fmap_size
+
+        for ind in range(depth):
+            
+            attn_class = partial(Attention, stable=stable)
+
+            attn = attn_class(
+                dim,
+                causal=causal,
+                seq_len=seq_len,
+                heads=heads,
+                dim_head=dim_head,
+                dropout=attn_dropout,
+            )
+
+            ff = FeedForward(dim, mult=ff_mult, dropout=ff_dropout)
+
+            layers.append(
+                nn.ModuleList(
+                    [
+                        LayerScale(
+                            dim, ind + 1, PreNorm(dim, attn)
+                        ),
+                        LayerScale(
+                            dim, ind + 1, PreNorm(dim, ff)
+                        ),
+                    ]
+                )
+            )
+
+        # pairs arguments with attention layer
+        route_attn = ((True, False),) * depth
+        attn_route_map = {
+            "mask": route_attn,
+            "rotary_pos_emb": route_attn,
+        }
+
+        self.layers = SequentialSequence(layers, args_route=attn_route_map)
+
+        # generate positional embeddings for rotary
+
+        pos_emb = None
+        if rotary_emb:
+            rot_dim = dim_head // 3
+            img_seq_len = ((image_fmap_size // num_images) ** 2) * num_images
+
+            text_len = seq_len - img_seq_len + 1
+
+            text_pos_emb = RotaryEmbedding(dim=rot_dim)
+
+            img_axial_pos_emb = RotaryEmbedding(dim=rot_dim, freqs_for="pixel")
+
+            text_freqs = text_pos_emb(torch.arange(text_len))
+
+            img_to_text_freqs = text_pos_emb(
+                torch.full((img_seq_len,), 8192)
+            )  # image is given a position far away from text
+
+            text_freqs = torch.cat((text_freqs, img_to_text_freqs), dim=0)
+
+            img_freqs_axial = img_axial_pos_emb(
+                torch.linspace(-1, 1, steps=image_fmap_size)
+            )
+
+            if num_images > 1:
+                split_img_freqs_axial = torch.split(
+                    img_freqs_axial, image_fmap_size // num_images, dim=0
+                )
+
+                split_img_freqs = [
+                    broadcat(
+                        (
+                            rearrange(img_freqs_axial_per_image, "i d -> i () d"),
+                            rearrange(img_freqs_axial_per_image, "j d -> () j d"),
+                        ),
+                        dim=-1,
+                    )
+                    for img_freqs_axial_per_image in split_img_freqs_axial
+                ]
+
+                split_img_freqs = [
+                    rearrange(img_freqs_per_image, "h w d -> (h w) d")
+                    for img_freqs_per_image in split_img_freqs
+                ]
+
+                # concat per image-image_freqs
+
+                img_freqs = torch.cat(split_img_freqs, dim=0)
+
+            elif num_images == 1:
+                img_freqs = broadcat(
+                    (
+                        rearrange(img_freqs_axial, "i d -> i () d"),
+                        rearrange(img_freqs_axial, "j d -> () j d"),
+                    ),
+                    dim=-1,
+                )
+
+                img_freqs = rearrange(img_freqs, "h w d -> (h w) d")
+
+            else:
+                assert False, "num_images must be int greater than 0"
+            self.img_axial_pos_emb = img_axial_pos_emb
+            self.text_pos_emb = text_pos_emb
+
+            text_axial_freqs = img_axial_pos_emb(
+                torch.full((text_len,), -10.0)
+            )  # text is given a position of -10 apart from the image axial positions, which is from range [-1, 1]
+
+            text_axial_freqs = torch.cat((text_axial_freqs, text_axial_freqs), dim=-1)
+
+            img_freqs = torch.cat((text_axial_freqs, img_freqs), dim=0)
+
+            pos_emb = torch.cat((text_freqs, img_freqs), dim=-1)
+
+            pos_emb = rearrange(pos_emb, "n d -> () n d")
+
+        self.register_buffer("pos_emb", pos_emb)
+
+    def forward(self, x, **kwargs):
+        return self.layers(x, rotary_pos_emb=self.pos_emb, **kwargs)

celle/utils.py

@@ -0,0 +1,193 @@
+import torch
+from math import pi
+
+# Define helper functions
+def exists(val):
+    """Check if a variable exists"""
+    return val is not None
+
+
+def uniq(arr):
+    return {el: True for el in arr}.keys()
+
+
+def default(val, d):
+    """If a value exists, return it; otherwise, return a default value"""
+    return val if exists(val) else d
+
+
+def max_neg_value(t):
+    return -torch.finfo(t.dtype).max
+
+
+def cast_tuple(val, depth=1):
+    if isinstance(val, list):
+        val = tuple(val)
+    return val if isinstance(val, tuple) else (val,) * depth
+
+
+def is_empty(t):
+    """Check if a tensor is empty"""
+    # Return True if the number of elements in the tensor is zero, else False
+    return t.nelement() == 0
+
+
+def masked_mean(t, mask, dim=1):
+    """
+    Compute the mean of a tensor, masked by a given mask
+
+    Args:
+        t (torch.Tensor): input tensor of shape (batch_size, seq_len, hidden_dim)
+        mask (torch.Tensor): mask tensor of shape (batch_size, seq_len)
+        dim (int): dimension along which to compute the mean (default=1)
+
+    Returns:
+        torch.Tensor: masked mean tensor of shape (batch_size, hidden_dim)
+    """
+    t = t.masked_fill(~mask[:, :, None], 0.0)
+    return t.sum(dim=1) / mask.sum(dim=1)[..., None]
+
+
+def set_requires_grad(model, value):
+    """
+    Set whether or not the model's parameters require gradients
+
+    Args:
+        model (torch.nn.Module): the PyTorch model to modify
+        value (bool): whether or not to require gradients
+    """
+    for param in model.parameters():
+        param.requires_grad = value
+
+
+def eval_decorator(fn):
+    """
+    Decorator function to evaluate a given function
+
+    Args:
+        fn (callable): function to evaluate
+
+    Returns:
+        callable: the decorated function
+    """
+
+    def inner(model, *args, **kwargs):
+        was_training = model.training
+        model.eval()
+        out = fn(model, *args, **kwargs)
+        model.train(was_training)
+        return out
+
+    return inner
+
+
+def log(t, eps=1e-20):
+    """
+    Compute the natural logarithm of a tensor
+
+    Args:
+        t (torch.Tensor): input tensor
+        eps (float): small value to add to prevent taking the log of 0 (default=1e-20)
+
+    Returns:
+        torch.Tensor: the natural logarithm of the input tensor
+    """
+    return torch.log(t + eps)
+
+
+def gumbel_noise(t):
+    """
+    Generate Gumbel noise
+
+    Args:
+        t (torch.Tensor): input tensor
+
+    Returns:
+        torch.Tensor: a tensor of Gumbel noise with the same shape as the input tensor
+    """
+    noise = torch.zeros_like(t).uniform_(0, 1)
+    return -log(-log(noise))
+
+
+def gumbel_sample(t, temperature=0.9, dim=-1):
+    """
+    Sample from a Gumbel-softmax distribution
+
+    Args:
+        t (torch.Tensor): input tensor of shape (batch_size, num_classes)
+        temperature (float): temperature for the Gumbel-softmax distribution (default=0.9)
+        dim (int): dimension along which to sample (default=-1)
+
+    Returns:
+        torch.Tensor: a tensor of samples from the Gumbel-softmax distribution with the same shape as the input tensor
+    """
+    return (t / max(temperature, 1e-10)) + gumbel_noise(t)
+
+
+def top_k(logits, thres=0.5):
+    """
+    Return a tensor where all but the top k values are set to negative infinity
+
+    Args:
+        logits (torch.Tensor): input tensor of shape (batch_size, num_classes)
+        thres (float): threshold for the top k values (default=0.5)
+
+    Returns:
+        torch.Tensor: a tensor with the same shape as the input tensor, where all but the top k values are set to negative infinity
+    """
+    num_logits = logits.shape[-1]
+    k = max(int((1 - thres) * num_logits), 1)
+    val, ind = torch.topk(logits, k)
+    probs = torch.full_like(logits, float("-inf"))
+    probs.scatter_(-1, ind, val)
+    return probs
+
+
+def gamma_func(mode="cosine", scale=0.15):
+    """Return a function that takes a single input r and returns a value based on the selected mode"""
+
+    # Define a different function based on the selected mode
+    if mode == "linear":
+        return lambda r: 1 - r
+    elif mode == "cosine":
+        return lambda r: torch.cos(r * pi / 2)
+    elif mode == "square":
+        return lambda r: 1 - r**2
+    elif mode == "cubic":
+        return lambda r: 1 - r**3
+    elif mode == "scaled-cosine":
+        return lambda r: scale * (torch.cos(r * pi / 2))
+    else:
+        # Raise an error if the selected mode is not implemented
+        raise NotImplementedError
+
+
+class always:
+    """Helper class to always return a given value"""
+
+    def __init__(self, val):
+        self.val = val
+
+    def __call__(self, x, *args, **kwargs):
+        return self.val
+
+
+class DivideMax(torch.nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        maxes = x.amax(dim=self.dim, keepdim=True).detach()
+        return x / maxes
+
+def process_image(image_path):
+    image = Image.open(image_path)
+    transform = transforms.Compose([
+        transforms.RandomCrop(256),
+        transforms.ToTensor()
+    ])
+    image_tensor = transform(image)
+    if image_tensor.shape[0] > 1:
+        image_tensor = torch.mean(image_tensor, dim=0, keepdim=True)
+    return image_tensor.unsqueeze(0)

celle/vae.py

@@ -0,0 +1,112 @@
+from math import sqrt, log
+from omegaconf import OmegaConf
+import importlib
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from einops import rearrange
+
+# helpers methods
+
+
+def load_model(path):
+    with open(path, "rb") as f:
+        return torch.load(f, map_location=torch.device("cpu"))
+
+
+def map_pixels(x, eps=0.1):
+    return (1 - 2 * eps) * x + eps
+
+
+def unmap_pixels(x, eps=0.1):
+    return torch.clamp((x - eps) / (1 - 2 * eps), 0, 1)
+
+
+def make_contiguous(module):
+    with torch.no_grad():
+        for param in module.parameters():
+            param.set_(param.contiguous())
+
+
+# VQGAN from Taming Transformers paper
+# https://arxiv.org/abs/2012.09841
+
+
+def get_obj_from_str(string, reload=False):
+    module, cls = string.rsplit(".", 1)
+    if reload:
+        module_imp = importlib.import_module(module)
+        importlib.reload(module_imp)
+    return getattr(importlib.import_module(module, package=None), cls)
+
+
+def instantiate_from_config(config):
+    if not "target" in config:
+        raise KeyError("Expected key `target` to instantiate.")
+    return get_obj_from_str(config["target"])(**config.get("params", dict()))
+
+
+class VQGanVAE(nn.Module):
+    def __init__(self, vqgan_model_path=None, vqgan_config_path=None, channels=1):
+        super().__init__()
+
+        assert vqgan_config_path is not None
+
+        model_path = vqgan_model_path
+        config_path = vqgan_config_path
+
+        config = OmegaConf.load(config_path)
+
+        model = instantiate_from_config(config["model"])
+
+        if vqgan_model_path:
+
+            state = torch.load(model_path, map_location="cpu")["state_dict"]
+            model.load_state_dict(state, strict=True)
+
+        print(f"Loaded VQGAN from {model_path} and {config_path}")
+
+        self.model = model
+
+        # f as used in https://github.com/CompVis/taming-transformers#overview-of-pretrained-models
+        f = (
+            config.model.params.ddconfig.resolution
+            / config.model.params.ddconfig.attn_resolutions[0]
+        )
+        self.num_layers = int(log(f) / log(2))
+        self.image_size = config.model.params.ddconfig.resolution
+        self.num_tokens = config.model.params.n_embed
+        # self.is_gumbel = isinstance(self.model, GumbelVQ)
+        self.is_gumbel = False
+        self.channels = config.model.params.ddconfig.in_channels
+
+    def encode(self, img):
+        return self.model.encode(img)
+
+    def get_codebook_indices(self, img):
+        b = img.shape[0]
+        # img = (2 * img) - 1
+        _, _, [_, _, indices] = self.encode(img)
+        if self.is_gumbel:
+            return rearrange(indices, "b h w -> b (h w)", b=b)
+        return rearrange(indices, "(b n) -> b n", b=b)
+
+    def decode(self, img_seq):
+        b, n = img_seq.shape
+        one_hot_indices = F.one_hot(img_seq, num_classes=self.num_tokens).float()
+        z = (
+            one_hot_indices @ self.model.quantize.embed.weight
+            if self.is_gumbel
+            else (one_hot_indices @ self.model.quantize.embedding.weight)
+        )
+
+        z = rearrange(z, "b (h w) c -> b c h w", h=int(sqrt(n)))
+        img = self.model.decode(z)
+
+        # img = (img.clamp(-1.0, 1.0) + 1) * 0.5
+        return img
+
+    def forward(self, img, optimizer_idx=1):
+        return self.model.training_step(img, optimizer_idx=optimizer_idx)

celle_main.py

@@ -0,0 +1,619 @@
+import os
+import numpy as np
+
+import torch
+import torch.random
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+import pytorch_lightning as pl
+from pytorch_lightning import seed_everything
+from pytorch_lightning.trainer import Trainer
+
+from dataloader import CellLoader
+from celle import VQGanVAE, CELLE
+from omegaconf import OmegaConf
+import argparse, os, sys, datetime, glob
+
+from celle.celle import gumbel_sample, top_k
+
+torch.random.manual_seed(42)
+np.random.seed(42)
+
+from celle_taming_main import (
+    instantiate_from_config,
+    nondefault_trainer_args,
+    get_parser,
+)
+
+
+class CellDataModule(pl.LightningDataModule):
+    def __init__(
+        self,
+        data_csv,
+        dataset,
+        sequence_mode="standard",
+        vocab="bert",
+        crop_size=256,
+        resize=600,
+        batch_size=1,
+        threshold="median",
+        text_seq_len=1000,
+        num_workers=1,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.data_csv = data_csv
+        self.dataset = dataset
+        self.protein_sequence_length = 0
+        self.image_folders = []
+        self.crop_size = crop_size
+        self.resize = resize
+        self.batch_size = batch_size
+        self.sequence_mode = sequence_mode
+        self.threshold = threshold
+        self.text_seq_len = int(text_seq_len)
+        self.vocab = vocab
+        self.num_workers = num_workers if num_workers is not None else batch_size * 2
+
+    def setup(self, stage=None):
+        # called on every GPU
+        self.cell_dataset_train = CellLoader(
+            data_csv=self.data_csv,
+            dataset=self.dataset,
+            crop_size=self.crop_size,
+            resize=self.resize,
+            split_key="train",
+            crop_method="random",
+            sequence_mode=self.sequence_mode,
+            vocab=self.vocab,
+            text_seq_len=self.text_seq_len,
+            threshold=self.threshold,
+        )
+
+        self.cell_dataset_val = CellLoader(
+            data_csv=self.data_csv,
+            dataset=self.dataset,
+            crop_size=self.crop_size,
+            resize=self.resize,
+            crop_method="center",
+            split_key="val",
+            sequence_mode=self.sequence_mode,
+            vocab=self.vocab,
+            text_seq_len=self.text_seq_len,
+            threshold=self.threshold,
+        )
+
+    def prepare_data(self):
+
+        pass
+
+    def train_dataloader(self):
+        return DataLoader(
+            self.cell_dataset_train,
+            num_workers=self.num_workers,
+            shuffle=True,
+            batch_size=self.batch_size,
+        )
+
+    def val_dataloader(self):
+        return DataLoader(
+            self.cell_dataset_val,
+            num_workers=self.num_workers,
+            batch_size=self.batch_size,
+        )
+
+    # def test_dataloader(self):
+    #    transforms = ...
+    #    return DataLoader(self.test, batch_size=64)
+
+
+class CELLE_trainer(pl.LightningModule):
+    def __init__(
+        self,
+        vqgan_model_path,
+        vqgan_config_path,
+        ckpt_path=None,
+        image_key="threshold",
+        condition_model_path=None,
+        condition_config_path=None,
+        num_images=2,
+        dim=2,
+        num_text_tokens=30,
+        text_seq_len=1000,
+        depth=16,
+        heads=16,
+        dim_head=64,
+        attn_dropout=0.1,
+        ff_dropout=0.1,
+        attn_types="full",
+        loss_img_weight=7,
+        stable=False,
+        rotary_emb=True,
+        text_embedding="bert",
+        fixed_embedding=True,
+        loss_cond_weight=1,
+        learning_rate=3e-4,
+        monitor="val_loss",
+    ):
+        super().__init__()
+
+        vae = VQGanVAE(
+            vqgan_model_path=vqgan_model_path, vqgan_config_path=vqgan_config_path
+        )
+
+        self.image_key = image_key
+
+        if condition_config_path:
+            condition_vae = VQGanVAE(
+                vqgan_model_path=condition_model_path,
+                vqgan_config_path=condition_config_path,
+            )
+        else:
+            condition_vae = None
+
+        self.celle = CELLE(
+            dim=dim,
+            vae=vae,  # automatically infer (1) image sequence length and (2) number of image tokens
+            condition_vae=condition_vae,
+            num_images=num_images,
+            num_text_tokens=num_text_tokens,  # vocab size for text
+            text_seq_len=text_seq_len,  # text sequence length
+            depth=depth,  # should aim to be 64
+            heads=heads,  # attention heads
+            dim_head=dim_head,  # attention head dimension
+            attn_dropout=attn_dropout,  # attention dropout
+            ff_dropout=ff_dropout,  # feedforward dropout
+            loss_img_weight=loss_img_weight,
+            stable=stable,
+            rotary_emb=rotary_emb,
+            text_embedding=text_embedding,
+            fixed_embedding=fixed_embedding,
+            loss_cond_weight=loss_cond_weight,
+        )
+
+        self.learning_rate = learning_rate
+        self.num_text_tokens = num_text_tokens
+        self.num_images = num_images
+
+        if monitor is not None:
+            self.monitor = monitor
+
+        ignore_keys = []
+
+        if condition_model_path:
+            ignore_keys.append("celle.condition_vae")
+
+        if vqgan_model_path:
+            ignore_keys.append("celle.vae")
+
+        if ckpt_path is not None:
+            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
+
+    def init_from_ckpt(self, path, ignore_keys=list()):
+        sd = torch.load(path, map_location="cpu")["state_dict"]
+        ckpt = sd.copy()
+        for k in sd.keys():
+            for ik in ignore_keys:
+                if k.startswith(ik):
+                    # print("Deleting key {} from state_dict.".format(k))
+                    del ckpt[k]
+        self.load_state_dict(ckpt, strict=True)
+        print(f"Restored from {path}")
+
+    def forward(self, text, condition, target, return_loss=True):
+
+        return self.celle(
+            text=text, condition=condition, image=target, return_loss=return_loss
+        )
+
+    def get_input(self, batch):
+        text = batch["sequence"].squeeze(1)
+        condition = batch["nucleus"]
+        target = batch[self.image_key]
+
+        return text, condition, target
+
+    def get_image_from_logits(self, logits, temperature=0.9):
+
+        filtered_logits = top_k(logits, thres=0.5)
+        sample = gumbel_sample(filtered_logits, temperature=temperature, dim=-1)
+
+        self.celle.vae.eval()
+        out = self.celle.vae.decode(
+            sample[:, self.celle.text_seq_len + self.celle.condition_seq_len :]
+            - (self.celle.num_text_tokens + self.celle.num_condition_tokens)
+        )
+
+        return out
+
+    def get_loss(self, text, condition, target):
+
+        loss_dict = {}
+
+        loss, loss_dict, logits = self(text, condition, target, return_loss=True)
+
+        return loss, loss_dict
+
+    def total_loss(
+        self,
+        loss,
+        loss_dict,
+        mode="train",
+    ):
+
+        loss_dict = {f"{mode}/{key}": value for key, value in loss_dict.items()}
+
+        for key, value in loss_dict.items():
+            self.log(
+                key,
+                value,
+                prog_bar=True,
+                logger=True,
+                on_step=True,
+                on_epoch=True,
+                sync_dist=True,
+            )
+
+        return loss
+
+    def training_step(self, batch, batch_idx):
+
+        text, condition, target = self.get_input(batch)
+        loss, log_dict = self.get_loss(text, condition, target)
+
+        loss = self.total_loss(loss, log_dict, mode="train")
+
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+
+        with torch.no_grad():
+
+            text, condition, target = self.get_input(batch)
+            loss, log_dict = self.get_loss(text, condition, target)
+
+            loss = self.total_loss(loss, log_dict, mode="val")
+
+        return loss
+
+    def configure_optimizers(self):
+
+        optimizer = AdamW(self.parameters(), lr=self.learning_rate, betas=(0.9, 0.95))
+
+        return optimizer
+
+    def scale_image(self, image):
+
+        for tensor in image:
+            if torch.min(tensor) < 0:
+                tensor += -torch.min(tensor)
+            else:
+                tensor -= torch.min(tensor)
+
+            tensor /= torch.max(tensor)
+
+        return image
+
+    @torch.no_grad()
+    def log_images(self, batch, **kwargs):
+        
+        log = []
+
+        text, condition, target = self.get_input(batch)
+        text = text.squeeze(1).to(self.device)
+        condition = condition.to(self.device)
+
+        out = self.celle.generate_images(text=text, condition=condition)
+
+        log["condition"] = self.scale_image(condition)
+        log["output"] = self.scale_image(out)
+        if self.image_key == "threshold":
+            log["threshold"] = self.scale_image(target)
+            log["target"] = self.scale_image(batch["target"])
+        else:
+            log["target"] = self.scale_image(target)
+
+        return log
+
+
+# from https://github.com/CompVis/taming-transformers/blob/master/celle_main.py
+
+if __name__ == "__main__":
+    # custom parser to specify config files, train, test and debug mode,
+    # postfix, resume.
+    # `--key value` arguments are interpreted as arguments to the trainer.
+    # `nested.key=value` arguments are interpreted as config parameters.
+    # configs are merged from left-to-right followed by command line parameters.
+
+    # model:
+    #   learning_rate: float
+    #   target: path to lightning module
+    #   params:
+    #       key: value
+    # data:
+    #   target: celle_main.DataModuleFromConfig
+    #   params:
+    #      batch_size: int
+    #      wrap: bool
+    #      train:
+    #          target: path to train dataset
+    #          params:
+    #              key: value
+    #      validation:
+    #          target: path to validation dataset
+    #          params:
+    #              key: value
+    #      test:
+    #          target: path to test dataset
+    #          params:
+    #              key: value
+    # lightning: (optional, has sane defaults and can be specified on cmdline)
+    #   trainer:
+    #       additional arguments to trainer
+    #   logger:
+    #       logger to instantiate
+    #   modelcheckpoint:
+    #       modelcheckpoint to instantiate
+    #   callbacks:
+    #       callback1:
+    #           target: importpath
+    #           params:
+    #               key: value
+
+    now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
+
+    # add cwd for convenience and to make classes in this file available when
+    # running as `python celle_main.py`
+    # (in particular `celle_main.DataModuleFromConfig`)
+    sys.path.append(os.getcwd())
+
+    parser = get_parser()
+    parser = Trainer.add_argparse_args(parser)
+
+    opt, unknown = parser.parse_known_args()
+    if opt.name and opt.resume:
+        raise ValueError(
+            "-n/--name and -r/--resume cannot be specified both."
+            "If you want to resume training in a new log folder, "
+            "use -n/--name in combination with --resume_from_checkpoint"
+        )
+    if opt.resume:
+        if not os.path.exists(opt.resume):
+            raise ValueError("Cannot find {}".format(opt.resume))
+        if os.path.isfile(opt.resume):
+            paths = opt.resume.split("/")
+            idx = len(paths) - paths[::-1].index("logs") + 1
+            logdir = "/".join(paths[:idx])
+            ckpt = opt.resume
+        else:
+            assert os.path.isdir(opt.resume), opt.resume
+            logdir = opt.resume.rstrip("/")
+            ckpt = os.path.join(logdir, "checkpoints", "last.ckpt")
+
+        opt.resume_from_checkpoint = ckpt
+        base_configs = sorted(glob.glob(os.path.join(logdir, "configs/*.yaml")))
+        opt.base = base_configs + opt.base
+        _tmp = logdir.split("/")
+        nowname = _tmp[_tmp.index("logs") + 1]
+    else:
+        if opt.name:
+            name = "_" + opt.name
+        elif opt.base:
+            cfg_fname = os.path.split(opt.base[0])[-1]
+            cfg_name = os.path.splitext(cfg_fname)[0]
+            name = "_" + cfg_name
+        else:
+            name = ""
+        nowname = now + name + opt.postfix
+        logdir = os.path.join("logs", nowname)
+
+    ckptdir = os.path.join(logdir, "checkpoints")
+    cfgdir = os.path.join(logdir, "configs")
+    seed_everything(opt.seed)
+
+    try:
+        # init and save configs
+        configs = [OmegaConf.load(cfg) for cfg in opt.base]
+        cli = OmegaConf.from_dotlist(unknown)
+        config = OmegaConf.merge(*configs, cli)
+        lightning_config = config.pop("lightning", OmegaConf.create())
+        # merge trainer cli with config
+        trainer_config = lightning_config.get("trainer", OmegaConf.create())
+        # default to ddp
+        # trainer_config["distributed_backend"] = "ddp"
+        for k in nondefault_trainer_args(opt):
+            trainer_config[k] = getattr(opt, k)
+        if not "gpus" in trainer_config:
+            del trainer_config["distributed_backend"]
+            cpu = True
+        else:
+            gpuinfo = trainer_config["gpus"]
+            print(f"Running on GPUs {gpuinfo}")
+            cpu = False
+        trainer_opt = argparse.Namespace(**trainer_config)
+        lightning_config.trainer = trainer_config
+
+        # model
+        # model = instantiate_from_config(config.model)
+        model = instantiate_from_config(config.model)
+        # trainer and callbacks
+        trainer_kwargs = dict()
+
+        # default logger configs
+        # NOTE wandb < 0.10.0 interferes with shutdown
+        # wandb >= 0.10.0 seems to fix it but still interferes with pudb
+        # debugging (wrongly sized pudb ui)
+        # thus prefer testtube for now
+        default_logger_cfgs = {
+            "wandb": {
+                "target": "pytorch_lightning.loggers.WandbLogger",
+                "params": {
+                    "name": nowname,
+                    "save_dir": logdir,
+                    "offline": opt.debug,
+                    "id": nowname,
+                },
+            },
+            "testtube": {
+                # "target": "pytorch_lightning.loggers.TestTubeLogger",
+                "target": "pytorch_lightning.loggers.TensorBoardLogger",
+                "params": {
+                    "name": "testtube",
+                    "save_dir": logdir,
+                },
+            },
+        }
+        default_logger_cfg = default_logger_cfgs["testtube"]
+        # logger_cfg = lightning_config.logger or OmegaConf.create()
+        try:
+            logger_cfg = lightning_config.logger
+        except:
+            logger_cfg = OmegaConf.create()
+        logger_cfg = OmegaConf.merge(default_logger_cfg, logger_cfg)
+        trainer_kwargs["logger"] = instantiate_from_config(logger_cfg)
+
+        # modelcheckpoint - use TrainResult/EvalResult(checkpoint_on=metric) to
+        # specify which metric is used to determine best models
+        default_modelckpt_cfg = {
+            "checkpoint_callback": {
+                "target": "pytorch_lightning.callbacks.ModelCheckpoint",
+                "params": {
+                    "dirpath": ckptdir,
+                    "filename": "{epoch:06}",
+                    "verbose": True,
+                    "save_last": True,
+                },
+            }
+        }
+        if hasattr(model, "monitor"):
+            print(f"Monitoring {model.monitor} as checkpoint metric.")
+            default_modelckpt_cfg["checkpoint_callback"]["params"][
+                "monitor"
+            ] = model.monitor
+            default_modelckpt_cfg["checkpoint_callback"]["params"]["save_top_k"] = 3
+        try:
+            modelckpt_cfg = lightning_config.modelcheckpoint
+        except:
+            modelckpt_cfg = OmegaConf.create()
+        modelckpt_cfg = OmegaConf.merge(default_modelckpt_cfg, modelckpt_cfg)
+        # trainer_kwargs["checkpoint_callback"] = instantiate_from_config(modelckpt_cfg)
+
+        # add callback which sets up log directory
+        default_callbacks_cfg = {
+            "setup_callback": {
+                "target": "celle_taming_main.SetupCallback",
+                "params": {
+                    "resume": opt.resume,
+                    "now": now,
+                    "logdir": logdir,
+                    "ckptdir": ckptdir,
+                    "cfgdir": cfgdir,
+                    "config": config,
+                    "lightning_config": lightning_config,
+                },
+            },
+            # "image_logger": {
+            #     "target": "celle_taming_main.ImageLogger",
+            #     "params": {
+            #         "batch_frequency": 0,
+            #         "max_images": 0,
+            #         "clamp": False,
+            #         "increase_log_steps": False,
+            #     },
+            # },
+            # "learning_rate_logger": {
+            #     "target": "celle_taming_main.LearningRateMonitor",
+            #     "params": {
+            #         "logging_interval": "step",
+            #         # "log_momentum": True
+            # },
+            # },
+        }
+        try:
+            callbacks_cfg = lightning_config.callbacks
+        except:
+            callbacks_cfg = OmegaConf.create()
+        callbacks_cfg = OmegaConf.merge(default_callbacks_cfg, callbacks_cfg)
+        callbacks_cfg = OmegaConf.merge(modelckpt_cfg, callbacks_cfg)
+        trainer_kwargs["callbacks"] = [
+            instantiate_from_config(callbacks_cfg[k]) for k in callbacks_cfg
+        ]
+
+        trainer = Trainer.from_argparse_args(
+            trainer_opt, **trainer_kwargs, profiler="simple"
+        )
+
+        # data
+        data = instantiate_from_config(config.data)
+        # NOTE according to https://pytorch-lightning.readthedocs.io/en/latest/datamodules.html
+        # calling these ourselves should not be necessary but it is.
+        # lightning still takes care of proper multiprocessing though
+        data.setup()
+        data.prepare_data()
+
+        # configure learning rate
+        bs, lr = config.data.params.batch_size, config.model.learning_rate
+
+        if not cpu:
+            ngpu = len(lightning_config.trainer.gpus.strip(",").split(","))
+        else:
+            ngpu = 1
+        try:
+            accumulate_grad_batches = lightning_config.trainer.accumulate_grad_batches
+        except:
+            accumulate_grad_batches = 1
+        print(f"accumulate_grad_batches = {accumulate_grad_batches}")
+        lightning_config.trainer.accumulate_grad_batches = accumulate_grad_batches
+        model.learning_rate = accumulate_grad_batches * ngpu * bs * lr
+
+        print(
+            "Setting learning rate to {:.2e} = {} (accumulate_grad_batches) * {} (num_gpus) * {} (batchsize) * {:.2e} (lr)".format(
+                model.learning_rate, accumulate_grad_batches, ngpu, bs, lr
+            )
+        )
+
+        # allow checkpointing via USR1
+        def melk(*args, **kwargs):
+            # run all checkpoint hooks
+            if trainer.global_rank == 0:
+                print("Summoning checkpoint.")
+                ckpt_path = os.path.join(ckptdir, "last.ckpt")
+                trainer.save_checkpoint(ckpt_path)
+
+        def divein(*args, **kwargs):
+            if trainer.global_rank == 0:
+                import pudb
+
+                pudb.set_trace()
+
+        import signal
+
+        signal.signal(signal.SIGUSR1, melk)
+        signal.signal(signal.SIGUSR2, divein)
+
+        # run
+        if opt.train:
+            try:
+                # model = torch.compile(model, mode="reduce_overhead")
+                torch.compile(trainer.fit(model, data), mode="max-autotune")
+            except Exception:
+                melk()
+                raise
+        if not opt.no_test and not trainer.interrupted:
+            trainer.test(model, data)
+    except Exception:
+        if opt.debug and trainer.global_rank == 0:
+            try:
+                import pudb as debugger
+            except ImportError:
+                import pdb as debugger
+            debugger.post_mortem()
+        raise
+    finally:
+        # move newly created debug project to debug_runs
+        if opt.debug and not opt.resume and trainer.global_rank == 0:
+            dst, name = os.path.split(logdir)
+            dst = os.path.join(dst, "debug_runs", name)
+            os.makedirs(os.path.split(dst)[0], exist_ok=True)
+            os.rename(logdir, dst)

celle_taming_main.py

@@ -0,0 +1,695 @@
+import argparse, os, sys, datetime, glob, importlib
+from omegaconf import OmegaConf
+import numpy as np
+from PIL import Image
+import torch
+import torchvision
+from torch.utils.data import DataLoader, Dataset
+from dataloader import CellLoader
+from pytorch_lightning.callbacks import ModelCheckpoint, Callback, LearningRateMonitor
+import pytorch_lightning as pl
+from pytorch_lightning import seed_everything
+from pytorch_lightning.trainer import Trainer
+from pytorch_lightning.callbacks import Callback
+from pytorch_lightning.utilities import rank_zero_only
+
+
+def get_obj_from_str(string, reload=False):
+    module, cls = string.rsplit(".", 1)
+    if reload:
+        module_imp = importlib.import_module(module)
+        importlib.reload(module_imp)
+    return getattr(importlib.import_module(module, package=None), cls)
+
+
+def get_parser(**parser_kwargs):
+    def str2bool(v):
+        if isinstance(v, bool):
+            return v
+        if v.lower() in ("yes", "true", "t", "y", "1"):
+            return True
+        elif v.lower() in ("no", "false", "f", "n", "0"):
+            return False
+        else:
+            raise argparse.ArgumentTypeError("Boolean value expected.")
+
+    parser = argparse.ArgumentParser(**parser_kwargs)
+    parser.add_argument(
+        "-n",
+        "--name",
+        type=str,
+        const=True,
+        default="",
+        nargs="?",
+        help="postfix for logdir",
+    )
+    parser.add_argument(
+        "-r",
+        "--resume",
+        type=str,
+        const=True,
+        default="",
+        nargs="?",
+        help="resume from logdir or checkpoint in logdir",
+    )
+    parser.add_argument(
+        "-b",
+        "--base",
+        nargs="*",
+        metavar="base_config.yaml",
+        help="paths to base configs. Loaded from left-to-right. "
+        "Parameters can be overwritten or added with command-line options of the form `--key value`.",
+        default=list(),
+    )
+    parser.add_argument(
+        "-t",
+        "--train",
+        type=str2bool,
+        const=True,
+        default=False,
+        nargs="?",
+        help="train",
+    )
+    parser.add_argument(
+        "--no-test",
+        type=str2bool,
+        const=True,
+        default=False,
+        nargs="?",
+        help="disable test",
+    )
+    parser.add_argument(
+        "-p", "--project", help="name of new or path to existing project"
+    )
+    parser.add_argument(
+        "-d",
+        "--debug",
+        type=str2bool,
+        nargs="?",
+        const=True,
+        default=False,
+        help="enable post-mortem debugging",
+    )
+    parser.add_argument(
+        "-s",
+        "--seed",
+        type=int,
+        default=42,
+        help="seed for seed_everything",
+    )
+    parser.add_argument(
+        "-f",
+        "--postfix",
+        type=str,
+        default="",
+        help="post-postfix for default name",
+    )
+
+    return parser
+
+
+def nondefault_trainer_args(opt):
+    parser = argparse.ArgumentParser()
+    parser = Trainer.add_argparse_args(parser)
+    args = parser.parse_args([])
+    return sorted(k for k in vars(args) if getattr(opt, k) != getattr(args, k))
+
+
+def instantiate_from_config(config):
+    if not "target" in config:
+        raise KeyError("Expected key `target` to instantiate.")
+    return get_obj_from_str(config["target"])(**config.get("params", dict()))
+
+
+class WrappedDataset(Dataset):
+    """Wraps an arbitrary object with __len__ and __getitem__ into a pytorch dataset"""
+
+    def __init__(self, dataset):
+        self.data = dataset
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        return self.data[idx]
+
+
+class DataModuleFromConfig(pl.LightningDataModule):
+    def __init__(
+        self,
+        data_csv,
+        dataset,
+        crop_size=256,
+        resize=600,
+        batch_size=1,
+        sequence_mode="latent",
+        vocab="bert",
+        text_seq_len=0,
+        num_workers=1,
+        threshold=False,
+        train=True,
+        validation=True,
+        test=None,
+        wrap=False,
+        **kwargs,
+    ):
+        super().__init__()
+        self.data_csv = data_csv
+        self.dataset = dataset
+        self.image_folders = []
+        self.crop_size = crop_size
+        self.resize = resize
+        self.batch_size = batch_size
+        self.sequence_mode = sequence_mode
+        self.threshold = threshold
+        self.text_seq_len = int(text_seq_len)
+        self.vocab = vocab
+        self.dataset_configs = dict()
+        self.num_workers = num_workers if num_workers is not None else batch_size * 2
+        if train is not None:
+            self.dataset_configs["train"] = train
+            self.train_dataloader = self._train_dataloader
+        if validation is not None:
+            self.dataset_configs["validation"] = validation
+            self.val_dataloader = self._val_dataloader
+        if test is not None:
+            self.dataset_configs["test"] = test
+            self.test_dataloader = self._test_dataloader
+        self.wrap = wrap
+
+    def prepare_data(self):
+        pass
+
+    def setup(self, stage=None):
+        # called on every GPU
+        self.cell_dataset_train = CellLoader(
+            data_csv=self.data_csv,
+            dataset=self.dataset,
+            crop_size=self.crop_size,
+            split_key="train",
+            crop_method="random",
+            sequence_mode=None,
+            vocab=self.vocab,
+            text_seq_len=self.text_seq_len,
+            threshold=self.threshold,
+        )
+
+        self.cell_dataset_val = CellLoader(
+            data_csv=self.data_csv,
+            dataset=self.dataset,
+            crop_size=self.crop_size,
+            split_key="val",
+            crop_method="center",
+            sequence_mode=None,
+            vocab=self.vocab,
+            text_seq_len=self.text_seq_len,
+            threshold=self.threshold,
+        )
+
+    def _train_dataloader(self):
+        return DataLoader(
+            self.cell_dataset_train,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            shuffle=True,
+            batch_size=self.batch_size,
+        )
+
+    def _val_dataloader(self):
+        return DataLoader(
+            self.cell_dataset_val,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            batch_size=self.batch_size,
+        )
+
+    # def _test_dataloader(self):
+    #    return DataLoader(self.datasets["test"], batch_size=self.batch_size,
+    #                     num_workers=self.num_workers)
+
+
+class SetupCallback(Callback):
+    def __init__(self, resume, now, logdir, ckptdir, cfgdir, config, lightning_config):
+        super().__init__()
+        self.resume = resume
+        self.now = now
+        self.logdir = logdir
+        self.ckptdir = ckptdir
+        self.cfgdir = cfgdir
+        self.config = config
+        self.lightning_config = lightning_config
+
+    def on_fit_start(self, trainer, pl_module):
+        if trainer.global_rank == 0:
+            # Create logdirs and save configs
+            os.makedirs(self.logdir, exist_ok=True)
+            os.makedirs(self.ckptdir, exist_ok=True)
+            os.makedirs(self.cfgdir, exist_ok=True)
+
+            print("Project config")
+            print(OmegaConf.to_yaml(self.config))
+            OmegaConf.save(
+                self.config,
+                os.path.join(self.cfgdir, "{}-project.yaml".format(self.now)),
+            )
+
+            print("Lightning config")
+            print(OmegaConf.to_yaml(self.lightning_config))
+            OmegaConf.save(
+                OmegaConf.create({"lightning": self.lightning_config}),
+                os.path.join(self.cfgdir, "{}-lightning.yaml".format(self.now)),
+            )
+
+        else:
+            # ModelCheckpoint callback created log directory --- remove it
+            if not self.resume and os.path.exists(self.logdir):
+                dst, name = os.path.split(self.logdir)
+                dst = os.path.join(dst, "child_runs", name)
+                os.makedirs(os.path.split(dst)[0], exist_ok=True)
+                try:
+                    os.rename(self.logdir, dst)
+                except FileNotFoundError:
+                    pass
+
+
+class ImageLogger(Callback):
+    def __init__(
+        self, batch_frequency, max_images, clamp=True, increase_log_steps=True
+    ):
+        super().__init__()
+        self.batch_freq = batch_frequency
+        self.max_images = max_images
+        self.logger_log_images = {
+            pl.loggers.WandbLogger: self._wandb,
+            # pl.loggers.TestTubeLogger: self._testtube,
+            pl.loggers.TensorBoardLogger: self._testtube,
+        }
+        self.log_steps = [2**n for n in range(int(np.log2(self.batch_freq)) + 1)]
+        if not increase_log_steps:
+            self.log_steps = [self.batch_freq]
+        self.clamp = clamp
+
+    @rank_zero_only
+    def _wandb(self, pl_module, images, batch_idx, split):
+        raise ValueError("No way wandb")
+        grids = dict()
+        for k in images:
+            grid = torchvision.utils.make_grid(images[k])
+            grids[f"{split}/{k}"] = wandb.Image(grid)
+        pl_module.logger.experiment.log(grids)
+
+    @rank_zero_only
+    def _testtube(self, pl_module, images, batch_idx, split):
+        for k in images:
+            images[k] -= torch.min(images[k])
+            images[k] /= torch.max(images[k])
+            grid = torchvision.utils.make_grid(images[k])
+            # grid = (grid + 1.0) / 2.0  # -1,1 -> 0,1; c,h,w
+
+            tag = f"{split}/{k}"
+            pl_module.logger.experiment.add_image(
+                tag, grid, global_step=pl_module.global_step
+            )
+
+    @rank_zero_only
+    def log_local(self, save_dir, split, images, global_step, current_epoch, batch_idx):
+        root = os.path.join(save_dir, "images", split)
+        for k in images:
+            images[k] -= torch.min(images[k])
+            images[k] /= torch.max(images[k])
+            grid = torchvision.utils.make_grid(images[k], nrow=4)
+
+            # grid = (grid + 1.0) / 2.0  # -1,1 -> 0,1; c,h,w
+            grid = grid.transpose(0, 1).transpose(1, 2).squeeze(-1)
+            grid = grid.numpy()
+            grid = (grid * 255).astype(np.uint8)
+            filename = "{}_gs-{:06}_e-{:06}_b-{:06}.png".format(
+                k, global_step, current_epoch, batch_idx
+            )
+            path = os.path.join(root, filename)
+            os.makedirs(os.path.split(path)[0], exist_ok=True)
+            Image.fromarray(grid).save(path)
+
+    def log_img(self, pl_module, batch, batch_idx, split="train"):
+        if (
+            self.check_frequency(batch_idx)
+            and hasattr(pl_module, "log_images")  # batch_idx % self.batch_freq == 0
+            and callable(pl_module.log_images)
+            and self.max_images > 0
+        ):
+            logger = type(pl_module.logger)
+
+            is_train = pl_module.training
+            if is_train:
+                pl_module.eval()
+
+            with torch.no_grad():
+                images = pl_module.log_images(batch, split=split)
+
+            for k in images:
+                N = min(images[k].shape[0], self.max_images)
+                images[k] = images[k][:N]
+                if isinstance(images[k], torch.Tensor):
+                    images[k] = images[k].detach().cpu()
+                    if self.clamp:
+                        images[k] = torch.clamp(images[k], -1.0, 1.0)
+
+            self.log_local(
+                pl_module.logger.save_dir,
+                split,
+                images,
+                pl_module.global_step,
+                pl_module.current_epoch,
+                batch_idx,
+            )
+
+            logger_log_images = self.logger_log_images.get(
+                logger, lambda *args, **kwargs: None
+            )
+            logger_log_images(pl_module, images, pl_module.global_step, split)
+
+            if is_train:
+                pl_module.train()
+
+    def check_frequency(self, batch_idx):
+        if (batch_idx % self.batch_freq) == 0 or (batch_idx in self.log_steps):
+            try:
+                self.log_steps.pop(0)
+            except IndexError:
+                pass
+            return True
+        return False
+
+    # def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
+    # def on_train_batch_end(self, *args, **kwargs):
+    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
+        self.log_img(pl_module, batch, batch_idx, split="train")
+
+    def on_validation_batch_end(
+        self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx
+    ):
+        self.log_img(pl_module, batch, batch_idx, split="val")
+
+
+if __name__ == "__main__":
+    # custom parser to specify config files, train, test and debug mode,
+    # postfix, resume.
+    # `--key value` arguments are interpreted as arguments to the trainer.
+    # `nested.key=value` arguments are interpreted as config parameters.
+    # configs are merged from left-to-right followed by command line parameters.
+
+    # model:
+    #   base_learning_rate: float
+    #   target: path to lightning module
+    #   params:
+    #       key: value
+    # data:
+    #   target: main.DataModuleFromConfig
+    #   params:
+    #      batch_size: int
+    #      wrap: bool
+    #      train:
+    #          target: path to train dataset
+    #          params:
+    #              key: value
+    #      validation:
+    #          target: path to validation dataset
+    #          params:
+    #              key: value
+    #      test:
+    #          target: path to test dataset
+    #          params:
+    #              key: value
+    # lightning: (optional, has sane defaults and can be specified on cmdline)
+    #   trainer:
+    #       additional arguments to trainer
+    #   logger:
+    #       logger to instantiate
+    #   modelcheckpoint:
+    #       modelcheckpoint to instantiate
+    #   callbacks:
+    #       callback1:
+    #           target: importpath
+    #           params:
+    #               key: value
+
+    now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
+
+    # add cwd for convenience and to make classes in this file available when
+    # running as `python main.py`
+    # (in particular `main.DataModuleFromConfig`)
+    sys.path.append(os.getcwd())
+
+    parser = get_parser()
+    parser = Trainer.add_argparse_args(parser)
+
+    opt, unknown = parser.parse_known_args()
+    if opt.name and opt.resume:
+        raise ValueError(
+            "-n/--name and -r/--resume cannot be specified both."
+            "If you want to resume training in a new log folder, "
+            "use -n/--name in combination with --resume_from_checkpoint"
+        )
+    if opt.resume:
+        if not os.path.exists(opt.resume):
+            raise ValueError("Cannot find {}".format(opt.resume))
+        if os.path.isfile(opt.resume):
+            paths = opt.resume.split("/")
+            idx = len(paths) - paths[::-1].index("logs") + 1
+            logdir = "/".join(paths[:idx])
+            ckpt = opt.resume
+        else:
+            assert os.path.isdir(opt.resume), opt.resume
+            logdir = opt.resume.rstrip("/")
+            ckpt = os.path.join(logdir, "checkpoints", "last.ckpt")
+
+        opt.resume_from_checkpoint = ckpt
+        base_configs = sorted(glob.glob(os.path.join(logdir, "configs/*.yaml")))
+        opt.base = base_configs + opt.base
+        _tmp = logdir.split("/")
+        nowname = _tmp[_tmp.index("logs") + 1]
+    else:
+        if opt.name:
+            name = "_" + opt.name
+        elif opt.base:
+            cfg_fname = os.path.split(opt.base[0])[-1]
+            cfg_name = os.path.splitext(cfg_fname)[0]
+            name = "_" + cfg_name
+        else:
+            name = ""
+        nowname = now + name + opt.postfix
+        logdir = os.path.join("logs", nowname)
+
+    ckptdir = os.path.join(logdir, "checkpoints")
+    cfgdir = os.path.join(logdir, "configs")
+    seed_everything(opt.seed)
+
+    try:
+        # init and save configs
+        configs = [OmegaConf.load(cfg) for cfg in opt.base]
+        cli = OmegaConf.from_dotlist(unknown)
+        config = OmegaConf.merge(*configs, cli)
+        lightning_config = config.pop("lightning", OmegaConf.create())
+        # merge trainer cli with config
+        trainer_config = lightning_config.get("trainer", OmegaConf.create())
+        # default to ddp
+        trainer_config["distributed_backend"] = "ddp"
+        trainer_config["replace_sampler_ddp"] = False
+        trainer_config["strategy"] = "ddp"
+        trainer_config["persistent_workers"] = True
+        for k in nondefault_trainer_args(opt):
+            trainer_config[k] = getattr(opt, k)
+        if not "gpus" in trainer_config:
+            del trainer_config["distributed_backend"]
+            cpu = True
+        else:
+            gpuinfo = trainer_config["gpus"]
+            print(f"Running on GPUs {gpuinfo}")
+            cpu = False
+        trainer_opt = argparse.Namespace(**trainer_config)
+        lightning_config.trainer = trainer_config
+
+        # model
+        model = instantiate_from_config(config.model)
+        # trainer and callbacks
+        trainer_kwargs = dict()
+
+        # default logger configs
+        # NOTE wandb < 0.10.0 interferes with shutdown
+        # wandb >= 0.10.0 seems to fix it but still interferes with pudb
+        # debugging (wrongly sized pudb ui)
+        # thus prefer testtube for now
+        default_logger_cfgs = {
+            "wandb": {
+                "target": "pytorch_lightning.loggers.WandbLogger",
+                "params": {
+                    "name": nowname,
+                    "save_dir": logdir,
+                    "offline": opt.debug,
+                    "id": nowname,
+                },
+            },
+            "testtube": {
+                # "target": "pytorch_lightning.loggers.TestTubeLogger",
+                "target": "pytorch_lightning.loggers.TensorBoardLogger",
+                "params": {
+                    "name": "testtube",
+                    "save_dir": logdir,
+                },
+            },
+        }
+        default_logger_cfg = default_logger_cfgs["testtube"]
+        try:
+            logger_cfg = lightning_config.logger
+        except:
+            logger_cfg = OmegaConf.create()
+        logger_cfg = OmegaConf.merge(default_logger_cfg, logger_cfg)
+        trainer_kwargs["logger"] = instantiate_from_config(logger_cfg)
+
+        # modelcheckpoint - use TrainResult/EvalResult(checkpoint_on=metric) to
+        # specify which metric is used to determine best models
+        default_modelckpt_cfg = {
+            "checkpoint_callback": {
+                "target": "pytorch_lightning.callbacks.ModelCheckpoint",
+                "params": {
+                    "dirpath": ckptdir,
+                    "filename": "{epoch:06}",
+                    "verbose": True,
+                    "save_last": True,
+                },
+            }
+        }
+        if hasattr(model, "monitor"):
+            print(f"Monitoring {model.monitor} as checkpoint metric.")
+            default_modelckpt_cfg["checkpoint_callback"]["params"][
+                "monitor"
+            ] = model.monitor
+            default_modelckpt_cfg["checkpoint_callback"]["params"]["save_top_k"] = 3
+        try:
+            modelckpt_cfg = lightning_config.modelcheckpoint
+        except:
+            modelckpt_cfg = OmegaConf.create()
+
+        modelckpt_cfg = OmegaConf.merge(default_modelckpt_cfg, modelckpt_cfg)
+        # trainer_kwargs["checkpoint_callback"] = instantiate_from_config(modelckpt_cfg)
+
+        # loaded_model_callbacks = instantiate_from_config(modelckpt_cfg)
+
+        # add callback which sets up log directory
+        default_callbacks_cfg = {
+            "setup_callback": {
+                "target": "celle_taming_main.SetupCallback",
+                "params": {
+                    "resume": opt.resume,
+                    "now": now,
+                    "logdir": logdir,
+                    "ckptdir": ckptdir,
+                    "cfgdir": cfgdir,
+                    "config": config,
+                    "lightning_config": lightning_config,
+                },
+            },
+            "image_logger": {
+                "target": "celle_taming_main.ImageLogger",
+                "params": {
+                    "batch_frequency": 2000,
+                    "max_images": 10,
+                    "clamp": True,
+                    "increase_log_steps": False,
+                },
+            },
+            "learning_rate_logger": {
+                "target": "celle_taming_main.LearningRateMonitor",
+                "params": {
+                    "logging_interval": "step",
+                    # "log_momentum": True
+                },
+            },
+        }
+        try:
+            callbacks_cfg = lightning_config.callbacks
+        except:
+            callbacks_cfg = OmegaConf.create()
+        callbacks_cfg = OmegaConf.merge(default_callbacks_cfg, callbacks_cfg)
+        callbacks_cfg = OmegaConf.merge(modelckpt_cfg, callbacks_cfg)
+        trainer_kwargs["callbacks"] = [
+            instantiate_from_config(callbacks_cfg[k]) for k in callbacks_cfg
+        ]
+        # loaded_callbacks = [
+        #     instantiate_from_config(callbacks_cfg[k]) for k in callbacks_cfg
+        # ]
+
+        # trainer_kwargs["callbacks"] = loaded_callbacks.append(loaded_model_callbacks)
+
+        trainer = Trainer.from_argparse_args(trainer_opt, **trainer_kwargs)
+
+        # data
+        data = instantiate_from_config(config.data)
+        # NOTE according to https://pytorch-lightning.readthedocs.io/en/latest/datamodules.html
+        # calling these ourselves should not be necessary but it is.
+        # lightning still takes care of proper multiprocessing though
+        data.prepare_data()
+        data.setup()
+
+        # configure learning rate
+        bs, base_lr = config.data.params.batch_size, config.model.base_learning_rate
+        if not cpu:
+            ngpu = len(lightning_config.trainer.gpus.strip(",").split(","))
+        else:
+            ngpu = 1
+        try:
+            accumulate_grad_batches = lightning_config.trainer.accumulate_grad_batches
+        except:
+            accumulate_grad_batches = 1
+        print(f"accumulate_grad_batches = {accumulate_grad_batches}")
+        lightning_config.trainer.accumulate_grad_batches = accumulate_grad_batches
+        model.learning_rate = accumulate_grad_batches * ngpu * bs * base_lr
+        print(
+            "Setting learning rate to {:.2e} = {} (accumulate_grad_batches) * {} (num_gpus) * {} (batchsize) * {:.2e} (base_lr)".format(
+                model.learning_rate, accumulate_grad_batches, ngpu, bs, base_lr
+            )
+        )
+
+        # allow checkpointing via USR1
+        def melk(*args, **kwargs):
+            # run all checkpoint hooks
+            if trainer.global_rank == 0:
+                print("Summoning checkpoint.")
+                ckpt_path = os.path.join(ckptdir, "last.ckpt")
+                trainer.save_checkpoint(ckpt_path)
+
+        def divein(*args, **kwargs):
+            if trainer.global_rank == 0:
+                import pudb
+
+                pudb.set_trace()
+
+        import signal
+
+        signal.signal(signal.SIGUSR1, melk)
+        signal.signal(signal.SIGUSR2, divein)
+        # model = torch.compile(model)
+        # run
+        if opt.train:
+            try:
+                torch.compile(trainer.fit(model, data))
+            except Exception:
+                melk()
+                raise
+        if not opt.no_test and not trainer.interrupted:
+            trainer.test(model, data)
+    except Exception:
+        if opt.debug and trainer.global_rank == 0:
+            try:
+                import pudb as debugger
+            except ImportError:
+                import pdb as debugger
+            debugger.post_mortem()
+        raise
+    finally:
+        # move newly created debug project to debug_runs
+        if opt.debug and not opt.resume and trainer.global_rank == 0:
+            dst, name = os.path.split(logdir)
+            dst = os.path.join(dst, "debug_runs", name)
+            os.makedirs(os.path.split(dst)[0], exist_ok=True)
+            os.rename(logdir, dst)

dataloader.py

@@ -0,0 +1,321 @@
+import os
+import numpy as np
+from PIL import Image, ImageSequence
+import json
+import pandas as pd
+
+import torch
+from torch.utils.data import Dataset
+from torchvision import transforms
+import torchvision.transforms.functional as TF
+
+
+def simple_conversion(seq):
+    """Create 26-dim embedding"""
+    chars = [
+        "-",
+        "M",
+        "R",
+        "H",
+        "K",
+        "D",
+        "E",
+        "S",
+        "T",
+        "N",
+        "Q",
+        "C",
+        "U",
+        "G",
+        "P",
+        "A",
+        "V",
+        "I",
+        "F",
+        "Y",
+        "W",
+        "L",
+        "O",
+        "X",
+        "Z",
+        "B",
+        "J",
+    ]
+
+    nums = range(len(chars))
+
+    seqs_x = np.zeros(len(seq))
+
+    for idx, char in enumerate(seq):
+
+        lui = chars.index(char)
+
+        seqs_x[idx] = nums[lui]
+
+    return torch.tensor([seqs_x]).long()
+
+
+def replace_outliers(image, percentile=0.0001):
+
+    lower_bound, upper_bound = torch.quantile(image, percentile), torch.quantile(
+        image, 1 - percentile
+    )
+    mask = (image <= upper_bound) & (image >= lower_bound)
+
+    valid_pixels = image[mask]
+
+    image[~mask] = torch.clip(image[~mask], min(valid_pixels), max(valid_pixels))
+
+    return image
+
+
+class CellLoader(Dataset):
+    """imports mined opencell images with protein sequence"""
+
+    def __init__(
+        self,
+        data_csv=None,
+        dataset=None,
+        split_key=None,
+        resize=600,
+        crop_size=600,
+        crop_method="random",
+        sequence_mode="simple",
+        vocab="bert",
+        threshold="median",
+        text_seq_len=0,
+        pad_mode="random",
+    ):
+        self.data_csv = data_csv
+        self.dataset = dataset
+        self.image_folders = []
+        self.crop_method = crop_method
+        self.resize = resize
+        self.crop_size = crop_size
+        self.sequence_mode = sequence_mode
+        self.threshold = threshold
+        self.text_seq_len = int(text_seq_len)
+        self.vocab = vocab
+        self.pad_mode = pad_mode
+
+        if self.sequence_mode == "embedding" or self.sequence_mode == "onehot":
+
+
+            if self.vocab == "esm1b" or self.vocab == "esm2":
+                from esm import Alphabet
+
+                self.tokenizer = Alphabet.from_architecture(
+                    "ESM-1b"
+                ).get_batch_converter()
+                self.text_seq_len += 2
+
+        if data_csv:
+
+            data = pd.read_csv(data_csv)
+
+            self.parent_path = os.path.dirname(data_csv).split(data_csv)[0]
+
+            if split_key == "train":
+                self.data = data[data["split"] == "train"]
+            elif split_key == "val":
+                self.data = data[data["split"] == "val"]
+            else:
+                self.data = data
+
+            self.data = self.data.reset_index(drop=True)
+            
+                
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(
+        self,
+        idx,
+        get_sequence=True,
+        get_images=True,
+    ):
+        if get_sequence and self.text_seq_len > 0:
+
+            protein_vector = self.get_protein_vector(idx)
+
+        else:
+            protein_vector = torch.zeros((1, 1))
+
+        if get_images:
+
+            nucleus, target, threshold = self.get_images(idx, self.dataset)
+        else:
+            nucleus, target, threshold = torch.zeros((3, 1))
+
+        data_dict = {
+            "nucleus": nucleus.float(),
+            "target": target.float(),
+            "threshold": threshold.float(),
+            "sequence": protein_vector.long(),
+        }
+
+        return data_dict
+
+    def get_protein_vector(self, idx):
+
+        if "protein_sequence" not in self.data.columns:
+
+            metadata = self.retrieve_metadata(idx)
+            protein_sequence = metadata["sequence"]
+        else:
+            protein_sequence = self.data.iloc[idx]["protein_sequence"]
+
+        protein_vector = self.tokenize_sequence(protein_sequence)
+
+        return protein_vector
+
+    def get_images(self, idx, dataset):
+
+        if dataset == "HPA":
+
+            nucleus = Image.open(
+                os.path.join(
+                    self.parent_path, self.data.iloc[idx]["nucleus_image_path"]
+                )
+            )
+
+            target = Image.open(
+                os.path.join(self.parent_path, self.data.iloc[idx]["target_image_path"])
+            )
+
+            nucleus = TF.to_tensor(nucleus)[0]
+            target = TF.to_tensor(target)[0]
+
+            image = torch.stack([nucleus, target], axis=0)
+
+            normalize = (0.0655, 0.0650), (0.1732, 0.1208)
+
+        elif dataset == "OpenCell":
+            image = Image.open(
+                os.path.join(self.parent_path, self.data.iloc[idx]["image_path"])
+            )
+            nucleus, target = [page.copy() for page in ImageSequence.Iterator(image)]
+
+            nucleus = replace_outliers(torch.divide(TF.to_tensor(nucleus), 65536))[0]
+            target = replace_outliers(torch.divide(TF.to_tensor(target), 65536))[0]
+
+            image = torch.stack([nucleus, target], axis=0)
+
+            normalize = (
+                (0.0272, 0.0244),
+                (0.0486, 0.0671),
+            )
+
+        # # from https://discuss.pytorch.org/t/how-to-apply-same-transform-on-a-pair-of-picture/14914
+
+        t_forms = [transforms.Resize(self.resize, antialias=None)]
+
+        if self.crop_method == "random":
+
+            t_forms.append(transforms.RandomCrop(self.crop_size))
+            t_forms.append(transforms.RandomHorizontalFlip(p=0.5))
+            t_forms.append(transforms.RandomVerticalFlip(p=0.5))
+
+        elif self.crop_method == "center":
+
+            t_forms.append(transforms.CenterCrop(self.crop_size))
+
+        t_forms.append(transforms.Normalize(normalize[0], normalize[1]))
+
+        image = transforms.Compose(t_forms)(image)
+
+        nucleus, target = image
+
+        nucleus /= torch.abs(nucleus).max()
+        target -= target.min()
+        target /= target.max()
+
+        nucleus = nucleus.unsqueeze(0)
+        target = target.unsqueeze(0)
+
+        threshold = target
+
+        if self.threshold == "mean":
+
+            threshold = 1.0 * (threshold > (torch.mean(threshold)))
+
+        elif self.threshold == "median":
+
+            threshold = 1.0 * (threshold > (torch.median(threshold)))
+
+        elif self.threshold == "1090_IQR":
+
+            p10 = torch.quantile(threshold, 0.1, None)
+            p90 = torch.quantile(threshold, 0.9, None)
+            threshold = torch.clip(threshold, p10, p90)
+
+        nucleus = torch.nan_to_num(nucleus, 0.0, 1.0, 0.0)
+        target = torch.nan_to_num(target, 0.0, 1.0, 0.0)
+        threshold = torch.nan_to_num(threshold, 0.0, 1.0, 0.0)
+
+        return nucleus, target, threshold
+
+    def retrieve_metadata(self, idx):
+        with open(
+            os.path.join(self.parent_path, self.data.iloc[idx]["metadata_path"])
+        ) as f:
+            metadata = json.load(f)
+
+        return metadata
+
+    def tokenize_sequence(self, protein_sequence):
+
+        pad_token = 0
+
+        if self.sequence_mode == "simple":
+            protein_vector = simple_conversion(protein_sequence)
+
+        elif self.sequence_mode == "center":
+            protein_sequence = protein_sequence.center(self.text_seq_length, "-")
+            protein_vector = simple_conversion(protein_sequence)
+
+        elif self.sequence_mode == "alternating":
+            protein_sequence = protein_sequence.center(self.text_seq_length, "-")
+            protein_sequence = protein_sequence[::18]
+            protein_sequence = protein_sequence.center(
+                int(self.text_seq_length / 18) + 1, "-"
+            )
+            protein_vector = simple_conversion(protein_sequence)
+
+
+        elif self.sequence_mode == "embedding":
+
+            if self.vocab == "esm1b" or self.vocab == "esm2":
+                pad_token = 1
+                protein_vector = self.tokenizer([("", protein_sequence)])[-1]
+
+        if protein_vector.shape[-1] < self.text_seq_len:
+
+            diff = self.text_seq_len - protein_vector.shape[-1]
+
+            if self.pad_mode == "end":
+                protein_vector = torch.nn.functional.pad(
+                    protein_vector, (0, diff), "constant", pad_token
+                )
+            elif self.pad_mode == "random":
+                split = diff - np.random.randint(0, diff + 1)
+
+                protein_vector = torch.cat(
+                    [torch.ones(1, split) * 0, protein_vector], dim=1
+                )
+
+                protein_vector = torch.nn.functional.pad(
+                    protein_vector, (0, diff - split), "constant", pad_token
+                )
+
+        elif protein_vector.shape[-1] > self.text_seq_len:
+            start_int = np.random.randint(
+                0, protein_vector.shape[-1] - self.text_seq_len
+            )
+
+            protein_vector = protein_vector[
+                :, start_int : start_int + self.text_seq_len
+            ]
+
+        return protein_vector.long()

prediction.py

@@ -0,0 +1,267 @@
+import argparse
+import torch
+import os
+os.chdir('..')
+from dataloader import CellLoader
+from matplotlib import pyplot as plt
+from celle_main import instantiate_from_config
+from omegaconf import OmegaConf
+from celle.utils import process_image
+
+def run_model(mode, sequence,
+            nucleus_image_path,
+            protein_image_path,
+            model_ckpt_path,
+            model_config_path,
+            device):
+    if mode == "image":
+        run_image_prediction(
+            sequence,
+            nucleus_image_path,
+            protein_image_path,
+            model_ckpt_path,
+            model_config_path,
+            device
+        )
+    elif mode == "sequence":
+        run_sequence_prediction(
+            sequence,
+            nucleus_image_path,
+            protein_image_path,
+            model_ckpt_path,
+            model_config_path,
+            device
+        )
+
+def run_sequence_prediction(
+    sequence_input,
+    nucleus_image_path,
+    protein_image_path,
+    model_ckpt_path,
+    model_config_path,
+    device
+):
+    """
+    Run Celle model with provided inputs and display results.
+
+    :param sequence: Path to sequence file
+    :param nucleus_image_path: Path to nucleus image
+    :param protein_image_path: Path to protein image (optional)
+    :param model_ckpt_path: Path to model checkpoint
+    :param model_config_path: Path to model config
+    """
+    
+    # Instantiate dataset object
+    dataset = CellLoader(
+        sequence_mode="embedding",
+        vocab="esm2",
+        split_key="val",
+        crop_method="center",
+        resize=600,
+        crop_size=256,
+        text_seq_len=1000,
+        pad_mode="end",
+        threshold="median",
+    )
+
+    # Check if sequence is provided and valid
+    if len(sequence_input) == 0:
+        raise ValueError("Sequence must be provided.")
+
+    if "<mask>" not in sequence_input:
+        print("Warning: Sequence does not contain any masked positions to predict.")
+
+    # Convert SEQUENCE to sequence using dataset.tokenize_sequence()
+    sequence = dataset.tokenize_sequence(sequence_input)
+
+    # Check if nucleus image path is provided and valid
+    if not os.path.exists(nucleus_image_path):
+        # Use default nucleus image from dataset and print warning
+        nucleus_image_path = 'images/nucleus.jpg'
+        print(
+            "Warning: No nucleus image provided. Using default nucleus image from dataset."
+        )
+    else:
+        # Load nucleus image from provided path
+        nucleus_image = process_image(nucleus_image_path)
+        
+    # Check if protein image path is provided and valid
+    if not os.path.exists(protein_image_path):
+        # Use default nucleus image from dataset and print warning
+        protein_image_path = 'images/protein.jpg'
+        print(
+            "Warning: No nucleus image provided. Using default protein image from dataset."
+        )
+    else:
+        # Load protein image from provided path
+        protein_image = process_image(protein_image_path)
+        protein_image = (protein_image > torch.median(protein_image,dim=0))*1.0
+
+    # Load model config and set ckpt_path if not provided in config
+    config = OmegaConf.load(model_config_path)
+    if config["model"]["params"]["ckpt_path"] is None:
+        config["model"]["params"]["ckpt_path"] = model_ckpt_path
+
+    # Set condition_model_path and vqgan_model_path to None
+    config["model"]["params"]["condition_model_path"] = None
+    config["model"]["params"]["vqgan_model_path"] = None
+
+    # Instantiate model from config and move to device
+    model = instantiate_from_config(config).to(device)
+
+    # Sample from model using provided sequence and nucleus image
+    _, predicted_sequence, _ = model.celle.sample_text(
+        text=sequence,
+        condition=nucleus_image,
+        image=protein_image,
+        force_aas=True,
+        timesteps=1,
+        temperature=1,
+        progress=True,
+    )
+
+    formatted_predicted_sequence = ""
+
+    for i in range(min(len(predicted_sequence), len(sequence))):
+        if predicted_sequence[i] != sequence[i]:
+            formatted_predicted_sequence += f"**{predicted_sequence[i]}**"
+        else:
+            formatted_predicted_sequence += predicted_sequence[i]
+
+    if len(predicted_sequence) > len(sequence):
+        formatted_predicted_sequence += f"**{predicted_sequence[len(sequence):]}**"
+
+    print("predicted_sequence:", formatted_predicted_sequence)
+
+
+def run_image_prediction(
+    sequence_input,
+    nucleus_image_path,
+    protein_image_path,
+    model_ckpt_path,
+    model_config_path,
+    device
+):
+    """
+    Run Celle model with provided inputs and display results.
+
+    :param sequence: Path to sequence file
+    :param nucleus_image_path: Path to nucleus image
+    :param protein_image_path: Path to protein image (optional)
+    :param model_ckpt_path: Path to model checkpoint
+    :param model_config_path: Path to model config
+    """
+    # Instantiate dataset object
+    dataset = CellLoader(
+        sequence_mode="embedding",
+        vocab="esm2",
+        split_key="val",
+        crop_method="center",
+        resize=600,
+        crop_size=256,
+        text_seq_len=1000,
+        pad_mode="end",
+        threshold="median",
+    )
+
+    # Check if sequence is provided and valid
+    if len(sequence_input) == 0:
+        sequence = "MSKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTFSYGVQCFSRYPDHMKQHDFFKSAMPEGYVQERTIFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK"
+        # Use default sequence for GFP and print warning
+        print("Warning: No sequence provided. Using default sequence for GFP.")
+
+    # Convert SEQUENCE to sequence using dataset.tokenize_sequence()
+    sequence = dataset.tokenize_sequence(sequence_input)
+
+    # Check if nucleus image path is provided and valid
+    if not os.path.exists(nucleus_image_path):
+        # Use default nucleus image from dataset and print warning
+        nucleus_image = dataset[0]["nucleus"]
+        print(
+            "Warning: No nucleus image provided. Using default nucleus image from dataset."
+        )
+    else:
+        # Load nucleus image from provided path
+        nucleus_image = process_image(nucleus_image_path)
+
+    # Load model config and set ckpt_path if not provided in config
+    config = OmegaConf.load(model_config_path)
+    if config["model"]["params"]["ckpt_path"] is None:
+        config["model"]["params"]["ckpt_path"] = model_ckpt_path
+
+    # Set condition_model_path and vqgan_model_path to None
+    config["model"]["params"]["condition_model_path"] = None
+    config["model"]["params"]["vqgan_model_path"] = None
+
+    # Instantiate model from config and move to device
+    model = instantiate_from_config(config).to(device)
+
+    # Sample from model using provided sequence and nucleus image
+    _, _, _, predicted_threshold, predicted_heatmap = model.celle.sample(
+        text=sequence,
+        condition=nucleus_image,
+        timesteps=1,
+        temperature=1,
+        progress=True,
+    )
+
+    # Move predicted_threshold and predicted_heatmap to CPU and select first element of batch
+    predicted_threshold = predicted_threshold.cpu()[0, 0]
+    predicted_heatmap = predicted_heatmap.cpu()[0, 0]
+
+    # Create 3 or 4 panel plot depending on whether protein image path is provided
+    fig, axs = plt.subplots(1, 3 if protein_image_path is None else 4)
+    axs[0].imshow(nucleus_image)
+    axs[0].set_title("Nucleus Input")
+    axs[1].imshow(predicted_threshold)
+    axs[1].set_title("Predicted Threshold")
+    if protein_image_path is not None:
+        protein_image = process_image(protein_image_path)
+        axs[2].imshow(protein_image)
+        axs[2].set_title("Protein Image")
+    axs[-1].imshow(predicted_heatmap)
+    axs[-1].set_title("Predicted Heatmap")
+    plt.show()
+
+
+if __name__ == "__main__":
+    # Parse command line arguments for input parameters
+    parser = argparse.ArgumentParser(
+        description="Run Celle model with provided inputs."
+    )
+    parser.add_argument("--mode", type=str, default="", help="Sequence or Image")
+    parser.add_argument(
+        "--sequence", type=str, default="", help="Path to sequence file"
+    )
+    parser.add_argument(
+        "--nucleus_image_path",
+        type=str,
+        default="images/nucleus.jpg",
+        help="Path to nucleus image",
+    )
+    parser.add_argument(
+        "--protein_image_path",
+        type=str,
+        default=None,
+        help="Path to protein image (optional)",
+    )
+    parser.add_argument(
+        "--model_ckpt_path", type=str, required=True, help="Path to model checkpoint"
+    )
+    parser.add_argument(
+        "--model_config_path", type=str, required=True, help="Path to model config"
+    )
+    parser.add_argument(
+        "--device", type=str, default="cpu", required=True, help="device"
+    )
+    args = parser.parse_args()
+
+    run_model(
+        args.mode,
+        args.sequence,
+        args.nucleus_image_path,
+        args.protein_image_path,
+        args.model_ckpt_path,
+        args.model_config_path,
+        args.device
+    )

requirements.txt

@@ -0,0 +1,160 @@
+aiohttp==3.8.4
+aiosignal==1.3.1
+antlr4-python3-runtime==4.9.3
+anyio==3.6.2
+argon2-cffi==21.3.0
+argon2-cffi-bindings==21.2.0
+arrow==1.2.3
+asttokens==2.2.1
+async-timeout==4.0.2
+attrs==23.1.0
+axial-positional-embedding==0.2.1
+backcall==0.2.0
+beautifulsoup4==4.12.2
+bleach==6.0.0
+blessed==1.20.0
+certifi==2023.5.7
+cffi==1.15.1
+charset-normalizer==3.1.0
+click==8.1.3
+cmake==3.26.3
+comm==0.1.3
+contourpy==1.0.7
+croniter==1.3.14
+cycler==0.11.0
+dateutils==0.6.12
+debugpy==1.6.7
+decorator==5.1.1
+deepdiff==6.3.0
+defusedxml==0.7.1
+einops==0.6.1
+executing==1.2.0
+fair-esm==2.0.0
+fastapi==0.88.0
+fastjsonschema==2.16.3
+filelock==3.12.0
+fonttools==4.39.4
+fqdn==1.5.1
+frozenlist==1.3.3
+fsspec==2023.5.0
+h11==0.14.0
+idna==3.4
+inquirer==3.1.3
+ipykernel==6.23.0
+ipython==8.13.2
+ipython-genutils==0.2.0
+ipywidgets==8.0.6
+isoduration==20.11.0
+itsdangerous==2.1.2
+jedi==0.18.2
+Jinja2==3.1.2
+jsonpointer==2.3
+jsonschema==4.17.3
+jupyter==1.0.0
+jupyter-console==6.6.3
+jupyter-events==0.6.3
+jupyter_client==8.2.0
+jupyter_core==5.3.0
+jupyter_server==2.5.0
+jupyter_server_terminals==0.4.4
+jupyterlab-pygments==0.2.2
+jupyterlab-widgets==3.0.7
+kiwisolver==1.4.4
+lightning==2.0.2
+lightning-cloud==0.5.34
+lightning-utilities==0.8.0
+lit==16.0.3
+markdown-it-py==2.2.0
+MarkupSafe==2.1.2
+matplotlib==3.7.1
+matplotlib-inline==0.1.6
+mdurl==0.1.2
+mistune==2.0.5
+mpmath==1.3.0
+multidict==6.0.4
+nbclassic==1.0.0
+nbclient==0.7.4
+nbconvert==7.4.0
+nbformat==5.8.0
+nest-asyncio==1.5.6
+networkx==3.1
+notebook==6.5.4
+notebook_shim==0.2.3
+numpy==1.24.3
+nvidia-cublas-cu11==11.10.3.66
+nvidia-cuda-cupti-cu11==11.7.101
+nvidia-cuda-nvrtc-cu11==11.7.99
+nvidia-cuda-runtime-cu11==11.7.99
+nvidia-cudnn-cu11==8.5.0.96
+nvidia-cufft-cu11==10.9.0.58
+nvidia-curand-cu11==10.2.10.91
+nvidia-cusolver-cu11==11.4.0.1
+nvidia-cusparse-cu11==11.7.4.91
+nvidia-nccl-cu11==2.14.3
+nvidia-nvtx-cu11==11.7.91
+omegaconf==2.3.0
+ordered-set==4.1.0
+packaging==23.1
+pandas==2.0.1
+pandocfilters==1.5.0
+parso==0.8.3
+pexpect==4.8.0
+pickleshare==0.7.5
+Pillow==9.5.0
+platformdirs==3.5.1
+prometheus-client==0.16.0
+prompt-toolkit==3.0.38
+psutil==5.9.5
+ptyprocess==0.7.0
+pure-eval==0.2.2
+pycparser==2.21
+pydantic==1.10.7
+Pygments==2.15.1
+PyJWT==2.7.0
+pyparsing==3.0.9
+pyrsistent==0.19.3
+python-dateutil==2.8.2
+python-editor==1.0.4
+python-json-logger==2.0.7
+python-multipart==0.0.6
+pytorch-lightning==1.9.0
+pytz==2023.3
+PyYAML==6.0
+pyzmq==25.0.2
+qtconsole==5.4.3
+QtPy==2.3.1
+readchar==4.0.5
+requests==2.30.0
+rfc3339-validator==0.1.4
+rfc3986-validator==0.1.1
+rich==13.3.5
+rotary-embedding-torch==0.2.3
+Send2Trash==1.8.2
+six==1.16.0
+sniffio==1.3.0
+soupsieve==2.4.1
+stack-data==0.6.2
+starlette==0.22.0
+starsessions==1.3.0
+sympy==1.12
+terminado==0.17.1
+tinycss2==1.2.1
+torch==2.0.0
+torchmetrics==0.11.4
+torchvision @ https://download.pytorch.org/whl/cpu/torchvision-0.15.1%2Bcpu-cp311-cp311-linux_x86_64.whl
+tornado==6.3.1
+tqdm==4.65.0
+traitlets==5.9.0
+triton==2.0.0
+typing_extensions==4.5.0
+tzdata==2023.3
+uri-template==1.2.0
+urllib3==2.0.2
+uvicorn==0.22.0
+wcwidth==0.2.6
+webcolors==1.13
+webencodings==0.5.1
+websocket-client==1.5.1
+websockets==11.0.3
+widgetsnbextension==4.0.7
+yarl==1.9.2

taming/lr_scheduler.py

@@ -0,0 +1,34 @@
+import numpy as np
+
+
+class LambdaWarmUpCosineScheduler:
+    """
+    note: use with a base_lr of 1.0
+    """
+    def __init__(self, warm_up_steps, lr_min, lr_max, lr_start, max_decay_steps, verbosity_interval=0):
+        self.lr_warm_up_steps = warm_up_steps
+        self.lr_start = lr_start
+        self.lr_min = lr_min
+        self.lr_max = lr_max
+        self.lr_max_decay_steps = max_decay_steps
+        self.last_lr = 0.
+        self.verbosity_interval = verbosity_interval
+
+    def schedule(self, n):
+        if self.verbosity_interval > 0:
+            if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_lr}")
+        if n < self.lr_warm_up_steps:
+            lr = (self.lr_max - self.lr_start) / self.lr_warm_up_steps * n + self.lr_start
+            self.last_lr = lr
+            return lr
+        else:
+            t = (n - self.lr_warm_up_steps) / (self.lr_max_decay_steps - self.lr_warm_up_steps)
+            t = min(t, 1.0)
+            lr = self.lr_min + 0.5 * (self.lr_max - self.lr_min) * (
+                    1 + np.cos(t * np.pi))
+            self.last_lr = lr
+            return lr
+
+    def __call__(self, n):
+        return self.schedule(n)
+

taming/models/cond_transformer.py

@@ -0,0 +1,349 @@
+import os, math
+import torch
+import torch.nn.functional as F
+import pytorch_lightning as pl
+
+from main import instantiate_from_config
+from taming.modules.util import SOSProvider
+
+
+def disabled_train(self, mode=True):
+    """Overwrite model.train with this function to make sure train/eval mode
+    does not change anymore."""
+    return self
+
+
+class Net2NetTransformer(pl.LightningModule):
+    def __init__(self,
+                 transformer_config,
+                 first_stage_config,
+                 cond_stage_config,
+                 permuter_config=None,
+                 ckpt_path=None,
+                 ignore_keys=[],
+                 first_stage_key="image",
+                 cond_stage_key="depth",
+                 downsample_cond_size=-1,
+                 pkeep=1.0,
+                 sos_token=0,
+                 unconditional=False,
+                 ):
+        super().__init__()
+        self.be_unconditional = unconditional
+        self.sos_token = sos_token
+        self.first_stage_key = first_stage_key
+        self.cond_stage_key = cond_stage_key
+        self.init_first_stage_from_ckpt(first_stage_config)
+        self.init_cond_stage_from_ckpt(cond_stage_config)
+        if permuter_config is None:
+            permuter_config = {"target": "taming.modules.transformer.permuter.Identity"}
+        self.permuter = instantiate_from_config(config=permuter_config)
+        self.transformer = instantiate_from_config(config=transformer_config)
+
+        if ckpt_path is not None:
+            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
+        self.downsample_cond_size = downsample_cond_size
+        self.pkeep = pkeep
+
+    def init_from_ckpt(self, path, ignore_keys=list()):
+        sd = torch.load(path, map_location="cpu")["state_dict"]
+        for k in sd.keys():
+            for ik in ignore_keys:
+                if k.startswith(ik):
+                    self.print("Deleting key {} from state_dict.".format(k))
+                    del sd[k]
+        self.load_state_dict(sd, strict=False)
+        print(f"Restored from {path}")
+
+    def init_first_stage_from_ckpt(self, config):
+        model = instantiate_from_config(config)
+        model = model.eval()
+        model.train = disabled_train
+        self.first_stage_model = model
+
+    def init_cond_stage_from_ckpt(self, config):
+        if config == "__is_first_stage__":
+            print("Using first stage also as cond stage.")
+            self.cond_stage_model = self.first_stage_model
+        elif config == "__is_unconditional__" or self.be_unconditional:
+            print(f"Using no cond stage. Assuming the training is intended to be unconditional. "
+                  f"Prepending {self.sos_token} as a sos token.")
+            self.be_unconditional = True
+            self.cond_stage_key = self.first_stage_key
+            self.cond_stage_model = SOSProvider(self.sos_token)
+        else:
+            model = instantiate_from_config(config)
+            model = model.eval()
+            model.train = disabled_train
+            self.cond_stage_model = model
+
+    def forward(self, x, c):
+        # one step to produce the logits
+        # x = target
+        # c = nucleus
+        _, z_indices = self.encode_to_z(x)
+        _, c_indices = self.encode_to_c(c) 
+        
+        if self.training and self.pkeep < 1.0:
+            mask = torch.bernoulli(self.pkeep*torch.ones(z_indices.shape,
+                                                         device=z_indices.device))
+            mask = mask.round().to(dtype=torch.int64)
+            r_indices = torch.randint_like(z_indices, self.transformer.config.vocab_size)
+            a_indices = mask*z_indices+(1-mask)*r_indices
+        else:
+            a_indices = z_indices
+        
+        cz_indices = torch.cat((c_indices, a_indices), dim=1)
+
+        # target includes all sequence elements (no need to handle first one
+        # differently because we are conditioning)
+        target = z_indices
+        # make the prediction
+        logits, _ = self.transformer(cz_indices[:, :-1])
+        # cut off conditioning outputs - output i corresponds to p(z_i | z_{<i}, c)
+        logits = logits[:, c_indices.shape[1]-1:]
+
+        return logits, target
+
+    def top_k_logits(self, logits, k):
+        v, ix = torch.topk(logits, k)
+        out = logits.clone()
+        out[out < v[..., [-1]]] = -float('Inf')
+        return out
+
+    @torch.no_grad()
+    def sample(self, x, c, steps, temperature=1.0, sample=False, top_k=None,
+               callback=lambda k: None):
+        x = torch.cat((c,x),dim=1)
+        block_size = self.transformer.get_block_size()
+        assert not self.transformer.training
+        if self.pkeep <= 0.0:
+            # one pass suffices since input is pure noise anyway
+            assert len(x.shape)==2
+            noise_shape = (x.shape[0], steps-1)
+            #noise = torch.randint(self.transformer.config.vocab_size, noise_shape).to(x)
+            noise = c.clone()[:,x.shape[1]-c.shape[1]:-1]
+            x = torch.cat((x,noise),dim=1)
+            logits, _ = self.transformer(x)
+            # take all logits for now and scale by temp
+            logits = logits / temperature
+            # optionally crop probabilities to only the top k options
+            if top_k is not None:
+                logits = self.top_k_logits(logits, top_k)
+            # apply softmax to convert to probabilities
+            probs = F.softmax(logits, dim=-1)
+            # sample from the distribution or take the most likely
+            if sample:
+                shape = probs.shape
+                probs = probs.reshape(shape[0]*shape[1],shape[2])
+                ix = torch.multinomial(probs, num_samples=1)
+                probs = probs.reshape(shape[0],shape[1],shape[2])
+                ix = ix.reshape(shape[0],shape[1])
+            else:
+                _, ix = torch.topk(probs, k=1, dim=-1)
+            # cut off conditioning
+            x = ix[:, c.shape[1]-1:]
+        else:
+            for k in range(steps):
+                callback(k)
+                assert x.size(1) <= block_size # make sure model can see conditioning
+                x_cond = x if x.size(1) <= block_size else x[:, -block_size:]  # crop context if needed
+                logits, _ = self.transformer(x_cond)
+                # pluck the logits at the final step and scale by temperature
+                logits = logits[:, -1, :] / temperature
+                # optionally crop probabilities to only the top k options
+                if top_k is not None:
+                    logits = self.top_k_logits(logits, top_k)
+                # apply softmax to convert to probabilities
+                probs = F.softmax(logits, dim=-1)
+                # sample from the distribution or take the most likely
+                if sample:
+                    ix = torch.multinomial(probs, num_samples=1)
+                else:
+                    _, ix = torch.topk(probs, k=1, dim=-1)
+                # append to the sequence and continue
+                x = torch.cat((x, ix), dim=1)
+            # cut off conditioning
+            x = x[:, c.shape[1]:]
+        return x
+
+    @torch.no_grad()
+    def encode_to_z(self, x):
+        quant_z, _, info = self.first_stage_model.encode(x)
+        indices = info[2].view(quant_z.shape[0], -1)
+        indices = self.permuter(indices)
+        return quant_z, indices
+
+    @torch.no_grad()
+    def encode_to_c(self, c):
+        if self.downsample_cond_size > -1:
+            c = F.interpolate(c, size=(self.downsample_cond_size, self.downsample_cond_size))
+            
+        #quant_c, _, info = self.cond_stage_model.encode(x)
+        #indices = info[2].view(quant_c.shape[0], -1)
+        #indices = self.permuter(indices)    
+        quant_c, _, [_,_,indices] = self.cond_stage_model.encode(c)
+        if len(indices.shape) != 2:
+            indices = indices.view(c.shape[0], -1)
+        return quant_c, indices
+    
+    @torch.no_grad()
+    def decode_to_img(self, index, zshape):
+        index = self.permuter(index, reverse=True)
+        bhwc = (zshape[0],zshape[2],zshape[3],zshape[1])
+        quant_z = self.first_stage_model.quantize.get_codebook_entry(
+            index.reshape(-1), shape=bhwc)
+        x = self.first_stage_model.decode(quant_z)
+        return x
+
+    @torch.no_grad()
+    def log_images(self, batch, temperature=None, top_k=None, callback=None, lr_interface=False, **kwargs):
+        log = dict()
+
+        N = 4
+        if lr_interface:
+            x, c = self.get_xc(batch, N, diffuse=False, upsample_factor=8)
+        else:
+            x, c = self.get_xc(batch, N)
+        x = x.to(device=self.device)
+        c = c.to(device=self.device)
+
+        quant_z, z_indices = self.encode_to_z(x)
+        quant_c, c_indices = self.encode_to_c(c)
+
+        # create a "half"" sample
+        z_start_indices = z_indices[:,:z_indices.shape[1]//2]
+        index_sample = self.sample(z_start_indices, c_indices,
+                                   steps=z_indices.shape[1]-z_start_indices.shape[1],
+                                   temperature=temperature if temperature is not None else 1.0,
+                                   sample=True,
+                                   top_k=top_k if top_k is not None else 100,
+                                   callback=callback if callback is not None else lambda k: None)
+        x_sample = self.decode_to_img(index_sample, quant_z.shape)
+
+        # sample
+        z_start_indices = z_indices[:, :0]
+        index_sample = self.sample(z_start_indices, c_indices,
+                                   steps=z_indices.shape[1],
+                                   temperature=temperature if temperature is not None else 1.0,
+                                   sample=True,
+                                   top_k=top_k if top_k is not None else 100,
+                                   callback=callback if callback is not None else lambda k: None)
+        x_sample_nopix = self.decode_to_img(index_sample, quant_z.shape)
+
+        # det sample
+        z_start_indices = z_indices[:, :0]
+        index_sample = self.sample(z_start_indices, c_indices,
+                                   steps=z_indices.shape[1],
+                                   sample=False,
+                                   callback=callback if callback is not None else lambda k: None)
+        x_sample_det = self.decode_to_img(index_sample, quant_z.shape)
+
+        # reconstruction
+        x_rec = self.decode_to_img(z_indices, quant_z.shape)
+
+        log["inputs"] = x
+        log["reconstructions"] = x_rec
+
+        if self.cond_stage_key != "image" or self.cond_stage_key != "nucleus" or self.cond_stage_key != "target":
+            cond_rec = self.cond_stage_model.decode(quant_c)
+            if self.cond_stage_key == "segmentation":
+                # get image from segmentation mask
+                num_classes = cond_rec.shape[1]
+
+                c = torch.argmax(c, dim=1, keepdim=True)
+                c = F.one_hot(c, num_classes=num_classes)
+                c = c.squeeze(1).permute(0, 3, 1, 2).float()
+                c = self.cond_stage_model.to_rgb(c)
+
+                cond_rec = torch.argmax(cond_rec, dim=1, keepdim=True)
+                cond_rec = F.one_hot(cond_rec, num_classes=num_classes)
+                cond_rec = cond_rec.squeeze(1).permute(0, 3, 1, 2).float()
+                cond_rec = self.cond_stage_model.to_rgb(cond_rec)
+            log["conditioning_rec"] = cond_rec
+            log["conditioning"] = c
+
+        log["samples_half"] = x_sample
+        log["samples_nopix"] = x_sample_nopix
+        log["samples_det"] = x_sample_det
+        return log
+
+    def get_input(self, key, batch):
+        x = batch[key]
+        if len(x.shape) == 3:
+            x = x[..., None]
+        #if len(x.shape) == 4:
+        #    x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format)
+        if x.dtype == torch.double:
+            x = x.float()
+        return x
+
+    def get_xc(self, batch, N=None):
+        x = self.get_input(self.first_stage_key, batch)
+        c = self.get_input(self.cond_stage_key, batch)
+        if N is not None:
+            x = x[:N]
+            c = c[:N]
+        return x, c
+
+    def shared_step(self, batch):
+        x, c = self.get_xc(batch)
+        logits, target = self(x, c)
+        loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), target.reshape(-1))
+        return loss
+
+    def training_step(self, batch, batch_idx):
+        loss = self.shared_step(batch)
+        self.log("train/loss", loss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        loss = self.shared_step(batch)
+        self.log("val/loss", loss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
+        return loss
+
+    def configure_optimizers(self):
+        """
+        Following minGPT:
+        This long function is unfortunately doing something very simple and is being very defensive:
+        We are separating out all parameters of the model into two buckets: those that will experience
+        weight decay for regularization and those that won't (biases, and layernorm/embedding weights).
+        We are then returning the PyTorch optimizer object.
+        """
+        # separate out all parameters to those that will and won't experience regularizing weight decay
+        decay = set()
+        no_decay = set()
+        whitelist_weight_modules = (torch.nn.Linear, )
+        blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding)
+        for mn, m in self.transformer.named_modules():
+            for pn, p in m.named_parameters():
+                fpn = '%s.%s' % (mn, pn) if mn else pn # full param name
+
+                if pn.endswith('bias'):
+                    # all biases will not be decayed
+                    no_decay.add(fpn)
+                elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules):
+                    # weights of whitelist modules will be weight decayed
+                    decay.add(fpn)
+                elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules):
+                    # weights of blacklist modules will NOT be weight decayed
+                    no_decay.add(fpn)
+
+        # special case the position embedding parameter in the root GPT module as not decayed
+        no_decay.add('pos_emb')
+
+        # validate that we considered every parameter
+        param_dict = {pn: p for pn, p in self.transformer.named_parameters()}
+        inter_params = decay & no_decay
+        union_params = decay | no_decay
+        assert len(inter_params) == 0, "parameters %s made it into both decay/no_decay sets!" % (str(inter_params), )
+        assert len(param_dict.keys() - union_params) == 0, "parameters %s were not separated into either decay/no_decay set!" \
+                                                    % (str(param_dict.keys() - union_params), )
+
+        # create the pytorch optimizer object
+        optim_groups = [
+            {"params": [param_dict[pn] for pn in sorted(list(decay))], "weight_decay": 0.01},
+            {"params": [param_dict[pn] for pn in sorted(list(no_decay))], "weight_decay": 0.0},
+        ]
+        optimizer = torch.optim.AdamW(optim_groups, lr=self.learning_rate, betas=(0.9, 0.95))
+        return optimizer

taming/models/dummy_cond_stage.py

@@ -0,0 +1,22 @@
+from torch import Tensor
+
+
+class DummyCondStage:
+    def __init__(self, conditional_key):
+        self.conditional_key = conditional_key
+        self.train = None
+
+    def eval(self):
+        return self
+
+    @staticmethod
+    def encode(c: Tensor):
+        return c, None, (None, None, c)
+
+    @staticmethod
+    def decode(c: Tensor):
+        return c
+
+    @staticmethod
+    def to_rgb(c: Tensor):
+        return c

taming/models/vqgan.py

@@ -0,0 +1,649 @@
+import torch
+import torch.nn.functional as F
+import pytorch_lightning as pl
+
+from celle_taming_main import instantiate_from_config
+
+from taming.modules.diffusionmodules.model import Encoder, Decoder
+from taming.modules.vqvae.quantize import VectorQuantizer2 as VectorQuantizer
+from taming.modules.vqvae.quantize import GumbelQuantize
+from taming.modules.vqvae.quantize import EMAVectorQuantizer
+
+
+class VQModel(pl.LightningModule):
+    def __init__(
+        self,
+        ddconfig,
+        lossconfig,
+        n_embed,
+        embed_dim,
+        ckpt_path=None,
+        ignore_keys=[],
+        image_key="image",
+        colorize_nlabels=None,
+        monitor=None,
+        remap=None,
+        sane_index_shape=False,  # tell vector quantizer to return indices as bhw
+    ):
+        super().__init__()
+        self.image_key = image_key
+        self.encoder = Encoder(**ddconfig)
+        self.decoder = Decoder(**ddconfig)
+        self.loss = instantiate_from_config(lossconfig)
+        self.quantize = VectorQuantizer(
+            n_embed,
+            embed_dim,
+            beta=0.25,
+            remap=remap,
+            sane_index_shape=sane_index_shape,
+        )
+        self.quant_conv = torch.nn.Conv2d(ddconfig["z_channels"], embed_dim, 1)
+        self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
+        if ckpt_path is not None:
+            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
+        self.image_key = image_key
+        if colorize_nlabels is not None:
+            assert type(colorize_nlabels) == int
+            self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
+        if monitor is not None:
+            self.monitor = monitor
+
+    def init_from_ckpt(self, path, ignore_keys=list()):
+        sd = torch.load(path, map_location="cpu")["state_dict"]
+        keys = list(sd.keys())
+        for k in keys:
+            for ik in ignore_keys:
+                if k.startswith(ik):
+                    print("Deleting key {} from state_dict.".format(k))
+                    del sd[k]
+        self.load_state_dict(sd, strict=False)
+        print(f"Restored from {path}")
+
+    def encode(self, x):
+        h = self.encoder(x)
+        h = self.quant_conv(h)
+        quant, emb_loss, info = self.quantize(h)
+        return quant, emb_loss, info
+
+    def decode(self, quant):
+        quant = self.post_quant_conv(quant)
+        dec = self.decoder(quant)
+        return dec
+
+    def decode_code(self, code_b):
+        quant_b = self.quantize.embed_code(code_b)
+        dec = self.decode(quant_b)
+        return dec
+
+    def forward(self, input):
+        quant, diff, _ = self.encode(input)
+        dec = self.decode(quant)
+        return dec, diff
+
+    def get_input(self, batch, k):
+
+        if k == "mixed":
+            keys = ["nucleus", "target"]
+            index = torch.randint(low=0, high=2, size=(1,), dtype=int).item()
+            k = keys[index]
+
+        x = batch[k]
+        if len(x.shape) == 3:
+            x = x[..., None]
+
+        # x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format)
+        return x
+
+    def training_step(self, batch, batch_idx=None, optimizer_idx=0):
+
+        if type(batch) == dict:
+
+            x = self.get_input(batch, self.image_key)
+
+        else:
+            x = batch
+
+        xrec, qloss = self(
+            x,
+        )
+
+        if optimizer_idx == 0:
+            # autoencode
+            aeloss, log_dict_ae = self.loss(
+                qloss,
+                x,
+                xrec,
+                optimizer_idx,
+                self.global_step,
+                last_layer=self.get_last_layer(),
+                split="train",
+            )
+
+            self.log(
+                "train/aeloss",
+                aeloss,
+                prog_bar=True,
+                logger=True,
+                on_step=True,
+                on_epoch=True,
+                sync_dist=True,
+            )
+            self.log_dict(
+                log_dict_ae,
+                prog_bar=False,
+                logger=True,
+                on_step=True,
+                on_epoch=True,
+                sync_dist=True,
+            )
+            return aeloss
+
+        if optimizer_idx == 1:
+            # discriminator
+            discloss, log_dict_disc = self.loss(
+                qloss,
+                x,
+                xrec,
+                optimizer_idx,
+                self.global_step,
+                last_layer=self.get_last_layer(),
+                split="train",
+            )
+            self.log(
+                "train/discloss",
+                discloss,
+                prog_bar=True,
+                logger=True,
+                on_step=True,
+                on_epoch=True,
+                sync_dist=True,
+            )
+            self.log_dict(
+                log_dict_disc,
+                prog_bar=False,
+                logger=True,
+                on_step=True,
+                on_epoch=True,
+                sync_dist=True,
+            )
+            return discloss
+
+    def validation_step(self, batch, batch_idx):
+
+        if type(batch) == dict:
+
+            x = self.get_input(batch, self.image_key)
+
+        else:
+            x = batch
+
+        xrec, qloss = self(x)
+        aeloss, log_dict_ae = self.loss(
+            qloss,
+            x,
+            xrec,
+            0,
+            self.global_step,
+            last_layer=self.get_last_layer(),
+            split="val",
+        )
+
+        discloss, log_dict_disc = self.loss(
+            qloss,
+            x,
+            xrec,
+            1,
+            self.global_step,
+            last_layer=self.get_last_layer(),
+            split="val",
+        )
+        # rec_loss = log_dict_ae["val/rec_loss"]
+        # self.log(
+        #     "val/rec_loss",
+        #     rec_loss,
+        #     prog_bar=True,
+        #     logger=True,
+        #     on_step=True,
+        #     on_epoch=True,
+        #     sync_dist=True,
+        # )
+        # self.log(
+        #     "val/aeloss",
+        #     aeloss,
+        #     prog_bar=True,
+        #     logger=True,
+        #     on_step=True,
+        #     on_epoch=True,
+        #     sync_dist=True,
+        # )
+
+        for key, value in log_dict_disc.items():
+            if key in log_dict_ae:
+                log_dict_ae[key].extend(value)
+            else:
+                log_dict_ae[key] = value
+
+        self.log_dict(log_dict_ae, sync_dist=True)
+        return self.log_dict
+
+    def configure_optimizers(self):
+        lr = self.learning_rate
+        opt_ae = torch.optim.Adam(
+            list(self.encoder.parameters())
+            + list(self.decoder.parameters())
+            + list(self.quantize.parameters())
+            + list(self.quant_conv.parameters())
+            + list(self.post_quant_conv.parameters()),
+            lr=lr,
+            betas=(0.5, 0.9),
+        )
+        opt_disc = torch.optim.Adam(
+            self.loss.discriminator.parameters(), lr=lr, betas=(0.5, 0.9)
+        )
+        return [opt_ae, opt_disc], []
+
+    def get_last_layer(self):
+        return self.decoder.conv_out.weight
+
+    def log_images(self, batch, **kwargs):
+        log = dict()
+        x = self.get_input(batch, self.image_key)
+        x = x.to(self.device)
+        xrec, _ = self(x)
+        if x.shape[1] > 3:
+            # colorize with random projection
+            assert xrec.shape[1] > 3
+            x = self.to_rgb(x)
+            xrec = self.to_rgb(xrec)
+        log["inputs"] = x
+        log["reconstructions"] = xrec
+        return log
+
+    def to_rgb(self, x):
+        assert self.image_key == "segmentation"
+        if not hasattr(self, "colorize"):
+            self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x))
+        x = F.conv2d(x, weight=self.colorize)
+        x = 2.0 * (x - x.min()) / (x.max() - x.min()) - 1.0
+        return x
+
+
+class VQSegmentationModel(VQModel):
+    def __init__(self, n_labels, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.register_buffer("colorize", torch.randn(3, n_labels, 1, 1))
+
+    def configure_optimizers(self):
+        lr = self.learning_rate
+        opt_ae = torch.optim.Adam(
+            list(self.encoder.parameters())
+            + list(self.decoder.parameters())
+            + list(self.quantize.parameters())
+            + list(self.quant_conv.parameters())
+            + list(self.post_quant_conv.parameters()),
+            lr=lr,
+            betas=(0.5, 0.9),
+        )
+        return opt_ae
+
+    def training_step(self, batch, batch_idx):
+        x = self.get_input(batch, self.image_key)
+        xrec, qloss = self(x)
+        aeloss, log_dict_ae = self.loss(qloss, x, xrec, split="train")
+        self.log_dict(
+            log_dict_ae,
+            prog_bar=False,
+            logger=True,
+            on_step=True,
+            on_epoch=True,
+            sync_dist=True,
+        )
+        return aeloss
+
+    def validation_step(self, batch, batch_idx):
+        x = self.get_input(batch, self.image_key)
+        xrec, qloss = self(x)
+        aeloss, log_dict_ae = self.loss(qloss, x, xrec, split="val")
+        self.log_dict(
+            log_dict_ae,
+            prog_bar=False,
+            logger=True,
+            on_step=True,
+            on_epoch=True,
+            sync_dist=True,
+        )
+        total_loss = log_dict_ae["val/total_loss"]
+        self.log(
+            "val/total_loss",
+            total_loss,
+            prog_bar=True,
+            logger=True,
+            on_step=True,
+            on_epoch=True,
+            sync_dist=True,
+        )
+        return aeloss
+
+    @torch.no_grad()
+    def log_images(self, batch, **kwargs):
+        log = dict()
+        x = self.get_input(batch, self.image_key)
+        x = x.to(self.device)
+        xrec, _ = self(x)
+        if x.shape[1] > 3:
+            # colorize with random projection
+            assert xrec.shape[1] > 3
+            # convert logits to indices
+            xrec = torch.argmax(xrec, dim=1, keepdim=True)
+            xrec = F.one_hot(xrec, num_classes=x.shape[1])
+            xrec = xrec.squeeze(1).permute(0, 3, 1, 2).float()
+            x = self.to_rgb(x)
+            xrec = self.to_rgb(xrec)
+        log["inputs"] = x
+        log["reconstructions"] = xrec
+        return log
+
+
+class VQNoDiscModel(VQModel):
+    def __init__(
+        self,
+        ddconfig,
+        lossconfig,
+        n_embed,
+        embed_dim,
+        ckpt_path=None,
+        ignore_keys=[],
+        image_key="image",
+        colorize_nlabels=None,
+    ):
+        super().__init__(
+            ddconfig=ddconfig,
+            lossconfig=lossconfig,
+            n_embed=n_embed,
+            embed_dim=embed_dim,
+            ckpt_path=ckpt_path,
+            ignore_keys=ignore_keys,
+            image_key=image_key,
+            colorize_nlabels=colorize_nlabels,
+        )
+
+    def training_step(self, batch, batch_idx):
+        x = self.get_input(batch, self.image_key)
+        xrec, qloss = self(x)
+        # autoencode
+        aeloss, log_dict_ae = self.loss(qloss, x, xrec, self.global_step, split="train")
+        output = pl.TrainResult(minimize=aeloss)
+        output.log(
+            "train/aeloss",
+            aeloss,
+            prog_bar=True,
+            logger=True,
+            on_step=True,
+            on_epoch=True,
+        )
+        output.log_dict(
+            log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True
+        )
+        return output
+
+    def validation_step(self, batch, batch_idx):
+        x = self.get_input(batch, self.image_key)
+        xrec, qloss = self(x)
+        aeloss, log_dict_ae = self.loss(qloss, x, xrec, self.global_step, split="val")
+        rec_loss = log_dict_ae["val/rec_loss"]
+        output = pl.EvalResult(checkpoint_on=rec_loss)
+        output.log(
+            "val/rec_loss",
+            rec_loss,
+            prog_bar=True,
+            logger=True,
+            on_step=True,
+            on_epoch=True,
+        )
+        output.log(
+            "val/aeloss",
+            aeloss,
+            prog_bar=True,
+            logger=True,
+            on_step=True,
+            on_epoch=True,
+        )
+        output.log_dict(log_dict_ae)
+
+        return output
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(
+            list(self.encoder.parameters())
+            + list(self.decoder.parameters())
+            + list(self.quantize.parameters())
+            + list(self.quant_conv.parameters())
+            + list(self.post_quant_conv.parameters()),
+            lr=self.learning_rate,
+            betas=(0.5, 0.9),
+        )
+        return optimizer
+
+
+class GumbelVQ(VQModel):
+    def __init__(
+        self,
+        ddconfig,
+        lossconfig,
+        n_embed,
+        embed_dim,
+        temperature_scheduler_config,
+        ckpt_path=None,
+        ignore_keys=[],
+        image_key="image",
+        colorize_nlabels=None,
+        monitor=None,
+        kl_weight=1e-8,
+        remap=None,
+    ):
+
+        z_channels = ddconfig["z_channels"]
+        super().__init__(
+            ddconfig,
+            lossconfig,
+            n_embed,
+            embed_dim,
+            ckpt_path=None,
+            ignore_keys=ignore_keys,
+            image_key=image_key,
+            colorize_nlabels=colorize_nlabels,
+            monitor=monitor,
+        )
+
+        self.loss.n_classes = n_embed
+        self.vocab_size = n_embed
+
+        self.quantize = GumbelQuantize(
+            z_channels,
+            embed_dim,
+            n_embed=n_embed,
+            kl_weight=kl_weight,
+            temp_init=1.0,
+            remap=remap,
+        )
+
+        self.temperature_scheduler = instantiate_from_config(
+            temperature_scheduler_config
+        )  # annealing of temp
+
+        if ckpt_path is not None:
+            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
+
+    def temperature_scheduling(self):
+        self.quantize.temperature = self.temperature_scheduler(self.global_step)
+
+    def encode_to_prequant(self, x):
+        h = self.encoder(x)
+        h = self.quant_conv(h)
+        return h
+
+    def decode_code(self, code_b):
+        raise NotImplementedError
+
+    def training_step(self, batch, batch_idx=None, optimizer_idx=0):
+        self.temperature_scheduling()
+        x = self.get_input(batch, self.image_key)
+        xrec, qloss = self(x)
+
+        if optimizer_idx == 0:
+            # autoencode
+            aeloss, log_dict_ae = self.loss(
+                qloss,
+                x,
+                xrec,
+                optimizer_idx,
+                self.global_step,
+                last_layer=self.get_last_layer(),
+                split="train",
+            )
+
+            self.log_dict(
+                log_dict_ae,
+                prog_bar=False,
+                logger=True,
+                on_step=True,
+                on_epoch=True,
+                sync_dist=True,
+            )
+            self.log(
+                "temperature",
+                self.quantize.temperature,
+                prog_bar=False,
+                logger=True,
+                on_step=True,
+                on_epoch=True,
+                sync_dist=True,
+            )
+            return aeloss
+
+        if optimizer_idx == 1:
+            # discriminator
+            discloss, log_dict_disc = self.loss(
+                qloss,
+                x,
+                xrec,
+                optimizer_idx,
+                self.global_step,
+                last_layer=self.get_last_layer(),
+                split="train",
+            )
+            self.log_dict(
+                log_dict_disc,
+                prog_bar=False,
+                logger=True,
+                on_step=True,
+                on_epoch=True,
+                sync_dist=True,
+            )
+            return discloss
+
+    def validation_step(self, batch, batch_idx):
+        x = self.get_input(batch, self.image_key)
+        xrec, qloss = self(x)
+        aeloss, log_dict_ae = self.loss(
+            qloss,
+            x,
+            xrec,
+            0,
+            self.global_step,
+            last_layer=self.get_last_layer(),
+            split="val",
+        )
+
+        discloss, log_dict_disc = self.loss(
+            qloss,
+            x,
+            xrec,
+            1,
+            self.global_step,
+            last_layer=self.get_last_layer(),
+            split="val",
+        )
+        rec_loss = log_dict_ae["val/rec_loss"]
+        self.log(
+            "val/rec_loss",
+            rec_loss,
+            prog_bar=True,
+            logger=True,
+            on_step=False,
+            on_epoch=True,
+            sync_dist=True,
+        )
+        self.log(
+            "val/aeloss",
+            aeloss,
+            prog_bar=True,
+            logger=True,
+            on_step=False,
+            on_epoch=True,
+            sync_dist=True,
+        )
+        self.log_dict(log_dict_ae, sync_dist=True)
+        self.log_dict(log_dict_disc, sync_dist=True)
+        return self.log_dict
+
+    def log_images(self, batch, **kwargs):
+        log = dict()
+        x = self.get_input(batch, self.image_key)
+        x = x.to(self.device)
+        # encode
+        h = self.encoder(x)
+        h = self.quant_conv(h)
+        quant, _, _ = self.quantize(h)
+        # decode
+        x_rec = self.decode(quant)
+        log["inputs"] = x
+        log["reconstructions"] = x_rec
+        return log
+
+
+class EMAVQ(VQModel):
+    def __init__(
+        self,
+        ddconfig,
+        lossconfig,
+        n_embed,
+        embed_dim,
+        ckpt_path=None,
+        ignore_keys=[],
+        image_key="image",
+        colorize_nlabels=None,
+        monitor=None,
+        remap=None,
+        sane_index_shape=False,  # tell vector quantizer to return indices as bhw
+    ):
+        super().__init__(
+            ddconfig,
+            lossconfig,
+            n_embed,
+            embed_dim,
+            ckpt_path=None,
+            ignore_keys=ignore_keys,
+            image_key=image_key,
+            colorize_nlabels=colorize_nlabels,
+            monitor=monitor,
+        )
+        self.quantize = EMAVectorQuantizer(
+            n_embed=n_embed, embedding_dim=embed_dim, beta=0.25, remap=remap
+        )
+
+    def configure_optimizers(self):
+        lr = self.learning_rate
+        # Remove self.quantize from parameter list since it is updated via EMA
+        opt_ae = torch.optim.Adam(
+            list(self.encoder.parameters())
+            + list(self.decoder.parameters())
+            + list(self.quant_conv.parameters())
+            + list(self.post_quant_conv.parameters()),
+            lr=lr,
+            betas=(0.5, 0.9),
+        )
+        opt_disc = torch.optim.Adam(
+            self.loss.discriminator.parameters(), lr=lr, betas=(0.5, 0.9)
+        )
+        return [opt_ae, opt_disc], []

taming/modules/autoencoder/lpips/vgg.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a78928a0af1e5f0fcb1f3b9e8f8c3a2a5a3de244d830ad5c1feddc79b8432868
+size 7289

taming/modules/diffusionmodules/model.py

@@ -0,0 +1,776 @@
+# pytorch_diffusion + derived encoder decoder
+import math
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+def get_timestep_embedding(timesteps, embedding_dim):
+    """
+    This matches the implementation in Denoising Diffusion Probabilistic Models:
+    From Fairseq.
+    Build sinusoidal embeddings.
+    This matches the implementation in tensor2tensor, but differs slightly
+    from the description in Section 3.5 of "Attention Is All You Need".
+    """
+    assert len(timesteps.shape) == 1
+
+    half_dim = embedding_dim // 2
+    emb = math.log(10000) / (half_dim - 1)
+    emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
+    emb = emb.to(device=timesteps.device)
+    emb = timesteps.float()[:, None] * emb[None, :]
+    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
+    if embedding_dim % 2 == 1:  # zero pad
+        emb = torch.nn.functional.pad(emb, (0,1,0,0))
+    return emb
+
+
+def nonlinearity(x):
+    # swish
+    return x*torch.sigmoid(x)
+
+
+def Normalize(in_channels):
+    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+class Upsample(nn.Module):
+    def __init__(self, in_channels, with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            self.conv = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, x):
+        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        if self.with_conv:
+            x = self.conv(x)
+        return x
+
+
+class Downsample(nn.Module):
+    def __init__(self, in_channels, with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            # no asymmetric padding in torch conv, must do it ourselves
+            self.conv = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=3,
+                                        stride=2,
+                                        padding=0)
+
+    def forward(self, x):
+        if self.with_conv:
+            pad = (0,1,0,1)
+            x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
+            x = self.conv(x)
+        else:
+            x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
+        return x
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False,
+                 dropout, temb_channels=512):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        self.norm1 = Normalize(in_channels)
+        self.conv1 = torch.nn.Conv2d(in_channels,
+                                     out_channels,
+                                     kernel_size=3,
+                                     stride=1,
+                                     padding=1)
+        if temb_channels > 0:
+            self.temb_proj = torch.nn.Linear(temb_channels,
+                                             out_channels)
+        self.norm2 = Normalize(out_channels)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = torch.nn.Conv2d(out_channels,
+                                     out_channels,
+                                     kernel_size=3,
+                                     stride=1,
+                                     padding=1)
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = torch.nn.Conv2d(in_channels,
+                                                     out_channels,
+                                                     kernel_size=3,
+                                                     stride=1,
+                                                     padding=1)
+            else:
+                self.nin_shortcut = torch.nn.Conv2d(in_channels,
+                                                    out_channels,
+                                                    kernel_size=1,
+                                                    stride=1,
+                                                    padding=0)
+
+    def forward(self, x, temb):
+        h = x
+        h = self.norm1(h)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+
+        if temb is not None:
+            h = h + self.temb_proj(nonlinearity(temb))[:,:,None,None]
+
+        h = self.norm2(h)
+        h = nonlinearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+
+        return x+h
+
+
+class AttnBlock(nn.Module):
+    def __init__(self, in_channels):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = Normalize(in_channels)
+        self.q = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.k = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.v = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.proj_out = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=1,
+                                        stride=1,
+                                        padding=0)
+
+
+    def forward(self, x):
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        b,c,h,w = q.shape
+        q = q.reshape(b,c,h*w)
+        q = q.permute(0,2,1)   # b,hw,c
+        k = k.reshape(b,c,h*w) # b,c,hw
+        w_ = torch.bmm(q,k)     # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+        w_ = w_ * (int(c)**(-0.5))
+        w_ = torch.nn.functional.softmax(w_, dim=2)
+
+        # attend to values
+        v = v.reshape(b,c,h*w)
+        w_ = w_.permute(0,2,1)   # b,hw,hw (first hw of k, second of q)
+        h_ = torch.bmm(v,w_)     # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+        h_ = h_.reshape(b,c,h,w)
+
+        h_ = self.proj_out(h_)
+
+        return x+h_
+
+
+class Model(nn.Module):
+    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
+                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
+                 resolution, use_timestep=True):
+        super().__init__()
+        self.ch = ch
+        self.temb_ch = self.ch*4
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        self.use_timestep = use_timestep
+        if self.use_timestep:
+            # timestep embedding
+            self.temb = nn.Module()
+            self.temb.dense = nn.ModuleList([
+                torch.nn.Linear(self.ch,
+                                self.temb_ch),
+                torch.nn.Linear(self.temb_ch,
+                                self.temb_ch),
+            ])
+
+        # downsampling
+        self.conv_in = torch.nn.Conv2d(in_channels,
+                                       self.ch,
+                                       kernel_size=3,
+                                       stride=1,
+                                       padding=1)
+
+        curr_res = resolution
+        in_ch_mult = (1,)+tuple(ch_mult)
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch*in_ch_mult[i_level]
+            block_out = ch*ch_mult[i_level]
+            for i_block in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions-1:
+                down.downsample = Downsample(block_in, resamp_with_conv)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch*ch_mult[i_level]
+            skip_in = ch*ch_mult[i_level]
+            for i_block in range(self.num_res_blocks+1):
+                if i_block == self.num_res_blocks:
+                    skip_in = ch*in_ch_mult[i_level]
+                block.append(ResnetBlock(in_channels=block_in+skip_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in, resamp_with_conv)
+                curr_res = curr_res * 2
+            self.up.insert(0, up) # prepend to get consistent order
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in,
+                                        out_ch,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+
+    def forward(self, x, t=None):
+        #assert x.shape[2] == x.shape[3] == self.resolution
+
+        if self.use_timestep:
+            # timestep embedding
+            assert t is not None
+            temb = get_timestep_embedding(t, self.ch)
+            temb = self.temb.dense[0](temb)
+            temb = nonlinearity(temb)
+            temb = self.temb.dense[1](temb)
+        else:
+            temb = None
+
+        # downsampling
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1], temb)
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions-1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks+1):
+                h = self.up[i_level].block[i_block](
+                    torch.cat([h, hs.pop()], dim=1), temb)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+class Encoder(nn.Module):
+    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
+                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
+                 resolution, z_channels, double_z=True, **ignore_kwargs):
+        super().__init__()
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        # downsampling
+        self.conv_in = torch.nn.Conv2d(in_channels,
+                                       self.ch,
+                                       kernel_size=3,
+                                       stride=1,
+                                       padding=1)
+
+        curr_res = resolution
+        in_ch_mult = (1,)+tuple(ch_mult)
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch*in_ch_mult[i_level]
+            block_out = ch*ch_mult[i_level]
+            for i_block in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions-1:
+                down.downsample = Downsample(block_in, resamp_with_conv)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in,
+                                        2*z_channels if double_z else z_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+
+    def forward(self, x):
+        #assert x.shape[2] == x.shape[3] == self.resolution, "{}, {}, {}".format(x.shape[2], x.shape[3], self.resolution)
+
+        # timestep embedding
+        temb = None
+
+        # downsampling
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1], temb)
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions-1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+class Decoder(nn.Module):
+    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
+                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
+                 resolution, z_channels, give_pre_end=False, **ignorekwargs):
+        super().__init__()
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.give_pre_end = give_pre_end
+
+        # compute in_ch_mult, block_in and curr_res at lowest res
+        in_ch_mult = (1,)+tuple(ch_mult)
+        block_in = ch*ch_mult[self.num_resolutions-1]
+        curr_res = resolution // 2**(self.num_resolutions-1)
+        self.z_shape = (1,z_channels,curr_res,curr_res)
+        print("Working with z of shape {} = {} dimensions.".format(
+            self.z_shape, np.prod(self.z_shape)))
+
+        # z to block_in
+        self.conv_in = torch.nn.Conv2d(z_channels,
+                                       block_in,
+                                       kernel_size=3,
+                                       stride=1,
+                                       padding=1)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch*ch_mult[i_level]
+            for i_block in range(self.num_res_blocks+1):
+                block.append(ResnetBlock(in_channels=block_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in, resamp_with_conv)
+                curr_res = curr_res * 2
+            self.up.insert(0, up) # prepend to get consistent order
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in,
+                                        out_ch,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, z):
+        #assert z.shape[1:] == self.z_shape[1:]
+        self.last_z_shape = z.shape
+
+        # timestep embedding
+        temb = None
+
+        # z to block_in
+        h = self.conv_in(z)
+
+        # middle
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks+1):
+                h = self.up[i_level].block[i_block](h, temb)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        # end
+        if self.give_pre_end:
+            return h
+
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+class VUNet(nn.Module):
+    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
+                 attn_resolutions, dropout=0.0, resamp_with_conv=True,
+                 in_channels, c_channels,
+                 resolution, z_channels, use_timestep=False, **ignore_kwargs):
+        super().__init__()
+        self.ch = ch
+        self.temb_ch = self.ch*4
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+
+        self.use_timestep = use_timestep
+        if self.use_timestep:
+            # timestep embedding
+            self.temb = nn.Module()
+            self.temb.dense = nn.ModuleList([
+                torch.nn.Linear(self.ch,
+                                self.temb_ch),
+                torch.nn.Linear(self.temb_ch,
+                                self.temb_ch),
+            ])
+
+        # downsampling
+        self.conv_in = torch.nn.Conv2d(c_channels,
+                                       self.ch,
+                                       kernel_size=3,
+                                       stride=1,
+                                       padding=1)
+
+        curr_res = resolution
+        in_ch_mult = (1,)+tuple(ch_mult)
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch*in_ch_mult[i_level]
+            block_out = ch*ch_mult[i_level]
+            for i_block in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions-1:
+                down.downsample = Downsample(block_in, resamp_with_conv)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        self.z_in = torch.nn.Conv2d(z_channels,
+                                    block_in,
+                                    kernel_size=1,
+                                    stride=1,
+                                    padding=0)
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=2*block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch*ch_mult[i_level]
+            skip_in = ch*ch_mult[i_level]
+            for i_block in range(self.num_res_blocks+1):
+                if i_block == self.num_res_blocks:
+                    skip_in = ch*in_ch_mult[i_level]
+                block.append(ResnetBlock(in_channels=block_in+skip_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in, resamp_with_conv)
+                curr_res = curr_res * 2
+            self.up.insert(0, up) # prepend to get consistent order
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in,
+                                        out_ch,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+
+    def forward(self, x, z):
+        #assert x.shape[2] == x.shape[3] == self.resolution
+
+        if self.use_timestep:
+            # timestep embedding
+            assert t is not None
+            temb = get_timestep_embedding(t, self.ch)
+            temb = self.temb.dense[0](temb)
+            temb = nonlinearity(temb)
+            temb = self.temb.dense[1](temb)
+        else:
+            temb = None
+
+        # downsampling
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1], temb)
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions-1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        z = self.z_in(z)
+        h = torch.cat((h,z),dim=1)
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks+1):
+                h = self.up[i_level].block[i_block](
+                    torch.cat([h, hs.pop()], dim=1), temb)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+class SimpleDecoder(nn.Module):
+    def __init__(self, in_channels, out_channels, *args, **kwargs):
+        super().__init__()
+        self.model = nn.ModuleList([nn.Conv2d(in_channels, in_channels, 1),
+                                     ResnetBlock(in_channels=in_channels,
+                                                 out_channels=2 * in_channels,
+                                                 temb_channels=0, dropout=0.0),
+                                     ResnetBlock(in_channels=2 * in_channels,
+                                                out_channels=4 * in_channels,
+                                                temb_channels=0, dropout=0.0),
+                                     ResnetBlock(in_channels=4 * in_channels,
+                                                out_channels=2 * in_channels,
+                                                temb_channels=0, dropout=0.0),
+                                     nn.Conv2d(2*in_channels, in_channels, 1),
+                                     Upsample(in_channels, with_conv=True)])
+        # end
+        self.norm_out = Normalize(in_channels)
+        self.conv_out = torch.nn.Conv2d(in_channels,
+                                        out_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, x):
+        for i, layer in enumerate(self.model):
+            if i in [1,2,3]:
+                x = layer(x, None)
+            else:
+                x = layer(x)
+
+        h = self.norm_out(x)
+        h = nonlinearity(h)
+        x = self.conv_out(h)
+        return x
+
+
+class UpsampleDecoder(nn.Module):
+    def __init__(self, in_channels, out_channels, ch, num_res_blocks, resolution,
+                 ch_mult=(2,2), dropout=0.0):
+        super().__init__()
+        # upsampling
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        block_in = in_channels
+        curr_res = resolution // 2 ** (self.num_resolutions - 1)
+        self.res_blocks = nn.ModuleList()
+        self.upsample_blocks = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            res_block = []
+            block_out = ch * ch_mult[i_level]
+            for i_block in range(self.num_res_blocks + 1):
+                res_block.append(ResnetBlock(in_channels=block_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+            self.res_blocks.append(nn.ModuleList(res_block))
+            if i_level != self.num_resolutions - 1:
+                self.upsample_blocks.append(Upsample(block_in, True))
+                curr_res = curr_res * 2
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in,
+                                        out_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, x):
+        # upsampling
+        h = x
+        for k, i_level in enumerate(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.res_blocks[i_level][i_block](h, None)
+            if i_level != self.num_resolutions - 1:
+                h = self.upsample_blocks[k](h)
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+

taming/modules/discriminator/model.py

@@ -0,0 +1,67 @@
+import functools
+import torch.nn as nn
+
+
+from taming.modules.util import ActNorm
+
+
+def weights_init(m):
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1:
+        nn.init.normal_(m.weight.data, 0.0, 0.02)
+    elif classname.find('BatchNorm') != -1:
+        nn.init.normal_(m.weight.data, 1.0, 0.02)
+        nn.init.constant_(m.bias.data, 0)
+
+
+class NLayerDiscriminator(nn.Module):
+    """Defines a PatchGAN discriminator as in Pix2Pix
+        --> see https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py
+    """
+    def __init__(self, input_nc=3, ndf=64, n_layers=3, use_actnorm=False):
+        """Construct a PatchGAN discriminator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            ndf (int)       -- the number of filters in the last conv layer
+            n_layers (int)  -- the number of conv layers in the discriminator
+            norm_layer      -- normalization layer
+        """
+        super(NLayerDiscriminator, self).__init__()
+        if not use_actnorm:
+            norm_layer = nn.BatchNorm2d
+        else:
+            norm_layer = ActNorm
+        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters
+            use_bias = norm_layer.func != nn.BatchNorm2d
+        else:
+            use_bias = norm_layer != nn.BatchNorm2d
+
+        kw = 4
+        padw = 1
+        sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
+        nf_mult = 1
+        nf_mult_prev = 1
+        for n in range(1, n_layers):  # gradually increase the number of filters
+            nf_mult_prev = nf_mult
+            nf_mult = min(2 ** n, 8)
+            sequence += [
+                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),
+                norm_layer(ndf * nf_mult),
+                nn.LeakyReLU(0.2, True)
+            ]
+
+        nf_mult_prev = nf_mult
+        nf_mult = min(2 ** n_layers, 8)
+        sequence += [
+            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
+            norm_layer(ndf * nf_mult),
+            nn.LeakyReLU(0.2, True)
+        ]
+
+        sequence += [
+            nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]  # output 1 channel prediction map
+        self.main = nn.Sequential(*sequence)
+
+    def forward(self, input):
+        """Standard forward."""
+        return self.main(input)

taming/modules/losses/__init__.py

@@ -0,0 +1,2 @@
+from taming.modules.losses.vqperceptual import DummyLoss
+

taming/modules/losses/lpips.py

@@ -0,0 +1,123 @@
+"""Stripped version of https://github.com/richzhang/PerceptualSimilarity/tree/master/models"""
+
+import torch
+import torch.nn as nn
+from torchvision import models
+from collections import namedtuple
+
+from taming.util import get_ckpt_path
+
+
+class LPIPS(nn.Module):
+    # Learned perceptual metric
+    def __init__(self, use_dropout=True):
+        super().__init__()
+        self.scaling_layer = ScalingLayer()
+        self.chns = [64, 128, 256, 512, 512]  # vg16 features
+        self.net = vgg16(pretrained=True, requires_grad=False)
+        self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
+        self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
+        self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
+        self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
+        self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
+        self.load_from_pretrained()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def load_from_pretrained(self, name="vgg_lpips"):
+        ckpt = get_ckpt_path(name, "taming/modules/autoencoder/lpips")
+        self.load_state_dict(torch.load(ckpt, map_location=torch.device("cpu")), strict=False)
+        print("loaded pretrained LPIPS loss from {}".format(ckpt))
+
+    @classmethod
+    def from_pretrained(cls, name="vgg_lpips"):
+        if name != "vgg_lpips":
+            raise NotImplementedError
+        model = cls()
+        ckpt = get_ckpt_path(name)
+        model.load_state_dict(torch.load(ckpt, map_location=torch.device("cpu")), strict=False)
+        return model
+
+    def forward(self, input, target):
+        in0_input, in1_input = (self.scaling_layer(input), self.scaling_layer(target))
+        outs0, outs1 = self.net(in0_input), self.net(in1_input)
+        feats0, feats1, diffs = {}, {}, {}
+        lins = [self.lin0, self.lin1, self.lin2, self.lin3, self.lin4]
+        for kk in range(len(self.chns)):
+            feats0[kk], feats1[kk] = normalize_tensor(outs0[kk]), normalize_tensor(outs1[kk])
+            diffs[kk] = (feats0[kk] - feats1[kk]) ** 2
+
+        res = [spatial_average(lins[kk].model(diffs[kk]), keepdim=True) for kk in range(len(self.chns))]
+        val = res[0]
+        for l in range(1, len(self.chns)):
+            val += res[l]
+        return val
+
+
+class ScalingLayer(nn.Module):
+    def __init__(self):
+        super(ScalingLayer, self).__init__()
+        self.register_buffer('shift', torch.Tensor([-.030, -.088, -.188])[None, :, None, None])
+        self.register_buffer('scale', torch.Tensor([.458, .448, .450])[None, :, None, None])
+
+    def forward(self, inp):
+        return (inp - self.shift) / self.scale
+
+
+class NetLinLayer(nn.Module):
+    """ A single linear layer which does a 1x1 conv """
+    def __init__(self, chn_in, chn_out=1, use_dropout=False):
+        super(NetLinLayer, self).__init__()
+        layers = [nn.Dropout(), ] if (use_dropout) else []
+        layers += [nn.Conv2d(chn_in, chn_out, 1, stride=1, padding=0, bias=False), ]
+        self.model = nn.Sequential(*layers)
+
+
+class vgg16(torch.nn.Module):
+    def __init__(self, requires_grad=False, pretrained=True):
+        super(vgg16, self).__init__()
+        vgg_pretrained_features = models.vgg16(pretrained=pretrained).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        self.N_slices = 5
+        for x in range(4):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(4, 9):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(9, 16):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(16, 23):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(23, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h = self.slice1(X)
+        h_relu1_2 = h
+        h = self.slice2(h)
+        h_relu2_2 = h
+        h = self.slice3(h)
+        h_relu3_3 = h
+        h = self.slice4(h)
+        h_relu4_3 = h
+        h = self.slice5(h)
+        h_relu5_3 = h
+        vgg_outputs = namedtuple("VggOutputs", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3', 'relu5_3'])
+        out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3)
+        return out
+
+
+def normalize_tensor(x,eps=1e-10):
+    norm_factor = torch.sqrt(torch.sum(x**2,dim=1,keepdim=True))
+    return x/(norm_factor+eps)
+
+
+def spatial_average(x, keepdim=True):
+    return x.mean([2,3],keepdim=keepdim)
+

taming/modules/losses/segmentation.py

@@ -0,0 +1,22 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class BCELoss(nn.Module):
+    def forward(self, prediction, target):
+        loss = F.binary_cross_entropy_with_logits(prediction,target)
+        return loss, {}
+
+
+class BCELossWithQuant(nn.Module):
+    def __init__(self, codebook_weight=1.):
+        super().__init__()
+        self.codebook_weight = codebook_weight
+
+    def forward(self, qloss, target, prediction, split):
+        bce_loss = F.binary_cross_entropy_with_logits(prediction,target)
+        loss = bce_loss + self.codebook_weight*qloss
+        return loss, {"{}/total_loss".format(split): loss.clone().detach().mean(),
+                      "{}/bce_loss".format(split): bce_loss.detach().mean(),
+                      "{}/quant_loss".format(split): qloss.detach().mean()
+                      }

taming/modules/losses/vqperceptual.py

@@ -0,0 +1,182 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from taming.modules.losses.lpips import LPIPS
+from taming.modules.discriminator.model import NLayerDiscriminator, weights_init
+
+
+class DummyLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+
+def adopt_weight(weight, global_step, threshold=0, value=0.0):
+    if global_step < threshold:
+        weight = value
+    return weight
+
+
+def hinge_d_loss(logits_real, logits_fake):
+    loss_real = torch.mean(F.relu(1.0 - logits_real))
+    loss_fake = torch.mean(F.relu(1.0 + logits_fake))
+    d_loss = 0.5 * (loss_real + loss_fake)
+    return d_loss
+
+
+def vanilla_d_loss(logits_real, logits_fake):
+    d_loss = 0.5 * (
+        torch.mean(torch.nn.functional.softplus(-logits_real))
+        + torch.mean(torch.nn.functional.softplus(logits_fake))
+    )
+    return d_loss
+
+
+class VQLPIPSWithDiscriminator(nn.Module):
+    def __init__(
+        self,
+        disc_start,
+        codebook_weight=1.0,
+        pixelloss_weight=1.0,
+        disc_num_layers=3,
+        disc_in_channels=3,
+        disc_factor=1.0,
+        disc_weight=1.0,
+        perceptual_weight=1.0,
+        use_actnorm=False,
+        disc_conditional=False,
+        disc_ndf=64,
+        disc_loss="hinge",
+    ):
+        super().__init__()
+        assert disc_loss in ["hinge", "vanilla"]
+        self.codebook_weight = codebook_weight
+        self.pixel_weight = pixelloss_weight
+        self.perceptual_loss = LPIPS().eval()
+        self.perceptual_weight = perceptual_weight
+
+        self.discriminator = NLayerDiscriminator(
+            input_nc=disc_in_channels,
+            n_layers=disc_num_layers,
+            use_actnorm=use_actnorm,
+            ndf=disc_ndf,
+        ).apply(weights_init)
+        self.discriminator_iter_start = disc_start
+        if disc_loss == "hinge":
+            self.disc_loss = hinge_d_loss
+        elif disc_loss == "vanilla":
+            self.disc_loss = vanilla_d_loss
+        else:
+            raise ValueError(f"Unknown GAN loss '{disc_loss}'.")
+        print(f"VQLPIPSWithDiscriminator running with {disc_loss} loss.")
+        self.disc_factor = disc_factor
+        self.discriminator_weight = disc_weight
+        self.disc_conditional = disc_conditional
+
+    def calculate_adaptive_weight(self, nll_loss, g_loss, last_layer=None):
+        if last_layer is not None:
+            nll_grads = torch.autograd.grad(nll_loss, last_layer, retain_graph=True)[0]
+            g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0]
+        else:
+            nll_grads = torch.autograd.grad(
+                nll_loss, self.last_layer[0], retain_graph=True
+            )[0]
+            g_grads = torch.autograd.grad(
+                g_loss, self.last_layer[0], retain_graph=True
+            )[0]
+
+        d_weight = torch.norm(nll_grads) / (torch.norm(g_grads) + 1e-4)
+        d_weight = torch.clamp(d_weight, 0.0, 1e4).detach()
+        d_weight = d_weight * self.discriminator_weight
+        return d_weight
+
+    def forward(
+        self,
+        codebook_loss,
+        inputs,
+        reconstructions,
+        optimizer_idx,
+        global_step,
+        last_layer=None,
+        cond=None,
+        split="train",
+    ):
+        rec_loss = torch.abs(inputs.contiguous() - reconstructions.contiguous())
+        if self.perceptual_weight > 0:
+            p_loss = self.perceptual_loss(
+                inputs.contiguous(), reconstructions.contiguous()
+            )
+            rec_loss = rec_loss + self.perceptual_weight * p_loss
+        else:
+            p_loss = torch.tensor([0.0])
+
+        nll_loss = rec_loss
+        # nll_loss = torch.sum(nll_loss) / nll_loss.shape[0]
+        nll_loss = torch.mean(nll_loss)
+
+        # now the GAN part
+        if optimizer_idx == 0:
+            # generator update
+            if cond is None:
+                assert not self.disc_conditional
+                logits_fake = self.discriminator(reconstructions.contiguous())
+            else:
+                assert self.disc_conditional
+                logits_fake = self.discriminator(
+                    torch.cat((reconstructions.contiguous(), cond), dim=1)
+                )
+            g_loss = -torch.mean(logits_fake)
+
+            try:
+                d_weight = self.calculate_adaptive_weight(
+                    nll_loss, g_loss, last_layer=last_layer
+                )
+            except RuntimeError:
+                assert not self.training
+                d_weight = torch.tensor(0.0)
+
+            disc_factor = adopt_weight(
+                self.disc_factor, global_step, threshold=self.discriminator_iter_start
+            )
+            loss = (
+                nll_loss
+                + d_weight * disc_factor * g_loss
+                + self.codebook_weight * codebook_loss.mean()
+            )
+
+            log = {
+                "{}/total_loss".format(split): loss.clone().detach().mean(),
+                "{}/quant_loss".format(split): codebook_loss.detach().mean(),
+                "{}/nll_loss".format(split): nll_loss.detach().mean(),
+                "{}/rec_loss".format(split): rec_loss.detach().mean(),
+                "{}/p_loss".format(split): p_loss.detach().mean(),
+                "{}/d_weight".format(split): d_weight.detach(),
+                "{}/disc_factor".format(split): torch.tensor(disc_factor),
+                "{}/g_loss".format(split): g_loss.detach().mean(),
+            }
+            return loss, log
+
+        if optimizer_idx == 1:
+            # second pass for discriminator update
+            if cond is None:
+                logits_real = self.discriminator(inputs.contiguous().detach())
+                logits_fake = self.discriminator(reconstructions.contiguous().detach())
+            else:
+                logits_real = self.discriminator(
+                    torch.cat((inputs.contiguous().detach(), cond), dim=1)
+                )
+                logits_fake = self.discriminator(
+                    torch.cat((reconstructions.contiguous().detach(), cond), dim=1)
+                )
+
+            disc_factor = adopt_weight(
+                self.disc_factor, global_step, threshold=self.discriminator_iter_start
+            )
+            d_loss = disc_factor * self.disc_loss(logits_real, logits_fake)
+
+            log = {
+                "{}/disc_loss".format(split): d_loss.clone().detach().mean(),
+                "{}/logits_real".format(split): logits_real.detach().mean(),
+                "{}/logits_fake".format(split): logits_fake.detach().mean(),
+            }
+            return d_loss, log

taming/modules/misc/coord.py

@@ -0,0 +1,31 @@
+import torch
+
+class CoordStage(object):
+    def __init__(self, n_embed, down_factor):
+        self.n_embed = n_embed
+        self.down_factor = down_factor
+
+    def eval(self):
+        return self
+
+    def encode(self, c):
+        """fake vqmodel interface"""
+        assert 0.0 <= c.min() and c.max() <= 1.0
+        b,ch,h,w = c.shape
+        assert ch == 1
+
+        c = torch.nn.functional.interpolate(c, scale_factor=1/self.down_factor,
+                                            mode="area")
+        c = c.clamp(0.0, 1.0)
+        c = self.n_embed*c
+        c_quant = c.round()
+        c_ind = c_quant.to(dtype=torch.long)
+
+        info = None, None, c_ind
+        return c_quant, None, info
+
+    def decode(self, c):
+        c = c/self.n_embed
+        c = torch.nn.functional.interpolate(c, scale_factor=self.down_factor,
+                                            mode="nearest")
+        return c

taming/modules/transformer/mingpt.py

@@ -0,0 +1,415 @@
+"""
+taken from: https://github.com/karpathy/minGPT/
+GPT model:
+- the initial stem consists of a combination of token encoding and a positional encoding
+- the meat of it is a uniform sequence of Transformer blocks
+    - each Transformer is a sequential combination of a 1-hidden-layer MLP block and a self-attention block
+    - all blocks feed into a central residual pathway similar to resnets
+- the final decoder is a linear projection into a vanilla Softmax classifier
+"""
+
+import math
+import logging
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from transformers import top_k_top_p_filtering
+
+logger = logging.getLogger(__name__)
+
+
+class GPTConfig:
+    """ base GPT config, params common to all GPT versions """
+    embd_pdrop = 0.1
+    resid_pdrop = 0.1
+    attn_pdrop = 0.1
+
+    def __init__(self, vocab_size, block_size, **kwargs):
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        for k,v in kwargs.items():
+            setattr(self, k, v)
+
+
+class GPT1Config(GPTConfig):
+    """ GPT-1 like network roughly 125M params """
+    n_layer = 12
+    n_head = 12
+    n_embd = 768
+
+
+class CausalSelfAttention(nn.Module):
+    """
+    A vanilla multi-head masked self-attention layer with a projection at the end.
+    It is possible to use torch.nn.MultiheadAttention here but I am including an
+    explicit implementation here to show that there is nothing too scary here.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        # key, query, value projections for all heads
+        self.key = nn.Linear(config.n_embd, config.n_embd)
+        self.query = nn.Linear(config.n_embd, config.n_embd)
+        self.value = nn.Linear(config.n_embd, config.n_embd)
+        # regularization
+        self.attn_drop = nn.Dropout(config.attn_pdrop)
+        self.resid_drop = nn.Dropout(config.resid_pdrop)
+        # output projection
+        self.proj = nn.Linear(config.n_embd, config.n_embd)
+        # causal mask to ensure that attention is only applied to the left in the input sequence
+        mask = torch.tril(torch.ones(config.block_size,
+                                     config.block_size))
+        if hasattr(config, "n_unmasked"):
+            mask[:config.n_unmasked, :config.n_unmasked] = 1
+        self.register_buffer("mask", mask.view(1, 1, config.block_size, config.block_size))
+        self.n_head = config.n_head
+
+    def forward(self, x, layer_past=None):
+        B, T, C = x.size()
+
+        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
+        k = self.key(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        q = self.query(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        v = self.value(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+
+        present = torch.stack((k, v))
+        if layer_past is not None:
+            past_key, past_value = layer_past
+            k = torch.cat((past_key, k), dim=-2)
+            v = torch.cat((past_value, v), dim=-2)
+
+        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
+        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+        if layer_past is None:
+            att = att.masked_fill(self.mask[:,:,:T,:T] == 0, float('-inf'))
+
+        att = F.softmax(att, dim=-1)
+        att = self.attn_drop(att)
+        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
+        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
+
+        # output projection
+        y = self.resid_drop(self.proj(y))
+        return y, present   # TODO: check that this does not break anything
+
+
+class Block(nn.Module):
+    """ an unassuming Transformer block """
+    def __init__(self, config):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(config.n_embd)
+        self.ln2 = nn.LayerNorm(config.n_embd)
+        self.attn = CausalSelfAttention(config)
+        self.mlp = nn.Sequential(
+            nn.Linear(config.n_embd, 4 * config.n_embd),
+            nn.GELU(),  # nice
+            nn.Linear(4 * config.n_embd, config.n_embd),
+            nn.Dropout(config.resid_pdrop),
+        )
+
+    def forward(self, x, layer_past=None, return_present=False):
+        # TODO: check that training still works
+        if return_present: assert not self.training
+        # layer past: tuple of length two with B, nh, T, hs
+        attn, present = self.attn(self.ln1(x), layer_past=layer_past)
+
+        x = x + attn
+        x = x + self.mlp(self.ln2(x))
+        if layer_past is not None or return_present:
+            return x, present
+        return x
+
+
+class GPT(nn.Module):
+    """  the full GPT language model, with a context size of block_size """
+    def __init__(self, vocab_size, block_size, n_layer=12, n_head=8, n_embd=256,
+                 embd_pdrop=0., resid_pdrop=0., attn_pdrop=0., n_unmasked=0):
+        super().__init__()
+        config = GPTConfig(vocab_size=vocab_size, block_size=block_size,
+                           embd_pdrop=embd_pdrop, resid_pdrop=resid_pdrop, attn_pdrop=attn_pdrop,
+                           n_layer=n_layer, n_head=n_head, n_embd=n_embd,
+                           n_unmasked=n_unmasked)
+        # input embedding stem
+        self.tok_emb = nn.Embedding(config.vocab_size, config.n_embd)
+        self.pos_emb = nn.Parameter(torch.zeros(1, config.block_size, config.n_embd))
+        self.drop = nn.Dropout(config.embd_pdrop)
+        # transformer
+        self.blocks = nn.Sequential(*[Block(config) for _ in range(config.n_layer)])
+        # decoder head
+        self.ln_f = nn.LayerNorm(config.n_embd)
+        self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        self.block_size = config.block_size
+        self.apply(self._init_weights)
+        self.config = config
+        logger.info("number of parameters: %e", sum(p.numel() for p in self.parameters()))
+
+    def get_block_size(self):
+        return self.block_size
+
+    def _init_weights(self, module):
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+            if isinstance(module, nn.Linear) and module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+    def forward(self, idx, embeddings=None, targets=None):
+        # forward the GPT model
+        token_embeddings = self.tok_emb(idx) # each index maps to a (learnable) vector
+
+        if embeddings is not None: # prepend explicit embeddings
+            token_embeddings = torch.cat((embeddings, token_embeddings), dim=1)
+
+        t = token_embeddings.shape[1]
+        assert t <= self.block_size, "Cannot forward, model block size is exhausted."
+        position_embeddings = self.pos_emb[:, :t, :] # each position maps to a (learnable) vector
+        x = self.drop(token_embeddings + position_embeddings)
+        x = self.blocks(x)
+        x = self.ln_f(x)
+        logits = self.head(x)
+
+        # if we are given some desired targets also calculate the loss
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
+
+        return logits, loss
+
+    def forward_with_past(self, idx, embeddings=None, targets=None, past=None, past_length=None):
+        # inference only
+        assert not self.training
+        token_embeddings = self.tok_emb(idx)    # each index maps to a (learnable) vector
+        if embeddings is not None:              # prepend explicit embeddings
+            token_embeddings = torch.cat((embeddings, token_embeddings), dim=1)
+
+        if past is not None:
+            assert past_length is not None
+            past = torch.cat(past, dim=-2)   # n_layer, 2, b, nh, len_past, dim_head
+            past_shape = list(past.shape)
+            expected_shape = [self.config.n_layer, 2, idx.shape[0], self.config.n_head, past_length, self.config.n_embd//self.config.n_head]
+            assert past_shape == expected_shape, f"{past_shape} =/= {expected_shape}"
+            position_embeddings = self.pos_emb[:, past_length, :]  # each position maps to a (learnable) vector
+        else:
+            position_embeddings = self.pos_emb[:, :token_embeddings.shape[1], :]
+
+        x = self.drop(token_embeddings + position_embeddings)
+        presents = []  # accumulate over layers
+        for i, block in enumerate(self.blocks):
+            x, present = block(x, layer_past=past[i, ...] if past is not None else None, return_present=True)
+            presents.append(present)
+
+        x = self.ln_f(x)
+        logits = self.head(x)
+        # if we are given some desired targets also calculate the loss
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
+
+        return logits, loss, torch.stack(presents)  # _, _, n_layer, 2, b, nh, 1, dim_head
+
+
+class DummyGPT(nn.Module):
+    # for debugging
+    def __init__(self, add_value=1):
+        super().__init__()
+        self.add_value = add_value
+
+    def forward(self, idx):
+        return idx + self.add_value, None
+
+
+class CodeGPT(nn.Module):
+    """Takes in semi-embeddings"""
+    def __init__(self, vocab_size, block_size, in_channels, n_layer=12, n_head=8, n_embd=256,
+                 embd_pdrop=0., resid_pdrop=0., attn_pdrop=0., n_unmasked=0):
+        super().__init__()
+        config = GPTConfig(vocab_size=vocab_size, block_size=block_size,
+                           embd_pdrop=embd_pdrop, resid_pdrop=resid_pdrop, attn_pdrop=attn_pdrop,
+                           n_layer=n_layer, n_head=n_head, n_embd=n_embd,
+                           n_unmasked=n_unmasked)
+        # input embedding stem
+        self.tok_emb = nn.Linear(in_channels, config.n_embd)
+        self.pos_emb = nn.Parameter(torch.zeros(1, config.block_size, config.n_embd))
+        self.drop = nn.Dropout(config.embd_pdrop)
+        # transformer
+        self.blocks = nn.Sequential(*[Block(config) for _ in range(config.n_layer)])
+        # decoder head
+        self.ln_f = nn.LayerNorm(config.n_embd)
+        self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        self.block_size = config.block_size
+        self.apply(self._init_weights)
+        self.config = config
+        logger.info("number of parameters: %e", sum(p.numel() for p in self.parameters()))
+
+    def get_block_size(self):
+        return self.block_size
+
+    def _init_weights(self, module):
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+            if isinstance(module, nn.Linear) and module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+    def forward(self, idx, embeddings=None, targets=None):
+        # forward the GPT model
+        token_embeddings = self.tok_emb(idx) # each index maps to a (learnable) vector
+
+        if embeddings is not None: # prepend explicit embeddings
+            token_embeddings = torch.cat((embeddings, token_embeddings), dim=1)
+
+        t = token_embeddings.shape[1]
+        assert t <= self.block_size, "Cannot forward, model block size is exhausted."
+        position_embeddings = self.pos_emb[:, :t, :] # each position maps to a (learnable) vector
+        x = self.drop(token_embeddings + position_embeddings)
+        x = self.blocks(x)
+        x = self.taming_cinln_f(x)
+        logits = self.head(x)
+
+        # if we are given some desired targets also calculate the loss
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
+
+        return logits, loss
+
+
+
+#### sampling utils
+
+def top_k_logits(logits, k):
+    v, ix = torch.topk(logits, k)
+    out = logits.clone()
+    out[out < v[:, [-1]]] = -float('Inf')
+    return out
+
+@torch.no_grad()
+def sample(model, x, steps, temperature=1.0, sample=False, top_k=None):
+    """
+    take a conditioning sequence of indices in x (of shape (b,t)) and predict the next token in
+    the sequence, feeding the predictions back into the model each time. Clearly the sampling
+    has quadratic complexity unlike an RNN that is only linear, and has a finite context window
+    of block_size, unlike an RNN that has an infinite context window.
+    """
+    block_size = model.get_block_size()
+    model.eval()
+    for k in range(steps):
+        x_cond = x if x.size(1) <= block_size else x[:, -block_size:]  # crop context if needed
+        logits, _ = model(x_cond)
+        # pluck the logits at the final step and scale by temperature
+        logits = logits[:, -1, :] / temperature
+        # optionally crop probabilities to only the top k options
+        if top_k is not None:
+            logits = top_k_logits(logits, top_k)
+        # apply softmax to convert to probabilities
+        probs = F.softmax(logits, dim=-1)
+        # sample from the distribution or take the most likely
+        if sample:
+            ix = torch.multinomial(probs, num_samples=1)
+        else:
+            _, ix = torch.topk(probs, k=1, dim=-1)
+        # append to the sequence and continue
+        x = torch.cat((x, ix), dim=1)
+
+    return x
+
+
+@torch.no_grad()
+def sample_with_past(x, model, steps, temperature=1., sample_logits=True,
+                     top_k=None, top_p=None, callback=None):
+    # x is conditioning
+    sample = x
+    cond_len = x.shape[1]
+    past = None
+    for n in range(steps):
+        if callback is not None:
+            callback(n)
+        logits, _, present = model.forward_with_past(x, past=past, past_length=(n+cond_len-1))
+        if past is None:
+            past = [present]
+        else:
+            past.append(present)
+        logits = logits[:, -1, :] / temperature
+        if top_k is not None:
+            logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
+
+        probs = F.softmax(logits, dim=-1)
+        if not sample_logits:
+            _, x = torch.topk(probs, k=1, dim=-1)
+        else:
+            x = torch.multinomial(probs, num_samples=1)
+        # append to the sequence and continue
+        sample = torch.cat((sample, x), dim=1)
+    del past
+    sample = sample[:, cond_len:]  # cut conditioning off
+    return sample
+
+
+#### clustering utils
+
+class KMeans(nn.Module):
+    def __init__(self, ncluster=512, nc=3, niter=10):
+        super().__init__()
+        self.ncluster = ncluster
+        self.nc = nc
+        self.niter = niter
+        self.shape = (3,32,32)
+        self.register_buffer("C", torch.zeros(self.ncluster,nc))
+        self.register_buffer('initialized', torch.tensor(0, dtype=torch.uint8))
+
+    def is_initialized(self):
+        return self.initialized.item() == 1
+
+    @torch.no_grad()
+    def initialize(self, x):
+        N, D = x.shape
+        assert D == self.nc, D
+        c = x[torch.randperm(N)[:self.ncluster]] # init clusters at random
+        for i in range(self.niter):
+            # assign all pixels to the closest codebook element
+            a = ((x[:, None, :] - c[None, :, :])**2).sum(-1).argmin(1)
+            # move each codebook element to be the mean of the pixels that assigned to it
+            c = torch.stack([x[a==k].mean(0) for k in range(self.ncluster)])
+            # re-assign any poorly positioned codebook elements
+            nanix = torch.any(torch.isnan(c), dim=1)
+            ndead = nanix.sum().item()
+            print('done step %d/%d, re-initialized %d dead clusters' % (i+1, self.niter, ndead))
+            c[nanix] = x[torch.randperm(N)[:ndead]] # re-init dead clusters
+
+        self.C.copy_(c)
+        self.initialized.fill_(1)
+
+
+    def forward(self, x, reverse=False, shape=None):
+        if not reverse:
+            # flatten
+            bs,c,h,w = x.shape
+            assert c == self.nc
+            x = x.reshape(bs,c,h*w,1)
+            C = self.C.permute(1,0)
+            C = C.reshape(1,c,1,self.ncluster)
+            a = ((x-C)**2).sum(1).argmin(-1) # bs, h*w indices
+            return a
+        else:
+            # flatten
+            bs, HW = x.shape
+            """
+            c = self.C.reshape( 1, self.nc,  1, self.ncluster)
+            c = c[bs*[0],:,:,:]
+            c = c[:,:,HW*[0],:]
+            x =      x.reshape(bs,       1, HW,             1)
+            x = x[:,3*[0],:,:]
+            x = torch.gather(c, dim=3, index=x)
+            """
+            x = self.C[x]
+            x = x.permute(0,2,1)
+            shape = shape if shape is not None else self.shape
+            x = x.reshape(bs, *shape)
+
+            return x

taming/modules/transformer/permuter.py

@@ -0,0 +1,248 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+class AbstractPermuter(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+    def forward(self, x, reverse=False):
+        raise NotImplementedError
+
+
+class Identity(AbstractPermuter):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, reverse=False):
+        return x
+
+
+class Subsample(AbstractPermuter):
+    def __init__(self, H, W):
+        super().__init__()
+        C = 1
+        indices = np.arange(H*W).reshape(C,H,W)
+        while min(H, W) > 1:
+            indices = indices.reshape(C,H//2,2,W//2,2)
+            indices = indices.transpose(0,2,4,1,3)
+            indices = indices.reshape(C*4,H//2, W//2)
+            H = H//2
+            W = W//2
+            C = C*4
+        assert H == W == 1
+        idx = torch.tensor(indices.ravel())
+        self.register_buffer('forward_shuffle_idx',
+                             nn.Parameter(idx, requires_grad=False))
+        self.register_buffer('backward_shuffle_idx',
+                             nn.Parameter(torch.argsort(idx), requires_grad=False))
+
+    def forward(self, x, reverse=False):
+        if not reverse:
+            return x[:, self.forward_shuffle_idx]
+        else:
+            return x[:, self.backward_shuffle_idx]
+
+
+def mortonify(i, j):
+    """(i,j) index to linear morton code"""
+    i = np.uint64(i)
+    j = np.uint64(j)
+
+    z = np.uint(0)
+
+    for pos in range(32):
+        z = (z |
+             ((j & (np.uint64(1) << np.uint64(pos))) << np.uint64(pos)) |
+             ((i & (np.uint64(1) << np.uint64(pos))) << np.uint64(pos+1))
+             )
+    return z
+
+
+class ZCurve(AbstractPermuter):
+    def __init__(self, H, W):
+        super().__init__()
+        reverseidx = [np.int64(mortonify(i,j)) for i in range(H) for j in range(W)]
+        idx = np.argsort(reverseidx)
+        idx = torch.tensor(idx)
+        reverseidx = torch.tensor(reverseidx)
+        self.register_buffer('forward_shuffle_idx',
+                             idx)
+        self.register_buffer('backward_shuffle_idx',
+                             reverseidx)
+
+    def forward(self, x, reverse=False):
+        if not reverse:
+            return x[:, self.forward_shuffle_idx]
+        else:
+            return x[:, self.backward_shuffle_idx]
+
+
+class SpiralOut(AbstractPermuter):
+    def __init__(self, H, W):
+        super().__init__()
+        assert H == W
+        size = W
+        indices = np.arange(size*size).reshape(size,size)
+
+        i0 = size//2
+        j0 = size//2-1
+
+        i = i0
+        j = j0
+
+        idx = [indices[i0, j0]]
+        step_mult = 0
+        for c in range(1, size//2+1):
+            step_mult += 1
+            # steps left
+            for k in range(step_mult):
+                i = i - 1
+                j = j
+                idx.append(indices[i, j])
+
+            # step down
+            for k in range(step_mult):
+                i = i
+                j = j + 1
+                idx.append(indices[i, j])
+
+            step_mult += 1
+            if c < size//2:
+                # step right
+                for k in range(step_mult):
+                    i = i + 1
+                    j = j
+                    idx.append(indices[i, j])
+
+                # step up
+                for k in range(step_mult):
+                    i = i
+                    j = j - 1
+                    idx.append(indices[i, j])
+            else:
+                # end reached
+                for k in range(step_mult-1):
+                    i = i + 1
+                    idx.append(indices[i, j])
+
+        assert len(idx) == size*size
+        idx = torch.tensor(idx)
+        self.register_buffer('forward_shuffle_idx', idx)
+        self.register_buffer('backward_shuffle_idx', torch.argsort(idx))
+
+    def forward(self, x, reverse=False):
+        if not reverse:
+            return x[:, self.forward_shuffle_idx]
+        else:
+            return x[:, self.backward_shuffle_idx]
+
+
+class SpiralIn(AbstractPermuter):
+    def __init__(self, H, W):
+        super().__init__()
+        assert H == W
+        size = W
+        indices = np.arange(size*size).reshape(size,size)
+
+        i0 = size//2
+        j0 = size//2-1
+
+        i = i0
+        j = j0
+
+        idx = [indices[i0, j0]]
+        step_mult = 0
+        for c in range(1, size//2+1):
+            step_mult += 1
+            # steps left
+            for k in range(step_mult):
+                i = i - 1
+                j = j
+                idx.append(indices[i, j])
+
+            # step down
+            for k in range(step_mult):
+                i = i
+                j = j + 1
+                idx.append(indices[i, j])
+
+            step_mult += 1
+            if c < size//2:
+                # step right
+                for k in range(step_mult):
+                    i = i + 1
+                    j = j
+                    idx.append(indices[i, j])
+
+                # step up
+                for k in range(step_mult):
+                    i = i
+                    j = j - 1
+                    idx.append(indices[i, j])
+            else:
+                # end reached
+                for k in range(step_mult-1):
+                    i = i + 1
+                    idx.append(indices[i, j])
+
+        assert len(idx) == size*size
+        idx = idx[::-1]
+        idx = torch.tensor(idx)
+        self.register_buffer('forward_shuffle_idx', idx)
+        self.register_buffer('backward_shuffle_idx', torch.argsort(idx))
+
+    def forward(self, x, reverse=False):
+        if not reverse:
+            return x[:, self.forward_shuffle_idx]
+        else:
+            return x[:, self.backward_shuffle_idx]
+
+
+class Random(nn.Module):
+    def __init__(self, H, W):
+        super().__init__()
+        indices = np.random.RandomState(1).permutation(H*W)
+        idx = torch.tensor(indices.ravel())
+        self.register_buffer('forward_shuffle_idx', idx)
+        self.register_buffer('backward_shuffle_idx', torch.argsort(idx))
+
+    def forward(self, x, reverse=False):
+        if not reverse:
+            return x[:, self.forward_shuffle_idx]
+        else:
+            return x[:, self.backward_shuffle_idx]
+
+
+class AlternateParsing(AbstractPermuter):
+    def __init__(self, H, W):
+        super().__init__()
+        indices = np.arange(W*H).reshape(H,W)
+        for i in range(1, H, 2):
+            indices[i, :] = indices[i, ::-1]
+        idx = indices.flatten()
+        assert len(idx) == H*W
+        idx = torch.tensor(idx)
+        self.register_buffer('forward_shuffle_idx', idx)
+        self.register_buffer('backward_shuffle_idx', torch.argsort(idx))
+
+    def forward(self, x, reverse=False):
+        if not reverse:
+            return x[:, self.forward_shuffle_idx]
+        else:
+            return x[:, self.backward_shuffle_idx]
+
+
+if __name__ == "__main__":
+    p0 = AlternateParsing(16, 16)
+    print(p0.forward_shuffle_idx)
+    print(p0.backward_shuffle_idx)
+
+    x = torch.randint(0, 768, size=(11, 256))
+    y = p0(x)
+    xre = p0(y, reverse=True)
+    assert torch.equal(x, xre)
+
+    p1 = SpiralOut(2, 2)
+    print(p1.forward_shuffle_idx)
+    print(p1.backward_shuffle_idx)

taming/modules/util.py

@@ -0,0 +1,130 @@
+import torch
+import torch.nn as nn
+
+
+def count_params(model):
+    total_params = sum(p.numel() for p in model.parameters())
+    return total_params
+
+
+class ActNorm(nn.Module):
+    def __init__(self, num_features, logdet=False, affine=True,
+                 allow_reverse_init=False):
+        assert affine
+        super().__init__()
+        self.logdet = logdet
+        self.loc = nn.Parameter(torch.zeros(1, num_features, 1, 1))
+        self.scale = nn.Parameter(torch.ones(1, num_features, 1, 1))
+        self.allow_reverse_init = allow_reverse_init
+
+        self.register_buffer('initialized', torch.tensor(0, dtype=torch.uint8))
+
+    def initialize(self, input):
+        with torch.no_grad():
+            flatten = input.permute(1, 0, 2, 3).contiguous().view(input.shape[1], -1)
+            mean = (
+                flatten.mean(1)
+                .unsqueeze(1)
+                .unsqueeze(2)
+                .unsqueeze(3)
+                .permute(1, 0, 2, 3)
+            )
+            std = (
+                flatten.std(1)
+                .unsqueeze(1)
+                .unsqueeze(2)
+                .unsqueeze(3)
+                .permute(1, 0, 2, 3)
+            )
+
+            self.loc.data.copy_(-mean)
+            self.scale.data.copy_(1 / (std + 1e-6))
+
+    def forward(self, input, reverse=False):
+        if reverse:
+            return self.reverse(input)
+        if len(input.shape) == 2:
+            input = input[:,:,None,None]
+            squeeze = True
+        else:
+            squeeze = False
+
+        _, _, height, width = input.shape
+
+        if self.training and self.initialized.item() == 0:
+            self.initialize(input)
+            self.initialized.fill_(1)
+
+        h = self.scale * (input + self.loc)
+
+        if squeeze:
+            h = h.squeeze(-1).squeeze(-1)
+
+        if self.logdet:
+            log_abs = torch.log(torch.abs(self.scale))
+            logdet = height*width*torch.sum(log_abs)
+            logdet = logdet * torch.ones(input.shape[0]).to(input)
+            return h, logdet
+
+        return h
+
+    def reverse(self, output):
+        if self.training and self.initialized.item() == 0:
+            if not self.allow_reverse_init:
+                raise RuntimeError(
+                    "Initializing ActNorm in reverse direction is "
+                    "disabled by default. Use allow_reverse_init=True to enable."
+                )
+            else:
+                self.initialize(output)
+                self.initialized.fill_(1)
+
+        if len(output.shape) == 2:
+            output = output[:,:,None,None]
+            squeeze = True
+        else:
+            squeeze = False
+
+        h = output / self.scale - self.loc
+
+        if squeeze:
+            h = h.squeeze(-1).squeeze(-1)
+        return h
+
+
+class AbstractEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def encode(self, *args, **kwargs):
+        raise NotImplementedError
+
+
+class Labelator(AbstractEncoder):
+    """Net2Net Interface for Class-Conditional Model"""
+    def __init__(self, n_classes, quantize_interface=True):
+        super().__init__()
+        self.n_classes = n_classes
+        self.quantize_interface = quantize_interface
+
+    def encode(self, c):
+        c = c[:,None]
+        if self.quantize_interface:
+            return c, None, [None, None, c.long()]
+        return c
+
+
+class SOSProvider(AbstractEncoder):
+    # for unconditional training
+    def __init__(self, sos_token, quantize_interface=True):
+        super().__init__()
+        self.sos_token = sos_token
+        self.quantize_interface = quantize_interface
+
+    def encode(self, x):
+        # get batch size from data and replicate sos_token
+        c = torch.ones(x.shape[0], 1)*self.sos_token
+        c = c.long().to(x.device)
+        if self.quantize_interface:
+            return c, None, [None, None, c]
+        return c

taming/modules/vqvae/quantize.py

@@ -0,0 +1,445 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from torch import einsum
+from einops import rearrange
+
+
+class VectorQuantizer(nn.Module):
+    """
+    see https://github.com/MishaLaskin/vqvae/blob/d761a999e2267766400dc646d82d3ac3657771d4/models/quantizer.py
+    ____________________________________________
+    Discretization bottleneck part of the VQ-VAE.
+    Inputs:
+    - n_e : number of embeddings
+    - e_dim : dimension of embedding
+    - beta : commitment cost used in loss term, beta * ||z_e(x)-sg[e]||^2
+    _____________________________________________
+    """
+
+    # NOTE: this class contains a bug regarding beta; see VectorQuantizer2 for
+    # a fix and use legacy=False to apply that fix. VectorQuantizer2 can be
+    # used wherever VectorQuantizer has been used before and is additionally
+    # more efficient.
+    def __init__(self, n_e, e_dim, beta):
+        super(VectorQuantizer, self).__init__()
+        self.n_e = n_e
+        self.e_dim = e_dim
+        self.beta = beta
+
+        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)
+
+    def forward(self, z):
+        """
+        Inputs the output of the encoder network z and maps it to a discrete
+        one-hot vector that is the index of the closest embedding vector e_j
+        z (continuous) -> z_q (discrete)
+        z.shape = (batch, channel, height, width)
+        quantization pipeline:
+            1. get encoder input (B,C,H,W)
+            2. flatten input to (B*H*W,C)
+        """
+        # reshape z -> (batch, height, width, channel) and flatten
+        z = z.permute(0, 2, 3, 1).contiguous()
+        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
+
+        d = torch.sum(z_flattened ** 2, dim=1, keepdim=True) + \
+            torch.sum(self.embedding.weight**2, dim=1) - 2 * \
+            torch.matmul(z_flattened, self.embedding.weight.t())
+
+        ## could possible replace this here
+        # #\start...
+        # find closest encodings
+        min_encoding_indices = torch.argmin(d, dim=1).unsqueeze(1)
+
+        min_encodings = torch.zeros(
+            min_encoding_indices.shape[0], self.n_e).to(z)
+        min_encodings.scatter_(1, min_encoding_indices, 1)
+
+        # dtype min encodings: torch.float32
+        # min_encodings shape: torch.Size([2048, 512])
+        # min_encoding_indices.shape: torch.Size([2048, 1])
+
+        # get quantized latent vectors
+        z_q = torch.matmul(min_encodings, self.embedding.weight).view(z.shape)
+        #.........\end
+
+        # with:
+        # .........\start
+        #min_encoding_indices = torch.argmin(d, dim=1)
+        #z_q = self.embedding(min_encoding_indices)
+        # ......\end......... (TODO)
+
+        # compute loss for embedding
+        loss = torch.mean((z_q.detach()-z)**2) + self.beta * \
+            torch.mean((z_q - z.detach()) ** 2)
+
+        # preserve gradients
+        z_q = z + (z_q - z).detach()
+
+        # perplexity
+        e_mean = torch.mean(min_encodings, dim=0)
+        perplexity = torch.exp(-torch.sum(e_mean * torch.log(e_mean + 1e-10)))
+
+        # reshape back to match original input shape
+        z_q = z_q.permute(0, 3, 1, 2).contiguous()
+
+        return z_q, loss, (perplexity, min_encodings, min_encoding_indices)
+
+    def get_codebook_entry(self, indices, shape):
+        # shape specifying (batch, height, width, channel)
+        # TODO: check for more easy handling with nn.Embedding
+        min_encodings = torch.zeros(indices.shape[0], self.n_e).to(indices)
+        min_encodings.scatter_(1, indices[:,None], 1)
+
+        # get quantized latent vectors
+        z_q = torch.matmul(min_encodings.float(), self.embedding.weight)
+
+        if shape is not None:
+            z_q = z_q.view(shape)
+
+            # reshape back to match original input shape
+            z_q = z_q.permute(0, 3, 1, 2).contiguous()
+
+        return z_q
+
+
+class GumbelQuantize(nn.Module):
+    """
+    credit to @karpathy: https://github.com/karpathy/deep-vector-quantization/blob/main/model.py (thanks!)
+    Gumbel Softmax trick quantizer
+    Categorical Reparameterization with Gumbel-Softmax, Jang et al. 2016
+    https://arxiv.org/abs/1611.01144
+    """
+    def __init__(self, num_hiddens, embedding_dim, n_embed, straight_through=True,
+                 kl_weight=5e-4, temp_init=1.0, use_vqinterface=True,
+                 remap=None, unknown_index="random"):
+        super().__init__()
+
+        self.embedding_dim = embedding_dim
+        self.n_embed = n_embed
+
+        self.straight_through = straight_through
+        self.temperature = temp_init
+        self.kl_weight = kl_weight
+
+        self.proj = nn.Conv2d(num_hiddens, n_embed, 1)
+        self.embed = nn.Embedding(n_embed, embedding_dim)
+
+        self.use_vqinterface = use_vqinterface
+
+        self.remap = remap
+        if self.remap is not None:
+            self.register_buffer("used", torch.tensor(np.load(self.remap)))
+            self.re_embed = self.used.shape[0]
+            self.unknown_index = unknown_index # "random" or "extra" or integer
+            if self.unknown_index == "extra":
+                self.unknown_index = self.re_embed
+                self.re_embed = self.re_embed+1
+            print(f"Remapping {self.n_embed} indices to {self.re_embed} indices. "
+                  f"Using {self.unknown_index} for unknown indices.")
+        else:
+            self.re_embed = n_embed
+
+    def remap_to_used(self, inds):
+        ishape = inds.shape
+        assert len(ishape)>1
+        inds = inds.reshape(ishape[0],-1)
+        used = self.used.to(inds)
+        match = (inds[:,:,None]==used[None,None,...]).long()
+        new = match.argmax(-1)
+        unknown = match.sum(2)<1
+        if self.unknown_index == "random":
+            new[unknown]=torch.randint(0,self.re_embed,size=new[unknown].shape).to(device=new.device)
+        else:
+            new[unknown] = self.unknown_index
+        return new.reshape(ishape)
+
+    def unmap_to_all(self, inds):
+        ishape = inds.shape
+        assert len(ishape)>1
+        inds = inds.reshape(ishape[0],-1)
+        used = self.used.to(inds)
+        if self.re_embed > self.used.shape[0]: # extra token
+            inds[inds>=self.used.shape[0]] = 0 # simply set to zero
+        back=torch.gather(used[None,:][inds.shape[0]*[0],:], 1, inds)
+        return back.reshape(ishape)
+
+    def forward(self, z, temp=None, return_logits=False):
+        # force hard = True when we are in eval mode, as we must quantize. actually, always true seems to work
+        hard = self.straight_through if self.training else True
+        temp = self.temperature if temp is None else temp
+
+        logits = self.proj(z)
+        if self.remap is not None:
+            # continue only with used logits
+            full_zeros = torch.zeros_like(logits)
+            logits = logits[:,self.used,...]
+
+        soft_one_hot = F.gumbel_softmax(logits, tau=temp, dim=1, hard=hard)
+        if self.remap is not None:
+            # go back to all entries but unused set to zero
+            full_zeros[:,self.used,...] = soft_one_hot
+            soft_one_hot = full_zeros
+        z_q = einsum('b n h w, n d -> b d h w', soft_one_hot, self.embed.weight)
+
+        # + kl divergence to the prior loss
+        qy = F.softmax(logits, dim=1)
+        diff = self.kl_weight * torch.sum(qy * torch.log(qy * self.n_embed + 1e-10), dim=1).mean()
+
+        ind = soft_one_hot.argmax(dim=1)
+        if self.remap is not None:
+            ind = self.remap_to_used(ind)
+        if self.use_vqinterface:
+            if return_logits:
+                return z_q, diff, (None, None, ind), logits
+            return z_q, diff, (None, None, ind)
+        return z_q, diff, ind
+
+    def get_codebook_entry(self, indices, shape):
+        b, h, w, c = shape
+        assert b*h*w == indices.shape[0]
+        indices = rearrange(indices, '(b h w) -> b h w', b=b, h=h, w=w)
+        if self.remap is not None:
+            indices = self.unmap_to_all(indices)
+        one_hot = F.one_hot(indices, num_classes=self.n_embed).permute(0, 3, 1, 2).float()
+        z_q = einsum('b n h w, n d -> b d h w', one_hot, self.embed.weight)
+        return z_q
+
+
+class VectorQuantizer2(nn.Module):
+    """
+    Improved version over VectorQuantizer, can be used as a drop-in replacement. Mostly
+    avoids costly matrix multiplications and allows for post-hoc remapping of indices.
+    """
+    # NOTE: due to a bug the beta term was applied to the wrong term. for
+    # backwards compatibility we use the buggy version by default, but you can
+    # specify legacy=False to fix it.
+    def __init__(self, n_e, e_dim, beta, remap=None, unknown_index="random",
+                 sane_index_shape=False, legacy=True):
+        super().__init__()
+        self.n_e = n_e
+        self.e_dim = e_dim
+        self.beta = beta
+        self.legacy = legacy
+
+        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)
+
+        self.remap = remap
+        if self.remap is not None:
+            self.register_buffer("used", torch.tensor(np.load(self.remap)))
+            self.re_embed = self.used.shape[0]
+            self.unknown_index = unknown_index # "random" or "extra" or integer
+            if self.unknown_index == "extra":
+                self.unknown_index = self.re_embed
+                self.re_embed = self.re_embed+1
+            print(f"Remapping {self.n_e} indices to {self.re_embed} indices. "
+                  f"Using {self.unknown_index} for unknown indices.")
+        else:
+            self.re_embed = n_e
+
+        self.sane_index_shape = sane_index_shape
+
+    def remap_to_used(self, inds):
+        ishape = inds.shape
+        assert len(ishape)>1
+        inds = inds.reshape(ishape[0],-1)
+        used = self.used.to(inds)
+        match = (inds[:,:,None]==used[None,None,...]).long()
+        new = match.argmax(-1)
+        unknown = match.sum(2)<1
+        if self.unknown_index == "random":
+            new[unknown]=torch.randint(0,self.re_embed,size=new[unknown].shape).to(device=new.device)
+        else:
+            new[unknown] = self.unknown_index
+        return new.reshape(ishape)
+
+    def unmap_to_all(self, inds):
+        ishape = inds.shape
+        assert len(ishape)>1
+        inds = inds.reshape(ishape[0],-1)
+        used = self.used.to(inds)
+        if self.re_embed > self.used.shape[0]: # extra token
+            inds[inds>=self.used.shape[0]] = 0 # simply set to zero
+        back=torch.gather(used[None,:][inds.shape[0]*[0],:], 1, inds)
+        return back.reshape(ishape)
+
+    def forward(self, z, temp=None, rescale_logits=False, return_logits=False):
+        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()
+        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
+
+        d = torch.sum(z_flattened ** 2, dim=1, keepdim=True) + \
+            torch.sum(self.embedding.weight**2, dim=1) - 2 * \
+            torch.einsum('bd,dn->bn', z_flattened, rearrange(self.embedding.weight, 'n d -> d n'))
+
+        min_encoding_indices = torch.argmin(d, dim=1)
+        z_q = self.embedding(min_encoding_indices).view(z.shape)
+        perplexity = None
+        min_encodings = None
+
+        # compute loss for embedding
+        if not self.legacy:
+            loss = self.beta * torch.mean((z_q.detach()-z)**2) + \
+                   torch.mean((z_q - z.detach()) ** 2)
+        else:
+            loss = torch.mean((z_q.detach()-z)**2) + self.beta * \
+                   torch.mean((z_q - z.detach()) ** 2)
+
+        # preserve gradients
+        z_q = z + (z_q - z).detach()
+
+        # reshape back to match original input shape
+        z_q = rearrange(z_q, 'b h w c -> b c h w').contiguous()
+
+        if self.remap is not None:
+            min_encoding_indices = min_encoding_indices.reshape(z.shape[0],-1) # add batch axis
+            min_encoding_indices = self.remap_to_used(min_encoding_indices)
+            min_encoding_indices = min_encoding_indices.reshape(-1,1) # flatten
+
+        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, loss, (perplexity, min_encodings, min_encoding_indices)
+
+    def get_codebook_entry(self, indices, shape):
+        # shape specifying (batch, height, width, channel)
+        if self.remap is not None:
+            indices = indices.reshape(shape[0],-1) # add batch axis
+            indices = self.unmap_to_all(indices)
+            indices = indices.reshape(-1) # flatten again
+
+        # get quantized latent vectors
+        z_q = self.embedding(indices)
+
+        if shape is not None:
+            z_q = z_q.view(shape)
+            # reshape back to match original input shape
+            z_q = z_q.permute(0, 3, 1, 2).contiguous()
+
+        return z_q
+
+class EmbeddingEMA(nn.Module):
+    def __init__(self, num_tokens, codebook_dim, decay=0.99, eps=1e-5):
+        super().__init__()
+        self.decay = decay
+        self.eps = eps        
+        weight = torch.randn(num_tokens, codebook_dim)
+        self.weight = nn.Parameter(weight, requires_grad = False)
+        self.cluster_size = nn.Parameter(torch.zeros(num_tokens), requires_grad = False)
+        self.embed_avg = nn.Parameter(weight.clone(), requires_grad = False)
+        self.update = True
+
+    def forward(self, embed_id):
+        return F.embedding(embed_id, self.weight)
+
+    def cluster_size_ema_update(self, new_cluster_size):
+        self.cluster_size.data.mul_(self.decay).add_(new_cluster_size, alpha=1 - self.decay)
+
+    def embed_avg_ema_update(self, new_embed_avg): 
+        self.embed_avg.data.mul_(self.decay).add_(new_embed_avg, alpha=1 - self.decay)
+
+    def weight_update(self, num_tokens):
+        n = self.cluster_size.sum()
+        smoothed_cluster_size = (
+                (self.cluster_size + self.eps) / (n + num_tokens * self.eps) * n
+            )
+        #normalize embedding average with smoothed cluster size
+        embed_normalized = self.embed_avg / smoothed_cluster_size.unsqueeze(1)
+        self.weight.data.copy_(embed_normalized)   
+
+
+class EMAVectorQuantizer(nn.Module):
+    def __init__(self, n_embed, embedding_dim, beta, decay=0.99, eps=1e-5,
+                remap=None, unknown_index="random"):
+        super().__init__()
+        self.codebook_dim = codebook_dim
+        self.num_tokens = num_tokens
+        self.beta = beta
+        self.embedding = EmbeddingEMA(self.num_tokens, self.codebook_dim, decay, eps)
+
+        self.remap = remap
+        if self.remap is not None:
+            self.register_buffer("used", torch.tensor(np.load(self.remap)))
+            self.re_embed = self.used.shape[0]
+            self.unknown_index = unknown_index # "random" or "extra" or integer
+            if self.unknown_index == "extra":
+                self.unknown_index = self.re_embed
+                self.re_embed = self.re_embed+1
+            print(f"Remapping {self.n_embed} indices to {self.re_embed} indices. "
+                  f"Using {self.unknown_index} for unknown indices.")
+        else:
+            self.re_embed = n_embed
+
+    def remap_to_used(self, inds):
+        ishape = inds.shape
+        assert len(ishape)>1
+        inds = inds.reshape(ishape[0],-1)
+        used = self.used.to(inds)
+        match = (inds[:,:,None]==used[None,None,...]).long()
+        new = match.argmax(-1)
+        unknown = match.sum(2)<1
+        if self.unknown_index == "random":
+            new[unknown]=torch.randint(0,self.re_embed,size=new[unknown].shape).to(device=new.device)
+        else:
+            new[unknown] = self.unknown_index
+        return new.reshape(ishape)
+
+    def unmap_to_all(self, inds):
+        ishape = inds.shape
+        assert len(ishape)>1
+        inds = inds.reshape(ishape[0],-1)
+        used = self.used.to(inds)
+        if self.re_embed > self.used.shape[0]: # extra token
+            inds[inds>=self.used.shape[0]] = 0 # simply set to zero
+        back=torch.gather(used[None,:][inds.shape[0]*[0],:], 1, inds)
+        return back.reshape(ishape)
+
+    def forward(self, z):
+        # reshape z -> (batch, height, width, channel) and flatten
+        #z, 'b c h w -> b h w c'
+        z = rearrange(z, 'b c h w -> b h w c')
+        z_flattened = z.reshape(-1, self.codebook_dim)
+        
+        # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
+        d = z_flattened.pow(2).sum(dim=1, keepdim=True) + \
+            self.embedding.weight.pow(2).sum(dim=1) - 2 * \
+            torch.einsum('bd,nd->bn', z_flattened, self.embedding.weight) # 'n d -> d n'
+
+
+        encoding_indices = torch.argmin(d, dim=1)
+
+        z_q = self.embedding(encoding_indices).view(z.shape)
+        encodings = F.one_hot(encoding_indices, self.num_tokens).type(z.dtype)     
+        avg_probs = torch.mean(encodings, dim=0)
+        perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))
+
+        if self.training and self.embedding.update:
+            #EMA cluster size
+            encodings_sum = encodings.sum(0)            
+            self.embedding.cluster_size_ema_update(encodings_sum)
+            #EMA embedding average
+            embed_sum = encodings.transpose(0,1) @ z_flattened            
+            self.embedding.embed_avg_ema_update(embed_sum)
+            #normalize embed_avg and update weight
+            self.embedding.weight_update(self.num_tokens)
+
+        # compute loss for embedding
+        loss = self.beta * F.mse_loss(z_q.detach(), z) 
+
+        # preserve gradients
+        z_q = z + (z_q - z).detach()
+
+        # reshape back to match original input shape
+        #z_q, 'b h w c -> b c h w'
+        z_q = rearrange(z_q, 'b h w c -> b c h w')
+        return z_q, loss, (perplexity, encodings, encoding_indices)

taming/util.py

@@ -0,0 +1,157 @@
+import os, hashlib
+import requests
+from tqdm import tqdm
+
+URL_MAP = {
+    "vgg_lpips": "https://heibox.uni-heidelberg.de/f/607503859c864bc1b30b/?dl=1"
+}
+
+CKPT_MAP = {
+    "vgg_lpips": "vgg.pth"
+}
+
+MD5_MAP = {
+    "vgg_lpips": "d507d7349b931f0638a25a48a722f98a"
+}
+
+
+def download(url, local_path, chunk_size=1024):
+    os.makedirs(os.path.split(local_path)[0], exist_ok=True)
+    with requests.get(url, stream=True) as r:
+        total_size = int(r.headers.get("content-length", 0))
+        with tqdm(total=total_size, unit="B", unit_scale=True) as pbar:
+            with open(local_path, "wb") as f:
+                for data in r.iter_content(chunk_size=chunk_size):
+                    if data:
+                        f.write(data)
+                        pbar.update(chunk_size)
+
+
+def md5_hash(path):
+    with open(path, "rb") as f:
+        content = f.read()
+    return hashlib.md5(content).hexdigest()
+
+
+def get_ckpt_path(name, root, check=False):
+    assert name in URL_MAP
+    path = os.path.join(root, CKPT_MAP[name])
+    if not os.path.exists(path) or (check and not md5_hash(path) == MD5_MAP[name]):
+        print("Downloading {} model from {} to {}".format(name, URL_MAP[name], path))
+        download(URL_MAP[name], path)
+        md5 = md5_hash(path)
+        assert md5 == MD5_MAP[name], md5
+    return path
+
+
+class KeyNotFoundError(Exception):
+    def __init__(self, cause, keys=None, visited=None):
+        self.cause = cause
+        self.keys = keys
+        self.visited = visited
+        messages = list()
+        if keys is not None:
+            messages.append("Key not found: {}".format(keys))
+        if visited is not None:
+            messages.append("Visited: {}".format(visited))
+        messages.append("Cause:\n{}".format(cause))
+        message = "\n".join(messages)
+        super().__init__(message)
+
+
+def retrieve(
+    list_or_dict, key, splitval="/", default=None, expand=True, pass_success=False
+):
+    """Given a nested list or dict return the desired value at key expanding
+    callable nodes if necessary and :attr:`expand` is ``True``. The expansion
+    is done in-place.
+
+    Parameters
+    ----------
+        list_or_dict : list or dict
+            Possibly nested list or dictionary.
+        key : str
+            key/to/value, path like string describing all keys necessary to
+            consider to get to the desired value. List indices can also be
+            passed here.
+        splitval : str
+            String that defines the delimiter between keys of the
+            different depth levels in `key`.
+        default : obj
+            Value returned if :attr:`key` is not found.
+        expand : bool
+            Whether to expand callable nodes on the path or not.
+
+    Returns
+    -------
+        The desired value or if :attr:`default` is not ``None`` and the
+        :attr:`key` is not found returns ``default``.
+
+    Raises
+    ------
+        Exception if ``key`` not in ``list_or_dict`` and :attr:`default` is
+        ``None``.
+    """
+
+    keys = key.split(splitval)
+
+    success = True
+    try:
+        visited = []
+        parent = None
+        last_key = None
+        for key in keys:
+            if callable(list_or_dict):
+                if not expand:
+                    raise KeyNotFoundError(
+                        ValueError(
+                            "Trying to get past callable node with expand=False."
+                        ),
+                        keys=keys,
+                        visited=visited,
+                    )
+                list_or_dict = list_or_dict()
+                parent[last_key] = list_or_dict
+
+            last_key = key
+            parent = list_or_dict
+
+            try:
+                if isinstance(list_or_dict, dict):
+                    list_or_dict = list_or_dict[key]
+                else:
+                    list_or_dict = list_or_dict[int(key)]
+            except (KeyError, IndexError, ValueError) as e:
+                raise KeyNotFoundError(e, keys=keys, visited=visited)
+
+            visited += [key]
+        # final expansion of retrieved value
+        if expand and callable(list_or_dict):
+            list_or_dict = list_or_dict()
+            parent[last_key] = list_or_dict
+    except KeyNotFoundError as e:
+        if default is None:
+            raise e
+        else:
+            list_or_dict = default
+            success = False
+
+    if not pass_success:
+        return list_or_dict
+    else:
+        return list_or_dict, success
+
+
+if __name__ == "__main__":
+    config = {"keya": "a",
+              "keyb": "b",
+              "keyc":
+                  {"cc1": 1,
+                   "cc2": 2,
+                   }
+              }
+    from omegaconf import OmegaConf
+    config = OmegaConf.create(config)
+    print(config)
+    retrieve(config, "keya")
+