networks.py

# Copyright (c) 2024 Alibaba Inc (authors: Shengkui Zhao, Zexu Pan)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch
import torch.nn as nn
import soundfile as sf
import os
import subprocess
import librosa
from tqdm import tqdm
import numpy as np
from pydub import AudioSegment
from utils.decode import decode_one_audio
from dataloader.dataloader import DataReader

MAX_WAV_VALUE = 32768.0

class SpeechModel:
    """
    The SpeechModel class is a base class designed to handle speech processing tasks,
    such as loading, processing, and decoding audio data. It initializes the computational 
    device (CPU or GPU) and holds model-related attributes. The class is flexible and intended 
    to be extended by specific speech models for tasks like speech enhancement, speech separation, 
    target speaker extraction etc.

    Attributes:
    - args: Argument parser object that contains configuration settings.
    - device: The device (CPU or GPU) on which the model will run.
    - model: The actual model used for speech processing tasks (to be loaded by subclasses).
    - name: A placeholder for the model's name.
    - data: A dictionary to store any additional data related to the model, such as audio input.
    """

    def __init__(self, args):
        """
        Initializes the SpeechModel class by determining the computation device 
        (GPU or CPU) to be used for running the model, based on system availability.

        Args:
        - args: Argument parser object containing settings like whether to use CUDA (GPU) or not.
        """
        # Check if a GPU is available
        if torch.cuda.is_available():
            # Find the GPU with the most free memory using a custom method
            free_gpu_id = self.get_free_gpu()
            if free_gpu_id is not None:
                args.use_cuda = 1
                torch.cuda.set_device(free_gpu_id)
                self.device = torch.device('cuda')
            else:
                # If no GPU is detected, use the CPU
                #print("No GPU found. Using CPU.")
                args.use_cuda = 0
                self.device = torch.device('cpu')
        else:
            # If no GPU is detected, use the CPU
            args.use_cuda = 0
            self.device = torch.device('cpu')

        self.args = args
        self.model = None
        self.name = None
        self.data = {}
        self.print = False

    def get_free_gpu(self):
        """
        Identifies the GPU with the most free memory using 'nvidia-smi' and returns its index.

        This function queries the available GPUs on the system and determines which one has 
        the highest amount of free memory. It uses the `nvidia-smi` command-line tool to gather 
        GPU memory usage data. If successful, it returns the index of the GPU with the most free memory.
        If the query fails or an error occurs, it returns None.

        Returns:
        int: Index of the GPU with the most free memory, or None if no GPU is found or an error occurs.
        """
        try:
            # Run nvidia-smi to query GPU memory usage and free memory
            result = subprocess.run(['nvidia-smi', '--query-gpu=memory.used,memory.free', '--format=csv,nounits,noheader'], stdout=subprocess.PIPE)
            gpu_info = result.stdout.decode('utf-8').strip().split('\n')

            free_gpu = None
            max_free_memory = 0
            for i, info in enumerate(gpu_info):
                used, free = map(int, info.split(','))
                if free > max_free_memory:
                    max_free_memory = free
                    free_gpu = i
            return free_gpu
        except Exception as e:
            print(f"Error finding free GPU: {e}")
            return None

    def download_model(self, model_name):
        checkpoint_dir = self.args.checkpoint_dir
        from huggingface_hub import snapshot_download
        if not os.path.exists(checkpoint_dir):
            os.makedirs(checkpoint_dir)
        print(f'downloading checkpoint for {model_name}')
        try:
            snapshot_download(repo_id=f'alibabasglab/{model_name}', local_dir=checkpoint_dir)
            return True
        except:
            return False
            
    def load_model(self):
        """
        Loads a pre-trained model checkpoints from a specified directory. It checks for
        the best model ('last_best_checkpoint') in the checkpoint directory. If a model is
        found, it loads the model state into the current model instance.

        If no checkpoint is found, it will try to download the model from huggingface.
        If the downloading fails, it prints a warning message.

        Steps:
        - Search for the best model checkpoint or the most recent one.
        - Load the model's state dictionary from the checkpoint file.

        Raises:
        - FileNotFoundError: If neither 'last_best_checkpoint' nor 'last_checkpoint' files are found.
        """
        # Define paths for the best model and the last checkpoint
        best_name = os.path.join(self.args.checkpoint_dir, 'last_best_checkpoint')
        # Check if the last best checkpoint exists
        if not os.path.isfile(best_name):
            if not self.download_model(self.name):
                # If downloading is unsuccessful
                print(f'Warning: Downloading model {self.name} is not successful. Please try again or manually download from https://huggingface.co/alibabasglab/{self.name}/tree/main !')
                return

        if isinstance(self.model, nn.ModuleList):
            with open(best_name, 'r') as f:
                model_name = f.readline().strip()
                checkpoint_path = os.path.join(self.args.checkpoint_dir, model_name)
                self._load_model(self.model[0], checkpoint_path, model_key='mossformer')
                model_name = f.readline().strip()
                checkpoint_path = os.path.join(self.args.checkpoint_dir, model_name)
                self._load_model(self.model[1], checkpoint_path, model_key='generator')
        else:
            # Read the model's checkpoint name from the file
            with open(best_name, 'r') as f:
                model_name = f.readline().strip()
            # Form the full path to the model's checkpoint
            checkpoint_path = os.path.join(self.args.checkpoint_dir, model_name)
            self._load_model(self.model, checkpoint_path, model_key='model')

    def _load_model(self, model, checkpoint_path, model_key=None):
        # Load the checkpoint file into memory (map_location ensures compatibility with different devices)
        checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)
        # Load the model's state dictionary (weights and biases) into the current model
        if model_key in checkpoint:
            pretrained_model = checkpoint[model_key]
        else:
            pretrained_model = checkpoint
        state = model.state_dict()
        for key in state.keys():
            if key in pretrained_model and state[key].shape == pretrained_model[key].shape:
                state[key] = pretrained_model[key]
            elif key.replace('module.', '') in pretrained_model and state[key].shape == pretrained_model[key.replace('module.', '')].shape:
                 state[key] = pretrained_model[key.replace('module.', '')]
            elif 'module.'+key in pretrained_model and state[key].shape == pretrained_model['module.'+key].shape:
                 state[key] = pretrained_model['module.'+key]
            elif self.print: print(f'{key} not loaded')
        model.load_state_dict(state)

    def decode(self):
        """
        Decodes the input audio data using the loaded model and ensures the output matches the original audio length.

        This method processes the audio through a speech model (e.g., for enhancement, separation, etc.),
        and truncates the resulting audio to match the original input's length. The method supports multiple speakers 
        if the model handles multi-speaker audio.

        Returns:
        output_audio: The decoded audio after processing, truncated to the input audio length. 
                  If multi-speaker audio is processed, a list of truncated audio outputs per speaker is returned.
        """
        # Decode the audio using the loaded model on the given device (e.g., CPU or GPU)
        output_audios = []
        for i in range(len(self.data['audio'])):
            output_audio = decode_one_audio(self.model, self.device, self.data['audio'][i], self.args)
            # Ensure the decoded output matches the length of the input audio
            if isinstance(output_audio, list):
                # If multi-speaker audio (a list of outputs), truncate each speaker's audio to input length
                for spk in range(self.args.num_spks):
                    output_audio[spk] = output_audio[spk][:self.data['audio_len']]
            else:
                # Single output, truncate to input audio length
                output_audio = output_audio[:self.data['audio_len']]
            output_audios.append(output_audio)
            
        if isinstance(output_audios[0], list):
            output_audios_np = []
            for spk in range(self.args.num_spks):
                output_audio_buf = []
                for i in range(len(output_audios)):
                    output_audio_buf.append(output_audios[i][spk])
                    #output_audio_buf = np.vstack((output_audio_buf, output_audios[i][spk])).T
                output_audios_np.append(np.array(output_audio_buf))
        else:
            output_audios_np = np.array(output_audios)
        return output_audios_np

    def process(self, input_path, online_write=False, output_path=None):
        """
        Load and process audio files from the specified input path. Optionally, 
        write the output audio files to the specified output directory.
        
        Args:
            input_path (str): Path to the input audio files or folder.
            online_write (bool): Whether to write the processed audio to disk in real-time.
            output_path (str): Optional path for writing output files. If None, output 
                               will be stored in self.result.
        
        Returns:
            dict or ndarray: Processed audio results either as a dictionary or as a single array, 
                             depending on the number of audio files processed. 
                             Returns None if online_write is enabled.
        """
        
        self.result = {}
        self.args.input_path = input_path
        data_reader = DataReader(self.args)  # Initialize a data reader to load the audio files


        # Check if online writing is enabled
        if online_write:
            output_wave_dir = self.args.output_dir  # Set the default output directory
            if isinstance(output_path, str):  # If a specific output path is provided, use it
                output_wave_dir = os.path.join(output_path, self.name)
            # Create the output directory if it does not exist
            if not os.path.isdir(output_wave_dir):
                os.makedirs(output_wave_dir)
        
        num_samples = len(data_reader)  # Get the total number of samples to process
        print(f'Running {self.name} ...')  # Display the model being used

        if self.args.task == 'target_speaker_extraction':
            from utils.video_process import process_tse
            assert online_write == True
            process_tse(self.args, self.model, self.device, data_reader, output_wave_dir)
        else:
            # Disable gradient calculation for better efficiency during inference
            with torch.no_grad():
                for idx in tqdm(range(num_samples)):  # Loop over all audio samples
                    self.data = {}
                    # Read the audio, waveform ID, and audio length from the data reader
                    input_audio, wav_id, input_len, scalars, audio_info = data_reader[idx]
                    # Store the input audio and metadata in self.data
                    self.data['audio'] = input_audio
                    self.data['id'] = wav_id
                    self.data['audio_len'] = input_len
                    self.data.update(audio_info)
                    
                    # Perform the audio decoding/processing
                    output_audios = self.decode()

                    # Perform audio renormalization
                    if not isinstance(output_audios, list):
                        if len(scalars) > 1:
                            for i in range(len(scalars)):
                                output_audios[:,i] = output_audios[:,i] * scalars[i]
                        else:
                                output_audios = output_audios * scalars[0]
                        
                    if online_write:
                        # If online writing is enabled, save the output audio to files
                        if isinstance(output_audios, list):
                            # In case of multi-speaker output, save each speaker's output separately
                            for spk in range(self.args.num_spks):
                                output_file = os.path.join(output_wave_dir, wav_id.replace('.'+self.data['ext'], f'_s{spk+1}.'+self.data['ext']))
                                self.write_audio(output_file, key=None, spk=spk, audio=output_audios)
                        else:
                            # Single-speaker or standard output
                            output_file = os.path.join(output_wave_dir, wav_id)
                            self.write_audio(output_file, key=None, spk=None, audio=output_audios)
                    else:
                        # If not writing to disk, store the output in the result dictionary
                        self.result[wav_id] = output_audios
            
            # Return the processed results if not writing to disk
            if not online_write:
                if len(self.result) == 1:
                    # If there is only one result, return it directly
                    return next(iter(self.result.values()))
                else:
                    # Otherwise, return the entire result dictionary
                    return self.result

    def write_audio(self, output_path, key=None, spk=None, audio=None):
        """
        This function writes an audio signal to an output file, applying necessary transformations
        such as resampling, channel handling, and format conversion based on the provided parameters
        and the instance's internal settings.
        
        Args:
            output_path (str): The file path where the audio will be saved.
            key (str, optional): The key used to retrieve audio from the internal result dictionary
                                  if audio is not provided.
            spk (str, optional): A specific speaker identifier, used to extract a particular speaker's
                                 audio from a multi-speaker dataset or result.
            audio (numpy.ndarray, optional): A numpy array containing the audio data to be written.
                                 If provided, key and spk are ignored.
        """
        
        if audio is not None:
            if spk is not None:
                result_ = audio[spk]
            else:
                result_ = audio
        else:
            if spk is not None:
                result_ = self.result[key][spk]
            else:
                result_ = self.result[key]
                
        if self.data['sample_rate'] != self.args.sampling_rate:
            if self.data['channels'] == 2:
                left_channel = librosa.resample(result_[0,:], orig_sr=self.args.sampling_rate, target_sr=self.data['sample_rate'])
                right_channel = librosa.resample(result_[1,:], orig_sr=self.args.sampling_rate, target_sr=self.data['sample_rate'])
                result = np.vstack((left_channel, right_channel)).T
            else:
                result = librosa.resample(result_[0,:], orig_sr=self.args.sampling_rate, target_sr=self.data['sample_rate'])
        else:
            if self.data['channels'] == 2:
                left_channel = result_[0,:]
                right_channel = result_[1,:]
                result = np.vstack((left_channel, right_channel)).T
            else:
                result = result_[0,:]
                
        if self.data['sample_width'] == 4: ##32 bit float
            MAX_WAV_VALUE = 2147483648.0
            np_type = np.int32
        elif self.data['sample_width'] == 2: ##16 bit int
            MAX_WAV_VALUE = 32768.0
            np_type = np.int16
        else:
            self.data['sample_width'] = 2 ##16 bit int
            MAX_WAV_VALUE = 32768.0
            np_type = np.int16
                        
        result = result * MAX_WAV_VALUE
        result = result.astype(np_type)
        audio_segment = AudioSegment(
            result.tobytes(),  # Raw audio data as bytes
            frame_rate=self.data['sample_rate'],  # Sample rate
            sample_width=self.data['sample_width'],          # No. bytes per sample
            channels=self.data['channels']               # No. channels
        )
        audio_format = 'ipod' if self.data['ext'] in ['m4a', 'aac'] else self.data['ext']
        audio_segment.export(output_path, format=audio_format)
                    
    def write(self, output_path, add_subdir=False, use_key=False):
        """
        Write the processed audio results to the specified output path.

        Args:
            output_path (str): The directory or file path where processed audio will be saved. If not 
                               provided, defaults to self.args.output_dir.
            add_subdir (bool): If True, appends the model name as a subdirectory to the output path.
            use_key (bool): If True, uses the result dictionary's keys (audio file IDs) for filenames.

        Returns:
            None: Outputs are written to disk, no data is returned.
        """

        # Ensure the output path is a string. If not provided, use the default output directory
        if not isinstance(output_path, str):
            output_path = self.args.output_dir

        # If add_subdir is enabled, create a subdirectory for the model name
        if add_subdir:
            if os.path.isfile(output_path):
                print(f'File exists: {output_path}, remove it and try again!')
                return
            output_path = os.path.join(output_path, self.name)
            if not os.path.isdir(output_path):
                os.makedirs(output_path)

        # Ensure proper directory setup when using keys for filenames
        if use_key and not os.path.isdir(output_path):
            if os.path.exists(output_path):
                print(f'File exists: {output_path}, remove it and try again!')
                return
            os.makedirs(output_path)
        # If not using keys and output path is a directory, check for conflicts
        if not use_key and os.path.isdir(output_path):
            print(f'Directory exists: {output_path}, remove it and try again!')
            return

        # Iterate over the results dictionary to write the processed audio to disk
        for key in self.result:
            if use_key:
                # If using keys, format filenames based on the result dictionary's keys (audio IDs)
                if isinstance(self.result[key], list):  # For multi-speaker outputs
                    for spk in range(self.args.num_spks):
                        output_file = os.path.join(output_path, key.replace('.'+self.data['ext'], f'_s{spk+1}.'+self.data['ext']))
                        self.write_audio(output_file, key, spk)
                else:
                    output_file = os.path.join(output_path, key)
                    self.write_audio(output_path, key)
            else:
                # If not using keys, write audio to the specified output path directly
                if isinstance(self.result[key], list):  # For multi-speaker outputs
                    for spk in range(self.args.num_spks):
                        output_file = output_path.replace('.'+self.data['ext'], f'_s{spk+1}.'+self.data['ext'])
                        self.write_audio(output_file, key, spk)
                else:
                    self.write_audio(output_path, key)
                    
