271 lines
8.8 KiB
Python
271 lines
8.8 KiB
Python
import torch
|
|
import numpy as np
|
|
import folder_paths
|
|
import os
|
|
from .audio_segmenter import segment_audio
|
|
from .audio_paragraph_segmenter import segment_by_paragraph
|
|
|
|
# Whisper模型路径
|
|
stt_model_dir = os.path.join(folder_paths.models_dir, "stt", "whisper")
|
|
os.makedirs(stt_model_dir, exist_ok=True)
|
|
|
|
# 获取可用模型列表
|
|
local_models = []
|
|
if os.path.isdir(stt_model_dir):
|
|
for item in os.listdir(stt_model_dir):
|
|
# 如果是文件夹，直接添加
|
|
if os.path.isdir(os.path.join(stt_model_dir, item)):
|
|
local_models.append(item)
|
|
# 如果是 .pt 文件，去掉后缀添加
|
|
elif item.endswith(".pt"):
|
|
local_models.append(os.path.splitext(item)[0])
|
|
|
|
standard_models = ["tiny", "base", "small", "medium", "large-v1", "large-v2", "large-v3"]
|
|
available_models = sorted(list(set(local_models + standard_models)))
|
|
|
|
|
|
# 自定义类型，用于 ComfyUI 识别
|
|
class AudioClip:
|
|
@classmethod
|
|
def is_list(cls, item):
|
|
return isinstance(item, list)
|
|
|
|
# 在 ComfyUI 中注册自定义类型
|
|
# 注意：这部分可能需要根据 ComfyUI 的插件系统进行调整
|
|
# 通常，我们通过返回一个特殊的字符串或 Python 类型来处理
|
|
# 在这里，我们先用"*"作为通配符，然后在节点中处理具体逻辑
|
|
|
|
class AudioSegmentByParagraphNode:
|
|
AUX_ID = "dreamidea/comfyui-audio-segment"
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"audio": ("AUDIO",),
|
|
"text": ("STRING", {
|
|
"multiline": True,
|
|
"default": "在此处输入按段落分隔的文本..."
|
|
}),
|
|
"model": (available_models, ),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("*", "INT", "STRING",)
|
|
RETURN_NAMES = ("AudioClip list", "size", "srt",)
|
|
FUNCTION = "segment"
|
|
CATEGORY = "Audio/Segmentation"
|
|
|
|
def segment(self, audio, text, model):
|
|
waveform = audio["waveform"]
|
|
sample_rate = audio["sample_rate"]
|
|
|
|
if waveform.device.type != 'cpu':
|
|
waveform = waveform.to('cpu')
|
|
|
|
# 假设 batch size 为 1
|
|
waveform_single = waveform[0]
|
|
waveform_np = waveform_single.numpy()
|
|
waveform_int16 = (waveform_np * 32767).astype(np.int16)
|
|
|
|
channels = waveform_int16.shape[0]
|
|
audio_data = waveform_int16.T.tobytes()
|
|
|
|
# 调用新的分割函数
|
|
audio_clips_raw, count, srt_string = segment_by_paragraph(
|
|
audio_data=audio_data,
|
|
sample_rate=sample_rate,
|
|
channels=channels,
|
|
text=text,
|
|
model_name=model,
|
|
model_dir=stt_model_dir
|
|
)
|
|
|
|
if not audio_clips_raw:
|
|
return ([], 0, "")
|
|
|
|
# 将 pydub 音频块转换回 ComfyUI 的 AUDIO 格式
|
|
final_audio_clips = []
|
|
for clip in audio_clips_raw:
|
|
pydub_chunk = clip["audio"]
|
|
|
|
chunk_np = np.array(pydub_chunk.get_array_of_samples())
|
|
|
|
if pydub_chunk.channels > 1:
|
|
chunk_np = chunk_np.reshape((-1, pydub_chunk.channels)).T
|
|
else:
|
|
chunk_np = chunk_np.reshape((1, -1))
|
|
|
|
chunk_tensor = torch.from_numpy(chunk_np.astype(np.float32) / 32767.0)
|
|
|
|
comfy_audio = {
|
|
"waveform": chunk_tensor.unsqueeze(0),
|
|
"sample_rate": sample_rate
|
|
}
|
|
|
|
final_audio_clips.append({
|
|
"audio": comfy_audio,
|
|
"startTime": clip["startTime"],
|
|
"endTime": clip["endTime"],
|
|
"text": clip["text"]
|
|
})
|
|
|
|
return (final_audio_clips, count, srt_string)
|
|
|
|
class AudioSegmenterNode:
|
|
AUX_ID = "dreamidea/comfyui-audio-segment"
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"audio": ("AUDIO",),
|
|
"max_length_s": ("INT", {
|
|
"default": 10, "min": 1, "max": 3600, "step": 1
|
|
}),
|
|
"silence_thresh_db": ("INT", {
|
|
"default": -40, "min": -100, "max": 0, "step": 1
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("*", "INT",)
|
|
RETURN_NAMES = ("audio_list", "count",)
|
|
FUNCTION = "segment"
|
|
CATEGORY = "Audio/Segmentation"
|
|
|
|
def segment(self, audio, max_length_s, silence_thresh_db):
|
|
waveform = audio["waveform"]
|
|
sample_rate = audio["sample_rate"]
|
|
|
|
if waveform.device.type != 'cpu':
|
|
waveform = waveform.to('cpu')
|
|
|
|
waveform_single = waveform[0]
|
|
waveform_np = waveform_single.numpy()
|
|
waveform_int16 = (waveform_np * 32767).astype(np.int16)
|
|
|
|
channels = waveform_int16.shape[0]
|
|
# Transpose to interleave channels for pydub
|
|
audio_data = waveform_int16.T.tobytes()
|
|
|
|
chunks = segment_audio(
|
|
audio_data=audio_data,
|
|
frame_rate=sample_rate,
|
|
channels=channels,
|
|
max_length_s=max_length_s,
|
|
silence_thresh_db=silence_thresh_db
|
|
)
|
|
|
|
if not chunks:
|
|
return ([], 0)
|
|
|
|
audio_list = []
|
|
for chunk in chunks:
|
|
chunk_np = np.array(chunk.get_array_of_samples())
|
|
|
|
if chunk.channels > 1:
|
|
chunk_np = chunk_np.reshape((-1, chunk.channels)).T
|
|
else:
|
|
chunk_np = chunk_np.reshape((1, -1))
|
|
|
|
chunk_tensor = torch.from_numpy(chunk_np.astype(np.float32) / 32767.0)
|
|
|
|
# 每个音频片段都是一个独立的 AUDIO 对象 (batch size 1)
|
|
audio_list.append({
|
|
"waveform": chunk_tensor.unsqueeze(0),
|
|
"sample_rate": sample_rate
|
|
})
|
|
|
|
return (audio_list, len(audio_list))
|
|
|
|
class SelectAudioClipFromList:
|
|
AUX_ID = "dreamidea/comfyui-audio-segment"
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"audio_clip_list": ("*",),
|
|
"index": ("INT", {"default": 0, "min": 0, "step": 1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("AUDIO", "STRING",)
|
|
RETURN_NAMES = ("audio", "text",)
|
|
FUNCTION = "select_clip"
|
|
CATEGORY = "Audio/Segmentation"
|
|
|
|
def select_clip(self, audio_clip_list, index):
|
|
if not audio_clip_list or index >= len(audio_clip_list):
|
|
# Fallback to empty audio and text if list is empty or index is out of bounds
|
|
empty_waveform = torch.zeros((1, 1, 1))
|
|
empty_audio = {"waveform": empty_waveform, "sample_rate": 44100}
|
|
return (empty_audio, "")
|
|
|
|
selected_clip = audio_clip_list[index]
|
|
return (selected_clip["audio"], selected_clip["text"])
|
|
|
|
class SelectAudioFromList:
|
|
AUX_ID = "dreamidea/comfyui-audio-segment"
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"audio_list": ("*",),
|
|
"index": ("INT", {"default": 0, "min": 0, "step": 1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("AUDIO",)
|
|
FUNCTION = "select_audio"
|
|
CATEGORY = "Audio/Segmentation"
|
|
|
|
def select_audio(self, audio_list, index):
|
|
if not audio_list or index >= len(audio_list):
|
|
# 如果列表为空或索引越界，返回一个空的音频对象
|
|
sample_rate = 44100 # Fallback sample rate
|
|
if audio_list:
|
|
sample_rate = audio_list[0].get("sample_rate", sample_rate)
|
|
empty_waveform = torch.zeros((1, 1, 1))
|
|
return ({"waveform": empty_waveform, "sample_rate": sample_rate},)
|
|
|
|
return (audio_list[index],)
|
|
|
|
class AudioListToBatch:
|
|
AUX_ID = "dreamidea/comfyui-audio-segment"
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"audio_list": ("*",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("AUDIO",)
|
|
FUNCTION = "convert_to_batch"
|
|
CATEGORY = "Audio/Segmentation"
|
|
|
|
def convert_to_batch(self, audio_list):
|
|
if not audio_list:
|
|
empty_waveform = torch.zeros((1, 1, 1))
|
|
return ({"waveform": empty_waveform, "sample_rate": 44100},)
|
|
|
|
waveforms = [item["waveform"] for item in audio_list]
|
|
sample_rate = audio_list[0]["sample_rate"]
|
|
|
|
# 确保所有 waveform 的 batch size 都是 1
|
|
waveforms = [w[0] for w in waveforms]
|
|
|
|
max_samples = max(t.shape[1] for t in waveforms)
|
|
|
|
padded_tensors = []
|
|
for t in waveforms:
|
|
padding_needed = max_samples - t.shape[1]
|
|
if padding_needed > 0:
|
|
padded_tensor = torch.nn.functional.pad(t, (0, padding_needed))
|
|
padded_tensors.append(padded_tensor)
|
|
else:
|
|
padded_tensors.append(t)
|
|
|
|
batched_waveform = torch.stack(padded_tensors, dim=0)
|
|
|
|
return ({"waveform": batched_waveform, "sample_rate": sample_rate},)
|