mirror of
https://github.com/Nighthawk42/soprano-factory.git
synced 2026-08-30 04:30:21 +00:00
117 lines
3.3 KiB
Python
117 lines
3.3 KiB
Python
# model/decoder.py - Standardized for V3
|
|
|
|
import torch
|
|
from torch import nn
|
|
import torch.nn.functional as F
|
|
from .common import VocosBackbone
|
|
|
|
class ISTFTHead(nn.Module):
|
|
"""
|
|
Inverse STFT head for waveform generation.
|
|
Converts hidden states -> complex spectrogram -> waveform.
|
|
"""
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
n_fft: int = 2048,
|
|
hop_length: int = 512,
|
|
padding: str = "same",
|
|
):
|
|
super().__init__()
|
|
self.n_fft = n_fft
|
|
self.hop_length = hop_length
|
|
out_dim = n_fft + 2
|
|
self.out = nn.Linear(dim, out_dim)
|
|
self.padding = padding
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Args:
|
|
x: [B, T, C] - hidden states
|
|
Returns:
|
|
audio: [B, 1, L] - waveform
|
|
"""
|
|
x = x.float() # Ensure float32
|
|
x = self.out(x) # [B, T, n_fft + 2]
|
|
x = x.transpose(1, 2) # [B, n_fft + 2, T]
|
|
|
|
mag, phase = x.chunk(2, dim=1) # each [B, (n_fft+2)//2, T]
|
|
mag = torch.exp(mag)
|
|
phase = torch.sin(phase)
|
|
|
|
# Construct complex spectrogram
|
|
S = mag * torch.exp(1j * phase)
|
|
|
|
# Inverse STFT
|
|
audio = torch.istft(
|
|
S,
|
|
n_fft=self.n_fft,
|
|
hop_length=self.hop_length,
|
|
window=torch.hann_window(self.n_fft).to(x.device),
|
|
center=True,
|
|
)
|
|
|
|
return audio.unsqueeze(1) # [B, 1, L]
|
|
|
|
class Decoder(nn.Module):
|
|
def __init__(
|
|
self,
|
|
input_channels=512,
|
|
decoder_dim=512,
|
|
decoder_layers=8,
|
|
):
|
|
super().__init__()
|
|
|
|
# 1. The Adapter (Projection)
|
|
self.input_proj = nn.Linear(input_channels, decoder_dim)
|
|
|
|
# Standardization Layer
|
|
self.post_proj_norm = nn.LayerNorm(decoder_dim)
|
|
|
|
# 2. Upsampler (Match Encoder's slice::4)
|
|
# We need to upsample by 4x to match the Hop=512 resolution
|
|
# Encoder: Hop=512 -> Slice(4) -> Effective Hop=2048
|
|
# Decoder: Latent(2048) -> Upsample(4) -> Latent(512) -> ISTFT(512) -> Audio
|
|
self.upsampler = nn.Sequential(
|
|
nn.ConvTranspose1d(decoder_dim, decoder_dim, kernel_size=4, stride=4),
|
|
nn.GELU()
|
|
)
|
|
|
|
# 3. Backbone
|
|
self.backbone = VocosBackbone(
|
|
input_channels=decoder_dim,
|
|
dim=decoder_dim,
|
|
intermediate_dim=decoder_dim * 3,
|
|
num_layers=decoder_layers,
|
|
)
|
|
|
|
# 4. Output Head (ISTFT)
|
|
self.head = ISTFTHead(
|
|
dim=decoder_dim,
|
|
n_fft=2048,
|
|
hop_length=512
|
|
)
|
|
|
|
def forward(self, x):
|
|
# x shape: [Batch, Time, Channels]
|
|
x = x.float() # Ensure float32 everywhere
|
|
|
|
# Project and Normalize
|
|
x = self.input_proj(x)
|
|
x = self.post_proj_norm(x)
|
|
|
|
# Upsample [B, T, C] -> [B, C, T] -> [B, C, T*4]
|
|
x = x.transpose(1, 2)
|
|
x = self.upsampler(x)
|
|
|
|
# Backbone ResBlocks [B, C, T*4] -> [B, C, T*4]
|
|
# (VocosBackbone expects [B, C, T] and returns [B, C, T])
|
|
x = self.backbone(x)
|
|
|
|
# Prepare for ISTFT Head [B, C, T] -> [B, T, C]
|
|
x = x.transpose(1, 2)
|
|
|
|
# Generate Waveform
|
|
x = self.head(x)
|
|
|
|
return x |