GPT-SoVITS/GPT_SoVITS/module/VoiceChange.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import torchaudio
import math
from torchaudio.transforms import Resample
import VoiceSave
import uuid

def get_train_set(voice_file_path):
    if type(voice_file_path) == str:
        voice_file_path = [voice_file_path]
    ret = []
    for i in voice_file_path:
        tensors_ = VoiceSave.load_tensor(i,
                              f"get_{uuid.uuid4()}",
                              find_func=VoiceSave.__find_func__,
                              MySet=set())
        ret.append(tensors_)
    return ret

class MelSpectrogram(nn.Module):
    def __init__(self, hps):
        super().__init__()
        self.filter_length = hps.data.filter_length
        self.hop_length = hps.data.hop_length
        self.win_length = hps.data.win_length
        self.sampling_rate = hps.data.sampling_rate
        self.n_mel_channels = hps.data.n_mel_channels  
        self.mel_fmin = hps.data.mel_fmin if hasattr(hps.data, 'mel_fmin') else 0
        self.mel_fmax = hps.data.mel_fmax if hasattr(hps.data, 'mel_fmax') else None

        # 构建梅尔频谱变换
        self.mel_transform = torchaudio.transforms.MelSpectrogram(
            sample_rate=self.sampling_rate,
            n_fft=self.filter_length,
            hop_length=self.hop_length,
            win_length=self.win_length,
            f_min=self.mel_fmin,
            f_max=self.mel_fmax,
            n_mels=192,  # self.n_mel_channels,
            window_fn=torch.hann_window,
            center=False,
            power=1.0,
        )

    def forward(self, audio):
        """
        输入：audio [B, 1, T] 或 [1, T]（单声道音频）
        输出：mel_spec [B, n_mel_channels, T']
        """
        if len(audio.shape) == 2:
            audio = audio.unsqueeze(0)  # [1, T] → [1, 1, T]
        
        # 提取梅尔频谱
        mel_spec = self.mel_transform(audio.squeeze(1))  # [B, n_mel, T']
        
        # 对数缩放（TTS标准操作）
        mel_spec = torch.log(torch.clamp(mel_spec, min=1e-5))
        
        return mel_spec

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_seq_length=5000):
        super(PositionalEncoding, self).__init__()
        self.pe = torch.zeros(max_seq_length, d_model)  # 初始化位置编码矩阵
        position = torch.arange(0, max_seq_length, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model))
        self.pe[:, 0::2] = torch.sin(position * div_term)  # 偶数位置使用正弦函数
        self.pe[:, 1::2] = torch.cos(position * div_term)  # 奇数位置使用余弦函数
        self.register_buffer('pe', self.pe.unsqueeze(0))  # 注册为缓冲区
        
    def forward(self, x):
        # 将位置编码添加到输入中
        return x + self.pe[:, :x.size(1)]
    
class Spliter(nn.Module):
    '''output: z_p shape: torch.Size([1, 192, x]), y_mask shape: torch.Size([1, 1, x]), ge shape: torch.Size([1, 1024, 1])'''
    def __init__(self,
                 hps,
                 ge,
                 device):
        super().__init__()
        self.hps = hps

        self.ge = ge
        self.device = device
        #TODO: 将mel_spec与ge输入Transformer模型
        self.mel_dim = 192
        self.ge_dim = 1024                
        self.transformer_dim = 512
        self.ge_proj = nn.Linear(self.ge_dim, self.transformer_dim).to(self.device)
        self.mel_proj = nn.Linear(self.mel_dim, self.transformer_dim).to(self.device)
        self.pos_encoder = PositionalEncoding(self.transformer_dim).to(self.device)
        self.transformer = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(
                d_model=self.transformer_dim,
                nhead=hps.model.nhead,
                dim_feedforward=hps.model.ffn_dim,
                batch_first=False,
                dropout=0.1
            ),
            num_layers=hps.model.num_layers
        ).to(self.device)
    
        self.out_proj = nn.Linear(self.transformer_dim, self.mel_dim).to(self.device)

    @torch.no_grad()
    def mel_(self,audio_path, hps, device, dtype):
        sr_target = int(hps.data.sampling_rate)
        audio, sr_origin = torchaudio.load(audio_path)
        if audio.shape[0] > 1:
            audio = audio.mean(0, keepdim=True)
        if sr_origin != sr_target:
            resampler = Resample(sr_origin, sr_target).to(device)
            audio = resampler(audio.to(device))
        else:
            audio = audio.to(device)
        max_audio = audio.abs().max()
        if max_audio > 1.0:
            audio = audio / max_audio
        mel_extractor = MelSpectrogram(hps).to(device)
        mel_spec = mel_extractor(audio).to(dtype)
        return mel_spec
    
    def forward(self, audio_path, ge,device,dtype):
        # 输入：audio_path, ge
        # 输出：z_p, y_mask, ge
        ge_ = ge
        mel = self.mel_(audio_path, self.hps, device, dtype)

        mel = mel.permute(2, 0, 1)  
        # 梅尔谱投影到Transformer维度：[T, 1, 512]
        mel_feat = self.mel_proj(mel)  
        
        # 全局情感特征GE处理：[1,1024,1] → [1,1024] → [1,1,512]
        ge = ge.to(device, dtype=dtype)
        ge_squeeze = ge.squeeze(-1)  # [1, 1024]
        ge_feat = self.ge_proj(ge_squeeze).unsqueeze(0)  # [1, 1, 512]
        
        # ===================== 3. 特征融合与Transformer输入 =====================
        # 将GE特征拼接在梅尔谱序列开头：[T+1, 1, 512]
        self.transformer_input = torch.cat([ge_feat, mel_feat], dim=0)
        # 添加位置编码
        self.transformer_input = self.pos_encoder(self.transformer_input)
        
        # ===================== 4. Transformer编码 =====================
        transformer_out = self.transformer(self.transformer_input)  # [T+1, 1, 512]
        
        # ===================== 5. 输出特征重构 =====================
        # 去除GE开头，提取梅尔谱对应的输出：[T, 1, 512]
        mel_out = transformer_out[1:, :, :]
        # 投影回原始梅尔维度：[T, 1, 192]
        mel_out = self.out_proj(mel_out)
        # 转换为目标格式：[1, 192, T] → z_p
        z_p = mel_out.permute(1, 2, 0)
        
        # ===================== 6. 生成掩码 =====================
        T = z_p.shape[-1]  # 梅尔谱时间步
        y_mask = torch.ones(1, 1, T, device=device, dtype=dtype)  # [1,1,T] 全1掩码
        
        # ===================== 7. 输出（严格匹配注释格式） =====================
        return z_p, y_mask, ge_

class SpliterDataset(torch.utils.data.Dataset):
    def __init__(self, voice_file_paths):
        self.voice_file_paths = voice_file_paths
        self.datas = get_train_set(voice_file_paths)

    def __len__(self):
        return len(self.datas)
    
    def __getitem__(self, idx):
        return self.datas[idx]