[WIP] scheduler

Author: williamberman

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

src/diffusers/models/vq_diffusion_attention.py

@@ -2,10 +2,10 @@
 
 import torch
 from torch import nn
+import torch.nn.functional as F
 
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.modeling_utils import ModelMixin
-from diffusers.models.embeddings import DalleMaskImageEmbedding
 
 from .attention import CrossAttention
 
@@ -23,20 +23,27 @@ def __init__(
         width: int,
         diffusion_steps: int,
         dropout: float = 0.0,
+        min_logged_value: float = -70.0
     ):
         super().__init__()
+
         self.n_heads = n_heads
         self.d_head = d_head
-        inner_dim = n_heads * d_head
+        self.inner_dim = n_heads * d_head
+        self.min_logged_value = min_logged_value
+        self.num_embed = num_embed
+        self.height = height
+        self.width = width
+        self.num_latent_pixels = self.height * self.width
 
         self.latent_image_embedding = DalleMaskImageEmbedding(
-            num_embed=num_embed, embed_dim=inner_dim, height=height, width=width
+            num_embed=self.num_embed, embed_dim=self.inner_dim, height=height, width=width
         )
 
         self.transformer_blocks = nn.ModuleList(
             [
                 BasicTransformerBlock(
-                    inner_dim,
+                    self.inner_dim,
                     n_heads,
                     d_head,
                     dropout=dropout,
@@ -48,21 +55,80 @@ def __init__(
             ]
         )
 
-        self.norm_out = nn.LayerNorm(inner_dim)
-        self.out = nn.Linear(inner_dim, num_embed)
+        self.norm_out = nn.LayerNorm(self.inner_dim)
+
+        # The output from the transformer is the embedding indices for the 
+        # quantized codebook. the output dimension is `num_embed - 1` because 
+        # it does not include additional index for the masked value since 
+        # the transformer predicts the unnoised image which has no masks
+        self.out = nn.Linear(self.inner_dim, self.num_embed - 1)
 
     def forward(self, latent_images, cond_emb, t):
+        # bsz = latent_images.shape[0]
+
         embedded_latent_images = self.latent_image_embedding(latent_images)
         hidden_states = embedded_latent_images
 
         for block in self.transformer_blocks:
             hidden_states = block(hidden_states, cond_emb, t)
 
         logits = self.out(self.norm_out(hidden_states))
-        out = logits.permute(0, 2, 1)
+        # (batch, self.num_embed - 1, self.num_latent_pixels)
+        logits = logits.permute(0, 2, 1)
+
+        # TODO remove
+        torch.save(logits, f"/content/diffusers-out/transformer_logits_diffusers-{t[0]}.pt")
+
+        # TODO document why we append the zero vector
+        # equivalent to `torch.zeros((bsz, 1, self.num_latent_pixels)).log().clamp(self.min_logged_value)`
+        # log_zero_vector = torch.full((bsz, 1, self.num_latent_pixels), self.min_logged_value, device=logits.device)
+
+        log_p_x_0 = F.log_softmax(logits.double(), dim=1).float().clamp(self.min_logged_value)
+
+        # (batch, self.num_embed, self.inner_dim)
+        # log_p_x_0 = torch.cat((log_p_x_0, log_zero_vector), dim=1)
+
+        return log_p_x_0
+
+
+# TODO(will) - document this
+class DalleMaskImageEmbedding(nn.Module):
+    def __init__(
+        self,
+        num_embed,
+        height,
+        width,
+        embed_dim,
+    ):
+        super().__init__()
+
+        self.height = height
+        self.width = width
+        self.num_embed = num_embed
+        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):
+        emb = self.emb(index)
+
+        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
 
-        return out
+        emb = emb + pos_emb[:, : emb.shape[1], :]
 
+        return emb
 
 class BasicTransformerBlock(nn.Module):
     def __init__(