53 lines
1.7 KiB
Python
53 lines
1.7 KiB
Python
from models.mossformer2_sr.generator import Mossformer, Generator
|
|
import torch.nn as nn
|
|
|
|
class MossFormer2_SR_48K(nn.Module):
|
|
"""
|
|
The MossFormer2_SR_48K model for speech super-resolution.
|
|
|
|
This class encapsulates the functionality of the MossFormer2 and HiFi-Gan
|
|
Generator within a higher-level model. It processes input audio data to produce
|
|
higher-resolution outputs.
|
|
|
|
Arguments
|
|
---------
|
|
args : Namespace
|
|
Configuration arguments that may include hyperparameters
|
|
and model settings (not utilized in this implementation but
|
|
can be extended for flexibility).
|
|
|
|
Example
|
|
---------
|
|
>>> model = MossFormer2_SR_48K(args).model
|
|
>>> x = torch.randn(10, 180, 2000) # Example input
|
|
>>> outputs = model(x) # Forward pass
|
|
>>> outputs.shape, mask.shape # Check output shapes
|
|
"""
|
|
|
|
def __init__(self, args):
|
|
super(MossFormer2_SR_48K, self).__init__()
|
|
# Initialize the TestNet model, which contains the MossFormer MaskNet
|
|
self.model_m = Mossformer() # Instance of TestNet
|
|
self.model_g = Generator(args)
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Forward pass through the model.
|
|
|
|
Arguments
|
|
---------
|
|
x : torch.Tensor
|
|
Input tensor of dimension [B, N, S], where B is the batch size,
|
|
N is the number of mel bins (80 in this case), and S is the
|
|
sequence length (e.g., time frames).
|
|
|
|
Returns
|
|
-------
|
|
outputs : torch.Tensor
|
|
Bandwidth expanded audio output tensor from the model.
|
|
|
|
"""
|
|
x = self.model_m(x) # Get outputs and mask from TestNet
|
|
outpus = self.model_g(x)
|
|
return outputs # Return the outputs
|