105 lines
3.0 KiB
Python
105 lines
3.0 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
Author: C0ffymachyne
|
|
License: GPLv3
|
|
Version: 1.0.0
|
|
|
|
Description:
|
|
various compression methods
|
|
"""
|
|
|
|
|
|
import cupy as cp
|
|
import torch
|
|
import numpy as np
|
|
|
|
from typing import Tuple, Any
|
|
from ..core.utilitiescuda import read_kernel_by_name
|
|
|
|
compressor_kernel = read_kernel_by_name(
|
|
"compressor", kernel_class="compressor", kernel_identifier="compexp_kernel"
|
|
)
|
|
|
|
|
|
def compressor(
|
|
audio_in: torch.Tensor,
|
|
sample_rate: int,
|
|
comp: float = -0.3, # Compression/expansion factor
|
|
attack: float = 0.1, # Attack time in ms
|
|
release: float = 60.0, # Release time in ms
|
|
a: float = 0.3, # Filter parameter < 1
|
|
device: str = "cuda",
|
|
) -> Tuple[torch.Tensor, Any]:
|
|
"""
|
|
Compresses or expands stereo audio using an optimized CUDA kernel.
|
|
|
|
Parameters:
|
|
audio_in (torch.Tensor or np.ndarray): Input stereo audio signal with shape (n_samples, 2).
|
|
sample_rate (int): Sampling rate in Hz.
|
|
comp (float): Compression/expansion factor.
|
|
attack (float): Attack time in milliseconds.
|
|
release (float): Release time in milliseconds.
|
|
a (float): Filter parameter (< 1) for envelope smoothing.
|
|
device (str): Device to place the output tensor ('cuda' or 'cpu').
|
|
|
|
Returns:
|
|
torch.Tensor: Compressed stereo audio with shape (n_samples, 2).
|
|
"""
|
|
# Convert input to NumPy array if necessary
|
|
audio_in = audio_in.T
|
|
if isinstance(audio_in, torch.Tensor):
|
|
audio_in = audio_in.detach().cpu().numpy()
|
|
else:
|
|
audio_in = np.asarray(audio_in, dtype=np.float64)
|
|
|
|
# 2. Ensure the audio is in shape (n_samples, 2)
|
|
if audio_in.ndim != 2 or audio_in.shape[1] != 2:
|
|
raise ValueError(
|
|
f"Input audio must have shape (n_samples, 2), but got {audio_in.shape}"
|
|
)
|
|
|
|
n_samples, n_channels = audio_in.shape
|
|
|
|
# Flatten the audio for kernel processing
|
|
wav_in_flat = audio_in.flatten()
|
|
|
|
# Move data to GPU
|
|
wav_in_gpu = cp.asarray(wav_in_flat, dtype=cp.float64)
|
|
wav_out_gpu = cp.zeros_like(wav_in_gpu)
|
|
|
|
# Define grid and block dimensions
|
|
block_size = 256 # Number of threads per block
|
|
grid_size = n_channels # One block per channel
|
|
|
|
# Launch the CUDA kernel
|
|
compressor_kernel(
|
|
(grid_size,), # Grid dimensions
|
|
(block_size,), # Block dimensions
|
|
(
|
|
wav_in_gpu,
|
|
wav_out_gpu,
|
|
np.int32(n_channels),
|
|
np.int32(n_samples),
|
|
np.float64(comp),
|
|
np.float64(release),
|
|
np.float64(attack),
|
|
np.float64(a),
|
|
np.float64(sample_rate),
|
|
),
|
|
)
|
|
|
|
# Retrieve results from GPU
|
|
wav_out_host = wav_out_gpu.get().astype(np.float64)
|
|
|
|
# Reshape the output
|
|
wav_out_stereo = wav_out_host.reshape((n_samples, n_channels))
|
|
|
|
# Convert back to Torch tensor if desired
|
|
if device == "cuda":
|
|
out_tensor = torch.from_numpy(wav_out_stereo).to("cuda")
|
|
else:
|
|
out_tensor = torch.from_numpy(wav_out_stereo).to("cpu")
|
|
|
|
return out_tensor.T, None
|