ute2e/tools/bridge-introduction-narration-repair.py

"""Regenerate only ASR-flagged pages using shorter complete clauses; retain originals."""
import argparse
import importlib.util
import json
from pathlib import Path
import re
import shutil
import subprocess
import threading
import numpy as np
import soundfile as sf
spec=importlib.util.spec_from_file_location('generator',Path(__file__).with_name('bridge-introduction-narration-generate.py'))
g=importlib.util.module_from_spec(spec);spec.loader.exec_module(g)

def shorter(text,limit=85):
    result=[];current=''
    for part in re.findall(r'[^，。！？；\n]+[，。！？；\n]?',text):
        if len(current)+len(part)>limit and current.strip():result.append(current.strip());current=''
        current+=part
    if current.strip():result.append(current.strip())
    return result

def main():
    parser=argparse.ArgumentParser();parser.add_argument('--keys',nargs='+',required=True);args=parser.parse_args()
    path=g.REPORT/'generation.json';data=json.loads(path.read_text('utf-8'));profile=data['profile']
    assert len(data['clips'])==11,'Finish the initial generation before repairs'
    settings=g.get_settings().model_copy(update={'service_base_url':'http://1.14.103.234:8123'})
    engine=g.MossTtsService(settings);renderer=g.OpenGenerationWorker(None,settings,None,None)
    with g.get_session_factory()() as db:
        agent=db.scalar(g.select(g.Agent).where(g.Agent.slug==profile['agentSlug'],g.Agent.is_delete==0))
        assert agent.data_version==profile['agentVersion'],'Default configuration changed'
        av=db.scalar(g.select(g.AgentAvatar).where(g.AgentAvatar.agent_id==agent.id).order_by(g.AgentAvatar.is_default.desc(),g.AgentAvatar.sort_order))
        vo=db.scalar(g.select(g.AgentVoice).where(g.AgentVoice.agent_id==agent.id).order_by(g.AgentVoice.is_default.desc(),g.AgentVoice.sort_order))
        avatar_id,voice_id=str(av.avatar_id),str(vo.capability_id)
    originals=g.REPORT/'initial-media';originals.mkdir(exist_ok=True)
    try:
        for key in args.keys:
            clip=next(c for c in data['clips'] if c['key']==key)
            if clip.get('repair')=='short-clauses-v1':continue
            for kind in ['audioUrl','videoUrl','posterUrl']:
                src=g.PUBLIC/clip[kind].lstrip('/');dest=originals/src.name
                if not dest.exists():shutil.copyfile(src,dest)
            g.helpers.write_json(originals/(key+'.json'),clip)
            # Preserve all already reviewed speech. Only replace the chunk that
            # contains the long omission/repetition identified in quality.json.
            affected={'TASK-VIRTUAL-005-intro-00':{1},'TASK-VIRTUAL-005-intro-03':{0},
                'TASK-VIRTUAL-005-intro-04':{0},'TASK-VIRTUAL-006-intro-00':{1},'TASK-VIRTUAL-007-intro-00':{1}}
            plan=[]
            for i,meta in enumerate(sorted(g.REPORT.glob(key+'-chunk-*.json'))):
                part=json.loads(meta.read_text('utf-8'))['text']
                if i in affected.get(key,set()):
                    normalized=part.replace('（','，').replace('）','，').replace('M口','艾姆口')
                    plan.extend((p,None) for p in shorter(normalized,50 if key.endswith('005-intro-04') else 85))
                else:plan.append((part,meta.with_suffix('.wav')))
            assert plan and key in affected
            text=''.join(p for p,_ in plan)
            identical=next((c for c in data['clips'] if c.get('repair')=='short-clauses-v1' and c['text']==text),None)
            if identical:
                for kind in ['audioUrl','videoUrl','posterUrl']:shutil.copyfile(g.PUBLIC/identical[kind].lstrip('/'),g.PUBLIC/clip[kind].lstrip('/'))
                for name in ['duration','rawDuration','pauseCuts','videoDuration','text','audioSha256','videoSha256','repair']:clip[name]=identical[name]
                clip['reusedIdenticalSpeechFrom']=identical['key'];g.helpers.write_json(path,data);print(key,'REUSED IDENTICAL INPUT',flush=True);continue
            samples=[];rate=None
            for i,(part,reuse) in enumerate(plan):
                if reuse:
                    pcm,sr=sf.read(reuse,always_2d=True);assert rate in (None,sr);rate=sr;samples.append(pcm);continue
                raw=g.REPORT/f'{key}-repair-v1-{i:02}.wav';meta=raw.with_suffix('.json')
                if raw.exists():assert json.loads(meta.read_text('utf-8'))=={'text':part,'profile':profile}
                else:
                    print(key,'repair speech',i+1,'/',len(plan),flush=True)
                    result=engine.synthesize(part,voice_code=profile['speaker'],speed=profile['speed'])
                    shutil.copyfile(engine.output_path(result['id']),raw);g.helpers.write_json(meta,{'text':part,'profile':profile})
                pcm,sr=sf.read(raw,always_2d=True);assert rate in (None,sr);rate=sr;samples.append(pcm)
            raw=g.REPORT/(key+'-repair-v1-raw.wav');sf.write(raw,np.concatenate(samples),rate,subtype='PCM_16')
            prepared=g.REPORT/(key+'-repair-v1-prepared.wav');timing=g.helpers.prepare_audio(settings,raw,prepared)
            audio=g.PUBLIC/clip['audioUrl'].lstrip('/');video=g.PUBLIC/clip['videoUrl'].lstrip('/');poster=g.PUBLIC/clip['posterUrl'].lstrip('/')
            subprocess.run([settings.ffmpeg_binary,'-v','error','-y','-i',str(prepared),'-ac','1','-ar','16000','-c:a','pcm_s16le',str(audio)],check=True)
            capture=g.REPORT/(key+'-repair-v1.webm')
            print(key,'repair render',timing['duration'],flush=True)
            command=g.GenerationCommand(title=clip['title'],text=text,avatarId=avatar_id,voiceId=voice_id,speed=profile['speed'],width=480,height=640,subtitles=False)
            renderer.capture(command,profile['providerAvatarId'],audio,None,capture,threading.Event())
            subprocess.run([settings.ffmpeg_binary,'-v','error','-y','-i',str(capture),'-r','25','-c:v','libx264','-preset','veryfast','-crf','20','-pix_fmt','yuv420p','-c:a','aac','-b:a','128k','-movflags','+faststart',str(video)],check=True)
            info=g.helpers.probe(settings,video);v=next(s for s in info['streams'] if s['codec_type']=='video')
            assert (v['width'],v['height'],v['r_frame_rate'])==(480,640,'25/1')
            assert abs(float(info['format']['duration'])-timing['duration'])<.7
            subprocess.run([settings.ffmpeg_binary,'-v','error','-y','-ss','0.1','-i',str(video),'-frames:v','1',str(poster)],check=True)
            clip.update(**timing,text=text,videoDuration=float(info['format']['duration']),audioSha256=g.helpers.digest(audio),videoSha256=g.helpers.digest(video),repair='short-clauses-v1')
            g.helpers.write_json(path,data);print(key,'REPAIRED',flush=True)
    finally:engine.close()

if __name__=='__main__':main()