163 lines
6.4 KiB
Python
163 lines
6.4 KiB
Python
# ComfyUI WB Audio Duration Node
|
|
# Author: Your Name
|
|
# Description: A simple node to get audio duration from audio input
|
|
|
|
# 直接定义节点类，避免导入问题
|
|
import os
|
|
import sys
|
|
import folder_paths
|
|
|
|
# 尝试导入torch，如果失败则使用备用方法
|
|
try:
|
|
import torch
|
|
TORCH_AVAILABLE = True
|
|
except ImportError:
|
|
TORCH_AVAILABLE = False
|
|
print("警告: torch未安装，将使用备用方法处理音频数据")
|
|
|
|
class WB_AudioDurationNode:
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
self.type = "output"
|
|
self.output_node = True
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"audio": ("AUDIO",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("FLOAT", "STRING",)
|
|
RETURN_NAMES = ("duration_seconds", "duration_text",)
|
|
FUNCTION = "get_duration"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "audio"
|
|
|
|
def get_duration(self, audio):
|
|
"""
|
|
简化的音频时长获取函数
|
|
支持1D, 2D, 3D张量格式
|
|
"""
|
|
try:
|
|
print(f"=== 音频分析开始 ===")
|
|
print(f"音频输入类型: {type(audio)}")
|
|
|
|
# 处理字典格式
|
|
if isinstance(audio, dict):
|
|
print(f"音频字典键: {list(audio.keys())}")
|
|
|
|
# 检查是否有必要的键 - 支持两种格式
|
|
if ('waveform' in audio and 'sample_rate' in audio) or ('audio' in audio and 'sample_rate' in audio):
|
|
# 优先使用waveform，如果没有则使用audio
|
|
if 'waveform' in audio:
|
|
audio_data = audio['waveform']
|
|
print("使用waveform键")
|
|
else:
|
|
audio_data = audio['audio']
|
|
print("使用audio键")
|
|
|
|
sample_rate = audio['sample_rate']
|
|
|
|
print(f"音频数据类型: {type(audio_data)}")
|
|
print(f"采样率: {sample_rate}")
|
|
|
|
# 尝试获取音频长度
|
|
if hasattr(audio_data, 'shape'):
|
|
print(f"音频数据形状: {audio_data.shape}")
|
|
|
|
# 计算样本数 - 支持1D, 2D, 3D张量
|
|
if len(audio_data.shape) == 1:
|
|
# 1D: [samples]
|
|
num_samples = audio_data.shape[0]
|
|
print("处理1D张量")
|
|
elif len(audio_data.shape) == 2:
|
|
# 2D: [channels, samples] 或 [batch, samples]
|
|
num_samples = audio_data.shape[1]
|
|
print("处理2D张量")
|
|
elif len(audio_data.shape) == 3:
|
|
# 3D: [batch, channels, samples]
|
|
num_samples = audio_data.shape[2]
|
|
print("处理3D张量")
|
|
else:
|
|
return (0.0, f"不支持的形状: {audio_data.shape}")
|
|
|
|
# 计算时长
|
|
duration_seconds = num_samples / sample_rate
|
|
duration_text = self.format_duration(duration_seconds)
|
|
|
|
print(f"样本数: {num_samples}")
|
|
print(f"计算出的时长: {duration_seconds}秒")
|
|
print(f"格式化时长: {duration_text}")
|
|
|
|
return (duration_seconds, duration_text)
|
|
|
|
elif hasattr(audio_data, '__len__'):
|
|
# 如果数据有长度属性
|
|
num_samples = len(audio_data)
|
|
duration_seconds = num_samples / sample_rate
|
|
duration_text = self.format_duration(duration_seconds)
|
|
|
|
print(f"使用长度属性，样本数: {num_samples}")
|
|
print(f"计算出的时长: {duration_seconds}秒")
|
|
|
|
return (duration_seconds, duration_text)
|
|
|
|
else:
|
|
return (0.0, f"无法获取音频数据长度，类型: {type(audio_data)}")
|
|
|
|
else:
|
|
return (0.0, f"音频字典缺少必要键，当前键: {list(audio.keys())}")
|
|
|
|
# 处理字符串（文件路径）
|
|
elif isinstance(audio, str):
|
|
print(f"音频文件路径: {audio}")
|
|
if os.path.exists(audio):
|
|
# 简单估算
|
|
file_size = os.path.getsize(audio)
|
|
estimated_duration = (file_size * 8) / (128 * 1000) # 假设128kbps
|
|
duration_text = self.format_duration(estimated_duration)
|
|
return (estimated_duration, f"估算时长: {duration_text}")
|
|
else:
|
|
return (0.0, f"文件不存在: {audio}")
|
|
|
|
else:
|
|
return (0.0, f"不支持的音频输入类型: {type(audio)}")
|
|
|
|
except Exception as e:
|
|
import traceback
|
|
print(f"=== 错误详情 ===")
|
|
print(f"错误类型: {type(e)}")
|
|
print(f"错误信息: {str(e)}")
|
|
print(f"错误堆栈: {traceback.format_exc()}")
|
|
return (0.0, f"处理错误: {str(e)}")
|
|
|
|
def format_duration(self, seconds):
|
|
"""
|
|
格式化时长显示
|
|
"""
|
|
if seconds < 1:
|
|
return f"{seconds:.3f}秒"
|
|
elif seconds < 60:
|
|
return f"{seconds:.2f}秒"
|
|
elif seconds < 3600:
|
|
minutes = int(seconds // 60)
|
|
remaining_seconds = seconds % 60
|
|
return f"{minutes}分{remaining_seconds:.2f}秒"
|
|
else:
|
|
hours = int(seconds // 3600)
|
|
remaining_minutes = int((seconds % 3600) // 60)
|
|
remaining_seconds = seconds % 60
|
|
return f"{hours}小时{remaining_minutes}分{remaining_seconds:.2f}秒"
|
|
|
|
# 节点注册
|
|
NODE_CLASS_MAPPINGS = {
|
|
"WB_AudioDuration": WB_AudioDurationNode
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"WB_AudioDuration": "WB Audio Duration"
|
|
}
|
|
|
|
print("ComfyUI WB Audio Duration Node 已加载") |