ApoorvBrooklyn/stable-diffusion-implementation

1import torch2from torch import nn3from torch.nn import functional as F4import math5 6class SelfAttention(nn.Module):7    def __init__(self, n_heads, d_embed, in_proj_bias=True, out_proj_bias=True):8        super().__init__()9        # This combines the Wq, Wk and Wv matrices into one matrix10        self.in_proj = nn.Linear(d_embed, 3 * d_embed, bias=in_proj_bias)11        # This one represents the Wo matrix12        self.out_proj = nn.Linear(d_embed, d_embed, bias=out_proj_bias)13        self.n_heads = n_heads14        self.d_head = d_embed // n_heads15 16    def forward(self, x, causal_mask=False):17        # x: # (Batch_Size, Seq_Len, Dim)18 19        # (Batch_Size, Seq_Len, Dim)20        input_shape = x.shape 21        22        # (Batch_Size, Seq_Len, Dim)23        batch_size, sequence_length, d_embed = input_shape 24 25        # (Batch_Size, Seq_Len, H, Dim / H)26        interim_shape = (batch_size, sequence_length, self.n_heads, self.d_head) 27 28        # (Batch_Size, Seq_Len, Dim) -> (Batch_Size, Seq_Len, Dim * 3) -> 3 tensor of shape (Batch_Size, Seq_Len, Dim)29        q, k, v = self.in_proj(x).chunk(3, dim=-1)30        31        # (Batch_Size, Seq_Len, Dim) -> (Batch_Size, Seq_Len, H, Dim / H) -> (Batch_Size, H, Seq_Len, Dim / H)32        q = q.view(interim_shape).transpose(1, 2)33        k = k.view(interim_shape).transpose(1, 2)34        v = v.view(interim_shape).transpose(1, 2)35 36        # (Batch_Size, H, Seq_Len, Dim / H) @ (Batch_Size, H, Dim / H, Seq_Len) -> (Batch_Size, H, Seq_Len, Seq_Len)37        weight = q @ k.transpose(-1, -2)38        39        if causal_mask:40            # Mask where the upper triangle (above the principal diagonal) is 141            mask = torch.ones_like(weight, dtype=torch.bool).triu(1) 42            # Fill the upper triangle with -inf43            weight.masked_fill_(mask, -torch.inf) 44        45        # Divide by d_k (Dim / H). 46        # (Batch_Size, H, Seq_Len, Seq_Len) -> (Batch_Size, H, Seq_Len, Seq_Len)47        weight /= math.sqrt(self.d_head) 48 49        # (Batch_Size, H, Seq_Len, Seq_Len) -> (Batch_Size, H, Seq_Len, Seq_Len)50        weight = F.softmax(weight, dim=-1) 51 52        # (Batch_Size, H, Seq_Len, Seq_Len) @ (Batch_Size, H, Seq_Len, Dim / H) -> (Batch_Size, H, Seq_Len, Dim / H)53        output = weight @ v54 55        # (Batch_Size, H, Seq_Len, Dim / H) -> (Batch_Size, Seq_Len, H, Dim / H)56        output = output.transpose(1, 2) 57 58        # (Batch_Size, Seq_Len, H, Dim / H) -> (Batch_Size, Seq_Len, Dim)59        output = output.reshape(input_shape) 60 61        # (Batch_Size, Seq_Len, Dim) -> (Batch_Size, Seq_Len, Dim)62        output = self.out_proj(output) 63        64        # (Batch_Size, Seq_Len, Dim)65        return output66 67class CrossAttention(nn.Module):68    def __init__(self, n_heads, d_embed, d_cross, in_proj_bias=True, out_proj_bias=True):69        super().__init__()70        self.q_proj   = nn.Linear(d_embed, d_embed, bias=in_proj_bias)71        self.k_proj   = nn.Linear(d_cross, d_embed, bias=in_proj_bias)72        self.v_proj   = nn.Linear(d_cross, d_embed, bias=in_proj_bias)73        self.out_proj = nn.Linear(d_embed, d_embed, bias=out_proj_bias)74        self.n_heads = n_heads75        self.d_head = d_embed // n_heads76    77    def forward(self, x, y):78        # x (latent): # (Batch_Size, Seq_Len_Q, Dim_Q)79        # y (context): # (Batch_Size, Seq_Len_KV, Dim_KV) = (Batch_Size, 77, 768)80 81        input_shape = x.shape82        batch_size, sequence_length, d_embed = input_shape83        # Divide each embedding of Q into multiple heads such that d_heads * n_heads = Dim_Q84        interim_shape = (batch_size, -1, self.n_heads, self.d_head)85        86        # (Batch_Size, Seq_Len_Q, Dim_Q) -> (Batch_Size, Seq_Len_Q, Dim_Q)87        q = self.q_proj(x)88        # (Batch_Size, Seq_Len_KV, Dim_KV) -> (Batch_Size, Seq_Len_KV, Dim_Q)89        k = self.k_proj(y)90        # (Batch_Size, Seq_Len_KV, Dim_KV) -> (Batch_Size, Seq_Len_KV, Dim_Q)91        v = self.v_proj(y)92 93        # (Batch_Size, Seq_Len_Q, Dim_Q) -> (Batch_Size, Seq_Len_Q, H, Dim_Q / H) -> (Batch_Size, H, Seq_Len_Q, Dim_Q / H)94        q = q.view(interim_shape).transpose(1, 2) 95        # (Batch_Size, Seq_Len_KV, Dim_Q) -> (Batch_Size, Seq_Len_KV, H, Dim_Q / H) -> (Batch_Size, H, Seq_Len_KV, Dim_Q / H)96        k = k.view(interim_shape).transpose(1, 2) 97        # (Batch_Size, Seq_Len_KV, Dim_Q) -> (Batch_Size, Seq_Len_KV, H, Dim_Q / H) -> (Batch_Size, H, Seq_Len_KV, Dim_Q / H)98        v = v.view(interim_shape).transpose(1, 2) 99        100        # (Batch_Size, H, Seq_Len_Q, Dim_Q / H) @ (Batch_Size, H, Dim_Q / H, Seq_Len_KV) -> (Batch_Size, H, Seq_Len_Q, Seq_Len_KV)101        weight = q @ k.transpose(-1, -2)102        103        # (Batch_Size, H, Seq_Len_Q, Seq_Len_KV)104        weight /= math.sqrt(self.d_head)105        106        # (Batch_Size, H, Seq_Len_Q, Seq_Len_KV)107        weight = F.softmax(weight, dim=-1)108        109        # (Batch_Size, H, Seq_Len_Q, Seq_Len_KV) @ (Batch_Size, H, Seq_Len_KV, Dim_Q / H) -> (Batch_Size, H, Seq_Len_Q, Dim_Q / H)110        output = weight @ v111        112        # (Batch_Size, H, Seq_Len_Q, Dim_Q / H) -> (Batch_Size, Seq_Len_Q, H, Dim_Q / H)113        output = output.transpose(1, 2).contiguous()114        115        # (Batch_Size, Seq_Len_Q, H, Dim_Q / H) -> (Batch_Size, Seq_Len_Q, Dim_Q)116        output = output.view(input_shape)117        118        # (Batch_Size, Seq_Len_Q, Dim_Q) -> (Batch_Size, Seq_Len_Q, Dim_Q)119        output = self.out_proj(output)120 121        # (Batch_Size, Seq_Len_Q, Dim_Q)122        return output