| """Native CUDA runtimes are created inside the allocated ZeroGPU worker.""" |
|
|
| import spaces |
|
|
| import json |
| import math |
| import os |
| import time |
| from functools import lru_cache |
| from types import SimpleNamespace |
|
|
| import onnxruntime as ort |
| import torch |
| from huggingface_hub import snapshot_download |
| from transformers import WhisperProcessor |
|
|
| from onnx_runtime import OpenPAOnnx |
| from asr import transcribe_with_whisperx, prepare_asr_mel |
| from asr_cache import TranscriptCache |
|
|
| ASR_CACHE = TranscriptCache(max_entries=16, ttl_seconds=600) |
|
|
| def asr_beam_size(): |
| return max(1, min(5, int(os.getenv("OPENPA_ASR_BEAM_SIZE", "5")))) |
|
|
| if "CUDAExecutionProvider" not in ort.get_available_providers(): |
| raise RuntimeError("CUDA ONNX Runtime wheel was overwritten by a CPU-only installation") |
|
|
| MODEL_CACHE = os.getenv("HF_HOME", "/tmp/openpa-hf") |
| MODEL_REVISION = "13a4f4d027c66bf78c36f58f7cf77b06becb67c0" |
| ASR_REVISION = "edaa852ec7e145841d8ffdb056a99866b5f0a478" |
|
|
|
|
| @lru_cache(maxsize=2) |
| def prepare_artifacts(reference_free): |
| """Download on the CPU process, outside GPU quota usage.""" |
| model = os.getenv("OPENPA_MODEL_DIR") or snapshot_download( |
| "fuann/openpa-onnx", revision=MODEL_REVISION, |
| token=os.getenv("OPENPA_HF_TOKEN"), cache_dir=MODEL_CACHE, |
| allow_patterns=["model.onnx", "model.onnx.data", "*config.json", "tokenizer.json"], |
| ) |
| asr = None |
| if reference_free: |
| asr = snapshot_download( |
| "Systran/faster-whisper-large-v3", revision=ASR_REVISION, |
| token=False, cache_dir=MODEL_CACHE, |
| allow_patterns=["model.bin", "config.json", "tokenizer.json", |
| "vocabulary.json", "preprocessor_config.json"], |
| ) |
| return model, asr |
|
|
|
|
| @lru_cache(maxsize=1) |
| def get_processor(model_dir): |
| return WhisperProcessor.from_pretrained(model_dir, local_files_only=True) |
|
|
|
|
| def prepare_request(waveform, rate, text, model_dir, asr_dir): |
| """Cache CPU-only processing and imports before allocating a GPU.""" |
| processor = get_processor(model_dir) |
| features = processor.feature_extractor( |
| waveform, sampling_rate=rate, return_tensors="np" |
| ).input_features.astype("float32") |
| if asr_dir: |
| import whisperx |
| from whisperx.vads.vad import Vad |
| return processor, features |
|
|
|
|
| def gpu_duration(*args, **kwargs): |
| return int(os.getenv("OPENPA_GPU_SECONDS", "90")) |
|
|
|
|
| @lru_cache(maxsize=1) |
| def get_cpu_scorer(model_dir): |
| """Keep the CPU session alive in the parent; no CUDA resource is retained.""" |
| return OpenPAOnnx(model_dir, threads=int(os.getenv("OPENPA_THREADS", "2")), |
| provider="CPUExecutionProvider", processor=get_processor(model_dir)) |
|
|
|
|
| def asr_duration(waveform, asr_dir, mel_features=None): |
| return 12 + math.ceil(len(waveform) / 16000 / 5) |
|
|
|
|
| @spaces.GPU(duration=asr_duration) |
| def transcribe_gpu(waveform, asr_dir, mel_features=None): |
| started = time.monotonic() |
| if not torch.cuda.is_available(): |
| raise RuntimeError("ZeroGPU did not allocate a CUDA device") |
| options = SimpleNamespace(asr_model=asr_dir, asr_compute_type="float16", |
| asr_batch_size=1, asr_beam_size=asr_beam_size(), local_files_only=True, |
| mel_features=mel_features) |
| transcript, compute_type = transcribe_with_whisperx(waveform, options, torch.device("cuda")) |
| elapsed = round(time.monotonic() - started, 3) |
| print("OPENPA_ASR_GPU_SUCCESS " + json.dumps({"seconds": elapsed, |
| "device": torch.cuda.get_device_name(0), "model": "large-v3", "beam_size": asr_beam_size()}), flush=True) |
| return transcript, compute_type, elapsed |
|
|
|
|
| def run_assessment(waveform, rate, text, model_dir, asr_dir, processor, features, on_stage=None): |
| notify = on_stage or (lambda stage: None) |
| if os.getenv("OPENPA_SCORING_DEVICE", "cpu") == "cuda": |
| notify("transcribe" if text is None else "score") |
| return infer_gpu(waveform, rate, text, model_dir, asr_dir, processor, features) |
| started = time.monotonic() |
| scorer = get_cpu_scorer(model_dir) |
| loaded = time.monotonic() |
| reference_free = text is None |
| asr_seconds = 0 |
| asr_cache_hit = False |
| if reference_free: |
| notify("transcribe") |
| cache_key = ASR_CACHE.key(waveform, (ASR_REVISION, asr_dir, asr_beam_size(), "float16")) |
| cached = ASR_CACHE.get(cache_key) |
| if cached is None: |
| |
| mel_features = prepare_asr_mel(waveform) |
| text, compute_type, asr_seconds = transcribe_gpu(waveform, asr_dir, mel_features) |
| ASR_CACHE.put(cache_key, (text, compute_type)) |
| else: |
| text, compute_type = cached |
| asr_cache_hit = True |
| transcribed = time.monotonic() |
| notify("score") |
| inputs, transcript, words = scorer.prepare(waveform, rate, text, input_features=features) |
| result = scorer.format_scores(scorer.predict(inputs), transcript, words) |
| scored = time.monotonic() |
| result["mode"] = "reference-free" if reference_free else "reference-text" |
| if reference_free: |
| result["asr"] = {"backend": "whisperx", "model": "large-v3", |
| "compute_type": compute_type, "beam_size": asr_beam_size()} |
| result["execution"] = { |
| "providers": scorer.providers, "asr_cache_hit": asr_cache_hit, |
| "timings": {"session_load": round(loaded - started, 3), "asr": asr_seconds, |
| "asr_wall": round(transcribed - loaded, 3), |
| "score": round(scored - transcribed, 3)}, |
| "compute_seconds": round((loaded - started) + asr_seconds + (scored - transcribed), 3), |
| "wall_seconds": round(scored - started, 3), |
| } |
| print("OPENPA_HYBRID_SUCCESS " + json.dumps({"mode": result["mode"], **result["execution"]}), flush=True) |
| return result |
|
|
|
|
| @spaces.GPU(duration=gpu_duration) |
| def infer_gpu(waveform, rate, text, model_dir, asr_dir, processor, features): |
| started = time.monotonic() |
| if not torch.cuda.is_available(): |
| raise RuntimeError("ZeroGPU did not allocate a CUDA device") |
| ort.preload_dlls() |
| ready = time.monotonic() |
| scorer = OpenPAOnnx(model_dir, threads=2, provider="CUDAExecutionProvider", |
| processor=processor, cuda_options={"cudnn_conv_algo_search": "HEURISTIC"}) |
| loaded = time.monotonic() |
| reference_free = text is None |
| if reference_free: |
| options = SimpleNamespace(asr_model=asr_dir, asr_compute_type="float16", |
| asr_batch_size=1, asr_beam_size=asr_beam_size(), local_files_only=True) |
| text, compute_type = transcribe_with_whisperx(waveform, options, torch.device("cuda")) |
| transcribed = time.monotonic() |
| inputs, transcript, words = scorer.prepare(waveform, rate, text, input_features=features) |
| prepared = time.monotonic() |
| predictions = scorer.predict(inputs) |
| predicted = time.monotonic() |
| result = scorer.format_scores(predictions, transcript, words) |
| result["mode"] = "reference-free" if reference_free else "reference-text" |
| if reference_free: |
| result["asr"] = {"backend": "whisperx", "model": "large-v3", |
| "compute_type": compute_type, "beam_size": asr_beam_size()} |
| result["execution"] = { |
| "providers": scorer.providers, "device": torch.cuda.get_device_name(0), |
| "seconds": round(time.monotonic() - started, 3), |
| "timings": {"cuda_ready": round(ready - started, 3), |
| "session_load": round(loaded - ready, 3), |
| "asr": round(transcribed - loaded, 3), |
| "token_preparation": round(prepared - transcribed, 3), |
| "forward": round(predicted - prepared, 3)}, |
| "cuda_options": scorer.session.get_provider_options().get("CUDAExecutionProvider", {}), |
| } |
| print("OPENPA_GPU_SUCCESS " + json.dumps({ |
| "mode": result["mode"], **result["execution"] |
| }), flush=True) |
| return result |
|
|
