Spaces:
Running on Zero
Running on Zero
File size: 4,358 Bytes
1b2da16 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | """Known-transcript alignment, separate from synthesis and its GPU lease."""
import logging
import math
from pathlib import Path
import unicodedata
import spaces # Must precede torch.
import numpy as np
import torch
import torchaudio.functional as AF
from transformers import AutoModelForTokenClassification, AutoProcessor
from config import ALIGNER_ID, ALIGNER_REVISION, SAMPLE_RATE
from lifecycle import check_cancelled
log = logging.getLogger("zonos.alignment")
ALIGNER = None
def canonical(text):
# Keep a source-index map so punctuation and typography survive unchanged.
chars, positions = [], []
for i, char in enumerate(text):
for normalized in unicodedata.normalize("NFKC", char).casefold():
if normalized.isalnum():
chars.append(normalized)
positions.append(i)
return "".join(chars), positions
def attach_timings(text, words, duration):
"""Validate measured timestamps and map them back to the exact input text."""
stream, positions = canonical(text)
if not stream or not words or not math.isfinite(duration) or duration <= 0:
raise ValueError("Missing alignment")
cursor, previous, segments = 0, 0.0, []
for word in words:
token, _ = canonical(str(word["text"]))
start, end = float(word["start_time"]), float(word["end_time"])
if not token or not stream.startswith(token, cursor):
raise ValueError("Alignment transcript mismatch")
if not all(math.isfinite(v) for v in (start,end)) or start < previous-.001 or end < start or start < -.08 or end > duration+.16:
raise ValueError("Invalid alignment timestamps")
start, end = max(0,start), min(duration,end)
if start > end:
raise ValueError("Alignment outside audio")
offset = positions[cursor]
segments.append({"offset":offset,"start":round(start,3),"end":round(end,3)})
cursor += len(token)
previous = end
if cursor != len(stream) or not any(w["end"] > w["start"] for w in segments):
raise ValueError("Incomplete alignment")
# Text strings, not offsets, cross the Python/JS boundary (Unicode-safe).
result=[]
for i, segment in enumerate(segments):
left = 0 if i == 0 else segment["offset"]
right = segments[i+1]["offset"] if i+1<len(segments) else len(text)
result.append({"text":text[left:right],"start":segment["start"],"end":segment["end"]})
return result
class Aligner:
def __init__(self):
log.info("alignment_loading revision=%s", ALIGNER_REVISION)
self.processor=AutoProcessor.from_pretrained(ALIGNER_ID,revision=ALIGNER_REVISION)
self.model=AutoModelForTokenClassification.from_pretrained(
ALIGNER_ID, revision=ALIGNER_REVISION, dtype=torch.bfloat16,
).to("cuda:0").eval()
log.info("alignment_loaded")
def prepare(self, pcm, text):
audio=AF.resample(torch.from_numpy(np.asarray(pcm,dtype=np.float32).copy()), SAMPLE_RATE,16000).numpy()
# Documented default segmentation: whitespace words + individual CJK
# characters. Avoid substituting an ASR transcript for the user's text.
return self.processor.prepare_forced_aligner_inputs(audio=audio,transcript=text,return_tensors="pt")
def run(self, prepared, marker, progress):
return _align_gpu(prepared,str(marker),progress)
@spaces.GPU(duration=20)
def _align_gpu(prepared, marker, progress):
log.info("alignment_gpu_enter request=%s",Path(marker).stem)
inputs = outputs = None
try:
check_cancelled(marker)
progress(None,desc="Synchronizing words")
cpu_inputs,word_lists=prepared
inputs=cpu_inputs.to("cuda:0",torch.bfloat16)
with torch.inference_mode():
outputs=ALIGNER.model(**inputs)
check_cancelled(marker)
return ALIGNER.processor.decode_forced_alignment(
logits=outputs.logits,input_ids=inputs["input_ids"],word_lists=word_lists,
timestamp_token_id=ALIGNER.model.config.timestamp_token_id,
)[0]
finally:
del outputs, inputs
torch.cuda.synchronize()
log.info("alignment_gpu_exit request=%s",Path(marker).stem)
def load_aligner():
global ALIGNER
if ALIGNER is None:
ALIGNER=Aligner()
return ALIGNER
|