130 lines
3.8 KiB
Python
130 lines
3.8 KiB
Python
import io
|
|
import os
|
|
import uuid
|
|
import wave
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
from aiohttp import web
|
|
from server import PromptServer
|
|
|
|
WEB_DIRECTORY = "./js"
|
|
|
|
TEMP_SUBDIR = "auto_audio_player"
|
|
|
|
|
|
def _get_temp_root() -> Path:
|
|
base = Path(tempfile.gettempdir()) / "ComfyUI" / TEMP_SUBDIR
|
|
base.mkdir(parents=True, exist_ok=True)
|
|
return base
|
|
|
|
|
|
def _to_numpy_audio(audio):
|
|
if not isinstance(audio, dict):
|
|
raise TypeError("Expected ComfyUI AUDIO input as a dict with waveform and sample_rate.")
|
|
|
|
waveform = audio.get("waveform")
|
|
sample_rate = int(audio.get("sample_rate", 44100))
|
|
|
|
if waveform is None:
|
|
raise ValueError("Audio input is missing 'waveform'.")
|
|
|
|
if isinstance(waveform, torch.Tensor):
|
|
data = waveform.detach().cpu().float().numpy()
|
|
else:
|
|
data = np.asarray(waveform, dtype=np.float32)
|
|
|
|
# Common ComfyUI audio shapes include:
|
|
# [batch, channels, samples], [channels, samples], [samples]
|
|
if data.ndim == 3:
|
|
data = data[0]
|
|
if data.ndim == 1:
|
|
data = data[np.newaxis, :]
|
|
elif data.ndim == 2:
|
|
# If someone hands in [samples, channels], rotate it.
|
|
if data.shape[0] > data.shape[1] and data.shape[1] <= 8:
|
|
data = data.T
|
|
else:
|
|
raise ValueError(f"Unsupported waveform shape: {data.shape}")
|
|
|
|
# Final shape: [channels, samples]
|
|
data = np.clip(data, -1.0, 1.0)
|
|
return data, sample_rate
|
|
|
|
|
|
def _write_wav(audio_path: Path, waveform: np.ndarray, sample_rate: int):
|
|
# Convert float32 [-1, 1] to 16-bit PCM.
|
|
pcm = (waveform * 32767.0).astype(np.int16)
|
|
pcm_interleaved = pcm.T # [samples, channels]
|
|
|
|
with wave.open(str(audio_path), "wb") as wav_file:
|
|
wav_file.setnchannels(pcm.shape[0])
|
|
wav_file.setsampwidth(2)
|
|
wav_file.setframerate(sample_rate)
|
|
wav_file.writeframes(pcm_interleaved.tobytes())
|
|
|
|
|
|
@PromptServer.instance.routes.get("/auto_audio_player/{filename}")
|
|
async def serve_auto_audio_player_file(request):
|
|
filename = request.match_info["filename"]
|
|
safe_name = os.path.basename(filename)
|
|
path = _get_temp_root() / safe_name
|
|
|
|
if not path.exists() or not path.is_file():
|
|
return web.Response(status=404, text="Audio file not found.")
|
|
|
|
return web.FileResponse(path, headers={"Cache-Control": "no-store"})
|
|
|
|
|
|
class AutoAudioPlayer:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"audio": ("AUDIO",),
|
|
"autoplay": ("BOOLEAN", {"default": True}),
|
|
"default_volume": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"loop": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("AUDIO",)
|
|
RETURN_NAMES = ("audio",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "audio/utils"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, audio, autoplay=True, default_volume=1.0, loop=False):
|
|
waveform, sample_rate = _to_numpy_audio(audio)
|
|
filename = f"auto_audio_{uuid.uuid4().hex}.wav"
|
|
out_path = _get_temp_root() / filename
|
|
_write_wav(out_path, waveform, sample_rate)
|
|
|
|
duration = float(waveform.shape[1]) / float(sample_rate) if sample_rate > 0 else 0.0
|
|
|
|
ui_payload = {
|
|
"filename": filename,
|
|
"url": f"/auto_audio_player/{filename}",
|
|
"sample_rate": sample_rate,
|
|
"channels": int(waveform.shape[0]),
|
|
"duration": duration,
|
|
"autoplay": bool(autoplay),
|
|
"default_volume": float(default_volume),
|
|
"loop": bool(loop),
|
|
}
|
|
|
|
return {"ui": {"audio": [ui_payload]}, "result": (audio,)}
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AutoAudioPlayer": AutoAudioPlayer,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AutoAudioPlayer": "Auto Audio Player",
|
|
}
|
|
|
|
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
|