# The model classes for specific sub-tasks

class CLS_FRCRN_SE_16K(SpeechModel):
    """
    A subclass of SpeechModel that implements a speech enhancement model using 
    the FRCRN architecture for 16 kHz speech enhancement.
    
    Args:
        args (Namespace): The argument parser containing model configurations and paths.
    """

    def __init__(self, args):
        # Initialize the parent SpeechModel class
        super(CLS_FRCRN_SE_16K, self).__init__(args)
        
        # Import the FRCRN speech enhancement model for 16 kHz
        from models.frcrn_se.frcrn import FRCRN_SE_16K
        
        # Initialize the model
        self.model = FRCRN_SE_16K(args).model
        self.name = 'FRCRN_SE_16K'
        
        # Load pre-trained model checkpoint
        self.load_model()
        
        # Move model to the appropriate device (GPU/CPU)
        if args.use_cuda == 1:
            self.model.to(self.device)
        
        # Set the model to evaluation mode (no gradient calculation)
        self.model.eval()

class CLS_MossFormer2_SE_48K(SpeechModel):
    """
    A subclass of SpeechModel that implements the MossFormer2 architecture for 
    48 kHz speech enhancement.
    
    Args:
        args (Namespace): The argument parser containing model configurations and paths.
    """

    def __init__(self, args):
        # Initialize the parent SpeechModel class
        super(CLS_MossFormer2_SE_48K, self).__init__(args)
        
        # Import the MossFormer2 speech enhancement model for 48 kHz
        from models.mossformer2_se.mossformer2_se_wrapper import MossFormer2_SE_48K
        
        # Initialize the model
        self.model = MossFormer2_SE_48K(args).model
        self.name = 'MossFormer2_SE_48K'
        
        # Load pre-trained model checkpoint
        self.load_model()
        
        # Move model to the appropriate device (GPU/CPU)
        if args.use_cuda == 1:
            self.model.to(self.device)
        
        # Set the model to evaluation mode (no gradient calculation)
        self.model.eval()

