414 lines
12 KiB
Python
414 lines
12 KiB
Python
import hashlib
|
|
import json
|
|
import os
|
|
import tempfile
|
|
from functools import lru_cache
|
|
|
|
import folder_paths
|
|
|
|
|
|
MODEL_OPTIONS = ["large-v3", "large-v3-turbo"]
|
|
DEVICE_OPTIONS = ["auto", "cuda", "cpu"]
|
|
COMPUTE_OPTIONS = ["auto", "float16", "int8_float16", "int8", "float32"]
|
|
TASK_OPTIONS = ["transcribe", "translate"]
|
|
|
|
|
|
def _audio_files():
|
|
input_dir = folder_paths.get_input_directory()
|
|
os.makedirs(input_dir, exist_ok=True)
|
|
files = [
|
|
name
|
|
for name in os.listdir(input_dir)
|
|
if os.path.isfile(os.path.join(input_dir, name))
|
|
]
|
|
files = folder_paths.filter_files_content_types(files, ["audio", "video"])
|
|
return sorted(files) or [""]
|
|
|
|
|
|
def _model_cache_dir():
|
|
models_dir = getattr(folder_paths, "models_dir", None)
|
|
if not models_dir:
|
|
models_dir = os.path.join(os.path.dirname(__file__), "models")
|
|
path = os.path.join(models_dir, "whisper-large-v3-repack")
|
|
os.makedirs(path, exist_ok=True)
|
|
return path
|
|
|
|
|
|
def _resolve_device(device):
|
|
if device != "auto":
|
|
return device
|
|
try:
|
|
import torch
|
|
|
|
return "cuda" if torch.cuda.is_available() else "cpu"
|
|
except Exception:
|
|
return "cpu"
|
|
|
|
|
|
def _resolve_compute_type(compute_type, device):
|
|
if compute_type != "auto":
|
|
return compute_type
|
|
return "float16" if device == "cuda" else "int8"
|
|
|
|
|
|
@lru_cache(maxsize=4)
|
|
def _load_model(model_name, device, compute_type, cache_dir):
|
|
try:
|
|
from faster_whisper import WhisperModel
|
|
except ImportError as exc:
|
|
raise ImportError(
|
|
"ComfyUI-WhisperLargeV3-Repack requires faster-whisper. "
|
|
"Install this node's requirements.txt, then restart ComfyUI."
|
|
) from exc
|
|
|
|
return WhisperModel(
|
|
model_name,
|
|
device=device,
|
|
compute_type=compute_type,
|
|
download_root=cache_dir,
|
|
)
|
|
|
|
|
|
def _words_to_json(words):
|
|
if not words:
|
|
return []
|
|
return [
|
|
{
|
|
"start": float(word.start) if word.start is not None else None,
|
|
"end": float(word.end) if word.end is not None else None,
|
|
"word": word.word,
|
|
"probability": float(word.probability)
|
|
if word.probability is not None
|
|
else None,
|
|
}
|
|
for word in words
|
|
]
|
|
|
|
|
|
def _transcribe_path(
|
|
audio_path,
|
|
model_name,
|
|
device,
|
|
compute_type,
|
|
language,
|
|
task,
|
|
beam_size,
|
|
vad_filter,
|
|
word_timestamps,
|
|
initial_prompt,
|
|
temperature,
|
|
):
|
|
resolved_device = _resolve_device(device)
|
|
resolved_compute_type = _resolve_compute_type(compute_type, resolved_device)
|
|
segments, info = _transcribe_with_fallback(
|
|
audio_path,
|
|
model_name,
|
|
requested_device=device,
|
|
resolved_device=resolved_device,
|
|
resolved_compute_type=resolved_compute_type,
|
|
original_compute_type=compute_type,
|
|
language=language,
|
|
task=task,
|
|
beam_size=beam_size,
|
|
vad_filter=vad_filter,
|
|
word_timestamps=word_timestamps,
|
|
initial_prompt=initial_prompt,
|
|
temperature=temperature,
|
|
)
|
|
|
|
segment_rows = []
|
|
text_parts = []
|
|
for segment in segments:
|
|
text = segment.text or ""
|
|
text_parts.append(text)
|
|
segment_rows.append(
|
|
{
|
|
"id": int(segment.id),
|
|
"start": float(segment.start),
|
|
"end": float(segment.end),
|
|
"text": text,
|
|
"words": _words_to_json(getattr(segment, "words", None)),
|
|
}
|
|
)
|
|
|
|
transcript = "".join(text_parts).strip()
|
|
detected_language = getattr(info, "language", "") or ""
|
|
duration = float(getattr(info, "duration", 0.0) or 0.0)
|
|
segments_json = json.dumps(segment_rows, ensure_ascii=False)
|
|
return transcript, detected_language, duration, segments_json
|
|
|
|
|
|
def _transcribe_with_fallback(
|
|
audio_path,
|
|
model_name,
|
|
requested_device,
|
|
resolved_device,
|
|
resolved_compute_type,
|
|
original_compute_type,
|
|
language,
|
|
task,
|
|
beam_size,
|
|
vad_filter,
|
|
word_timestamps,
|
|
initial_prompt,
|
|
temperature,
|
|
):
|
|
kwargs = {
|
|
"language": (language.strip() or None),
|
|
"task": task,
|
|
"beam_size": int(beam_size),
|
|
"vad_filter": bool(vad_filter),
|
|
"word_timestamps": bool(word_timestamps),
|
|
"initial_prompt": (initial_prompt.strip() or None),
|
|
"temperature": float(temperature),
|
|
}
|
|
|
|
try:
|
|
model = _load_model(
|
|
model_name,
|
|
resolved_device,
|
|
resolved_compute_type,
|
|
_model_cache_dir(),
|
|
)
|
|
return model.transcribe(audio_path, **kwargs)
|
|
except RuntimeError as exc:
|
|
if not _should_cpu_fallback(exc, requested_device, resolved_device):
|
|
raise
|
|
|
|
_load_model.cache_clear()
|
|
fallback_compute = (
|
|
original_compute_type
|
|
if original_compute_type in {"int8", "float32"}
|
|
else "int8"
|
|
)
|
|
model = _load_model(
|
|
model_name,
|
|
"cpu",
|
|
fallback_compute,
|
|
_model_cache_dir(),
|
|
)
|
|
return model.transcribe(audio_path, **kwargs)
|
|
|
|
|
|
def _should_cpu_fallback(exc, requested_device, resolved_device):
|
|
if requested_device != "auto" or resolved_device != "cuda":
|
|
return False
|
|
message = str(exc).lower()
|
|
cuda_markers = (
|
|
"cublas",
|
|
"cudnn",
|
|
"cuda",
|
|
"could not load library",
|
|
"is not found or cannot be loaded",
|
|
)
|
|
return any(marker in message for marker in cuda_markers)
|
|
|
|
|
|
def _audio_to_temp_wav(audio):
|
|
try:
|
|
import soundfile as sf
|
|
except ImportError as exc:
|
|
raise ImportError(
|
|
"ComfyUI-WhisperLargeV3-Repack requires soundfile for AUDIO inputs."
|
|
) from exc
|
|
|
|
waveform = audio["waveform"]
|
|
sample_rate = int(audio["sample_rate"])
|
|
samples = waveform[0].detach().cpu().float().permute(1, 0).numpy()
|
|
|
|
handle = tempfile.NamedTemporaryFile(
|
|
suffix=".wav",
|
|
prefix="whisper_large_v3_repack_",
|
|
delete=False,
|
|
)
|
|
handle.close()
|
|
sf.write(handle.name, samples, sample_rate, format="WAV", subtype="PCM_16")
|
|
return handle.name
|
|
|
|
|
|
def _ui_result(transcript, language, duration, segments_json):
|
|
preview = transcript or "[No speech detected]"
|
|
return {
|
|
"ui": {"text": (preview,)},
|
|
"result": (transcript, language, duration, segments_json),
|
|
}
|
|
|
|
|
|
class WhisperLargeV3RepackTranscribeFile:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"audio": (_audio_files(),),
|
|
"model_name": (MODEL_OPTIONS, {"default": "large-v3"}),
|
|
"device": (DEVICE_OPTIONS, {"default": "auto"}),
|
|
"compute_type": (COMPUTE_OPTIONS, {"default": "auto"}),
|
|
"language": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"placeholder": "Optional ISO code, empty = auto",
|
|
},
|
|
),
|
|
"task": (TASK_OPTIONS, {"default": "transcribe"}),
|
|
"beam_size": ("INT", {"default": 5, "min": 1, "max": 16}),
|
|
"vad_filter": ("BOOLEAN", {"default": True}),
|
|
"word_timestamps": ("BOOLEAN", {"default": False}),
|
|
"initial_prompt": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"multiline": True,
|
|
"placeholder": "Optional vocabulary or context",
|
|
},
|
|
),
|
|
"temperature": (
|
|
"FLOAT",
|
|
{"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.1},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING", "FLOAT", "STRING")
|
|
RETURN_NAMES = ("transcript", "language", "duration_seconds", "segments_json")
|
|
FUNCTION = "transcribe"
|
|
CATEGORY = "audio/whisper"
|
|
OUTPUT_NODE = True
|
|
|
|
def transcribe(
|
|
self,
|
|
audio,
|
|
model_name,
|
|
device,
|
|
compute_type,
|
|
language,
|
|
task,
|
|
beam_size,
|
|
vad_filter,
|
|
word_timestamps,
|
|
initial_prompt,
|
|
temperature,
|
|
):
|
|
audio_path = folder_paths.get_annotated_filepath(audio)
|
|
result = _transcribe_path(
|
|
audio_path,
|
|
model_name,
|
|
device,
|
|
compute_type,
|
|
language,
|
|
task,
|
|
beam_size,
|
|
vad_filter,
|
|
word_timestamps,
|
|
initial_prompt,
|
|
temperature,
|
|
)
|
|
return _ui_result(*result)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, audio, **kwargs):
|
|
if not folder_paths.exists_annotated_filepath(audio):
|
|
return audio
|
|
audio_path = folder_paths.get_annotated_filepath(audio)
|
|
digest = hashlib.sha256()
|
|
with open(audio_path, "rb") as handle:
|
|
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
digest.update(json.dumps(kwargs, sort_keys=True).encode("utf-8"))
|
|
return digest.hexdigest()
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(cls, audio, **kwargs):
|
|
if not audio or not folder_paths.exists_annotated_filepath(audio):
|
|
return f"Invalid audio file: {audio}"
|
|
return True
|
|
|
|
|
|
class WhisperLargeV3RepackTranscribeAudio:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"audio": ("AUDIO",),
|
|
"model_name": (MODEL_OPTIONS, {"default": "large-v3"}),
|
|
"device": (DEVICE_OPTIONS, {"default": "auto"}),
|
|
"compute_type": (COMPUTE_OPTIONS, {"default": "auto"}),
|
|
"language": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"placeholder": "Optional ISO code, empty = auto",
|
|
},
|
|
),
|
|
"task": (TASK_OPTIONS, {"default": "transcribe"}),
|
|
"beam_size": ("INT", {"default": 5, "min": 1, "max": 16}),
|
|
"vad_filter": ("BOOLEAN", {"default": True}),
|
|
"word_timestamps": ("BOOLEAN", {"default": False}),
|
|
"initial_prompt": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"multiline": True,
|
|
"placeholder": "Optional vocabulary or context",
|
|
},
|
|
),
|
|
"temperature": (
|
|
"FLOAT",
|
|
{"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.1},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING", "FLOAT", "STRING")
|
|
RETURN_NAMES = ("transcript", "language", "duration_seconds", "segments_json")
|
|
FUNCTION = "transcribe"
|
|
CATEGORY = "audio/whisper"
|
|
OUTPUT_NODE = True
|
|
|
|
def transcribe(
|
|
self,
|
|
audio,
|
|
model_name,
|
|
device,
|
|
compute_type,
|
|
language,
|
|
task,
|
|
beam_size,
|
|
vad_filter,
|
|
word_timestamps,
|
|
initial_prompt,
|
|
temperature,
|
|
):
|
|
temp_path = _audio_to_temp_wav(audio)
|
|
try:
|
|
result = _transcribe_path(
|
|
temp_path,
|
|
model_name,
|
|
device,
|
|
compute_type,
|
|
language,
|
|
task,
|
|
beam_size,
|
|
vad_filter,
|
|
word_timestamps,
|
|
initial_prompt,
|
|
temperature,
|
|
)
|
|
finally:
|
|
try:
|
|
os.remove(temp_path)
|
|
except OSError:
|
|
pass
|
|
return _ui_result(*result)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"WhisperLargeV3RepackTranscribeFile": WhisperLargeV3RepackTranscribeFile,
|
|
"WhisperLargeV3RepackTranscribeAudio": WhisperLargeV3RepackTranscribeAudio,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"WhisperLargeV3RepackTranscribeFile": "Whisper Large V3 Repack - Transcribe File",
|
|
"WhisperLargeV3RepackTranscribeAudio": "Whisper Large V3 Repack - Transcribe Audio",
|
|
}
|