inference.py

import torch
import soundfile as sf
import numpy as np
import argparse
import os
import yaml
import julius

import sys
currentdir = os.path.dirname(os.path.realpath(__file__))
sys.path.append(os.path.dirname(currentdir))
from networks import Dasp_Mastering_Style_Transfer, Effects_Encoder
from modules.loss import AudioFeatureLoss, Loss, CLAPFeatureLoss
from modules.data_normalization import Audio_Effects_Normalizer


def convert_audio(wav: torch.Tensor, from_rate: float,
                  to_rate: float, to_channels: int) -> torch.Tensor:
    """Convert audio to new sample rate and number of audio channels.
    """
    wav = julius.resample_frac(wav, int(from_rate), int(to_rate))
    wav = convert_audio_channels(wav, to_channels)
    return wav

class MasteringStyleTransfer:
    def __init__(self, args):
        self.args = args
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # Load models
        self.effects_encoder = self.load_effects_encoder()
        self.mastering_converter = self.load_mastering_converter()

        self.fx_normalizer = Audio_Effects_Normalizer(precomputed_feature_path=args.fx_norm_feature_path, \
                                                        STEMS=['mixture'], \
                                                        EFFECTS=['eq', 'imager', 'loudness'])
        # Loss functions
        self.clap_loss = CLAPFeatureLoss()

    def load_effects_encoder(self):
        effects_encoder = Effects_Encoder(self.args.cfg_enc)
        reload_weights(effects_encoder, self.args.encoder_path, self.device)
        effects_encoder.to(self.device)
        effects_encoder.eval()
        return effects_encoder

    def load_mastering_converter(self):
        mastering_converter = Dasp_Mastering_Style_Transfer(num_features=2048,
                                                            sample_rate=self.args.sample_rate,
                                                            tgt_fx_names=['eq', 'distortion', 'multiband_comp', 'gain', 'imager', 'limiter'],
                                                            model_type='tcn',
                                                            config=self.args.cfg_converter,
                                                            batch_size=1)
        reload_weights(mastering_converter, self.args.model_path, self.device)
        mastering_converter.to(self.device)
        mastering_converter.eval()
        return mastering_converter

    def get_reference_embedding(self, reference_tensor):
        with torch.no_grad():
            reference_feature = self.effects_encoder(reference_tensor)
        return reference_feature

    def mastering_style_transfer(self, input_tensor, reference_feature):
        with torch.no_grad():
            output_audio = self.mastering_converter(input_tensor, reference_feature)
            predicted_params = self.mastering_converter.get_last_predicted_params()
        return output_audio, predicted_params

    def inference_time_optimization(self, input_tensor, reference_tensor, ito_config, initial_reference_feature):
        fit_embedding = torch.nn.Parameter(initial_reference_feature, requires_grad=True)
        optimizer = getattr(torch.optim, ito_config['optimizer'])([fit_embedding], lr=ito_config['learning_rate'])

        min_loss = float('inf')
        min_loss_step = 0
        all_results = []

        af_loss = AudioFeatureLoss(
            weights=ito_config['af_weights'],
            sample_rate=ito_config['sample_rate'],
            stem_separation=False,
            use_clap=False
        )

        for step in range(ito_config['num_steps']):
            optimizer.zero_grad()

            output_audio = self.mastering_converter(input_tensor, fit_embedding)
            current_params = self.mastering_converter.get_last_predicted_params()

            # Compute loss
            if ito_config['loss_function'] == 'AudioFeatureLoss':
                losses = af_loss(output_audio, reference_tensor)
                total_loss = sum(losses.values())
            elif ito_config['loss_function'] == 'CLAPFeatureLoss':
                if ito_config['clap_target_type'] == 'Audio':
                    target = reference_tensor
                else:
                    target = ito_config['clap_text_prompt']
                total_loss = self.clap_loss(output_audio, target, self.args.sample_rate, distance_fn=ito_config['clap_distance_fn'])

            if total_loss < min_loss:
                min_loss = total_loss.item()
                min_loss_step = step

            # Log top 5 parameter differences
            if step == 0:
                initial_params = current_params
            top_5_diff = self.get_top_n_diff_string(initial_params, current_params, top_n=5)
            log_entry = f"Step {step + 1}\n   Loss: {total_loss.item():.4f}\n{top_5_diff}\n"

            all_results.append({
                'step': step + 1,
                'loss': total_loss.item(),
                'audio': output_audio.detach().cpu().numpy(),
                'params': current_params,
                'log': log_entry
            })

            total_loss.backward()
            optimizer.step()

            gc.collect()
            torch.cuda.empty_cache()

        return all_results, min_loss_step

    def preprocess_audio(self, audio, target_sample_rate=44100, normalize=False):
        sample_rate, data = audio

        # Normalize audio to -1 to 1 range
        if data.dtype == np.int16:
            data = data.astype(np.float32) / 32768.0
        elif data.dtype == np.float32:
            data = np.clip(data, -1.0, 1.0)
        else:
            raise ValueError(f"Unsupported audio data type: {data.dtype}")

        # Ensure stereo channels
        if data.ndim == 1:
            data = np.stack([data, data])
        elif data.ndim == 2:
            if data.shape[0] == 2:
                pass  # Already in correct shape
            elif data.shape[1] == 2:
                data = data.T
            else:
                data = np.stack([data[:, 0], data[:, 0]])  # Duplicate mono channel
        else:
            raise ValueError(f"Unsupported audio shape: {data.shape}")

        # Resample if necessary
        if sample_rate != target_sample_rate:
            data = julius.resample_frac(torch.from_numpy(data), sample_rate, target_sample_rate).numpy()

        # Apply fx normalization for input audio during mastering style transfer
        if normalize:
            data = self.fx_normalizer.normalize_audio(data.T, 'mixture').T

        # Convert to torch tensor
        data_tensor = torch.FloatTensor(data).unsqueeze(0)

        return data_tensor.to(self.device)

    def process_audio(self, input_audio, reference_audio):
        input_tensor = self.preprocess_audio(input_audio, self.args.sample_rate, normalize=True)
        reference_tensor = self.preprocess_audio(reference_audio, self.args.sample_rate)

        reference_feature = self.get_reference_embedding(reference_tensor)

        output_audio, predicted_params = self.mastering_style_transfer(input_tensor, reference_feature)

        return output_audio, predicted_params, self.args.sample_rate, input_tensor

    def get_param_output_string(self, params):
        if params is None:
            return "No parameters available"
        
        param_mapper = {
            'eq': {
                'low_shelf_gain_db': ('Low Shelf Gain', 'dB', -20, 20),
                'low_shelf_cutoff_freq': ('Low Shelf Cutoff', 'Hz', 20, 2000),
                'low_shelf_q_factor': ('Low Shelf Q', '', 0.1, 5.0),
                'band0_gain_db': ('Low-Mid Band Gain', 'dB', -20, 20),
                'band0_cutoff_freq': ('Low-Mid Band Frequency', 'Hz', 80, 2000),
                'band0_q_factor': ('Low-Mid Band Q', '', 0.1, 5.0),
                'band1_gain_db': ('Mid Band Gain', 'dB', -20, 20),
                'band1_cutoff_freq': ('Mid Band Frequency', 'Hz', 2000, 8000),
                'band1_q_factor': ('Mid Band Q', '', 0.1, 5.0),
                'band2_gain_db': ('High-Mid Band Gain', 'dB', -20, 20),
                'band2_cutoff_freq': ('High-Mid Band Frequency', 'Hz', 8000, 12000),
                'band2_q_factor': ('High-Mid Band Q', '', 0.1, 5.0),
                'band3_gain_db': ('High Band Gain', 'dB', -20, 20),
                'band3_cutoff_freq': ('High Band Frequency', 'Hz', 12000, 20000),
                'band3_q_factor': ('High Band Q', '', 0.1, 5.0),
                'high_shelf_gain_db': ('High Shelf Gain', 'dB', -20, 20),
                'high_shelf_cutoff_freq': ('High Shelf Cutoff', 'Hz', 4000, 20000),
                'high_shelf_q_factor': ('High Shelf Q', '', 0.1, 5.0),
            },
            'distortion': {
                'drive_db': ('Drive', 'dB', 0, 8),
                'parallel_weight_factor': ('Dry/Wet Mix', '%', 0, 100),
            },
            'multiband_comp': {
                'low_cutoff': ('Low/Mid Crossover', 'Hz', 20, 1000),
                'high_cutoff': ('Mid/High Crossover', 'Hz', 1000, 20000),
                'parallel_weight_factor': ('Dry/Wet Mix', '%', 0, 100),
                'low_shelf_comp_thresh': ('Low Band Comp Threshold', 'dB', -60, 0),
                'low_shelf_comp_ratio': ('Low Band Comp Ratio', ': 1', 1, 20),
                'low_shelf_exp_thresh': ('Low Band Exp Threshold', 'dB', -60, 0),
                'low_shelf_exp_ratio': ('Low Band Exp Ratio', ': 1', 1, 20),
                'low_shelf_at': ('Low Band Attack Time', 'ms', 5, 100),
                'low_shelf_rt': ('Low Band Release Time', 'ms', 5, 100),
                'mid_band_comp_thresh': ('Mid Band Comp Threshold', 'dB', -60, 0),
                'mid_band_comp_ratio': ('Mid Band Comp Ratio', ': 1', 1, 20),
                'mid_band_exp_thresh': ('Mid Band Exp Threshold', 'dB', -60, 0),
                'mid_band_exp_ratio': ('Mid Band Exp Ratio', ': 1', 0, 1),
                'mid_band_at': ('Mid Band Attack Time', 'ms', 5, 100),
                'mid_band_rt': ('Mid Band Release Time', 'ms', 5, 100),
                'high_shelf_comp_thresh': ('High Band Comp Threshold', 'dB', -60, 0),
                'high_shelf_comp_ratio': ('High Band Comp Ratio', ': 1', 1, 20),
                'high_shelf_exp_thresh': ('High Band Exp Threshold', 'dB', -60, 0),
                'high_shelf_exp_ratio': ('High Band Exp Ratio', ': 1', 1, 20),
                'high_shelf_at': ('High Band Attack Time', 'ms', 5, 100),
                'high_shelf_rt': ('High Band Release Time', 'ms', 5, 100),
            },
            'gain': {
                'gain_db': ('Output Gain', 'dB', -24, 24),
            },
            'imager': {
                'width': ('Stereo Width', '', 0, 1),
            },
            'limiter': {
                'threshold': ('Threshold', 'dB', -60, 0),
                'at': ('Attack Time', 'ms', 5, 100),
                'rt': ('Release Time', 'ms', 5, 100),
            },
        }
        
        output = []
        for fx_name, fx_params in params.items():
            output.append(f"{fx_name.upper()}:")
            if isinstance(fx_params, dict):
                for param_name, param_value in fx_params.items():
                    if isinstance(param_value, torch.Tensor):
                        param_value = param_value.item()
                    
                    if fx_name in param_mapper and param_name in param_mapper[fx_name]:
                        friendly_name, unit, min_val, max_val = param_mapper[fx_name][param_name]
                        if unit=='%':
                            param_value = param_value * 100
                        current_content = f"  {friendly_name}: {param_value:.2f} {unit}"
                        if param_name=='mid_band_exp_ratio':
                            current_content += f" (Range: {min_val}-{max_val})"
                        output.append(current_content)
                    else:
                        output.append(f"  {param_name}: {param_value:.2f}")
            else:
                # stereo imager
                width_percentage = fx_params.item() * 200
                output.append(f"  Stereo Width: {width_percentage:.2f}% (Range: 0-200%)")
    
        return "\n".join(output)

    def get_top_n_diff_string(self, initial_params, ito_params, top_n=5):
        if initial_params is None or ito_params is None:
            return "Cannot compare parameters"
        
        all_diffs = []
        for fx_name in initial_params.keys():
            if isinstance(initial_params[fx_name], dict):
                for param_name in initial_params[fx_name].keys():
                    initial_value = initial_params[fx_name][param_name]
                    ito_value = ito_params[fx_name][param_name]
                    
                    param_range = self.mastering_converter.fx_processors[fx_name].param_ranges[param_name]
                    normalized_diff = abs((ito_value - initial_value) / (param_range[1] - param_range[0]))
                    
                    all_diffs.append((fx_name, param_name, initial_value.item(), ito_value.item(), normalized_diff.item()))
            else:
                initial_value = initial_params[fx_name]
                ito_value = ito_params[fx_name]
                normalized_diff = abs(ito_value - initial_value)
                all_diffs.append((fx_name, 'width', initial_value.item(), ito_value.item(), normalized_diff.item()))

        top_diffs = sorted(all_diffs, key=lambda x: x[4], reverse=True)[:top_n]
        
        output = [f"   Top {top_n} parameter differences (initial / ITO / normalized diff):"]
        for fx_name, param_name, initial_value, ito_value, normalized_diff in top_diffs:
            output.append(f"      {fx_name.upper()} - {param_name}: {initial_value:.2f} / {ito_value:.2f} / {normalized_diff:.2f}")
        
        return "\n".join(output)

def reload_weights(model, ckpt_path, device):
    checkpoint = torch.load(ckpt_path, map_location=device)

    from collections import OrderedDict
    new_state_dict = OrderedDict()
    for k, v in checkpoint["model"].items():
        name = k[7:] # remove `module.`
        new_state_dict[name] = v
    model.load_state_dict(new_state_dict, strict=False)