[WIP] scheduler

Author: williamberman

orig_scheduler.py

@@ -0,0 +1,155 @@
+import torch
+import numpy as np
+
+
+def log_1_min_a(a):
+    return torch.log(1 - a.exp() + 1e-40)
+
+def log_add_exp(a, b):
+    maximum = torch.max(a, b)
+    return maximum + torch.log(torch.exp(a - maximum) + torch.exp(b - maximum))
+
+def extract(a, t, x_shape):
+    b, *_ = t.shape
+    out = a.gather(-1, t)
+    return out.reshape(b, *((1,) * (len(x_shape) - 1)))
+
+def log_categorical(log_x_start, log_prob):
+    return (log_x_start.exp() * log_prob).sum(dim=1)
+
+def index_to_log_onehot(x, num_classes):
+    assert x.max().item() < num_classes, \
+        f'Error: {x.max().item()} >= {num_classes}'
+    x_onehot = F.one_hot(x, num_classes)
+    permute_order = (0, -1) + tuple(range(1, len(x.size())))
+    x_onehot = x_onehot.permute(permute_order)
+    log_x = torch.log(x_onehot.float().clamp(min=1e-30))
+    return log_x
+
+def log_onehot_to_index(log_x):
+    return log_x.argmax(1)
+
+def alpha_schedule(time_step, N=100, att_1 = 0.99999, att_T = 0.000009, ctt_1 = 0.000009, ctt_T = 0.99999):
+    att = np.arange(0, time_step)/(time_step-1)*(att_T - att_1) + att_1
+    att = np.concatenate(([1], att))
+    at = att[1:]/att[:-1]
+    ctt = np.arange(0, time_step)/(time_step-1)*(ctt_T - ctt_1) + ctt_1
+    ctt = np.concatenate(([0], ctt))
+    one_minus_ctt = 1 - ctt
+    one_minus_ct = one_minus_ctt[1:] / one_minus_ctt[:-1]
+    ct = 1-one_minus_ct
+    bt = (1-at-ct)/N
+    att = np.concatenate((att[1:], [1]))
+    ctt = np.concatenate((ctt[1:], [0]))
+    btt = (1-att-ctt)/N
+    return at, bt, ct, att, btt, ctt
+
+
+class OrigScheduler:
+    def __init__(self, *, num_classes, content_seq_len, num_timesteps=100):
+        self.num_timesteps = num_timesteps
+        self.num_classes = num_classes
+        self.content_seq_len = content_seq_len
+
+        at, bt, ct, att, btt, ctt = alpha_schedule(self.num_timesteps, N=self.num_classes-1)
+
+        at = torch.tensor(at.astype('float64'))
+        bt = torch.tensor(bt.astype('float64'))
+        ct = torch.tensor(ct.astype('float64'))
+        log_at = torch.log(at)
+        log_bt = torch.log(bt)
+        log_ct = torch.log(ct)
+        att = torch.tensor(att.astype('float64'))
+        btt = torch.tensor(btt.astype('float64'))
+        ctt = torch.tensor(ctt.astype('float64'))
+        log_cumprod_at = torch.log(att)
+        log_cumprod_bt = torch.log(btt)
+        log_cumprod_ct = torch.log(ctt)
+
+        log_1_min_ct = log_1_min_a(log_ct)
+        log_1_min_cumprod_ct = log_1_min_a(log_cumprod_ct)
+
+        assert log_add_exp(log_ct, log_1_min_ct).abs().sum().item() < 1.e-5
+        assert log_add_exp(log_cumprod_ct, log_1_min_cumprod_ct).abs().sum().item() < 1.e-5
+
+        # Convert to float32 and register buffers.
+        self.log_at = log_at.float()
+        self.log_bt = log_bt.float()
+        self.log_ct = log_ct.float()
+        self.log_cumprod_at = log_cumprod_at.float()
+        self.log_cumprod_bt = log_cumprod_bt.float()
+        self.log_cumprod_ct = log_cumprod_ct.float()
+        self.log_1_min_ct = log_1_min_ct.float()
+        self.log_1_min_cumprod_ct = log_1_min_cumprod_ct.float()
+
+
+
+    def q_posterior(self, log_x_start, log_x_t, t):            # p_theta(xt_1|xt) = sum(q(xt-1|xt,x0')*p(x0'))
+        # notice that log_x_t is onehot
+        assert t.min().item() >= 0 and t.max().item() < self.num_timesteps
+        batch_size = log_x_start.size()[0]
+        onehot_x_t = log_onehot_to_index(log_x_t)
+        mask = (onehot_x_t == self.num_classes-1).unsqueeze(1) 
+        log_one_vector = torch.zeros(batch_size, 1, 1).type_as(log_x_t)
+        log_zero_vector = torch.log(log_one_vector+1.0e-30).expand(-1, -1, self.content_seq_len)
+
+        log_qt = self.q_pred(log_x_t, t)                                  # q(xt|x0)
+        # log_qt = torch.cat((log_qt[:,:-1,:], log_zero_vector), dim=1)
+        log_qt = log_qt[:,:-1,:]
+        log_cumprod_ct = extract(self.log_cumprod_ct, t, log_x_start.shape)         # ct~
+        ct_cumprod_vector = log_cumprod_ct.expand(-1, self.num_classes-1, -1)
+        # ct_cumprod_vector = torch.cat((ct_cumprod_vector, log_one_vector), dim=1)
+        log_qt = (~mask)*log_qt + mask*ct_cumprod_vector
+        
+
+        log_qt_one_timestep = self.q_pred_one_timestep(log_x_t, t)        # q(xt|xt_1)
+        log_qt_one_timestep = torch.cat((log_qt_one_timestep[:,:-1,:], log_zero_vector), dim=1)
+        log_ct = extract(self.log_ct, t, log_x_start.shape)         # ct
+        ct_vector = log_ct.expand(-1, self.num_classes-1, -1)
+        ct_vector = torch.cat((ct_vector, log_one_vector), dim=1)
+        log_qt_one_timestep = (~mask)*log_qt_one_timestep + mask*ct_vector
+        
+        # log_x_start = torch.cat((log_x_start, log_zero_vector), dim=1)
+        # q = log_x_start - log_qt
+        q = log_x_start[:,:-1,:] - log_qt
+        q = torch.cat((q, log_zero_vector), dim=1)
+        q_log_sum_exp = torch.logsumexp(q, dim=1, keepdim=True)
+        q = q - q_log_sum_exp
+        log_EV_xtmin_given_xt_given_xstart = self.q_pred(q, t-1) + log_qt_one_timestep + q_log_sum_exp
+        return torch.clamp(log_EV_xtmin_given_xt_given_xstart, -70, 0)
+
+
+    def q_pred_one_timestep(self, log_x_t, t):         # q(xt|xt_1)
+        log_at = extract(self.log_at, t, log_x_t.shape)             # at
+        log_bt = extract(self.log_bt, t, log_x_t.shape)             # bt
+        log_ct = extract(self.log_ct, t, log_x_t.shape)             # ct
+        log_1_min_ct = extract(self.log_1_min_ct, t, log_x_t.shape)          # 1-ct
+
+        log_probs = torch.cat(
+            [
+                log_add_exp(log_x_t[:,:-1,:]+log_at, log_bt),
+                log_add_exp(log_x_t[:, -1:, :] + log_1_min_ct, log_ct)
+            ],
+            dim=1
+        )
+
+        return log_probs
+
+    def q_pred(self, log_x_start, t):           # q(xt|x0)
+        # log_x_start can be onehot or not
+        t = (t + (self.num_timesteps + 1))%(self.num_timesteps + 1)
+        log_cumprod_at = extract(self.log_cumprod_at, t, log_x_start.shape)         # at~
+        log_cumprod_bt = extract(self.log_cumprod_bt, t, log_x_start.shape)         # bt~
+        log_cumprod_ct = extract(self.log_cumprod_ct, t, log_x_start.shape)         # ct~
+        log_1_min_cumprod_ct = extract(self.log_1_min_cumprod_ct, t, log_x_start.shape)       # 1-ct~
+        
+
+        log_probs = torch.cat(
+            [
+                log_add_exp(log_x_start[:,:-1,:]+log_cumprod_at, log_cumprod_bt),
+                log_add_exp(log_x_start[:,-1:,:]+log_1_min_cumprod_ct, log_cumprod_ct)
+            ],
+            dim=1
+        )
+
+        return log_probs

