Infinity AIandClaude Opus 4.8
844aa39622
Aligns generated speech against known text; outputs per-word timings + char offsets for the Synthy/iRead boris_tts_* provider. Pairs with an n8n/FastAPI mapping layer that converts to UTF-16 offsets on the original text. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
219 lines
8.2 KiB
Python
219 lines
8.2 KiB
Python
"""
|
||
ComfyUI-ForcedAlignTimings — word-level forced alignment for TTS markup.
|
||
|
||
Aligns generated speech audio against a KNOWN text using torchaudio's
|
||
MMS_FA (wav2vec2 multilingual forced aligner) and returns per-word timings +
|
||
character offsets into the text we aligned against.
|
||
|
||
Pipeline role (Synthy/iRead "boris_tts_*" provider PoC):
|
||
CosyVoice (text -> audio)
|
||
-> ForcedAlignTimings (audio + known text -> raw word timings)
|
||
-> n8n Code node (raw timings + ORIGINAL text -> X-Voice-Markup)
|
||
|
||
This node deliberately does NOT compute UTF-16 offsets and does NOT map text
|
||
normalization back to the original — that is pure string logic and lives in the
|
||
n8n Code node (see 1/tts-poc/map_timings.mjs). Here we only emit, per word:
|
||
value, char_start, char_end (Python codepoint offsets), time(ms), end(ms), score
|
||
|
||
Author: iMind Studio
|
||
License: MIT
|
||
"""
|
||
|
||
import json
|
||
import re
|
||
|
||
import torch
|
||
import torchaudio
|
||
|
||
|
||
# ── word tokenizer ───────────────────────────────────
|
||
# One token per run of letters (with internal apostrophes) OR run of digits.
|
||
# Surrounding punctuation/space is excluded, so char_start/char_end point at the
|
||
# word core only (matches SY-A-39: commas / "!" are not emitted as marks).
|
||
_WORD_RE = re.compile(r"[0-9]+|[^\W\d_]+(?:['’][^\W\d_]+)*", re.UNICODE)
|
||
|
||
# Compact Cyrillic -> Latin fallback, used only if the `uroman` package is absent.
|
||
_RU_MAP = {
|
||
"а": "a", "б": "b", "в": "v", "г": "g", "д": "d", "е": "e", "ё": "yo",
|
||
"ж": "zh", "з": "z", "и": "i", "й": "y", "к": "k", "л": "l", "м": "m",
|
||
"н": "n", "о": "o", "п": "p", "р": "r", "с": "s", "т": "t", "у": "u",
|
||
"ф": "f", "х": "kh", "ц": "ts", "ч": "ch", "ш": "sh", "щ": "shch",
|
||
"ъ": "", "ы": "y", "ь": "", "э": "e", "ю": "yu", "я": "ya",
|
||
}
|
||
|
||
_uroman = None
|
||
_uroman_tried = False
|
||
|
||
|
||
def _romanize(word):
|
||
"""Romanize a single word to lowercase [a-z'] for the MMS_FA dictionary."""
|
||
global _uroman, _uroman_tried
|
||
if not _uroman_tried:
|
||
_uroman_tried = True
|
||
try:
|
||
import uroman as _ur
|
||
_uroman = _ur.Uroman()
|
||
except Exception:
|
||
_uroman = None
|
||
|
||
if _uroman is not None:
|
||
try:
|
||
text = _uroman.romanize_string(word)
|
||
except Exception:
|
||
text = word
|
||
else:
|
||
out = []
|
||
for ch in word:
|
||
low = ch.lower()
|
||
out.append(_RU_MAP[low] if low in _RU_MAP else ch)
|
||
text = "".join(out)
|
||
|
||
text = text.lower()
|
||
text = re.sub(r"[^a-z']", "", text)
|
||
return text
|
||
|
||
|
||
class ForcedAlignTimings:
|
||
"""Forced-align an AUDIO clip against a KNOWN text; return raw word timings."""
|
||
|
||
TARGET_SR = 16000 # MMS_FA expects 16 kHz mono
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"audio": ("AUDIO",),
|
||
"text": ("STRING", {"default": "", "multiline": True}),
|
||
},
|
||
"optional": {
|
||
"device": (["auto", "cuda", "cpu"],),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("STRING", "INT", "INT")
|
||
RETURN_NAMES = ("timings_json", "audio_ms", "word_count")
|
||
FUNCTION = "align"
|
||
CATEGORY = "audio/tts"
|
||
DESCRIPTION = ("Forced alignment (MMS_FA) of audio against known text. "
|
||
"Outputs per-word timings + char offsets as JSON for TTS markup.")
|
||
OUTPUT_NODE = True # surface timings_json into /history outputs so n8n can read it via API
|
||
|
||
def _wrap(self, timings_json, audio_ms, count):
|
||
# Downstream sockets receive `result`; n8n reads `ui` from /history[pid].outputs[nodeId].
|
||
return {
|
||
"ui": {"timings_json": [timings_json], "audio_ms": [audio_ms], "word_count": [count]},
|
||
"result": (timings_json, audio_ms, count),
|
||
}
|
||
|
||
# ── helpers ──────────────────────────────────────
|
||
def _prep_audio(self, audio):
|
||
"""ComfyUI AUDIO dict -> mono 16 kHz [1, T] tensor."""
|
||
wav = audio["waveform"] # [B, C, T]
|
||
sr = int(audio["sample_rate"])
|
||
if wav.dim() == 3:
|
||
wav = wav[0] # [C, T]
|
||
if wav.dim() == 2 and wav.shape[0] > 1:
|
||
wav = wav.mean(dim=0, keepdim=True)
|
||
elif wav.dim() == 1:
|
||
wav = wav.unsqueeze(0)
|
||
if sr != self.TARGET_SR:
|
||
wav = torchaudio.functional.resample(wav, sr, self.TARGET_SR)
|
||
return wav.contiguous(), self.TARGET_SR
|
||
|
||
@staticmethod
|
||
def _get_dict(bundle):
|
||
try:
|
||
return bundle.get_dict(star="*")
|
||
except TypeError:
|
||
return bundle.get_dict()
|
||
|
||
@staticmethod
|
||
def _get_model(bundle, with_star):
|
||
try:
|
||
return bundle.get_model(with_star=with_star)
|
||
except TypeError:
|
||
return bundle.get_model()
|
||
|
||
# ── main ─────────────────────────────────────────
|
||
def align(self, audio, text, device="auto"):
|
||
from torchaudio.pipelines import MMS_FA as bundle
|
||
|
||
dev = ("cuda" if torch.cuda.is_available() else "cpu") if device == "auto" else device
|
||
|
||
waveform, sr = self._prep_audio(audio)
|
||
waveform = waveform.to(dev)
|
||
audio_ms = int(round(waveform.shape[-1] / sr * 1000))
|
||
|
||
# 1) tokenize text into words with codepoint char offsets
|
||
words = [{"value": m.group(0), "char_start": m.start(), "char_end": m.end()}
|
||
for m in _WORD_RE.finditer(text)]
|
||
|
||
if not words:
|
||
payload = {"sample_rate": sr, "audio_ms": audio_ms, "align_text": text, "words": []}
|
||
return self._wrap(json.dumps(payload, ensure_ascii=False), audio_ms, 0)
|
||
|
||
# 2) romanize words -> MMS_FA dictionary token ids (digits/unknown -> '*')
|
||
dictionary = self._get_dict(bundle)
|
||
star = dictionary.get("*")
|
||
|
||
tokens_flat, counts = [], []
|
||
for w in words:
|
||
roman = _romanize(w["value"])
|
||
ids = [dictionary[c] for c in roman if c in dictionary]
|
||
if not ids:
|
||
ids = [star if star is not None else 0]
|
||
tokens_flat.extend(ids)
|
||
counts.append(len(ids))
|
||
|
||
# 3) emission + forced alignment
|
||
model = self._get_model(bundle, with_star=(star is not None)).to(dev)
|
||
with torch.inference_mode():
|
||
emission, _ = model(waveform)
|
||
targets = torch.tensor([tokens_flat], dtype=torch.int32, device=dev)
|
||
aligned, scores = torchaudio.functional.forced_align(emission, targets, blank=0)
|
||
token_spans = torchaudio.functional.merge_tokens(aligned[0], scores[0].exp())
|
||
|
||
num_frames = emission.shape[1]
|
||
ratio = waveform.shape[-1] / num_frames / sr # seconds per emission frame
|
||
|
||
# 4) regroup token spans into words by token counts
|
||
out_words = []
|
||
idx = 0
|
||
for w, n in zip(words, counts):
|
||
span = token_spans[idx: idx + n]
|
||
idx += n
|
||
if span:
|
||
t0 = int(round(span[0].start * ratio * 1000))
|
||
t1 = int(round(span[-1].end * ratio * 1000))
|
||
sc = sum(s.score for s in span) / len(span)
|
||
else:
|
||
t0 = t1 = 0
|
||
sc = 0.0
|
||
out_words.append({
|
||
"value": w["value"],
|
||
"char_start": w["char_start"],
|
||
"char_end": w["char_end"],
|
||
"time": t0,
|
||
"end": t1,
|
||
"score": round(float(sc), 4),
|
||
})
|
||
|
||
payload = {"sample_rate": sr, "audio_ms": audio_ms, "align_text": text, "words": out_words}
|
||
return self._wrap(json.dumps(payload, ensure_ascii=False), audio_ms, len(out_words))
|
||
|
||
|
||
NODE_CLASS_MAPPINGS = {"ForcedAlignTimings": ForcedAlignTimings}
|
||
NODE_DISPLAY_NAME_MAPPINGS = {"ForcedAlignTimings": "Forced Align Timings 🎯⏱️"}
|
||
|
||
|
||
# ── GPU smoke test (run on the 5090 box BEFORE docker commit) ──
|
||
# Usage: python nodes.py path/to/audio.wav "Привет, мир"
|
||
if __name__ == "__main__":
|
||
import sys
|
||
|
||
wav_path, in_text = sys.argv[1], sys.argv[2]
|
||
wav, sr = torchaudio.load(wav_path)
|
||
audio_in = {"waveform": wav.unsqueeze(0), "sample_rate": sr}
|
||
out = ForcedAlignTimings().align(audio_in, in_text)
|
||
print(out["result"][0])
|