| import torch | |
| from diffusers import ModelMixin | |
| from einops import rearrange | |
| from torch import nn | |
| class IDProjConvModel(ModelMixin): | |
| def __init__(self, in_channels=2048, out_channels=1024): | |
| super().__init__() | |
| self.project1024 = torch.nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=False) | |
| self.final_norm = torch.nn.LayerNorm(out_channels) | |
| def forward(self, src_id_features_7_7_1024): | |
| c = self.project1024(src_id_features_7_7_1024) | |
| c = torch.flatten(c, 2) | |
| c = torch.transpose(c, 2, 1) | |
| c = self.final_norm(c) | |
| return c | |