scripts/convert_vq_diffusion_to_diffusers.py

@@ -32,9 +32,8 @@
 
 import yaml
 from accelerate import init_empty_weights, load_checkpoint_and_dispatch
-from diffusers import VQModel
+from diffusers import VQModel, VQDiffusionPipeline, VQDiffusionScheduler
 from diffusers.models.vq_diffusion_attention import VQDiffusionTransformer
-from diffusers.pipelines import VQDiffusionPipeline
 from transformers import CLIPTextModel, CLIPTokenizer
 from yaml.loader import FullLoader
 
@@ -492,7 +491,12 @@ def transformer_model_from_original_config(
 
     depth = original_transformer_config["n_layer"]
     context_dim = original_transformer_config["condition_dim"]
+
     num_embed = original_content_embedding_config["num_embed"]
+    # the number of embeddings in the transformer includes the mask embedding.
+    # the content embedding (the vqvae) does not include the mask embedding.
+    num_embed = num_embed + 1
+
     height = original_transformer_config["content_spatial_size"][0]
     width = original_transformer_config["content_spatial_size"][1]
     dropout = original_transformer_config["resid_pdrop"]
@@ -846,10 +850,23 @@ def read_config_file(filename):
 
     # done text encoder
 
+    # scheduler
+
+    scheduler_model = VQDiffusionScheduler(
+        # the scheduler has the same number of embeddings as the transformer
+        num_embed=transformer_model.num_embed
+    )
+
+    # done scheduler
+
     print(f"saving VQ diffusion model, path: {args.dump_path}")
 
     pipe = VQDiffusionPipeline(
-        vqvae=vqvae_model, transformer=transformer_model, tokenizer=tokenizer_model, text_encoder=text_encoder_model
+        vqvae=vqvae_model,
+        transformer=transformer_model,
+        tokenizer=tokenizer_model,
+        text_encoder=text_encoder_model,
+        scheduler=scheduler_model,
     )
     pipe.save_pretrained(args.dump_path)
 

src/diffusers/__init__.py

@@ -29,14 +29,15 @@
         get_scheduler,
     )
     from .pipeline_utils import DiffusionPipeline
