Files

121 lines
8.3 KiB
Python
Raw Permalink 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,select
from pathlib import Path
from tts_runtime import REPO,REVISION,SPEAKERS,LANGUAGES
from image_test import probe,cgroup_headroom
WORKER=Path(__file__).with_name('tts_worker.py')
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.gpu_uuid=None;self.cancel=threading.Event();self.acquire=lambda wait=False:lambda:None;self.policy='auto';self.warming=False
def status(self):
with self.lock:return dict(job=dict(self.job) if self.job else None,runtime=self.runtime.status(),loaded=bool(self.process and self.process.poll() is None),gpu_uuid=self.gpu_uuid,warming=self.warming,policy=self.policy)
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.warming:raise ValueError('TTS wird gerade vorgeladen. Gleich erneut versuchen.')
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.')
resident=self.process and self.process.poll() is None
gpu=next((g for g in probe() if g['uuid']==self.gpu_uuid),None) if resident else 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 or resident and gpu['processes']>1:raise ValueError('Die RTX 3060 ist nicht verfügbar oder wurde anderweitig belegt.')
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='Vorhandenes Sprachmodell wird verwendet' if resident else '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);process=self._ensure(gpu)
with self.lock:self.job['phase']='Sprache wird auf der RTX 3060 erzeugt'
process.stdin.write(json.dumps(dict(output=str(directory),text=text,speaker=speaker,language=language,speed=speed))+'\n');process.stdin.flush()
end=time.monotonic()+600
while not (directory/'result.json').is_file():
if self.cancel.wait(1):raise InterruptedError()
if time.monotonic()>end:raise ValueError('Zeitlimit der Sprachgenerierung erreicht.')
if process.poll() is not None:raise ValueError('TTS-Worker wurde unerwartet beendet.')
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 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' if self.policy=='per_request' else 'Modell bleibt geladen')
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 state!='complete' or self.policy=='per_request':self.unload_idle()
if state=='complete':(directory/'complete').touch()
with self.lock:self.job.update(state=state,phase=phase,finished_at=time.time())
release()
def _ensure(self,gpu):
python,model=self.runtime.paths()
with self.lock:
if self.cancel.is_set():raise InterruptedError()
process=self.process if self.process and self.process.poll() is None else None
if process:return process
self.root.mkdir(parents=True,exist_ok=True,mode=0o700)
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(self.root))
process=subprocess.Popen([str(python),str(WORKER),str(model)],stdin=subprocess.PIPE,stdout=subprocess.PIPE,stderr=subprocess.DEVNULL,start_new_session=True,env=env,text=True,bufsize=1);self.process=process;self.gpu_uuid=gpu['uuid']
if not select.select([process.stdout],[],[],180)[0]:raise ValueError('Zeitlimit beim Laden des TTS-Modells.')
if self.cancel.is_set():raise InterruptedError()
ready=json.loads(process.stdout.readline())
if not ready.get('ready'):raise ValueError(ready.get('error','TTS-Modell konnte nicht geladen werden.'))
return process
def warm(self,profile_id):
release=self.acquire(True)
try:
with self.lock:
if self.job and self.job.get('state')=='running':return False
if self.process and self.process.poll() is None:return True
if self.warming:return False
p=next((p for p in self.profiles.status()['profiles'] if p['id']==profile_id and p['kind']=='audio' and p['runnable']),None)
if not p:raise ValueError('Das ausgewählte TTS-Profil ist nicht ausführbar.')
self.warming=True;self.cancel.clear()
try:
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('Für dauerhaftes TTS sind mindestens 7 GiB freie RTX-3060-VRAM nötig.')
headroom=cgroup_headroom()
if headroom is not None and headroom<8*1024**3:raise ValueError('Zu wenig freier Deck-RAM für TTS.')
self._ensure(gpu);return True
except Exception:
self.unload_idle();raise
finally:
with self.lock:self.warming=False
finally:release()
def unload_idle(self):
with self.lock:
process=self.process;self.process=None;self.gpu_uuid=None
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 process:
if process.stdin:process.stdin.close()
if process.stdout:process.stdout.close()
def stop(self):
self.cancel.set()
self.unload_idle()
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()