File size: 5,853 Bytes
8c6b65f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | import torch
import torch.nn as nn
import torch.utils.checkpoint as checkpoint
from typing import Any, Tuple
from modules.waveform_branch import WaveformEncoder, WaveformDecoder
class Generator(nn.Module):
"""Hybrid GAN-BWE Generator with modular waveform and spectral branches."""
def __init__(self, config: Any):
super().__init__()
self.config = config
self.use_waveform = config.generator.use_waveform_branch
self.use_spectral = config.generator.use_spectral_branch
self.use_attention = config.generator.use_cross_attention
self.gradient_checkpointing = config.train.gradient_checkpointing
if not self.use_waveform and not self.use_spectral:
raise ValueError("At least one of the branches (waveform or spectral) must be enabled.")
# 1. Waveform Branch
if self.use_waveform:
self.waveform_encoder = WaveformEncoder(config.generator.waveform_encoder)
latent_channels = self.waveform_encoder.out_channels
else:
latent_channels = 0
# 2. Spectral Branch (Stub for baseline phase; implemented in Phase 8)
if self.use_spectral:
# We import dynamically to keep modules independent and prevent circular dependencies
try:
from modules.spectral_branch import SpectralEncoder
self.spectral_encoder = SpectralEncoder(config.generator.spectral_encoder)
spec_channels = self.spectral_encoder.out_channels
except ImportError:
# Fallback mock for early testing if spectral_branch is not yet created
self.spectral_encoder = None
spec_channels = 128
else:
spec_channels = 0
# 3. Fusion Branch
if self.use_waveform and self.use_spectral:
if self.use_attention:
try:
from modules.attention import CrossAttentionFusion
self.fusion = CrossAttentionFusion(
config.generator.attention,
wave_channels=latent_channels,
spec_channels=spec_channels
)
self.latent_dim = self.fusion.out_channels
except ImportError:
# Fallback simple concat if attention module not yet created
self.fusion = None
self.latent_dim = latent_channels + spec_channels
else:
self.latent_dim = latent_channels + spec_channels
elif self.use_waveform:
self.latent_dim = latent_channels
else:
self.latent_dim = spec_channels
# 4. Decoder
self.decoder = WaveformDecoder(config.generator.decoder, in_channels=self.latent_dim)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x: Degraded/Narrowband input waveform of shape (batch, 1, samples).
Returns:
Wideband reconstructed waveform of shape (batch, 1, samples).
"""
feat_wave = None
feat_spec = None
# Process Waveform branch
if self.use_waveform:
if self.gradient_checkpointing and self.training:
feat_wave = checkpoint.checkpoint(self.waveform_encoder, x, use_reentrant=False)
else:
feat_wave = self.waveform_encoder(x)
# Process Spectral branch
if self.use_spectral:
from utils.audio import wav_to_spec
# Compute STFT on input waveform
mag, phase = wav_to_spec(
x,
fft_size=self.config.generator.spectral_encoder.fft_size,
hop_length=self.config.generator.spectral_encoder.hop_length,
win_length=self.config.generator.spectral_encoder.win_length
)
if self.gradient_checkpointing and self.training:
feat_spec = checkpoint.checkpoint(self.spectral_encoder, mag, phase, use_reentrant=False)
else:
feat_spec = self.spectral_encoder(mag, phase)
# Fuse features
if self.use_waveform and self.use_spectral:
if self.use_attention and self.fusion is not None:
if self.gradient_checkpointing and self.training:
latent = checkpoint.checkpoint(self.fusion, feat_wave, feat_spec, use_reentrant=False)
else:
latent = self.fusion(feat_wave, feat_spec)
else:
# If sizes differ along the time axis, interpolate spectral features to match waveform features
if feat_wave.shape[-1] != feat_spec.shape[-1]:
feat_spec = nn.functional.interpolate(
feat_spec, size=feat_wave.shape[-1], mode='linear', align_corners=False
)
latent = torch.cat([feat_wave, feat_spec], dim=1)
elif self.use_waveform:
latent = feat_wave
else:
latent = feat_spec
# Decode features back to waveform
if self.gradient_checkpointing and self.training:
out = checkpoint.checkpoint(self.decoder, latent, use_reentrant=False)
else:
out = self.decoder(latent)
# Optional residual skip connection to add low frequencies directly
# Since x is low-pass filtered, we can add it to output to preserve low frequency phase
# The generator learns to add high-frequency details.
return out
def load_state_dict(self, state_dict, strict=True):
"""Override to load baseline and lightweight models dynamically by disabling strict checking."""
return super().load_state_dict(state_dict, strict=False)
|