155 lines
11 KiB
Python
155 lines
11 KiB
Python
"""Isolated, one-shot Qwen image prompt enhancement owned by Deck."""
|
||
import json,os,shutil,signal,subprocess,tempfile,threading,time,uuid
|
||
from pathlib import Path
|
||
from image_test import probe,cgroup_headroom,decode_references
|
||
from image_runtime import BUNDLED
|
||
from profiles import PROMPT_ENHANCER_REPOS
|
||
|
||
WORKER=Path(__file__).with_name('prompt_enhancer_worker.py')
|
||
GIB=1024**3
|
||
|
||
class PromptEnhancer:
|
||
def __init__(self,root,catalog,profiles,runtime):
|
||
self.root=Path(root);self.catalog=catalog;self.profiles=profiles;self.runtime=runtime
|
||
self.lock=threading.RLock();self.cancel=threading.Event();self.process=None;self.job=None;self.preview=None;self.acquire=lambda wait=False:lambda:None
|
||
def directory(self,task):return self.root/task
|
||
def python(self):return BUNDLED[0] if BUNDLED[0].is_file() else self.runtime.paths()[0]
|
||
def installed(self,task):
|
||
path=self.directory(task);marker=path/'installed.json'
|
||
if not marker.is_file():return None
|
||
try:
|
||
data=json.loads(marker.read_text());return data if data['repo']==PROMPT_ENHANCER_REPOS[task] and (path/'model.safetensors.index.json').is_file() else None
|
||
except (OSError,ValueError,KeyError):return None
|
||
def status(self):
|
||
with self.lock:
|
||
return dict(models={task:dict(repo=repo,installed=bool(self.installed(task)),revision=(self.installed(task) or {}).get('revision')) for task,repo in PROMPT_ENHANCER_REPOS.items()},job=dict(self.job) if self.job else None,preview=dict(self.preview) if self.preview else None)
|
||
def install(self,task,revision=None):
|
||
if task not in PROMPT_ENHANCER_REPOS:raise ValueError('Unbekannter Prompt-Aufwerter.')
|
||
with self.lock:
|
||
if self.job and self.job['state']=='running':raise ValueError('Ein Aufwerter-Download läuft bereits.')
|
||
if self.preview and self.preview['state']=='running':raise ValueError('Eine Prompt-Vorschau läuft; zuerst abschließen.')
|
||
if self.installed(task):raise ValueError('Dieser Aufwerter ist bereits installiert.')
|
||
if not self.runtime.status()['installed']:raise ValueError('Zuerst die Bildlaufzeit einrichten.')
|
||
if shutil.disk_usage(self.root if self.root.exists() else self.root.parent).free<30*GIB:raise ValueError('Mindestens 30 GiB freier Plattenspeicher erforderlich.')
|
||
source=self.catalog.files(PROMPT_ENHANCER_REPOS[task],revision) if revision else self.catalog.files(PROMPT_ENHANCER_REPOS[task]);revision=source['revision']
|
||
expected=sum(f['size'] for f in source['files'] if f['name'].endswith('.safetensors'))
|
||
if expected<15*GIB or expected>25*GIB:raise ValueError('Unerwartete Größe der offiziellen Aufwerter-Gewichte; Rezept prüfen.')
|
||
self.root.mkdir(parents=True,exist_ok=True,mode=0o700);self.cancel.clear()
|
||
self.job=dict(id=uuid.uuid4().hex,task=task,state='running',phase='Offizielle Dateien werden heruntergeladen',repo=source['repo'],revision=revision,total_bytes=expected,started_at=time.time())
|
||
threading.Thread(target=self._download,args=(task,revision),daemon=True).start();return self.status()
|
||
def _download(self,task,revision):
|
||
target=self.directory(task);target.mkdir(parents=True,exist_ok=True,mode=0o700)
|
||
env=dict(os.environ,HF_TOKEN=self.catalog.hub_auth.token() or '',HF_HUB_DISABLE_TELEMETRY='1',HOME=str(self.root),PYTHONDONTWRITEBYTECODE='1')
|
||
data=dict(repo=PROMPT_ENHANCER_REPOS[task],revision=revision,directory=str(target))
|
||
state='failed';phase='Aufwerter-Download fehlgeschlagen. Netzwerk, Hub-Zugang und freien Speicher prüfen.'
|
||
try:
|
||
python=self.python()
|
||
with self.lock:
|
||
if self.cancel.is_set():raise InterruptedError()
|
||
self.process=subprocess.Popen([str(python),str(WORKER),'download'],stdin=subprocess.PIPE,stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL,text=True,start_new_session=True,env=env)
|
||
process=self.process
|
||
process.stdin.write(json.dumps(data));process.stdin.close()
|
||
deadline=time.monotonic()+21600
|
||
while process.poll() is None:
|
||
if self.cancel.wait(1):raise InterruptedError()
|
||
if time.monotonic()>deadline:raise ValueError('Download-Zeitlimit erreicht.')
|
||
if shutil.disk_usage(target).free<10*GIB:raise ValueError('Plattenreserve unter 10 GiB.')
|
||
if self.cancel.is_set():raise InterruptedError()
|
||
if process.returncode:raise ValueError('Aufwerter-Dateien konnten nicht vollständig heruntergeladen werden.')
|
||
marker=target/'installed.tmp';marker.write_text(json.dumps(data));marker.replace(target/'installed.json')
|
||
state='complete';phase='Offizieller Aufwerter installiert. Es läuft kein Modell.'
|
||
except InterruptedError:state='cancelled';phase='Download abgebrochen.'
|
||
except ValueError as exc:phase=str(exc)
|
||
except Exception:pass
|
||
finally:
|
||
self._terminate()
|
||
with self.lock:self.job.update(state=state,phase=phase,finished_at=time.time())
|
||
def _terminate(self):
|
||
with self.lock:process=self.process;self.process=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
|
||
def stop(self):
|
||
self.cancel.set();self._terminate();return {'cancellation_requested':True}
|
||
def stop_preview(self):
|
||
with self.lock:active=bool(self.preview and self.preview['state']=='running')
|
||
if active:self.cancel.set();self._terminate()
|
||
return {'cancellation_requested':active}
|
||
def _devices(self,choice):
|
||
if choice=='cpu':return ''
|
||
rows=probe();names=['5080','3060'] if choice=='auto' else [choice];selected=[]
|
||
for name in names:
|
||
row=next((g for g in rows if ('RTX '+name) in g['name']),None)
|
||
if not row or row['processes'] or row['free_mib']<('5080'==name and 10000 or 7000):
|
||
if choice=='auto':continue
|
||
raise ValueError('Gewählte GPU für Prompt-Aufbereitung nicht frei oder mit zu wenig VRAM.')
|
||
selected.append(row['uuid'])
|
||
if not selected:raise ValueError('Keine geeignete freie GPU für den Prompt-Aufwerter.')
|
||
return ','.join(selected)
|
||
def rewrite(self,profile,prompt,references=(),cancel=None):
|
||
task='i2i' if references else 't2i';config=profile.get('prompt_enhancer') or {};repo=config.get(task)
|
||
if not repo:return dict(prompt=prompt,enhanced=False,wh_ratio='')
|
||
if repo!=PROMPT_ENHANCER_REPOS[task]:raise ValueError('Unbekannter Prompt-Aufwerter im Profil.')
|
||
if not self.installed(task):raise ValueError('Prompt-Aufwerter fehlt. Unter Bildgenerierung → Aufwerter installieren.')
|
||
if not isinstance(prompt,str) or not 1<=len(prompt.strip())<=4000:raise ValueError('Prompt mit 1–4000 Zeichen erforderlich.')
|
||
headroom=cgroup_headroom()
|
||
if headroom is not None and headroom<22*GIB:raise ValueError('Für den Prompt-Aufwerter sind mindestens 22 GiB freier Deck-RAM nötig.')
|
||
device=config.get('device','auto');visible=self._devices(device)
|
||
rows={g['uuid']:g for g in probe()}
|
||
gpu_limits=[min(12 if '5080' in rows[ident]['name'] else 9,max(4,int(rows[ident]['free_mib']/1024)-2)) for ident in visible.split(',') if ident]
|
||
self.root.mkdir(parents=True,exist_ok=True,mode=0o700)
|
||
with tempfile.TemporaryDirectory(prefix='rewrite-',dir=self.root) as folder:
|
||
paths=[]
|
||
for i,(raw,extension) in enumerate(references):
|
||
path=Path(folder)/f'image-{i}.{extension}';path.write_bytes(raw);paths.append(str(path))
|
||
env=dict(os.environ,CUDA_VISIBLE_DEVICES=visible,HF_HUB_OFFLINE='1',TRANSFORMERS_OFFLINE='1',HF_HUB_DISABLE_TELEMETRY='1',OMP_NUM_THREADS='2',HOME=folder,PYTHONDONTWRITEBYTECODE='1')
|
||
data=dict(directory=str(self.directory(task)),task=task,images=paths,prompt=prompt,device=device,gpu_limits=gpu_limits)
|
||
with self.lock:
|
||
if self.job and self.job['state']=='running':raise ValueError('Aufwerter-Download läuft; Bildauftrag danach erneut starten.')
|
||
if self.process and self.process.poll() is None:raise ValueError('Ein Prompt-Aufwerter läuft bereits.')
|
||
self.process=subprocess.Popen([str(self.python()),str(WORKER),'rewrite'],stdin=subprocess.PIPE,stdout=subprocess.PIPE,stderr=subprocess.DEVNULL,text=True,start_new_session=True,env=env)
|
||
process=self.process
|
||
try:
|
||
payload=json.dumps(data);deadline=time.monotonic()+900
|
||
while True:
|
||
if cancel and cancel.is_set():raise InterruptedError()
|
||
if time.monotonic()>deadline:raise ValueError('Prompt-Aufwerter hat das Zeitlimit erreicht.')
|
||
try:output,_=process.communicate(payload,timeout=1);break
|
||
except subprocess.TimeoutExpired:payload=None
|
||
if cancel and cancel.is_set():raise InterruptedError()
|
||
if process.returncode:raise ValueError('Prompt-Aufwerter fehlgeschlagen. GPU/RAM, Modell und Laufzeit prüfen; Bild wurde nicht gestartet.')
|
||
result=json.loads(output)
|
||
if not isinstance(result.get('prompt'),str) or not 1<=len(result['prompt'])<=4000:raise ValueError('Prompt-Aufwerter lieferte keinen gültigen Prompt.')
|
||
return dict(result,enhanced=True,source=repo)
|
||
finally:self._terminate()
|
||
def preview_start(self,profile_id,prompt,reference_images):
|
||
profile=next((p for p in self.profiles.status()['profiles'] if p['id']==profile_id and p['kind']=='image'),None)
|
||
if not profile:raise ValueError('Bildprofil nicht gefunden.')
|
||
if not isinstance(prompt,str) or not 1<=len(prompt.strip())<=4000:raise ValueError('Prompt mit 1–4000 Zeichen erforderlich.')
|
||
references=decode_references(reference_images,profile.get('capabilities',{}).get('reference_images',0))
|
||
task='i2i' if references else 't2i'
|
||
if not (profile.get('prompt_enhancer') or {}).get(task):raise ValueError('Für diese Bildaufgabe ist im Profil kein Prompt-Aufwerter aktiviert.')
|
||
if not self.installed(task):raise ValueError('Prompt-Aufwerter zuerst im Bildprofil installieren.')
|
||
release=self.acquire(False)
|
||
try:
|
||
with self.lock:
|
||
if self.preview and self.preview['state']=='running':raise ValueError('Eine Prompt-Vorschau läuft bereits.')
|
||
if self.job and self.job['state']=='running':raise ValueError('Aufwerter-Download läuft bereits.')
|
||
self.cancel.clear();self.preview=dict(id=uuid.uuid4().hex,state='running',phase='Prompt wird aufbereitet',profile_id=profile_id,task=task,started_at=time.time())
|
||
threading.Thread(target=self._preview,args=(profile,prompt,references,release),daemon=True).start();return dict(self.preview)
|
||
except Exception:release();raise
|
||
def _preview(self,profile,prompt,references,release):
|
||
try:
|
||
result=self.rewrite(profile,prompt,references,self.cancel)
|
||
with self.lock:self.preview.update(state='complete',phase='Prompt-Aufbereitung fertig',result=result)
|
||
except InterruptedError:
|
||
with self.lock:self.preview.update(state='cancelled',phase='Prompt-Aufbereitung abgebrochen')
|
||
except ValueError as exc:
|
||
with self.lock:self.preview.update(state='failed',phase=str(exc))
|
||
except Exception:
|
||
with self.lock:self.preview.update(state='failed',phase='Prompt-Aufwerter fehlgeschlagen. Ressourcen und Laufzeit prüfen.')
|
||
finally:release()
|