329 lines
12 KiB
Python
329 lines
12 KiB
Python
# MBM's Music Visualizer: Audio Feature Calculator
|
|
# Calculates relevant audio features from loaded audio.
|
|
|
|
# Imports
|
|
import librosa
|
|
import torch
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
from typing import Optional
|
|
from scipy.signal import resample
|
|
|
|
from .mbmMVShared import normalizeArray, chartData, renderChart
|
|
|
|
# Classes
|
|
class AudioFeatureCalculator:
|
|
"""
|
|
Calculates relevant audio features from loaded audio.
|
|
"""
|
|
# Class Constants
|
|
DEF_FEAT_MOD_MAX = 10000.0
|
|
DEF_FEAT_MOD_MIN = -10000.0
|
|
|
|
FEAT_MOD_CALC_V1_0_0 = "v1.0.0"
|
|
FEAT_MOD_CALC_V1_1_0 = "v1.1.0"
|
|
FEAT_MOD_CALC_V1_2_0 = "v1.2.0"
|
|
FEAT_MOD_CALCS = (FEAT_MOD_CALC_V1_2_0, FEAT_MOD_CALC_V1_1_0, FEAT_MOD_CALC_V1_0_0)
|
|
|
|
RETURN_TYPES = ("TENSOR_1D", "FLOAT", "FLOAT", "IMAGE")
|
|
RETURN_NAMES = ("FEAT_MODS", "FEAT_SECONDS", "FPS", "CHARTS")
|
|
FUNCTION = "process"
|
|
CATEGORY = "MBMnodes/Audio"
|
|
|
|
# Constructor
|
|
def __init__(self):
|
|
pass
|
|
|
|
# ComfyUI Functions
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"audio": ("AUDIO", ),
|
|
"intensity": ("FLOAT", {"default": 1.0}), # Muiltiplier for the audio features
|
|
"hop_length": ("INT", {"default": 512}),
|
|
"fps_target": ("FLOAT", {"default": 6, "min": -1, "max": 10000}), # Provide `<= 0` to use whatever audio sampling comes up with
|
|
"feat_mod_max": ("FLOAT", {"default": s.DEF_FEAT_MOD_MAX, "min": -10000.0, "max": 10000.0}), # The maximum value the feature modifier can be. 10,000 should be unattainable through normal usage.
|
|
"feat_mod_min": ("FLOAT", {"default": s.DEF_FEAT_MOD_MIN, "min": -10000.0, "max": 10000.0}), # The minimum value the feature modifier can be. -10,000 should be unattainable through normal usage.
|
|
"feat_mod_normalize": ([False, True], ), # If `True`, the feature modifier array will be normalized between 0 and the maximum value in the array.
|
|
"feat_mod_calc": (s.FEAT_MOD_CALCS, ), # The calculation to use for the feature modifier.
|
|
"beat_weight": ("FLOAT", {"default": 1.0}), # The weight to apply to the beat frames.
|
|
}
|
|
}
|
|
|
|
def process(self,
|
|
audio: tuple,
|
|
intensity: float,
|
|
hop_length: int,
|
|
fps_target: float,
|
|
feat_mod_max: float,
|
|
feat_mod_min: float,
|
|
feat_mod_normalize: bool,
|
|
feat_mod_calc: str,
|
|
beat_weight: float
|
|
):
|
|
## Validation
|
|
# Make sure the feature modifier values are valid
|
|
if feat_mod_max < feat_mod_min:
|
|
raise ValueError("The maximum feature modifier value must be greater than the minimum feature modifier value.")
|
|
|
|
if feat_mod_max == self.DEF_FEAT_MOD_MAX:
|
|
feat_mod_max = None
|
|
|
|
if feat_mod_min == self.DEF_FEAT_MOD_MIN:
|
|
feat_mod_min = None
|
|
|
|
## Calculation
|
|
# Unpack the audio
|
|
y, sr = audio
|
|
|
|
# Calculate the duration of the audio
|
|
duration = librosa.get_duration(y=y, sr=sr, hop_length=hop_length)
|
|
|
|
# # Calculate the duration of a single hop
|
|
# hopSeconds: float = (hop_length / sr)
|
|
|
|
# Calculate tempo
|
|
onset = librosa.onset.onset_strength(y=y, sr=sr)
|
|
tempo = normalizeArray(librosa.beat.tempo(onset_envelope=onset, sr=sr, hop_length=hop_length, aggregate=None))
|
|
|
|
# Calculate the spectrogram
|
|
spectro = librosa.feature.melspectrogram(
|
|
y=y,
|
|
sr=sr,
|
|
n_mels=128,
|
|
fmax=8000,
|
|
hop_length=hop_length
|
|
)
|
|
|
|
# Calculate normalized mean power per hop
|
|
spectroMean = np.mean(spectro, axis=0)
|
|
|
|
# Calculate the delta of the spectrogram mean
|
|
spectroMeanDelta = librosa.feature.delta(spectroMean)
|
|
|
|
# Normalize the spectro mean
|
|
spectroMean = normalizeArray(spectroMean)
|
|
|
|
# Calculate pitch chroma for hops
|
|
chroma = librosa.feature.chroma_cqt(
|
|
y=y,
|
|
sr=sr,
|
|
hop_length=hop_length
|
|
)
|
|
|
|
# Get the mean of the chroma for each step
|
|
chromaMean = np.mean(chroma, axis=0)
|
|
|
|
# Calculate the Harmonic and Percussive components
|
|
harmonics, percussives = librosa.effects.hpss(y)
|
|
|
|
# Calculate the output FPS
|
|
if fps_target <= 0:
|
|
# Calculate framerate based on audio
|
|
fps = len(tempo) / duration
|
|
else:
|
|
# Specific framerate to target
|
|
fps = fps_target
|
|
|
|
# Calculate desired frame count
|
|
desiredFrames = round(fps * duration)
|
|
|
|
# Calculate the beat frames
|
|
beatTempo, beats = librosa.beat.beat_track(y=y, sr=sr, hop_length=hop_length)
|
|
beatTimes = librosa.frames_to_time(beats, sr=sr, hop_length=hop_length) # Where each beat is in seconds
|
|
beatFrames = np.round(beatTimes * fps).astype(int) # Where each beat is in a Video frame
|
|
|
|
# Resample audio features to match desired frame count
|
|
tempo = resample(tempo, desiredFrames)
|
|
spectroMean = resample(spectroMean, desiredFrames)
|
|
spectroMeanDelta = resample(spectroMeanDelta, desiredFrames)
|
|
chromaMean = resample(chromaMean, desiredFrames)
|
|
harmonics = resample(harmonics, desiredFrames)
|
|
percussives = resample(percussives, desiredFrames)
|
|
|
|
# Calculate the feature modifier for each frame
|
|
featModifiers = torch.Tensor(
|
|
[self._calcFeatModifier(
|
|
i,
|
|
intensity,
|
|
tempo[i],
|
|
spectroMean[i],
|
|
spectroMeanDelta[i],
|
|
chromaMean[i],
|
|
harmonics[i],
|
|
percussives[i],
|
|
beatFrames,
|
|
beat_weight,
|
|
modMax=feat_mod_max,
|
|
modMin=feat_mod_min,
|
|
calc=feat_mod_calc
|
|
) for i in range(desiredFrames)]
|
|
)
|
|
|
|
# Normalize the feature modifiers if requested
|
|
if feat_mod_normalize:
|
|
# NOTE: The feat_mod_max and feat_mod_min values will be inaccurate realtive to the now normalized featModifiers array and must be adjusted if used later.
|
|
featModifiers = normalizeArray(featModifiers, minVal=0.0, maxVal=featModifiers.max().item())
|
|
|
|
# Build the charts
|
|
chartImages = torch.vstack([
|
|
chartData(tempo, "Tempo"),
|
|
chartData(spectroMean, "Spectrogram Mean"),
|
|
chartData(spectroMeanDelta, "Spectrogram Mean Delta"),
|
|
chartData(chromaMean, "Chroma Mean"),
|
|
chartData(harmonics, "Harmonic"),
|
|
chartData(percussives, "Percussive"),
|
|
self._chartBeatFrames(chromaMean, beatFrames),
|
|
self._chartFeatMod(featModifiers, intensity, feat_mod_normalize, modMax=feat_mod_max, modMin=feat_mod_min)
|
|
])
|
|
|
|
return (
|
|
featModifiers,
|
|
(duration / desiredFrames),
|
|
fps,
|
|
chartImages
|
|
)
|
|
|
|
# Internal Functions
|
|
def _calcFeatModifier(self,
|
|
frame: int,
|
|
intensity: float,
|
|
tempo: float,
|
|
spectroMean: float,
|
|
spectroMeanDelta: float,
|
|
chromaMean: float,
|
|
harmoic: float,
|
|
percussive: float,
|
|
beatFrames: np.ndarray,
|
|
beatWeight: float,
|
|
modMax: Optional[float] = None,
|
|
modMin: Optional[float] = None,
|
|
calc: str = FEAT_MOD_CALC_V1_1_0
|
|
) -> float:
|
|
"""
|
|
Calculates the overall feature modifier based on the provided audio features.
|
|
|
|
frame: The Video frame number to calculate the feature modifier for.
|
|
intensity: A modifier to increase (>1.0) or decrease (<1.0) the overall effect of the audio features.
|
|
tempo: The tempo for a single step of the audio.
|
|
spectroMean: The normalized mean power for a single step of the audio.
|
|
spectroMeanDelta: The delta of the normalized mean power for a single step of the audio.
|
|
chromaMean: The mean value of the chroma for a single step of the audio.
|
|
harmoic: The harmonic component of the audio.
|
|
percussive: The percussive component of the audio.
|
|
beatFrames: A Numpy array of Video frame numbers where each major beat element occurs.
|
|
beatWeight: The weight to apply to the beat frames.
|
|
modMax: The maximum value the feature modifier can be. Provide `None` to have no maximum.
|
|
modMin: The minimum value the feature modifier can be. Provide `None` to have no minimum.
|
|
calc: The calculation version to use.
|
|
|
|
Returns the calculated overall feature modifier.
|
|
"""
|
|
# Calculate the feature modifier
|
|
if calc == self.FEAT_MOD_CALC_V1_0_0:
|
|
# v1.0.0
|
|
modVal = (((tempo + 1.0) * (spectroMean + 1.0) * (chromaMean + 1.0)) * (intensity + spectroMeanDelta))
|
|
elif calc == self.FEAT_MOD_CALC_V1_1_0:
|
|
# v1.1.0
|
|
modVal = (((tempo + 1.0) * (spectroMean + 1.0) * (chromaMean + 1.0)) * (intensity + (spectroMeanDelta * (harmoic + percussive))))
|
|
else:
|
|
# v1.2.0
|
|
modVal = (((tempo + 1.0) * (spectroMean + 1.0) * (chromaMean + 1.0)) * (intensity + (spectroMeanDelta * (harmoic + percussive)))) + (beatWeight if frame in beatFrames else 0.0)
|
|
|
|
# Return step within bounds
|
|
if (modMax is not None) and (modVal > modMax):
|
|
return modMax
|
|
elif (modMin is not None) and (modVal < modMin):
|
|
return modMin
|
|
else:
|
|
return modVal
|
|
|
|
def _chartFeatMod(self,
|
|
featModifiers: torch.Tensor,
|
|
intensity: float,
|
|
isNormalized: bool,
|
|
modMax: Optional[float] = None,
|
|
modMin: Optional[float] = None
|
|
) -> torch.Tensor:
|
|
"""
|
|
Creates a chart representing the feature modifier.
|
|
|
|
featModifiers: The calculated feature modifiers.
|
|
intensity: The given intensity used in calculating the feature modifiers.
|
|
isNormalized: If the feature modifiers are normalized.
|
|
modMax: The maximum value the feature modifier can be. Provide `None` to display no maximum.
|
|
modMin: The minimum value the feature modifier can be. Provide `None` to display no minimum.
|
|
|
|
Returns a ComfyUI compatible Tensor image of the chart.
|
|
"""
|
|
# Build the chart
|
|
fig, ax = plt.subplots(figsize=(20, 4))
|
|
|
|
ax.plot(featModifiers, label="Modifiers")
|
|
|
|
modMaxStr = "none"
|
|
if modMax is not None:
|
|
if isNormalized:
|
|
# Use logical value
|
|
modMaxAlt = featModifiers.max().item()
|
|
modMaxStr = f"{modMaxAlt:.2f} ({modMax:.2f})"
|
|
plt.axhline(y=modMaxAlt, linestyle="dotted", color="yellow", label="Mod Max")
|
|
else:
|
|
# Use prescribed value
|
|
modMaxStr = f"{modMax:.2f}"
|
|
plt.axhline(y=modMax, linestyle="dotted", color="yellow", label="Mod Max")
|
|
|
|
modMinStr = "none"
|
|
if modMin is not None:
|
|
if isNormalized:
|
|
# Use logical value
|
|
modMinStr = f"0 ({modMin:.2f})"
|
|
plt.axhline(y=0, linestyle="dotted", color="green", label="Mod Min")
|
|
else:
|
|
# Use prescribed value
|
|
modMinStr = f"{modMin:.2f}"
|
|
plt.axhline(y=modMin, linestyle="dotted", color="green", label="Mod Min")
|
|
|
|
ax.legend()
|
|
ax.grid(True)
|
|
ax.set_title("Feature Modifiers")
|
|
ax.set_xlabel("Index")
|
|
ax.set_ylabel("Value")
|
|
|
|
# Add feature information
|
|
featureInfo = f"Normalized: {('yes' if isNormalized else 'no')}\n"
|
|
featureInfo += f"Max: {modMaxStr}\n"
|
|
featureInfo += f"Min: {modMinStr}\n"
|
|
featureInfo += f"Intensity: {intensity:.2f}"
|
|
ax.text(1.02, 0.5, featureInfo, transform=ax.transAxes, va="center")
|
|
|
|
# Render the chart
|
|
return renderChart(fig)
|
|
|
|
def _chartBeatFrames(self, chroma: np.ndarray, beatFrames: np.ndarray) -> torch.Tensor:
|
|
"""
|
|
Plots the beat frames over the chroma.
|
|
|
|
chroma: The chroma to plot.
|
|
beatFrames: The beat frames to plot.
|
|
|
|
Returns a ComfyUI compatible Tensor image of the chart.
|
|
"""
|
|
# Prepare figure
|
|
fig, ax = plt.subplots(figsize=(20, 4))
|
|
|
|
# Plot the data
|
|
ax.plot(chroma)
|
|
for beat in beatFrames:
|
|
ax.axvline(x=beat, color="r", linestyle="--")
|
|
|
|
# Add labels
|
|
ax.grid(True)
|
|
ax.set_xlabel("Frame Number")
|
|
ax.set_title("Beat Frames")
|
|
ax.legend()
|
|
|
|
# Render the chart
|
|
return renderChart(fig)
|