Spaces:
Running on Zero
Running on Zero
| import math | |
| from utils.feature_extractors.dsp_features import compute_log_rms_gated_max, compute_crest_factor, compute_stereo_width, compute_stereo_imbalance, compute_log_spread | |
| import torch | |
| class AF_fourier_embedding: | |
| def __init__(self, | |
| input_dim=8, | |
| output_dim=64, | |
| sigma=0.2, | |
| log_rms_shift=-26.5, #calculated as the mean from the dataset | |
| log_rms_scale=7.0, #calculated as the std from the dataset | |
| crest_shift=16.7, #calculated as the mean from the dataset | |
| crest_scale=6.3, | |
| log_spread_shift=-20.0, #calculated as the mean from the dataset | |
| log_spread_scale=20.0, #calculated as the std from the dataset | |
| stereo_width_shift=0.28, | |
| stereo_width_scale=0.39, | |
| stereo_imbalance_shift=0.0, | |
| stereo_imbalance_scale=0.35, | |
| device="cpu" | |
| ): | |
| """ | |
| Deterministic Fourier feature transformer using fixed cosine-based projection | |
| """ | |
| self.device = device | |
| # Ensure output_dim is even and >= 2 * input_dim | |
| self.output_dim = max(input_dim * 2, output_dim) | |
| if self.output_dim % 2 != 0: | |
| self.output_dim += 1 | |
| self.input_dim = input_dim | |
| self.sigma = sigma | |
| # Create deterministic projection matrix | |
| self.projection = self._create_deterministic_projection(input_dim, self.output_dim // 2, sigma) | |
| self.projection = self.projection.to(self.device) | |
| # Normalization factor | |
| self.scale_factor = math.sqrt(2.0 / self.output_dim) | |
| self.log_rms_shift = log_rms_shift | |
| self.log_rms_scale = log_rms_scale | |
| self.crest_shift = crest_shift | |
| self.crest_scale = crest_scale | |
| self.log_spread_shift = log_spread_shift | |
| self.log_spread_scale = log_spread_scale | |
| self.stereo_width_shift = stereo_width_shift | |
| self.stereo_width_scale = stereo_width_scale | |
| self.stereo_imbalance_shift = stereo_imbalance_shift | |
| self.stereo_imbalance_scale = stereo_imbalance_scale | |
| def _create_deterministic_projection(self, input_dim, proj_dim, sigma): | |
| """ | |
| Create a deterministic projection matrix using a cosine basis | |
| """ | |
| # Cosine-based matrix (like DCT type-II) | |
| projection = torch.zeros(input_dim, proj_dim) | |
| for i in range(input_dim): | |
| for j in range(proj_dim): | |
| projection[i, j] = math.cos(math.pi * (i + 0.5) * (j + 1) / proj_dim) | |
| return projection * sigma | |
| def encode(self, x): | |
| log_rms=compute_log_rms_gated_max(x) | |
| crest_factor= compute_crest_factor(x) | |
| log_spread= compute_log_spread(x) | |
| stereo_width= compute_stereo_width(x) | |
| stereo_imbalance= compute_stereo_imbalance(x) | |
| log_rms_std, crest_factor_std, log_spread_std, stereo_width_std, stereo_imbalance_std = self.standardize_features( | |
| log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance | |
| ) | |
| embedding= self.transform( | |
| log_rms_std, crest_factor_std, log_spread_std, stereo_width_std, stereo_imbalance_std | |
| ) | |
| return embedding, (log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance) | |
| def decode(self, fourier_features): | |
| """ | |
| Invert Fourier features back to original feature space | |
| (approximate due to phase-only reconstruction) | |
| """ | |
| reconstructed = self.inverse_transform(fourier_features) | |
| # Reshape back to original feature dimensions | |
| log_rms= reconstructed[:,0:2] | |
| crest_factor = reconstructed[:,2:4] | |
| log_spread= reconstructed[:,4:6] | |
| stereo_width = reconstructed[:,6:7] | |
| stereo_imbalance = reconstructed[:,7:8] | |
| log_rms, crest_factor, log_spread,stereo_width, stereo_imbalance = self.destandardize_features( | |
| log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance | |
| ) | |
| return log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance | |
| def standardize_features(self, log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance): | |
| """ | |
| Standardize features using pre-computed mean and std | |
| """ | |
| log_rms = (log_rms - self.log_rms_shift) / self.log_rms_scale | |
| crest_factor = (crest_factor - self.crest_shift) / self.crest_scale | |
| log_spread = (log_spread - self.log_spread_shift) / self.log_spread_scale | |
| stereo_width = (stereo_width - self.stereo_width_shift) / self.stereo_width_scale | |
| stereo_imbalance = (stereo_imbalance - self.stereo_imbalance_shift) / self.stereo_imbalance_scale | |
| return log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance | |
| def destandardize_features(self, log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance): | |
| """ | |
| Reverse standardization to get back to original feature space | |
| """ | |
| log_rms = log_rms * self.log_rms_scale + self.log_rms_shift | |
| crest_factor = crest_factor * self.crest_scale + self.crest_shift | |
| log_spread = log_spread * self.log_spread_scale + self.log_spread_shift | |
| stereo_width = stereo_width * self.stereo_width_scale + self.stereo_width_shift | |
| stereo_imbalance = stereo_imbalance * self.stereo_imbalance_scale + self.stereo_imbalance_shift | |
| return log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance | |
| def transform(self, log_rms, crest_factor,log_spread, stereo_width, stereo_imbalance): | |
| """ | |
| Transform features using the stored projection matrix | |
| """ | |
| flat_features=torch.cat([log_rms, crest_factor, log_spread, stereo_width.unsqueeze(-1), stereo_imbalance.unsqueeze(-1)], dim=-1) | |
| # Project and transform | |
| projected = flat_features @ self.projection | |
| cos_features = torch.cos(projected) | |
| sin_features = torch.sin(projected) | |
| # Concatenate and normalize | |
| return torch.cat([cos_features, sin_features], dim=-1) * self.scale_factor | |
| def inverse_transform(self, fourier_features): | |
| """ | |
| Invert Fourier features back to original feature space | |
| (approximate due to phase-only reconstruction) | |
| """ | |
| # Split into cosine and sine components | |
| feature_dim = fourier_features.shape[-1] // 2 | |
| cos_features = fourier_features[:, :feature_dim] | |
| sin_features = fourier_features[:, feature_dim:] | |
| # Denormalize | |
| cos_features = cos_features / self.scale_factor | |
| sin_features = sin_features / self.scale_factor | |
| # Compute phase angles | |
| phases = torch.atan2(sin_features, cos_features) | |
| # Use pseudo-inverse for approximate inversion | |
| projection_pinv = torch.pinverse(self.projection) | |
| reconstructed = phases @ projection_pinv | |
| return reconstructed | |