mirror of
https://github.com/AIGC-Audio/AudioGPT.git
synced 2026-07-10 04:20:09 +02:00
30 lines
888 B
Python
30 lines
888 B
Python
import torch
|
|
|
|
|
|
def squeeze(x, x_mask=None, n_sqz=2):
|
|
b, c, t = x.size()
|
|
|
|
t = (t // n_sqz) * n_sqz
|
|
x = x[:, :, :t]
|
|
x_sqz = x.view(b, c, t // n_sqz, n_sqz)
|
|
x_sqz = x_sqz.permute(0, 3, 1, 2).contiguous().view(b, c * n_sqz, t // n_sqz)
|
|
|
|
if x_mask is not None:
|
|
x_mask = x_mask[:, :, n_sqz - 1::n_sqz]
|
|
else:
|
|
x_mask = torch.ones(b, 1, t // n_sqz).to(device=x.device, dtype=x.dtype)
|
|
return x_sqz * x_mask, x_mask
|
|
|
|
|
|
def unsqueeze(x, x_mask=None, n_sqz=2):
|
|
b, c, t = x.size()
|
|
|
|
x_unsqz = x.view(b, n_sqz, c // n_sqz, t)
|
|
x_unsqz = x_unsqz.permute(0, 2, 3, 1).contiguous().view(b, c // n_sqz, t * n_sqz)
|
|
|
|
if x_mask is not None:
|
|
x_mask = x_mask.unsqueeze(-1).repeat(1, 1, 1, n_sqz).view(b, 1, t * n_sqz)
|
|
else:
|
|
x_mask = torch.ones(b, 1, t * n_sqz).to(device=x.device, dtype=x.dtype)
|
|
return x_unsqz * x_mask, x_mask
|