class CLS_MossFormer2_SR_48K(SpeechModel):
    """
    A subclass of SpeechModel that implements the MossFormer2 architecture for
    48 kHz speech super-resolution.
    
    Args:
        args (Namespace): The argument parser containing model configurations and paths.
    """

    def __init__(self, args):
        # Initialize the parent SpeechModel class
        super(CLS_MossFormer2_SR_48K, self).__init__(args)
        
        # Import the MossFormer2 speech enhancement model for 48 kHz
        from models.mossformer2_sr.mossformer2_sr_wrapper import MossFormer2_SR_48K
        
        # Initialize the model
        self.model = nn.ModuleList()
        self.model.append(MossFormer2_SR_48K(args).model_m)
        self.model.append(MossFormer2_SR_48K(args).model_g)
        self.name = 'MossFormer2_SR_48K'
        
        # Load pre-trained model checkpoint
        self.load_model()
        
        # Move model to the appropriate device (GPU/CPU)
        if args.use_cuda == 1:
            for model in self.model:
                model.to(self.device)
        
        # Set the model to evaluation mode (no gradient calculation)
        for model in self.model:
            model.eval()
        self.model[1].remove_weight_norm()

class CLS_MossFormerGAN_SE_16K(SpeechModel):
    """
    A subclass of SpeechModel that implements the MossFormerGAN architecture for 
    16 kHz speech enhancement, utilizing GAN-based speech processing.
    
    Args:
        args (Namespace): The argument parser containing model configurations and paths.
    """

    def __init__(self, args):
        # Initialize the parent SpeechModel class
        super(CLS_MossFormerGAN_SE_16K, self).__init__(args)
        
        # Import the MossFormerGAN speech enhancement model for 16 kHz
        from models.mossformer_gan_se.generator import MossFormerGAN_SE_16K
        
        # Initialize the model
        self.model = MossFormerGAN_SE_16K(args).model
        self.name = 'MossFormerGAN_SE_16K'
        
        # Load pre-trained model checkpoint
        self.load_model()
        
        # Move model to the appropriate device (GPU/CPU)
        if args.use_cuda == 1:
            self.model.to(self.device)
        
        # Set the model to evaluation mode (no gradient calculation)
        self.model.eval()

