119 lines
4.3 KiB
Python
119 lines
4.3 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
Author: C0ffymachyne
|
|
License: GPLv3
|
|
Version: 1.0.0
|
|
|
|
Description:
|
|
Classic Audio filter set
|
|
"""
|
|
from ast import literal_eval
|
|
import torch
|
|
import torchaudio
|
|
from typing import Dict, Any, List, Tuple, Union
|
|
|
|
from ..core.utilities import comfy_root_to_syspath
|
|
from ..core.io import audio_to_comfy_3d, audio_from_comfy_2d
|
|
from ..core.loudness import lufs_normalization, get_loudness
|
|
|
|
comfy_root_to_syspath() # add comfy to sys path for dev
|
|
|
|
|
|
class SignalProcessingHarmonicsEnhancer:
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
|
return {
|
|
"required": {
|
|
"audio_input": ("AUDIO",),
|
|
"harmonics": ("STRING", {"default": "1,3,5,7,9"}),
|
|
"mode": (["detect base frequency", "use base frequency"],),
|
|
"base_frequency": ("FLOAT", {"default": 440, "min": 0, "max": 20000}),
|
|
"gain_db": ("INT", {"default": 5, "min": 0, "max": 500, "step": 1}),
|
|
"Q": ("FLOAT", {"default": 0.707, "min": 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]],
|
|
harmonics: str = "1,3,5,7,9",
|
|
mode: str = "detect base frequency",
|
|
base_frequency: int = 440,
|
|
gain_db: int = 5,
|
|
Q: float = 0.707,
|
|
) -> Tuple[Dict[str, Union[torch.Tensor, int]]]:
|
|
waveform, sample_rate = audio_from_comfy_2d(audio_input, try_gpu=True)
|
|
loudness = get_loudness(waveform, sample_rate)
|
|
|
|
try:
|
|
harmonics_list: List[int] = [literal_eval(x) for x in harmonics.split(",")]
|
|
except Exception:
|
|
raise RuntimeWarning(
|
|
"Invalid Harmonics Format. Please delimit integers by a comma \
|
|
',' like this 1,,3,5,7,9 "
|
|
)
|
|
if mode == "detect base frequency":
|
|
filtered_waveform = self.enhance_harmonics(
|
|
waveform, sample_rate, harmonics=harmonics_list, gain_db=gain_db, Q=Q
|
|
)
|
|
elif mode == "use base frequency":
|
|
filtered_waveform = self.enhance_harmonics(
|
|
waveform,
|
|
sample_rate,
|
|
harmonics=harmonics_list,
|
|
gain_db=gain_db,
|
|
base_frequency=base_frequency,
|
|
Q=Q,
|
|
)
|
|
|
|
filtered_waveform = lufs_normalization(filtered_waveform, sample_rate, loudness)
|
|
return audio_to_comfy_3d(filtered_waveform, sample_rate)
|
|
|
|
def add_harmonics(self, audio: torch.Tensor, gain: float = 1.2) -> torch.Tensor:
|
|
# Apply saturation using a tanh curve
|
|
harmonic_audio = torch.tanh(audio * gain)
|
|
return harmonic_audio
|
|
|
|
def detect_fundamental(self, audio: torch.Tensor, sample_rate: int) -> torch.Tensor:
|
|
# Estimate the fundamental frequency using a pitch detection method
|
|
pitch = torchaudio.functional.detect_pitch_frequency(audio, sample_rate)
|
|
|
|
return pitch
|
|
|
|
def detect_fundamental_mean(self, audio: torch.Tensor, sample_rate: int) -> int:
|
|
# Estimate the fundamental frequency using a pitch detection method
|
|
pitch = torchaudio.functional.detect_pitch_frequency(audio, sample_rate)
|
|
|
|
return int(pitch.mean().item())
|
|
|
|
def enhance_harmonics(
|
|
self,
|
|
audio: torch.Tensor,
|
|
sample_rate: int,
|
|
harmonics: List[int] = [1, 3, 5, 7, 9, 11],
|
|
gain_db: float = 5,
|
|
base_frequency: float = 0,
|
|
Q: float = 0.707,
|
|
) -> torch.Tensor:
|
|
# Detect the base frequency
|
|
if base_frequency == 0:
|
|
base_frequency = self.detect_fundamental_mean(audio, sample_rate)
|
|
if base_frequency <= 0: # Fallback if pitch detection fails
|
|
base_frequency = 440 # Use a default base frequency
|
|
|
|
# Apply EQ boosts to specific harmonic frequencies
|
|
for harmonic in harmonics:
|
|
freq = base_frequency * harmonic
|
|
if freq < sample_rate / 2: # Ensure it's within the Nyquist frequency
|
|
audio = torchaudio.functional.equalizer_biquad(
|
|
audio, sample_rate, center_freq=freq, gain=gain_db, Q=Q
|
|
)
|
|
|
|
return audio
|