113 lines
4.4 KiB
Python
113 lines
4.4 KiB
Python
# Merge image encoder and fuse module to create an ID Encoder
|
|
# send multiple ID images, we can directly obtain the updated text encoder containing a stacked ID embedding
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from transformers.models.clip.modeling_clip import CLIPVisionModelWithProjection
|
|
from transformers.models.clip.configuration_clip import CLIPVisionConfig
|
|
from transformers import PretrainedConfig
|
|
|
|
VISION_CONFIG_DICT = {
|
|
"hidden_size": 1024,
|
|
"intermediate_size": 4096,
|
|
"num_attention_heads": 16,
|
|
"num_hidden_layers": 24,
|
|
"patch_size": 14,
|
|
"projection_dim": 768
|
|
}
|
|
|
|
class MLP(nn.Module):
|
|
def __init__(self, in_dim, out_dim, hidden_dim, use_residual=True):
|
|
super().__init__()
|
|
if use_residual:
|
|
assert in_dim == out_dim
|
|
self.layernorm = nn.LayerNorm(in_dim)
|
|
self.fc1 = nn.Linear(in_dim, hidden_dim)
|
|
self.fc2 = nn.Linear(hidden_dim, out_dim)
|
|
self.use_residual = use_residual
|
|
self.act_fn = nn.GELU()
|
|
|
|
def forward(self, x):
|
|
residual = x
|
|
x = self.layernorm(x)
|
|
x = self.fc1(x)
|
|
x = self.act_fn(x)
|
|
x = self.fc2(x)
|
|
if self.use_residual:
|
|
x = x + residual
|
|
return x
|
|
|
|
|
|
class FuseModule(nn.Module):
|
|
def __init__(self, embed_dim):
|
|
super().__init__()
|
|
self.mlp1 = MLP(embed_dim * 2, embed_dim, embed_dim, use_residual=False)
|
|
self.mlp2 = MLP(embed_dim, embed_dim, embed_dim, use_residual=True)
|
|
self.layer_norm = nn.LayerNorm(embed_dim)
|
|
|
|
def fuse_fn(self, prompt_embeds, id_embeds):
|
|
stacked_id_embeds = torch.cat([prompt_embeds, id_embeds], dim=-1)
|
|
stacked_id_embeds = self.mlp1(stacked_id_embeds) + prompt_embeds
|
|
stacked_id_embeds = self.mlp2(stacked_id_embeds)
|
|
stacked_id_embeds = self.layer_norm(stacked_id_embeds)
|
|
return stacked_id_embeds
|
|
|
|
def forward(
|
|
self,
|
|
prompt_embeds,
|
|
id_embeds,
|
|
class_tokens_mask,
|
|
) -> torch.Tensor:
|
|
# id_embeds shape: [b, max_num_inputs, 1, 2048]
|
|
id_embeds = id_embeds.to(prompt_embeds.dtype)
|
|
num_inputs = class_tokens_mask.sum().unsqueeze(0) # TODO: check for training case
|
|
batch_size, max_num_inputs = id_embeds.shape[:2]
|
|
# seq_length: 77
|
|
seq_length = prompt_embeds.shape[1]
|
|
# flat_id_embeds shape: [b*max_num_inputs, 1, 2048]
|
|
flat_id_embeds = id_embeds.view(
|
|
-1, id_embeds.shape[-2], id_embeds.shape[-1]
|
|
)
|
|
# valid_id_mask [b*max_num_inputs]
|
|
valid_id_mask = (
|
|
torch.arange(max_num_inputs, device=flat_id_embeds.device)[None, :]
|
|
< num_inputs[:, None]
|
|
)
|
|
valid_id_embeds = flat_id_embeds[valid_id_mask.flatten()]
|
|
|
|
prompt_embeds = prompt_embeds.view(-1, prompt_embeds.shape[-1])
|
|
class_tokens_mask = class_tokens_mask.view(-1)
|
|
valid_id_embeds = valid_id_embeds.view(-1, valid_id_embeds.shape[-1])
|
|
# slice out the image token embeddings
|
|
image_token_embeds = prompt_embeds[class_tokens_mask]
|
|
stacked_id_embeds = self.fuse_fn(image_token_embeds, valid_id_embeds)
|
|
assert class_tokens_mask.sum() == stacked_id_embeds.shape[0], f"{class_tokens_mask.sum()} != {stacked_id_embeds.shape[0]}"
|
|
prompt_embeds.masked_scatter_(class_tokens_mask[:, None], stacked_id_embeds.to(prompt_embeds.dtype))
|
|
updated_prompt_embeds = prompt_embeds.view(batch_size, seq_length, -1)
|
|
return updated_prompt_embeds
|
|
|
|
class PhotoMakerIDEncoder(CLIPVisionModelWithProjection):
|
|
def __init__(self):
|
|
super().__init__(CLIPVisionConfig(**VISION_CONFIG_DICT))
|
|
self.visual_projection_2 = nn.Linear(1024, 1280, bias=False)
|
|
self.fuse_module = FuseModule(2048)
|
|
|
|
def forward(self, id_pixel_values, prompt_embeds, class_tokens_mask):
|
|
b, num_inputs, c, h, w = id_pixel_values.shape
|
|
id_pixel_values = id_pixel_values.view(b * num_inputs, c, h, w)
|
|
|
|
shared_id_embeds = self.vision_model(id_pixel_values)[1]
|
|
id_embeds = self.visual_projection(shared_id_embeds)
|
|
id_embeds_2 = self.visual_projection_2(shared_id_embeds)
|
|
|
|
id_embeds = id_embeds.view(b, num_inputs, 1, -1)
|
|
id_embeds_2 = id_embeds_2.view(b, num_inputs, 1, -1)
|
|
|
|
id_embeds = torch.cat((id_embeds, id_embeds_2), dim=-1)
|
|
updated_prompt_embeds = self.fuse_module(prompt_embeds, id_embeds, class_tokens_mask)
|
|
|
|
return updated_prompt_embeds
|
|
|
|
|
|
if __name__ == "__main__":
|
|
PhotoMakerIDEncoder() |