UAI-Software/PhotoMaker

1# Merge image encoder and fuse module to create a ID Encoder2# send multiple ID images, we can directly obtain the updated text encoder containing a stacked ID embedding3 4import torch5import torch.nn as nn6from transformers.models.clip.modeling_clip import CLIPVisionModelWithProjection7from transformers.models.clip.configuration_clip import CLIPVisionConfig8from transformers import PretrainedConfig9 10VISION_CONFIG_DICT = {11    "hidden_size": 1024,12    "intermediate_size": 4096,13    "num_attention_heads": 16,14    "num_hidden_layers": 24,15    "patch_size": 14,16    "projection_dim": 76817}18 19class MLP(nn.Module):20    def __init__(self, in_dim, out_dim, hidden_dim, use_residual=True):21        super().__init__()22        if use_residual:23            assert in_dim == out_dim24        self.layernorm = nn.LayerNorm(in_dim)25        self.fc1 = nn.Linear(in_dim, hidden_dim)26        self.fc2 = nn.Linear(hidden_dim, out_dim)27        self.use_residual = use_residual28        self.act_fn = nn.GELU()29 30    def forward(self, x):31        residual = x32        x = self.layernorm(x)33        x = self.fc1(x)34        x = self.act_fn(x)35        x = self.fc2(x)36        if self.use_residual:37            x = x + residual38        return x39 40 41class FuseModule(nn.Module):42    def __init__(self, embed_dim):43        super().__init__()44        self.mlp1 = MLP(embed_dim * 2, embed_dim, embed_dim, use_residual=False)45        self.mlp2 = MLP(embed_dim, embed_dim, embed_dim, use_residual=True)46        self.layer_norm = nn.LayerNorm(embed_dim)47 48    def fuse_fn(self, prompt_embeds, id_embeds):49        print(prompt_embeds.shape, id_embeds.shape)50        stacked_id_embeds = torch.cat([prompt_embeds, id_embeds], dim=-1)51        stacked_id_embeds = self.mlp1(stacked_id_embeds) + prompt_embeds52        stacked_id_embeds = self.mlp2(stacked_id_embeds)53        stacked_id_embeds = self.layer_norm(stacked_id_embeds)54        return stacked_id_embeds55 56    def forward(57        self,58        prompt_embeds,59        id_embeds,60        class_tokens_mask,61    ) -> torch.Tensor:62        # id_embeds shape: [b, max_num_inputs, 1, 2048]63        id_embeds = id_embeds.to(prompt_embeds.dtype)64        num_inputs = class_tokens_mask.sum().unsqueeze(0) # TODO: check for training case65        batch_size, max_num_inputs = id_embeds.shape[:2]66        # seq_length: 7767        seq_length = prompt_embeds.shape[1]68        # flat_id_embeds shape: [b*max_num_inputs, 1, 2048]69        flat_id_embeds = id_embeds.view(70            -1, id_embeds.shape[-2], id_embeds.shape[-1]71        )72        # valid_id_mask [b*max_num_inputs]73        valid_id_mask = (74            torch.arange(max_num_inputs, device=flat_id_embeds.device)[None, :]75            < num_inputs[:, None]76        )77        valid_id_embeds = flat_id_embeds[valid_id_mask.flatten()]78 79        prompt_embeds = prompt_embeds.view(-1, prompt_embeds.shape[-1])80        class_tokens_mask = class_tokens_mask.view(-1)81        valid_id_embeds = valid_id_embeds.view(-1, valid_id_embeds.shape[-1])82        # slice out the image token embeddings83        image_token_embeds = prompt_embeds[class_tokens_mask]84        stacked_id_embeds = self.fuse_fn(image_token_embeds, valid_id_embeds)85        assert class_tokens_mask.sum() == stacked_id_embeds.shape[0], f"{class_tokens_mask.sum()} != {stacked_id_embeds.shape[0]}"86        prompt_embeds.masked_scatter_(class_tokens_mask[:, None], stacked_id_embeds.to(prompt_embeds.dtype))87        updated_prompt_embeds = prompt_embeds.view(batch_size, seq_length, -1)88        return updated_prompt_embeds89 90class PhotoMakerIDEncoder(CLIPVisionModelWithProjection):91    def __init__(self):92        super().__init__(CLIPVisionConfig(**VISION_CONFIG_DICT))93        self.visual_projection_2 = nn.Linear(1024, 1280, bias=False)94        self.fuse_module = FuseModule(2048)95 96    def forward(self, id_pixel_values, prompt_embeds, class_tokens_mask):97        b, num_inputs, c, h, w = id_pixel_values.shape98        id_pixel_values = id_pixel_values.view(b * num_inputs, c, h, w)99 100        shared_id_embeds = self.vision_model(id_pixel_values)[1]101        id_embeds = self.visual_projection(shared_id_embeds)102        id_embeds_2 = self.visual_projection_2(shared_id_embeds)103 104        id_embeds = id_embeds.view(b, num_inputs, 1, -1)105        id_embeds_2 = id_embeds_2.view(b, num_inputs, 1, -1)    106 107        id_embeds = torch.cat((id_embeds, id_embeds_2), dim=-1)108        updated_prompt_embeds = self.fuse_module(prompt_embeds, id_embeds, class_tokens_mask)109 110        return updated_prompt_embeds111 112 113if __name__ == "__main__":114    PhotoMakerIDEncoder()