Setup diffusion modules.

src/skelcast/models/diffusers/__init__.py

@@ -0,0 +1 @@
+from .embedding import DiffusionEmbedding

src/skelcast/models/diffusers/embedding.py

@@ -0,0 +1,58 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DiffusionEmbedding(nn.Module):
+    def __init__(self, num_steps, embedding_dim=128, projection_dim=None):
+        # Codes from CSDI
+        super().__init__()
+        if projection_dim is None:
+            projection_dim = embedding_dim
+        self.register_buffer(
+            "embedding",
+            self._build_embedding(num_steps, embedding_dim / 2),
+            persistent=False,
+        )
+        self.projection1 = nn.Linear(embedding_dim, projection_dim)
+        self.projection2 = nn.Linear(projection_dim, projection_dim)
+
+    def forward(self, diffusion_step):
+        x = self.embedding[diffusion_step]
+        x = self.projection1(x)
+        x = F.silu(x)
+        x = self.projection2(x)
+        x = F.silu(x)
+        return x
+
+    def _build_embedding(self, num_steps, dim=64):
+        steps = torch.arange(num_steps).unsqueeze(1)  # (T,1)
+        frequencies = 10.0 ** (torch.arange(dim) / (dim - 1) * 4.0).unsqueeze(
+            0
+        )  # (1,dim)
+        table = steps * frequencies  # (T,dim)
+        table = torch.cat([torch.sin(table), torch.cos(table)], dim=1)  # (T,dim*2)
+        return table
+
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 5000):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        position = torch.arange(max_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
+        pe = torch.zeros(max_len, 1, d_model)
+        pe[:, 0, 0::2] = torch.sin(position * div_term)
+        pe[:, 0, 1::2] = torch.cos(position * div_term)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        """
+        Args:
+            x: Tensor, shape [seq_len, batch_size, embedding_dim]
+        """
+        x = x + self.pe[:x.size(0)]
+        return self.dropout(x)

src/skelcast/models/diffusers/series_diffuser.py

@@ -0,0 +1,51 @@
+import torch
+import torch.nn as nn
+
+from skelcast.models.diffusers.embedding import DiffusionEmbedding, PositionalEncoding
+
+
+class Series_Denoiser(nn.Module):
+    def __init__(self, input_dim, qkv_dim, num_layers, num_heads, prefix_len, pred_len, diff_steps):
+        super().__init__()
+
+        self.input_dim = input_dim
+        self.qkv_dim = qkv_dim
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.prefix_len = prefix_len
+        self.pred_len = pred_len
+        self.diff_steps = diff_steps
+
+        self.spatial_layer = nn.TransformerEncoderLayer(d_model=qkv_dim, nhead=num_heads,
+                                                    dim_feedforward=qkv_dim, activation='gelu')
+        self.spatial_TF = nn.TransformerEncoder(self.spatial_layer, num_layers=num_layers)
+        self.spatial_inp_fc = nn.Linear(prefix_len+pred_len, qkv_dim)
+        self.spatial_out_fc = nn.Linear(qkv_dim, prefix_len+pred_len)
+
+        self.temporal_layer = nn.TransformerEncoderLayer(d_model=qkv_dim, nhead=num_heads,
+                                                    dim_feedforward=qkv_dim, activation='gelu')
+        self.temporal_TF = nn.TransformerEncoder(self.temporal_layer, num_layers=num_layers)
+        self.temporal_inp_fc = nn.Linear(input_dim, qkv_dim)
+        self.temporal_out_fc = nn.Linear(qkv_dim, input_dim)
+
+        self.pos_encoder = PositionalEncoding(qkv_dim)
+
+        self.step_temporal_encoder = DiffusionEmbedding(diff_steps, qkv_dim)
+        self.step_spatial_encoder = DiffusionEmbedding(diff_steps, qkv_dim)
+
+    def forward(self, noise, x, t):
+        step_spatial_embed = self.step_spatial_encoder(t)
+        step_temporal_embed = self.step_temporal_encoder(t)
+        window = torch.cat([x, noise], dim=0)
+
+        spatial = self.spatial_inp_fc(window.permute(2, 1, 0)) + step_spatial_embed # L B D -> D B L 
+        spatial = self.pos_encoder(spatial)
+        spatial = self.spatial_TF(spatial) 
+        spatial = self.spatial_out_fc(spatial).permute(2, 1, 0)
+
+        temporal = self.temporal_inp_fc(spatial) + step_temporal_embed
+        temporal = self.pos_encoder(temporal)
+        temporal = self.temporal_TF(temporal)
+        temporal = self.temporal_out_fc(temporal)
+
+        return temporal[x.shape[0]:]