class CLS_MossFormer2_SS_16K(SpeechModel):
    """
    A subclass of SpeechModel that implements the MossFormer2 architecture for 
    16 kHz speech separation.
    
    Args:
        args (Namespace): The argument parser containing model configurations and paths.
    """

    def __init__(self, args):
        # Initialize the parent SpeechModel class
        super(CLS_MossFormer2_SS_16K, self).__init__(args)
        
        # Import the MossFormer2 speech separation model for 16 kHz
        from models.mossformer2_ss.mossformer2 import MossFormer2_SS_16K
        
        # Initialize the model
        self.model = MossFormer2_SS_16K(args).model
        self.name = 'MossFormer2_SS_16K'
        
        # Load pre-trained model checkpoint
        self.load_model()
        
        # Move model to the appropriate device (GPU/CPU)
        if args.use_cuda == 1:
            self.model.to(self.device)
        
        # Set the model to evaluation mode (no gradient calculation)
        self.model.eval()


class CLS_AV_MossFormer2_TSE_16K(SpeechModel):
    """
    A subclass of SpeechModel that implements an audio-visual (AV) model using 
    the AV-MossFormer2 architecture for target speaker extraction (TSE) at 16 kHz. 
    This model leverages both audio and visual cues to perform speaker extraction.
    
    Args:
        args (Namespace): The argument parser containing model configurations and paths.
    """

    def __init__(self, args):
        # Initialize the parent SpeechModel class
        super(CLS_AV_MossFormer2_TSE_16K, self).__init__(args)
        
        # Import the AV-MossFormer2 model for 16 kHz target speech enhancement
        from models.av_mossformer2_tse.av_mossformer2 import AV_MossFormer2_TSE_16K
        
        # Initialize the model
        self.model = AV_MossFormer2_TSE_16K(args).model
        self.name = 'AV_MossFormer2_TSE_16K'
        
        # Load pre-trained model checkpoint
        self.load_model()
        
        # Move model to the appropriate device (GPU/CPU)
        if args.use_cuda == 1:
            self.model.to(self.device)
        
        # Set the model to evaluation mode (no gradient calculation)
        self.model.eval()