-    from .pipelines import DDIMPipeline, DDPMPipeline, KarrasVePipeline, LDMPipeline, PNDMPipeline, ScoreSdeVePipeline
+    from .pipelines import DDIMPipeline, DDPMPipeline, KarrasVePipeline, LDMPipeline, PNDMPipeline, ScoreSdeVePipeline, VQDiffusionPipeline
     from .schedulers import (
         DDIMScheduler,
         DDPMScheduler,
         KarrasVeScheduler,
         PNDMScheduler,
         SchedulerMixin,
         ScoreSdeVeScheduler,
+        VQDiffusionScheduler
     )
     from .training_utils import EMAModel
 else:

src/diffusers/models/embeddings.py

@@ -115,52 +115,3 @@ def forward(self, x):
         x_proj = x[:, None] * self.weight[None, :] * 2 * np.pi
         out = torch.cat([torch.sin(x_proj), torch.cos(x_proj)], dim=-1)
         return out
-
-
-# TODO(will) - document this. check if throwing errors internally is appropriate
-class DalleMaskImageEmbedding(nn.Module):
-    def __init__(
-        self,
-        num_embed,
-        height,
-        width,
-        embed_dim,
-    ):
-        super().__init__()
-
-        self.height = height
-        self.width = width
-        # TODO(will) add docs on why this is incremented by 1. (Has to do with mask?)
-        self.num_embed = num_embed + 1
-        self.embed_dim = embed_dim
-
-        self.emb = nn.Embedding(self.num_embed, embed_dim)
-        self.height_emb = nn.Embedding(self.height, embed_dim)
-        self.width_emb = nn.Embedding(self.width, embed_dim)
-
-    def forward(self, index):
-        assert index.dim() == 2  # B x L
-        try:
-            index[index < 0] = 0
-            emb = self.emb(index)
-        except:
-            raise RuntimeError(
-                "IndexError: index out of range in self, max index {}, num embed {}".format(
-                    index.max(), self.num_embed
-                )
-            )
-
-        # add col and row embedding
-        if emb.shape[1] > 0:
-            height_emb = self.height_emb(
-                torch.arange(self.height, device=index.device).view(1, self.height)
-            ).unsqueeze(
-                2
-            )  # 1 x H x D -> 1 x H x 1 x D
-            width_emb = self.width_emb(torch.arange(self.width, device=index.device).view(1, self.width)).unsqueeze(
-                1
-            )  # 1 x W x D -> 1 x 1 x W x D
-            pos_emb = (height_emb + width_emb).view(1, self.height * self.width, -1)  # 1 x H x W x D -> 1 x L xD
-            emb = emb + pos_emb[:, : emb.shape[1], :]
-
-        return emb