107 lines
4.1 KiB
Python
107 lines
4.1 KiB
Python
"""
|
|
VibeVoice ASR engine v4 — Pure VibeVoice source (model + processor).
|
|
Patched AutoModel.register(..., exist_ok=True) for Transformers 5.x compat.
|
|
"""
|
|
import torch
|
|
import sys
|
|
sys.path.insert(0, "/share/ymq/VibeVoice")
|
|
|
|
from longtasks.longtasks import LongTasks
|
|
from appPublic.worker import awaitify
|
|
from appPublic.jsonConfig import getConfig
|
|
from appPublic.log import debug
|
|
from ahserver.filestorage import FileStorage
|
|
|
|
from vibevoice.modular.modeling_vibevoice_asr import VibeVoiceASRForConditionalGeneration
|
|
from vibevoice.processor.vibevoice_asr_processor import VibeVoiceASRProcessor
|
|
|
|
|
|
class VibeVoiceASREngine(LongTasks):
|
|
def __init__(self):
|
|
self.config = getConfig()
|
|
super().__init__(self.config.redis_url, 'vibevoice-asr', worker_cnt=self.config.worker_cnt)
|
|
self.load_models()
|
|
|
|
def load_models(self):
|
|
model_path = self.config.model_path
|
|
debug('loading VibeVoice ASR on cuda:0...')
|
|
self.model = VibeVoiceASRForConditionalGeneration.from_pretrained(
|
|
model_path,
|
|
dtype=torch.bfloat16,
|
|
device_map="cuda:0",
|
|
attn_implementation="sdpa",
|
|
trust_remote_code=True,
|
|
ignore_mismatched_sizes=True,
|
|
)
|
|
self.model.eval()
|
|
self.processor = VibeVoiceASRProcessor.from_pretrained(
|
|
model_path,
|
|
language_model_pretrained_name="/share/models/Qwen2.5-7B",
|
|
)
|
|
debug('VibeVoice ASR loaded')
|
|
|
|
async def process_task(self, payload, workerid=None):
|
|
webpath = payload.get('audio_file')
|
|
if not webpath:
|
|
return {"task_status": "error", "message": "no audio_file"}
|
|
fs = FileStorage()
|
|
fpath = fs.realPath(webpath)
|
|
f = awaitify(self._transcribe)
|
|
return await f(fpath)
|
|
|
|
def _transcribe(self, fpath):
|
|
inputs = self.processor(
|
|
audio=fpath, sampling_rate=None, return_tensors="pt",
|
|
padding=True, add_generation_prompt=True)
|
|
device = torch.device("cuda:0")
|
|
inputs = {k: v.to(device) if isinstance(v, torch.Tensor) else v
|
|
for k, v in inputs.items()}
|
|
with torch.no_grad():
|
|
gen = self.model.generate(**inputs, max_new_tokens=4096,
|
|
temperature=0.0, do_sample=False)
|
|
text = self.processor.decode(gen[0], skip_special_tokens=True)
|
|
try:
|
|
segs = self.processor.post_process_transcription(text)
|
|
except Exception as e:
|
|
debug(f'parse failed: {e}')
|
|
segs = []
|
|
|
|
# Fallback: direct JSON extraction from model output
|
|
if not segs and text:
|
|
import re, json as _json
|
|
# Try full JSON first
|
|
m = re.search(r'assistant\n(.+?)(?:<\|im_end\|>)', text, re.DOTALL)
|
|
json_str = m.group(1).strip() if m else ''
|
|
if not json_str and 'assistant\n[' in text:
|
|
json_str = text[text.find('assistant\n[')+10:]
|
|
if json_str:
|
|
# Handle truncated JSON: find last complete object
|
|
try:
|
|
raw_segs = _json.loads(json_str)
|
|
except Exception:
|
|
last_good = json_str.rfind('"}')
|
|
if last_good > 0:
|
|
try:
|
|
raw_segs = _json.loads(json_str[:last_good+2] + ']')
|
|
except Exception:
|
|
raw_segs = []
|
|
else:
|
|
raw_segs = []
|
|
for s in raw_segs:
|
|
segs.append({
|
|
'start_time': s.get('Start', 0),
|
|
'end_time': s.get('End', 0),
|
|
'speaker_id': s.get('Speaker', 0),
|
|
'text': s.get('Content', ''),
|
|
})
|
|
|
|
ws, ft = [], []
|
|
for s in segs:
|
|
ws.append([s.get('start_time',0), s.get('end_time',0), s.get('text',''), []])
|
|
ft.append(s.get('text',''))
|
|
return {
|
|
'task_status': 'SUCCEEDED', 'language': 'auto',
|
|
'language_probability': 1.0, 'content': ' '.join(ft),
|
|
'segments': ws, 'raw_output': text, 'vibevoice_segments': segs,
|
|
}
|