Files
Athena-Deck/tts_test.py
T

81 lines
6.1 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""Serial own TTS workers with shared model reservation, WAV results and cancellation."""
import json,os,signal,subprocess,threading,time,uuid,hashlib
from pathlib import Path
from tts_runtime import REPO,REVISION,SPEAKERS,LANGUAGES
from image_test import probe,cgroup_headroom
class TTSTests:
def __init__(self,root,profiles,runtime):
self.root=Path(root);self.profiles=profiles;self.runtime=runtime;self.lock=threading.RLock();self.job=None;self.process=None;self.cancel=threading.Event();self.acquire=lambda wait=False:lambda:None
def status(self):
with self.lock:return dict(job=dict(self.job) if self.job else None,runtime=self.runtime.status())
def blockers(self,p):
if p['model']['repo']!=REPO or p['model']['file']!='model.safetensors' or p['model'].get('revision')!=REVISION:return ['Für diese TTS-Variante fehlt die Anbindung in Deck. Unter TTS → Einrichten kannst du ausdrücklich auf die unterstützte Variante mit eingebauten Stimmen wechseln. Die ursprüngliche Datei bleibt erhalten.']
return [] if self.runtime.status()['installed'] else ['TTS-Laufzeit unter TTS → Einrichten installieren.']
def assign(self,ident,revision):
if not self.runtime.status()['installed']:raise ValueError('Zuerst TTS einrichten.')
entry_id=self.runtime.ensure_library()
p=next((p for p in self.profiles.status()['profiles'] if p['id']==ident),None)
if not p or p['kind']!='audio':raise ValueError('TTS-Profil nicht gefunden.')
return self.profiles.save(dict(id=ident,revision=revision,name=p['name'],kind='audio',model_id=entry_id,parameters=p['parameters']))
def start(self,profile_id,text,speaker='Ryan',language='German',wait=False,speed=None):
if not isinstance(text,str) or not 1<=len(text.strip())<=1000:raise ValueError('Text mit 1–1000 Zeichen erforderlich.')
if speaker not in SPEAKERS or language not in LANGUAGES:raise ValueError('Ungültige Stimme oder Sprache.')
if speed is not None and (type(speed) not in (int,float) or not .25<=speed<=4):raise ValueError('Geschwindigkeit muss zwischen 0,25 und 4 liegen.')
release=self.acquire(wait)
try:
with self.lock:
if self.job and self.job['state']=='running':raise ValueError('TTS-Auftrag läuft bereits.')
p=next((p for p in self.profiles.status()['profiles'] if p['id']==profile_id and p['kind']=='audio'),None)
if not p or not p['runnable']:raise ValueError('TTS-Profil noch nicht eingerichtet.')
gpu=next((g for g in probe() if 'RTX 3060' in g['name'] and not g['processes'] and g['free_mib']>7000),None)
if not gpu:raise ValueError('Die RTX 3060 muss frei sein und mindestens 7 GiB freien Speicher haben.')
headroom=cgroup_headroom()
if headroom is not None and headroom<8*1024**3:raise ValueError('Zu wenig freier Deck-RAM für TTS.')
self.cancel.clear();ident=uuid.uuid4().hex;self.root.mkdir(parents=True,exist_ok=True,mode=0o700)
self.job=dict(id=ident,state='running',phase='Sprachmodell wird geladen',profile_id=profile_id,started_at=time.time())
threading.Thread(target=self._run,args=(ident,text,speaker,language,speed if speed is not None else p['parameters']['speed'],gpu,release),daemon=True).start();return dict(self.job)
except Exception:release();raise
def _run(self,ident,text,speaker,language,speed,gpu,release):
directory=self.root/ident;process=None;state='failed';phase='TTS-Ausführung fehlgeschlagen'
try:
directory.mkdir(mode=0o700);python,model=self.runtime.paths()
env={k:v for k,v in os.environ.items() if not k.startswith(('HF_','LLAMA_'))};env.update(CUDA_VISIBLE_DEVICES=gpu['uuid'],HF_HUB_OFFLINE='1',TRANSFORMERS_OFFLINE='1',OMP_NUM_THREADS='2',HOME=str(directory))
with self.lock:
if self.cancel.is_set():raise InterruptedError()
process=subprocess.Popen([str(python),str(Path(__file__).with_name('tts_worker.py'))],stdin=subprocess.PIPE,stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL,start_new_session=True,env=env);self.process=process;self.job['phase']='Sprache wird auf der RTX 3060 erzeugt'
process.stdin.write(json.dumps(dict(model=str(model),output=str(directory),text=text,speaker=speaker,language=language,speed=speed)).encode());process.stdin.close()
end=time.monotonic()+600
while process.poll() is None:
if self.cancel.wait(1):raise InterruptedError()
if time.monotonic()>end:raise ValueError('Zeitlimit der Sprachgenerierung erreicht.')
current=next((g for g in probe() if g['uuid']==gpu['uuid']),None)
if not current or current['processes']>1:raise ValueError('Die reservierte GPU wurde anderweitig belegt.')
result=json.loads((directory/'result.json').read_text()) if (directory/'result.json').exists() else dict(ok=False,error='Worker beendet: möglicher Speichermangel oder Laufzeitfehler.')
if self.cancel.is_set():raise InterruptedError()
if not result['ok']:raise ValueError(result['error'])
wav=directory/'result.wav'
if not wav.exists() or wav.stat().st_size>32*1024**2:raise ValueError('Ungültige Audioausgabe.')
state='complete';phase='Sprache fertig · Modell entladen'
except InterruptedError:state='cancelled';phase='Sprachgenerierung abgebrochen'
except Exception as exc:phase=str(exc) if isinstance(exc,ValueError) else 'Sprachgenerierung fehlgeschlagen. Laufzeit und Ressourcen prüfen.'
finally:
if process and process.poll() is None:
try:
os.killpg(process.pid,signal.SIGTERM)
try:process.wait(timeout=5)
except subprocess.TimeoutExpired:os.killpg(process.pid,signal.SIGKILL);process.wait()
except ProcessLookupError:pass
if state=='complete':(directory/'complete').touch()
with self.lock:self.process=None;self.job.update(state=state,phase=phase,finished_at=time.time())
release()
def stop(self):
self.cancel.set()
with self.lock:process=self.process
if process and process.poll() is None:
try:os.killpg(process.pid,signal.SIGTERM)
except ProcessLookupError:pass
return dict(cancellation_requested=True)
def audio(self,ident):
if not isinstance(ident,str) or len(ident)!=32 or any(c not in '0123456789abcdef' for c in ident) or not (self.root/ident/'complete').exists():raise ValueError('Audio noch nicht verfügbar.')
return (self.root/ident/'result.wav').read_bytes()