radames/Text2Human-API

1import torch2import torch.nn.functional as F3from torch import nn4 5 6class ShapeAttrEmbedding(nn.Module):7 8    def __init__(self, dim, out_dim, cls_num_list):9        super(ShapeAttrEmbedding, self).__init__()10 11        for idx, cls_num in enumerate(cls_num_list):12            setattr(13                self, f'attr_{idx}',14                nn.Sequential(15                    nn.Linear(cls_num, dim), nn.LeakyReLU(),16                    nn.Linear(dim, dim)))17        self.cls_num_list = cls_num_list18        self.attr_num = len(cls_num_list)19        self.fusion = nn.Sequential(20            nn.Linear(dim * self.attr_num, out_dim), nn.LeakyReLU(),21            nn.Linear(out_dim, out_dim))22 23    def forward(self, attr):24        attr_embedding_list = []25        for idx in range(self.attr_num):26            attr_embed_fc = getattr(self, f'attr_{idx}')27            attr_embedding_list.append(28                attr_embed_fc(29                    F.one_hot(30                        attr[:, idx],31                        num_classes=self.cls_num_list[idx]).to(torch.float32)))32        attr_embedding = torch.cat(attr_embedding_list, dim=1)33        attr_embedding = self.fusion(attr_embedding)34 35        return attr_embedding36