| 119 | |
| 120 | |
| 121 | class MultiProjModel(nn.Module): |
| 122 | def __init__(self, adapter_in_dim=1024, cross_attention_dim=1024): |
| 123 | super().__init__() |
| 124 | |
| 125 | self.generator = None |
| 126 | self.cross_attention_dim = cross_attention_dim |
| 127 | self.eye_proj = torch.nn.Linear(6, cross_attention_dim, bias=False) |
| 128 | self.emo_proj = torch.nn.Linear(30, cross_attention_dim, bias=False) |
| 129 | self.mouth_proj = torch.nn.Linear(512, cross_attention_dim, bias=False) |
| 130 | self.headpose_proj = torch.nn.Linear(6, cross_attention_dim, bias=False) |
| 131 | |
| 132 | self.norm = torch.nn.LayerNorm(cross_attention_dim) |
| 133 | |
| 134 | def forward(self, adapter_embeds): |
| 135 | B, num_frames, C = adapter_embeds.shape |
| 136 | embeds = adapter_embeds |
| 137 | split_sizes = [6, 6, 30, 512] |
| 138 | headpose, eye, emo, mouth = torch.split(embeds, split_sizes, dim=-1) |
| 139 | headpose = self.norm(self.headpose_proj(headpose)) |
| 140 | eye = self.norm(self.eye_proj(eye)) |
| 141 | emo = self.norm(self.emo_proj(emo)) |
| 142 | mouth = self.norm(self.mouth_proj(mouth)) |
| 143 | |
| 144 | all_features = torch.stack([headpose, eye, emo, mouth], dim=2) |
| 145 | result_final = all_features.view(B, num_frames * 4, self.cross_attention_dim) |
| 146 | |
| 147 | return result_final |
| 148 | |
| 149 | |
| 150 | class PortraitEncoder(ModelMixin, ConfigMixin, FromOriginalModelMixin): |