Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import math | |
| from transformers import CLIPProcessor, CLIPModel, CLIPTextModelWithProjection | |
| class JointAttention(nn.Module): | |
| def __init__(self, d_model=256, image_tokens=256, prompt_tokens=1, n_head=1): | |
| super().__init__() | |
| self.dim = d_model | |
| self.nhead = n_head | |
| self.image_q, self.image_v, self.image_k = nn.Linear(d_model, d_model), nn.Linear(d_model, d_model), nn.Linear(d_model, d_model) | |
| self.prompt_q, self.prompt_v, self.prompt_k = nn.Linear(d_model, d_model), nn.Linear(d_model, d_model), nn.Linear(d_model, d_model) | |
| self.mlp = nn.Sequential( | |
| nn.Linear(d_model, d_model*4), | |
| nn.ReLU(), | |
| nn.Linear(d_model*4, d_model) | |
| ) | |
| self.o_proj = nn.Linear(d_model, d_model) | |
| half_dim = d_model //2 | |
| freq = torch.exp( | |
| -math.log(10000) * torch.arange(0, half_dim, dtype=torch.float32) / half_dim | |
| ) | |
| self.register_buffer("timestep_freq", freq) | |
| self.image_emb = nn.Embedding(image_tokens,self.dim) | |
| self.register_buffer("pos_ids", torch.arange(image_tokens)) | |
| self.ln1 = nn.LayerNorm(self.dim) | |
| self.ln2 = nn.LayerNorm(self.dim) | |
| self.image_tokens = image_tokens | |
| # self.image_proj = nn.Linear(4*32*32//image_tokens,d_model) | |
| # self.prompt_proj = nn.Linear(512,d_model) | |
| # self.image_up = nn.Linear(d_model, 4*32*32//image_tokens) | |
| # self.prompt_up = nn.Linear(d_model, 512) | |
| self.prompt_tokens = prompt_tokens | |
| def timestep_embedding(self,t): | |
| freqs = self.timestep_freq.unsqueeze(0) | |
| # emb = t.unsqueeze(1) * freqs | |
| emb = t * freqs * 1000 | |
| emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) | |
| return emb | |
| def forward(self, image, prompt, timestep): | |
| # image shape: b,256,512 | |
| # prompt shape: b,1,512 | |
| timestep_emb = self.timestep_embedding(timestep) | |
| # timestep emb shape: b,1,512 (dim -1 follows d_model, no proj. required) | |
| # image = self.image_proj(image) | |
| # prompt = self.prompt_proj(prompt) | |
| # image: b,256,dim | |
| # propmt: b,1,dim | |
| # print(f"sanity: image: {image.shape}, prompt: {prompt.shape}, timestep_emb: {timestep_emb.shape}") | |
| image = image+timestep_emb | |
| image = image+self.image_emb(self.pos_ids) | |
| prompt = prompt+timestep_emb | |
| # print(f"sanity: image: {image.shape}, prompt: {prompt.shape}, timestep_emb: {timestep_emb.shape}") | |
| image_norm, prompt_norm = self.ln1(image), self.ln1(prompt) | |
| image_q, image_k, image_v = self.image_q(image_norm), self.image_k(image_norm), self.image_v(image_norm) | |
| prompt_q, prompt_k, prompt_v = self.prompt_q(prompt_norm), self.prompt_k(prompt_norm), self.prompt_v(prompt_norm) | |
| Q = torch.cat([image_q, prompt_q], dim=1) # B,257, dim | |
| K = torch.cat([image_k, prompt_k], dim=1) # "" | |
| V = torch.cat([image_v, prompt_v], dim=1) # "" | |
| out_resid = torch.cat([image, prompt], dim=1) # b,257,dim | |
| Q = Q.view(Q.size(0), Q.size(1), self.nhead, -1).transpose(1, 2) | |
| K = K.view(K.size(0), K.size(1), self.nhead, -1).transpose(1, 2) | |
| V = V.view(V.size(0), V.size(1), self.nhead, -1).transpose(1, 2) | |
| attn_out = F.scaled_dot_product_attention(Q, K, V) # B, H, L, hd | |
| attn_out = attn_out.transpose(1, 2).reshape(attn_out.size(0), -1, self.dim) | |
| # out = self.mlp(attn_out) + out_resid | |
| # out = self.ln2(out) | |
| out = out_resid + self.o_proj(attn_out) | |
| out = out + self.mlp(self.ln2(out)) | |
| out_image, out_prompt = out[:,:self.image_tokens,:], out[:,self.image_tokens:,:] | |
| # out_image = self.image_up(out_image) | |
| # out_prompt = self.prompt_up(out_prompt) | |
| return out_image, out_prompt | |
| class DiffusionModel(nn.Module): | |
| def __init__(self, dim=256, num_layers=4, image_tokens=256, nhead=1): | |
| prompt_tokens=1 | |
| super(DiffusionModel, self).__init__() | |
| self.attn_layers = nn.ModuleList([JointAttention(dim, image_tokens, prompt_tokens,n_head=nhead) for _ in range(num_layers)]) | |
| self.image_proj = nn.Linear(4*32*32//image_tokens,dim) | |
| self.prompt_proj = nn.Linear(512,dim) | |
| self.image_up = nn.Linear(dim, 4*32*32//image_tokens) | |
| # self.prompt_up = nn.Linear(dim, 512) | |
| self.ln_f = nn.LayerNorm(dim) | |
| def forward(self, image_enc, prompt_enc, timestep): | |
| image_enc = self.image_proj(image_enc) | |
| prompt_enc = self.prompt_proj(prompt_enc) | |
| # no need to take care of resid_image because layer(image) natively handles residual | |
| # resid_image_enc, resid_prompt_enc = image_enc, prompt_enc | |
| for layer in self.attn_layers: | |
| image_enc, prompt_enc = layer(image_enc, prompt_enc, timestep) | |
| # print(f"sanity shapes: {image_enc.shape}, {prompt_enc.shape}, {resid_image_enc.shape}, {resid_prompt_enc.shape}") | |
| # image_enc, prompt_enc = image_enc+resid_image_enc, prompt_enc+resid_prompt_enc | |
| # resid_image_enc, resid_prompt_enc = image_enc, prompt_enc | |
| # return image_enc, prompt_enc | |
| image_enc = self.image_up(self.ln_f(image_enc)) | |
| return image_enc # velocity pred | |