ShalomKing/infinitetalk

1import torch
2import torch.nn.functional as F
3
4
5def get_mask_from_lengths(lengths, max_len=None):
6    lengths = lengths.to(torch.long)
7    if max_len is None:
8        max_len = torch.max(lengths).item()
9
10    ids = torch.arange(0, max_len).unsqueeze(0).expand(lengths.shape[0], -1).to(lengths.device)
11    mask = ids < lengths.unsqueeze(1).expand(-1, max_len)
12
13    return mask
14
15
16def linear_interpolation(features, seq_len):
17    features = features.transpose(1, 2)
18    output_features = F.interpolate(features, size=seq_len, align_corners=True, mode='linear')
19    return output_features.transpose(1, 2)
20
21