harmonics enhancer, license, readme updates, blacked

This commit is contained in:
c0ffymachyne
2024-12-08 11:30:05 -08:00
parent b49d3f487d
commit 8442628888
29 changed files with 1430 additions and 1322 deletions
+44 -19
View File
@@ -22,21 +22,31 @@ import matplotlib.pyplot as plt
from ..core.io import audio_to_comfy_3d, audio_from_comfy_3d, audio_from_comfy_2d
from ..core.loudness import lufs_normalization, get_loudness
class SignalProcessingSpectrogram:
@classmethod
def INPUT_TYPES(cls):
cmaps = ['viridis', 'plasma', 'inferno', 'magma', 'cividis']
cmaps = ["viridis", "plasma", "inferno", "magma", "cividis"]
return {
"required": {
"audio_input": ("AUDIO",),
"color_map": (cmaps,),
},
"optional": {
"n_fft": ("INT", {"default": 4096, "min": 512, "max": 8192, "step": 256}),
"hop_length": ("INT", {"default": 128, "min": 64, "max": 4096, "step": 128}),
"n_fft": (
"INT",
{"default": 4096, "min": 512, "max": 8192, "step": 256},
),
"hop_length": (
"INT",
{"default": 128, "min": 64, "max": 4096, "step": 128},
),
"n_mels": ("INT", {"default": 512, "min": 32, "max": 2048, "step": 32}),
"top_db": ("FLOAT", {"default": 80.0, "min": 10.0, "max": 100.0, "step": 5.0}),
}
"top_db": (
"FLOAT",
{"default": 80.0, "min": 10.0, "max": 100.0, "step": 5.0},
),
},
}
RETURN_TYPES = ("IMAGE",)
@@ -44,11 +54,21 @@ class SignalProcessingSpectrogram:
CATEGORY = "Audio Processing"
FUNCTION = "process"
def process(self, audio_input, color_map, n_fft=2048, hop_length=512, n_mels=128, top_db=80.0):
waveform = audio_input.get('waveform') # [channels, samples] or [batch, channels, samples]
sample_rate = audio_input.get('sample_rate')
def process(
self,
audio_input,
color_map,
n_fft=2048,
hop_length=512,
n_mels=128,
top_db=80.0,
):
waveform = audio_input.get(
"waveform"
) # [channels, samples] or [batch, channels, samples]
sample_rate = audio_input.get("sample_rate")
#waveform, sample_rate = audio_from_comfy_2d(audio_input)
# waveform, sample_rate = audio_from_comfy_2d(audio_input)
# Convert to mono by averaging channels
if waveform.ndim == 3:
@@ -74,9 +94,9 @@ class SignalProcessingSpectrogram:
hop_length=hop_length,
n_mels=n_mels,
power=2.0,
norm='slaney',
mel_scale='htk',
).to(waveform.device, dtype = waveform.dtype)
norm="slaney",
mel_scale="htk",
).to(waveform.device, dtype=waveform.dtype)
spectrogram = spectrogram_transform(waveform) # [1, n_mels, time_frames]
# Convert to decibel scale
@@ -84,7 +104,9 @@ class SignalProcessingSpectrogram:
spectrogram_db = amplitude_to_db(spectrogram) # [1, n_mels, time_frames]
# Convert to numpy
spectrogram_db = spectrogram_db.squeeze().detach().cpu().numpy() # [n_mels, time_frames]
spectrogram_db = (
spectrogram_db.squeeze().detach().cpu().numpy()
) # [n_mels, time_frames]
# Clip spectrogram to a range for better contrast
spectrogram_db = np.clip(spectrogram_db, -top_db, 0.0)
@@ -94,10 +116,14 @@ class SignalProcessingSpectrogram:
# Apply a colormap (e.g., 'inferno') using matplotlib
cmap = plt.get_cmap(color_map)
spectrogram_colored = cmap(spectrogram_normalized) # [n_mels, time_frames, 4] RGBA
spectrogram_colored = cmap(
spectrogram_normalized
) # [n_mels, time_frames, 4] RGBA
# Convert to RGB by removing alpha channel
spectrogram_rgb = (spectrogram_colored[:, :, :3] * 255).astype(np.uint8) # [n_mels, time_frames, 3]
spectrogram_rgb = (spectrogram_colored[:, :, :3] * 255).astype(
np.uint8
) # [n_mels, time_frames, 3]
spectrogram_rgb = np.squeeze(spectrogram_rgb)
# Check the shape and adjust if necessary
@@ -112,16 +138,15 @@ class SignalProcessingSpectrogram:
# Optionally resize for better resolution
spectrogram_image = spectrogram_image.resize(
(spectrogram_image.width * 2, spectrogram_image.height * 2),
Image.BILINEAR
(spectrogram_image.width * 2, spectrogram_image.height * 2), Image.BILINEAR
)
# Convert to numpy array and normalize to [0,1]
image_np = np.array(spectrogram_image).astype(np.float32) / 255.0 # [H, W, 3]
# Convert to torch tensor and add batch dimension
#image_tensor = torch.from_numpy(image_np).permute(2, 0, 1).unsqueeze(0) # [1, 3, H, W]
# image_tensor = torch.from_numpy(image_np).permute(2, 0, 1).unsqueeze(0) # [1, 3, H, W]
image = torch.from_numpy(image_np)[None,]
return (image, )
return (image,)