232 lines
7.3 KiB
Python
232 lines
7.3 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
Author: C0ffymachyne
|
|
License: GPLv3
|
|
Version: 1.0.0
|
|
|
|
Description:
|
|
Waveform image rendering node
|
|
"""
|
|
|
|
import torch
|
|
from io import BytesIO
|
|
from PIL import Image
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
|
|
from ..core.utilities import comfy_root_to_syspath
|
|
|
|
comfy_root_to_syspath() # add comfy to sys path for dev
|
|
|
|
from ..core.io import audio_to_comfy_3d
|
|
|
|
|
|
class SignalProcessingWaveform:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"audio_input": ("AUDIO",),
|
|
},
|
|
"optional": {
|
|
"color": ("STRING", {"default": "black"}),
|
|
"background_color": ("STRING", {"default": "white"}),
|
|
"width": (
|
|
"INT",
|
|
{"default": 800, "min": 100, "max": 4000, "step": 100},
|
|
),
|
|
"height": ("INT", {"default": 200, "min": 50, "max": 1000, "step": 50}),
|
|
"line_width": (
|
|
"FLOAT",
|
|
{"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("waveform_image",)
|
|
CATEGORY = "Audio Processing"
|
|
FUNCTION = "process"
|
|
|
|
def process(
|
|
self,
|
|
audio_input,
|
|
color="white",
|
|
background_color="black",
|
|
width=800,
|
|
height=200,
|
|
line_width=1.0,
|
|
):
|
|
waveform = audio_input.get(
|
|
"waveform"
|
|
) # [channels, samples] or [batch, channels, samples]
|
|
sample_rate = audio_input.get("sample_rate")
|
|
|
|
# Convert to mono by averaging channels
|
|
if waveform.ndim == 3:
|
|
# [batch, channels, samples]
|
|
waveform = waveform.mean(dim=1, keepdim=True) # [batch, 1, samples]
|
|
waveform = waveform.squeeze(0) # [1, samples]
|
|
elif waveform.ndim == 2:
|
|
# [channels, samples]
|
|
if waveform.shape[0] > 1:
|
|
waveform = waveform.mean(dim=0, keepdim=True) # [1, samples]
|
|
else:
|
|
waveform = waveform.unsqueeze(0) # [1, samples]
|
|
elif waveform.ndim == 1:
|
|
# [samples]
|
|
waveform = waveform.unsqueeze(0) # [1, samples]
|
|
else:
|
|
raise ValueError(f"Unsupported waveform shape: {waveform.shape}")
|
|
|
|
# Convert waveform to numpy
|
|
waveform = waveform.to(dtype=torch.float32)
|
|
waveform_np = waveform.squeeze().detach().cpu().numpy() # [samples]
|
|
|
|
# Create a matplotlib figure without axes
|
|
plt.figure(figsize=(width / 100, height / 100), dpi=96)
|
|
plt.axis("off")
|
|
plt.margins(0, 0)
|
|
plt.gca().set_facecolor(background_color)
|
|
plt.gca().set_position([0, 0, 1, 1])
|
|
|
|
# Plot the waveform
|
|
plt.plot(waveform_np, color=color, linewidth=line_width)
|
|
plt.ylim(-1.1, 1.1) # Set y-axis limits to -1 and 1
|
|
plt.tight_layout(pad=0)
|
|
|
|
# Save the plot to a buffer
|
|
buf = BytesIO()
|
|
plt.savefig(buf, format="png", bbox_inches="tight", pad_inches=0)
|
|
plt.close()
|
|
|
|
# Load the image from the buffer
|
|
buf.seek(0)
|
|
waveform_image = Image.open(buf).convert("RGB")
|
|
|
|
# Resize if necessary
|
|
waveform_image = waveform_image.resize((width, height), Image.BILINEAR)
|
|
|
|
# Convert to numpy array and normalize to [0,1]
|
|
image_np = np.array(waveform_image).astype(np.float32) / 255.0 # [H, W, 3]
|
|
|
|
# Convert to torch tensor and add batch dimension
|
|
image = torch.from_numpy(image_np)[None,]
|
|
|
|
return (image,)
|
|
|
|
|
|
class SignalProcessingWaveform2:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"audio_input": ("AUDIO",),
|
|
},
|
|
"optional": {
|
|
"color": ("STRING", {"default": "black"}),
|
|
"background_color": ("STRING", {"default": "white"}),
|
|
"width": (
|
|
"INT",
|
|
{"default": 800, "min": 100, "max": 4000, "step": 100},
|
|
),
|
|
"height": ("INT", {"default": 200, "min": 50, "max": 1000, "step": 50}),
|
|
"line_width": (
|
|
"FLOAT",
|
|
{"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("waveform_image",)
|
|
CATEGORY = "Audio Processing"
|
|
FUNCTION = "process"
|
|
|
|
def process(
|
|
self,
|
|
audio_input,
|
|
color="black",
|
|
background_color="white",
|
|
width=800,
|
|
height=200,
|
|
line_width=1.0,
|
|
):
|
|
waveform = audio_input.get(
|
|
"waveform"
|
|
) # [channels, samples] or [batch, channels, samples]
|
|
sample_rate = audio_input.get("sample_rate")
|
|
|
|
# Convert to mono by averaging channels
|
|
if waveform.ndim == 3:
|
|
# [batch, channels, samples]
|
|
waveform = waveform.mean(dim=1, keepdim=True) # [batch, 1, samples]
|
|
waveform = waveform.squeeze(0) # [1, samples]
|
|
elif waveform.ndim == 2:
|
|
# [channels, samples]
|
|
if waveform.shape[0] > 1:
|
|
waveform = waveform.mean(dim=0, keepdim=True) # [1, samples]
|
|
else:
|
|
waveform = waveform.unsqueeze(0) # [1, samples]
|
|
elif waveform.ndim == 1:
|
|
# [samples]
|
|
waveform = waveform.unsqueeze(0) # [1, samples]
|
|
else:
|
|
raise ValueError(f"Unsupported waveform shape: {waveform.shape}")
|
|
|
|
# Ensure waveform is in float32
|
|
waveform = waveform.to(dtype=torch.float32)
|
|
waveform_np = waveform.squeeze().detach().cpu().numpy() # [samples]
|
|
|
|
# Create a matplotlib figure without axes
|
|
plt.figure(figsize=(width / 100, height / 100), dpi=96)
|
|
plt.axis("off")
|
|
plt.margins(0, 0)
|
|
ax = plt.gca()
|
|
ax.set_facecolor(background_color)
|
|
ax.set_position([0, 0, 1, 1])
|
|
|
|
# Plot the waveform with fixed y-axis limits
|
|
plt.plot(waveform_np, color=color, linewidth=line_width)
|
|
plt.ylim(-1, 1) # Set y-axis limits to -1 and 1
|
|
plt.tight_layout(pad=0)
|
|
|
|
# Save the plot to a buffer
|
|
buf = BytesIO()
|
|
plt.savefig(buf, format="png", bbox_inches="tight", pad_inches=0)
|
|
plt.close()
|
|
|
|
# Load the image from the buffer
|
|
buf.seek(0)
|
|
waveform_image = Image.open(buf).convert("RGB")
|
|
|
|
# Resize if necessary
|
|
waveform_image = waveform_image.resize((width, height), Image.BILINEAR)
|
|
|
|
# Convert to numpy array and normalize to [0,1]
|
|
image_np = np.array(waveform_image).astype(np.float32) / 255.0 # [H, W, 3]
|
|
|
|
# Convert to torch tensor and add batch dimension
|
|
image = torch.from_numpy(image_np).permute(2, 0, 1).unsqueeze(0) # [1, 3, H, W]
|
|
|
|
return (image,)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
from pathlib import Path
|
|
from ..core.io import from_disk_as_raw_2d, audio_to_comfy_3d
|
|
|
|
node = SignalProcessingWaveform()
|
|
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]
|
|
|
|
result = node.process(input)[0]
|
|
|
|
print("result", result)
|