| 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) |