96 lines
3.1 KiB
Python
96 lines
3.1 KiB
Python
import os
|
|
import shutil
|
|
import subprocess
|
|
from pathlib import Path
|
|
import numpy as np
|
|
import torch
|
|
import folder_paths
|
|
from scipy.io.wavfile import write
|
|
|
|
class LikeSpiderUINode:
|
|
"""
|
|
Base class for declarative UI-driven ComfyUI nodes.
|
|
Extend this class and define `UI_CONFIG` and `run(...)`.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
config = cls.UI_CONFIG
|
|
required = {}
|
|
|
|
for field in config["inputs"]:
|
|
name = field["name"]
|
|
type_ = field["type"]
|
|
default = field.get("default")
|
|
|
|
if type_ == "string":
|
|
required[name] = ("STRING", {"default": default or ""})
|
|
elif type_ == "select":
|
|
required[name] = (field["options"], {"default": default})
|
|
elif type_ == "audio":
|
|
required[name] = ("AUDIO",)
|
|
else:
|
|
raise ValueError(f"Unsupported field type: {type_}")
|
|
|
|
return {"required": required}
|
|
|
|
RETURN_TYPES = ("AUDIO",)
|
|
RETURN_NAMES = ("audio",)
|
|
FUNCTION = "interface_entry"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "LikeSpiderAI"
|
|
|
|
def interface_entry(self, **kwargs):
|
|
return self.run(**kwargs)
|
|
|
|
def save_audio_with_ffmpeg(self, audio_data, filename_base, format="mp3", bitrate=192):
|
|
samples = audio_data["waveform"]
|
|
if isinstance(samples, torch.Tensor):
|
|
samples = samples.detach().cpu().numpy()
|
|
if samples.ndim == 3 and samples.shape[0] == 1:
|
|
samples = samples[0]
|
|
|
|
max_val = np.max(np.abs(samples))
|
|
if max_val > 1.0:
|
|
samples = samples / max_val
|
|
samples = np.clip(samples, -1.0, 1.0)
|
|
audio_int16 = (samples * 32767).astype(np.int16)
|
|
|
|
if audio_int16.ndim == 2 and audio_int16.shape[0] in [1, 2]:
|
|
wav_data = audio_int16.T
|
|
elif audio_int16.ndim == 1:
|
|
wav_data = audio_int16
|
|
else:
|
|
raise Exception(f"❌ Unsupported audio shape: {audio_int16.shape}")
|
|
|
|
output_dir = Path(folder_paths.get_output_directory()) / "audio"
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
existing = sorted(output_dir.glob(f"{filename_base}_*.{format}"))
|
|
index = len(existing)
|
|
file_path = output_dir / f"{filename_base}_{index:05d}.{format}"
|
|
|
|
temp_wav = output_dir / "temp_input.wav"
|
|
write(temp_wav, 44100, wav_data)
|
|
|
|
if format == "mp3":
|
|
subprocess.run([
|
|
"ffmpeg", "-y", "-i", str(temp_wav),
|
|
"-codec:a", "libmp3lame", "-b:a", f"{bitrate}k",
|
|
str(file_path)
|
|
], check=True)
|
|
elif format == "wav":
|
|
shutil.copy(temp_wav, file_path)
|
|
elif format == "flac":
|
|
subprocess.run([
|
|
"ffmpeg", "-y", "-i", str(temp_wav),
|
|
"-codec:a", "flac",
|
|
str(file_path)
|
|
], check=True)
|
|
else:
|
|
raise Exception(f"Unsupported format: {format}")
|
|
|
|
os.remove(temp_wav)
|
|
print(f"[LikeSpiderUINode] ✅ Saved: {file_path}")
|
|
return file_path
|