Delete networks.py

networks.py

@@ -1,490 +0,0 @@
-"""
-Authors: Shengkui Zhao, Zexu Pan
-"""
-
-import torch
-import soundfile as sf
-import os
-import subprocess
-from tqdm import tqdm
-from utils.decode import decode_one_audio
-from dataloader.dataloader import DataReader
-
-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)
-                print(f'use GPU: {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')
-        """    
-        if torch.cuda.is_available():
-            print('GPU is found and used!')
-            self.device = torch.device('cuda')
-        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 = {}
-
-    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 load_model(self):
-        """
-        Loads a pre-trained model checkpoint from a specified directory. It checks for 
-        the best model ('last_best_checkpoint') or the most recent checkpoint ('last_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 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')
-        ckpt_name = os.path.join(self.args.checkpoint_dir, 'last_checkpoint')
-
-        # Check if the best checkpoint or last checkpoint exists
-        if os.path.isfile(best_name):
-            name = best_name  # Prioritize loading the best model
-        elif os.path.isfile(ckpt_name):
-            name = ckpt_name  # Otherwise, load the last saved checkpoint
-        else:
-            # If no checkpoint exists, print a warning and exit the function
-            print('Warning: No existing checkpoint or best model found!')
-            return
-
-        # Read the model's checkpoint name from the file
-        with open(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)
-        
-        # 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' in checkpoint:
-            # If the checkpoint contains a 'model' key, load the corresponding state dictionary
-            if self.args.task =='target_speaker_extraction':
-                pretrained_model = checkpoint['model']
-                state = self.model.state_dict()
-                for key in state.keys():
-                    pretrain_key = 'module.' + key
-                    state[key] = pretrained_model[pretrain_key]
-                self.model.load_state_dict(state, strict=True)
-            else:
-                self.model.load_state_dict(checkpoint['model'], strict=False)
-        else:
-            # If the checkpoint is a plain state dictionary, load it directly
-            self.model.load_state_dict(checkpoint, strict=False)
-        '''
-        if 'model' in checkpoint:
-            pretrained_model = checkpoint['model']
-        else:
-            pretrained_model = checkpoint
-        state = self.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')
-        self.model.load_state_dict(state)
-        print(f'Successfully loaded {model_name} for decoding')
-
-    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_audio = decode_one_audio(self.model, self.device, self.data['audio'], 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']]
-    
-        return output_audio
-
-    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, scalar = 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
-                    
-                    # Perform the audio decoding/processing
-                    output_audio = self.decode()
-                    
-                    #if isinstance(output_audio, list):
-                    #    for spk in range(self.args.num_spks):
-                    #        output_audio[spk] = output_audio[spk] * scalar
-                    #else:
-                    #if not isinstance(output_audio, list):
-                    if self.args.network == 'FRCRN_SE_16K':
-                        output_audio = output_audio * scalar
-                    
-                    if online_write:
-                        # If online writing is enabled, save the output audio to files
-                        if isinstance(output_audio, 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('.wav', f'_s{spk+1}.wav'))
-                                sf.write(output_file, output_audio[spk], self.args.sampling_rate)
-                        else:
-                            # Single-speaker or standard output
-                            output_file = os.path.join(output_wave_dir, wav_id)
-                            sf.write(output_file, output_audio, self.args.sampling_rate)
-                    else:
-                        # If not writing to disk, store the output in the result dictionary
-                        self.result[wav_id] = output_audio
-            
-            # 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(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):
-                        sf.write(os.path.join(output_path, key.replace('.wav', f'_s{spk+1}.wav')),
-                                 self.result[key][spk], self.args.sampling_rate)
-                else:
-                    sf.write(os.path.join(output_path, key), self.result[key], self.args.sampling_rate)
-            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):
-                        sf.write(output_path.replace('.wav', f'_s{spk+1}.wav'),
-                                 self.result[key][spk], self.args.sampling_rate)
-                else:
-                    sf.write(output_path, self.result[key], self.args.sampling_rate)
-
-# 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)
-        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)
-        self.model.to(self.device)
-        
-        # Set the model to evaluation mode (no gradient calculation)
-        self.model.eval()
-
-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)
-        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)
-        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)
-        self.model.to(self.device)
-        
-        # Set the model to evaluation mode (no gradient calculation)
-        self.model.eval()
-
-