92 lines
3.0 KiB
Python
92 lines
3.0 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
Author: C0ffymachyne
|
|
License: GPLv3
|
|
Version: 1.0.0
|
|
|
|
Description:
|
|
This file contains a description of a compressor node that utilizes a CUDA-optimized kernel.
|
|
"""
|
|
import torch
|
|
from typing import Dict, Any, Tuple, Union
|
|
|
|
from ..core.io import audio_to_comfy_3d, audio_from_comfy_2d
|
|
from ..core.compression import compressor
|
|
from ..core.loudness import get_loudness, lufs_normalization
|
|
|
|
|
|
class SignalProcessingCompressor:
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
|
return {
|
|
"required": {
|
|
"audio_input": ("AUDIO",),
|
|
"comp": (
|
|
"FLOAT",
|
|
{"default": -0.3, "min": -2.0, "max": 2.0, "step": 0.01},
|
|
),
|
|
"attack": (
|
|
"FLOAT",
|
|
{"default": 0.1, "min": 0.01, "max": 100.0, "step": 0.01},
|
|
),
|
|
"release": (
|
|
"FLOAT",
|
|
{"default": 60.0, "min": 0.01, "max": 1000.0, "step": 0.1},
|
|
),
|
|
"filter_param": (
|
|
"FLOAT",
|
|
{"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("AUDIO",)
|
|
RETURN_NAMES = ("audio",)
|
|
CATEGORY = "Signal Processing"
|
|
FUNCTION = "process"
|
|
|
|
def process(
|
|
self,
|
|
audio_input: Dict[str, Union[torch.Tensor, int]],
|
|
comp: float = -0.3, # Compression/expansion factor
|
|
attack: float = 0.1, # Attack time in ms
|
|
release: float = 60.0, # Release time in ms
|
|
filter_param: float = 0.3, # Filter parameter < 1
|
|
) -> Tuple[Dict[str, Union[torch.Tensor, int]]]:
|
|
"""
|
|
Apply compression or expansion to the audio input using CUDA.
|
|
|
|
Parameters:
|
|
audio_input (Dict[str, Union[torch.Tensor, int]]): Input audio waveform and sample rate.
|
|
comp (float): Compression/expansion factor.
|
|
attack (float): Attack time in milliseconds.
|
|
release (float): Release time in milliseconds.
|
|
filter_param (float): Filter parameter for envelope smoothing.
|
|
|
|
Returns:
|
|
Tuple[Dict[str, Union[torch.Tensor, int]]]: Compressed audio and sample rate.
|
|
"""
|
|
# Extract waveform and sample rate
|
|
waveform, sample_rate = audio_from_comfy_2d(audio_input, try_gpu=True)
|
|
|
|
loudness = get_loudness(waveform, sample_rate=sample_rate)
|
|
|
|
# Apply the compressor kernel
|
|
filtered_waveform, _ = compressor(
|
|
waveform,
|
|
sample_rate,
|
|
comp=comp,
|
|
attack=attack,
|
|
release=release,
|
|
a=filter_param,
|
|
device="cuda" if torch.cuda.is_available() else "cpu",
|
|
)
|
|
|
|
filtered_waveform = lufs_normalization(
|
|
filtered_waveform, sample_rate=sample_rate, target_lufs=loudness
|
|
)
|
|
|
|
# Return the processed audio
|
|
return audio_to_comfy_3d(filtered_waveform, sample_rate)
|