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)