Spaces:
Running on Zero
Running on Zero
| """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) | |
| 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 | |