| 148 | |
| 149 | |
| 150 | class PortraitEncoder(ModelMixin, ConfigMixin, FromOriginalModelMixin): |
| 151 | def __init__(self, adapter_in_dim: int, adapter_proj_dim: int): |
| 152 | super().__init__() |
| 153 | |
| 154 | self.adapter_in_dim = adapter_in_dim |
| 155 | self.adapter_proj_dim = adapter_proj_dim |
| 156 | self.proj_model = self.init_proj(self.adapter_proj_dim) |
| 157 | |
| 158 | self.mouth_proj_model = Resampler( |
| 159 | dim=1280, |
| 160 | depth=4, |
| 161 | dim_head=64, |
| 162 | heads=20, |
| 163 | num_queries=16, |
| 164 | embedding_dim=512, |
| 165 | output_dim=2048, |
| 166 | ff_mult=4, |
| 167 | ) |
| 168 | |
| 169 | self.emo_proj_model = Resampler( |
| 170 | dim=1280, |
| 171 | depth=4, |
| 172 | dim_head=64, |
| 173 | heads=20, |
| 174 | num_queries=4, |
| 175 | embedding_dim=30, |
| 176 | output_dim=2048, |
| 177 | ff_mult=4, |
| 178 | ) |
| 179 | |
| 180 | def init_proj(self, cross_attention_dim=5120): |
| 181 | proj_model = MultiProjModel(adapter_in_dim=self.adapter_in_dim, cross_attention_dim=cross_attention_dim) |
| 182 | return proj_model |
| 183 | |
| 184 | def get_adapter_proj(self, adapter_fea=None): |
| 185 | split_sizes = [6, 6, 30, 512] |
| 186 | headpose, eye, emo, mouth = torch.split( |
| 187 | adapter_fea, split_sizes, dim=-1 |
| 188 | ) |
| 189 | B, frames, dim = mouth.shape |
| 190 | mouth = mouth.view(B * frames, 1, 512) |
| 191 | emo = emo.view(B * frames, 1, 30) |
| 192 | |
| 193 | mouth_fea = self.mouth_proj_model(mouth) |
| 194 | emo_fea = self.emo_proj_model(emo) |
| 195 | |
| 196 | mouth_fea = mouth_fea.view(B, frames, 16, 2048) |
| 197 | emo_fea = emo_fea.view(B, frames, 4, 2048) |
| 198 | |
| 199 | adapter_fea = self.proj_model(adapter_fea) |
| 200 | |
| 201 | adapter_fea = adapter_fea.view(B, frames, 4, 2048) |
| 202 | |
| 203 | all_fea = torch.cat([adapter_fea, mouth_fea, emo_fea], dim=2) |
| 204 | |
| 205 | result_final = all_fea.view(B, frames * 24, 2048) |
| 206 | |
| 207 | return result_final |
no outgoing calls
no test coverage detected