Files
c0ffymachyne-ComfyUI_Signal…/processors/SignalProcessingBaxandallEQ.py
T
2024-12-07 20:16:47 -08:00

225 lines
8.9 KiB
Python

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Author: C0ffymachyne
License: GPLv3
Version: 1.0.0
Description:
Baxandall shelf EQ
references used : # reference https://webaudio.github.io/Audio-EQ-Cookbook/Audio-EQ-Cookbook.txt
"""
import torch
import torchaudio
import math
from ..core.io import audio_to_comfy_3d, audio_from_comfy_3d
from ..core.loudness import lufs_normalization, get_loudness
import torch
import math
class SignalProcessingBaxandallEQ:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"audio_input": ("AUDIO",),
"bass_gain_db": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.1}),
"treble_gain_db": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.1}),
},
}
RETURN_TYPES = ("AUDIO",)
RETURN_NAMES = ("processed_audio",)
CATEGORY = "Audio Processing"
FUNCTION = "process"
def process(self, audio_input, bass_gain_db=0.0, treble_gain_db=0.0):
waveform, sample_rate = audio_from_comfy_3d(audio_input, try_gpu=True)
loudness = get_loudness(waveform, sample_rate)
# Apply Bass Shelf (low shelf) using RBJ formula
b_bass, a_bass = self.design_rbj_shelf(sample_rate, freq=100.0, gain_db=bass_gain_db, shelf_type='low')
waveform = torchaudio.functional.lfilter(waveform, a_bass.to(waveform.device), b_bass.to(waveform.device),clamp=False)
# Apply Treble Shelf (high shelf) using RBJ formula
b_treble, a_treble = self.design_rbj_shelf(sample_rate, freq=10000.0, gain_db=treble_gain_db, shelf_type='high')
waveform = torchaudio.functional.lfilter(waveform, a_treble.to(waveform.device), b_treble.to(waveform.device),clamp=False)
waveform = lufs_normalization(waveform, sample_rate, loudness)
return audio_to_comfy_3d(waveform, sample_rate)
def design_rbj_shelf(self, sr, freq, gain_db, shelf_type='low'):
# RBJ Audio EQ Cookbook shelf filters
A = 10.0**(gain_db/40.0)
w0 = 2*math.pi*freq/sr
alpha = math.sin(w0)/2.0 * math.sqrt((A + 1/A)*(1.0/1.0 - 1) + 2.0) # S=1.0
cosw0 = math.cos(w0)
if shelf_type == 'low':
b0 = A*((A+1) - (A-1)*cosw0 + 2*math.sqrt(A)*alpha)
b1 = 2*A*((A-1) - (A+1)*cosw0)
b2 = A*((A+1) - (A-1)*cosw0 - 2*math.sqrt(A)*alpha)
a0 = (A+1) + (A-1)*cosw0 + 2*math.sqrt(A)*alpha
a1 = -2*((A-1) + (A+1)*cosw0)
a2 = (A+1) + (A-1)*cosw0 - 2*math.sqrt(A)*alpha
else: # high shelf
b0 = A*((A+1) + (A-1)*cosw0 + 2*math.sqrt(A)*alpha)
b1 = -2*A*((A-1) + (A+1)*cosw0)
b2 = A*((A+1) + (A-1)*cosw0 - 2*math.sqrt(A)*alpha)
a0 = (A+1) - (A-1)*cosw0 + 2*math.sqrt(A)*alpha
a1 = 2*((A-1) - (A+1)*cosw0)
a2 = (A+1) - (A-1)*cosw0 - 2*math.sqrt(A)*alpha
b = torch.tensor([b0/a0, b1/a0, b2/a0], dtype=torch.float64)
a = torch.tensor([1.0, a1/a0, a2/a0], dtype=torch.float64)
return b, a
class SignalProcessingBaxandall3BandEQ:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"audio_input": ("AUDIO",),
"bass_gain_db": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.1}),
"mid_gain_db": ("FLOAT", {"default": 0.0, "min": -20.0, "max": 20.0, "step": 0.1}),
"treble_gain_db": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.1}),
"low_freq": ("FLOAT", {"default": 100.0, "min": 20.0, "max": 500.0, "step": 1.0}),
"mid_freq": ("FLOAT", {"default": 1000.0, "min": 200.0, "max": 5000.0, "step": 10.0}),
"high_freq": ("FLOAT", {"default": 10000.0, "min": 2000.0, "max": 20000.0, "step": 100.0}),
"mid_q": ("FLOAT", {"default": 0.7, "min": 0.1, "max": 10.0, "step": 0.1}),
},
}
RETURN_TYPES = ("AUDIO",)
RETURN_NAMES = ("processed_audio",)
CATEGORY = "Audio Processing"
FUNCTION = "process"
def process(
self,
audio_input,
bass_gain_db: float = 0.0,
mid_gain_db: float = 0.0,
treble_gain_db: float = 0.0,
low_freq: float = 100.0,
mid_freq: float = 1000.0,
high_freq: float = 10000.0,
mid_q: float = 0.7
):
waveform, sample_rate = audio_from_comfy_3d(audio_input, try_gpu=True)
device = waveform.device
dtype = waveform.dtype
loudness = get_loudness(waveform, sample_rate)
# Low shelf filter
b_low, a_low = self.design_rbj_shelf(sample_rate, low_freq, bass_gain_db, shelf_type='low')
b_low = b_low.to(device=device, dtype=dtype)
a_low = a_low.to(device=device, dtype=dtype)
waveform = torchaudio.functional.lfilter(waveform, a_low, b_low,clamp=False)
# Mid peaking filter
b_mid, a_mid = self.design_rbj_peak(sample_rate, mid_freq, mid_gain_db, Q=mid_q)
b_mid = b_mid.to(device=device, dtype=dtype)
a_mid = a_mid.to(device=device, dtype=dtype)
waveform = torchaudio.functional.lfilter(waveform, a_mid, b_mid,clamp=False)
# High shelf filter
b_high, a_high = self.design_rbj_shelf(sample_rate, high_freq, treble_gain_db, shelf_type='high')
b_high = b_high.to(device=device, dtype=dtype)
a_high = a_high.to(device=device, dtype=dtype)
waveform = torchaudio.functional.lfilter(waveform, a_high, b_high,clamp=False)
# Normalize loudness after EQ
waveform = lufs_normalization(waveform, sample_rate, loudness)
return audio_to_comfy_3d(waveform, sample_rate)
def design_rbj_shelf(self, sr, freq, gain_db, shelf_type='low'):
# RBJ audio EQ cookbook formula for shelving filters
A = 10.0**(gain_db/40.0)
w0 = 2.0*math.pi*freq/sr
# Slope S=1.0 (Baxandall-like gentle slope)
S = 1.0
alpha = math.sin(w0)/2.0 * math.sqrt((A + 1.0/A)*(1.0/S - 1.0) + 2.0)
cosw0 = math.cos(w0)
if shelf_type == 'low':
b0 = A*((A+1.0)-(A-1.0)*cosw0+2.0*math.sqrt(A)*alpha)
b1 = 2.0*A*((A-1.0)-(A+1.0)*cosw0)
b2 = A*((A+1.0)-(A-1.0)*cosw0-2.0*math.sqrt(A)*alpha)
a0 = (A+1.0)+(A-1.0)*cosw0+2.0*math.sqrt(A)*alpha
a1 = -2.0*((A-1.0)+(A+1.0)*cosw0)
a2 = (A+1.0)+(A-1.0)*cosw0-2.0*math.sqrt(A)*alpha
else:
# high shelf
b0 = A*((A+1.0)+(A-1.0)*cosw0+2.0*math.sqrt(A)*alpha)
b1 = -2.0*A*((A-1.0)+(A+1.0)*cosw0)
b2 = A*((A+1.0)+(A-1.0)*cosw0-2.0*math.sqrt(A)*alpha)
a0 = (A+1.0)-(A-1.0)*cosw0+2.0*math.sqrt(A)*alpha
a1 = 2.0*((A-1.0)-(A+1.0)*cosw0)
a2 = (A+1.0)-(A-1.0)*cosw0-2.0*math.sqrt(A)*alpha
b = torch.tensor([b0/a0, b1/a0, b2/a0], dtype=torch.float64)
a = torch.tensor([1.0, a1/a0, a2/a0], dtype=torch.float64)
return b, a
def design_rbj_peak(self, sr, freq, gain_db, Q=0.7):
# RBJ audio EQ cookbook peak filter
A = 10.0**(gain_db/40.0)
w0 = 2.0*math.pi*freq/sr
alpha = math.sin(w0)/(2.0*Q)
cosw0 = math.cos(w0)
b0 = 1.0 + alpha*A
b1 = -2.0*cosw0
b2 = 1.0 - alpha*A
a0 = 1.0 + alpha/A
a1 = -2.0*cosw0
a2 = 1.0 - alpha/A
b = torch.tensor([b0/a0, b1/a0, b2/a0], dtype=torch.float64)
a = torch.tensor([1.0, a1/a0, a2/a0], dtype=torch.float64)
return b, a
if __name__ == "__main__":
import torchaudio
from pathlib import Path
from ..core.io import from_disk_as_raw_2d, audio_to_comfy_3d, audio_from_comfy_2d
node = SignalProcessingBaxandallEQ()
samples_path = Path('ComfyUI_SignalProcessing/audio/inputs/song.mp4')
source_path = samples_path.absolute()
source_audio, source_audio_sample_rate = from_disk_as_raw_2d(source_path)
input = audio_to_comfy_3d(source_audio, source_audio_sample_rate)[0]
# Test with some gain settings
result = node.process(input, bass_gain_db=5.0, treble_gain_db=-3.0)[0]
output_audio, sample_rate_audio = audio_from_comfy_2d(result)
# Save output for analysis
torchaudio.save('ComfyUI_SignalProcessing/audio/tests/baxandall_eq.wav', output_audio.cpu(), sample_rate_audio)
node = SignalProcessingBaxandall3BandEQ()
# Example usage:
# Provide a test audio file path
test_audio_path = Path('ComfyUI_SignalProcessing/audio/samples/test_audio.wav')
source_audio, source_audio_sample_rate = from_disk_as_raw_2d(test_audio_path)
input_tensors = audio_to_comfy_3d(source_audio, source_audio_sample_rate)[0]
# Apply EQ with some settings:
result = node.process(input_tensors, bass_gain_db=4.0, mid_gain_db=-2.0, treble_gain_db=3.0, low_freq=100, mid_freq=1000, high_freq=10000, mid_q=0.7)[0]
output_audio, out_sr = audio_from_comfy_2d(result)
torchaudio.save('ComfyUI_SignalProcessing/audio/tests/3band_baxandall_eq_output.wav', output_audio.cpu(), out_sr)