Spaces:
Running on Zero
Running on Zero
File size: 3,440 Bytes
338c6c4 ad316d6 338c6c4 | 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 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | """Thin wrapper around Qwen/Qwen3-ForcedAligner-0.6B.
Loads the model once (lazily, on first request) and exposes a single
`align()` function that both the Gradio UI and the Space API call into.
"""
from __future__ import annotations
import logging
import os
import threading
from dataclasses import dataclass
import torch
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("qwen3_force_aligner")
MODEL_ID = os.environ.get("ALIGNER_MODEL_ID", "Qwen/Qwen3-ForcedAligner-0.6B")
# Languages officially supported by Qwen3-ForcedAligner-0.6B, per the model
# card: Chinese, English, Cantonese, French, German, Italian, Japanese,
# Korean, Portuguese, Russian, Spanish. Vietnamese is included below because
# it was requested for this UI; the underlying model may not have been
# trained/evaluated on it, so alignment quality is not guaranteed.
SUPPORTED_LANGUAGES = [
"Chinese",
"English",
"Cantonese",
"French",
"German",
"Italian",
"Japanese",
"Korean",
"Portuguese",
"Russian",
"Spanish",
"Vietnamese",
"Auto",
]
MAX_AUDIO_SECONDS = 5 * 60 # model card: ~5 minutes per input
@dataclass(frozen=True)
class AlignedSpan:
index: int
text: str
start_time: float
end_time: float
_model = None
_model_lock = threading.Lock()
def _pick_device_and_dtype() -> tuple[str, torch.dtype]:
if torch.cuda.is_available():
return "cuda:0", torch.bfloat16
if torch.backends.mps.is_available():
# bf16 forced alignment on MPS is unreliable in practice; fp32 is safe.
return "mps", torch.float32
return "cpu", torch.float32
def _flash_attention_available() -> bool:
try:
import flash_attn # noqa: F401
return True
except ImportError:
return False
def get_model():
"""Lazily load and cache the Qwen3ForcedAligner model (thread-safe)."""
global _model
if _model is not None:
return _model
with _model_lock:
if _model is not None:
return _model
from qwen_asr import Qwen3ForcedAligner
device, dtype = _pick_device_and_dtype()
kwargs = dict(dtype=dtype, device_map=device)
if device.startswith("cuda") and _flash_attention_available():
kwargs["attn_implementation"] = "flash_attention_2"
logger.info("Loading %s on %s (%s)...", MODEL_ID, device, dtype)
_model = Qwen3ForcedAligner.from_pretrained(MODEL_ID, **kwargs)
logger.info("Model loaded.")
return _model
def align(audio, text: str, language: str) -> list[AlignedSpan]:
"""Run forced alignment for a single (audio, text, language) triple.
`audio` may be anything qwen_asr accepts: a local file path, a URL, a
base64 data URL, or an (np.ndarray, sample_rate) tuple.
"""
if not audio:
raise ValueError("An audio input is required.")
if not text or not text.strip():
raise ValueError("A text transcript is required.")
if language not in SUPPORTED_LANGUAGES:
raise ValueError(f"Unsupported language: {language!r}")
model = get_model()
results = model.align(audio=audio, text=text, language=language)
result = results[0]
return [
AlignedSpan(
index=i,
text=item.text,
start_time=float(item.start_time),
end_time=float(item.end_time),
)
for i, item in enumerate(result)
]
|