212 lines
6.9 KiB
Python
212 lines
6.9 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 typing import Any, Dict, Tuple
|
|
|
|
from ..core.utilities import comfy_root_to_syspath
|
|
|
|
comfy_root_to_syspath() # add comfy to sys path for dev
|
|
|
|
|
|
class SignalProcessingWaveform:
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
|
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 = "Signal Processing"
|
|
FUNCTION = "process"
|
|
|
|
def process(
|
|
self,
|
|
audio_input: torch.Tensor,
|
|
color: str = "white",
|
|
background_color: str = "black",
|
|
width: int = 800,
|
|
height: int = 200,
|
|
line_width: float = 1.0,
|
|
) -> Tuple[torch.Tensor]:
|
|
waveform = audio_input.get(
|
|
"waveform"
|
|
) # [channels, samples] or [batch, channels, samples]
|
|
|
|
# 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.3, 1.3) # 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) -> Dict[str, Any]:
|
|
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 = "Signal Processing"
|
|
FUNCTION = "process"
|
|
|
|
def process(
|
|
self,
|
|
audio_input: torch.Tensor,
|
|
color: str = "black",
|
|
background_color: str = "white",
|
|
width: int = 800,
|
|
height: int = 200,
|
|
line_width: float = 1.0,
|
|
) -> Tuple[torch.Tensor]:
|
|
waveform = audio_input.get(
|
|
"waveform"
|
|
) # [channels, samples] or [batch, channels, samples]
|
|
|
|
# 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,)
|