193 lines
14 KiB
Python
193 lines
14 KiB
Python
"""Owned, serial ComfyUI test jobs. Never talks to the production image router."""
|
||
import json
|
||
import os
|
||
from pathlib import Path
|
||
import secrets
|
||
import shutil
|
||
import signal
|
||
import socket
|
||
import subprocess
|
||
import threading
|
||
import time
|
||
import urllib.request
|
||
import uuid
|
||
from profiles import QWEN_REPO
|
||
|
||
PYTHON=Path('/opt/deck-image-python/bin/python')
|
||
COMFY=Path('/opt/deck-comfy')
|
||
GIB=1024**3
|
||
|
||
def cgroup_headroom():
|
||
"""Exclude reclaimable inactive file cache, not live anonymous allocations."""
|
||
try:
|
||
raw=Path('/sys/fs/cgroup/memory.max').read_text().strip()
|
||
if raw=='max':return None
|
||
used=int(Path('/sys/fs/cgroup/memory.current').read_text())
|
||
stat=dict(line.split() for line in Path('/sys/fs/cgroup/memory.stat').read_text().splitlines())
|
||
return int(raw)-max(0,used-int(stat.get('inactive_file',0)))
|
||
except FileNotFoundError:return None
|
||
|
||
def probe():
|
||
rows=subprocess.run(['nvidia-smi','--query-gpu=uuid,name,memory.total,memory.free','--format=csv,noheader,nounits'],capture_output=True,text=True,check=True,timeout=5).stdout
|
||
apps=subprocess.run(['nvidia-smi','--query-compute-apps=gpu_uuid,pid','--format=csv,noheader,nounits'],capture_output=True,text=True,check=True,timeout=5).stdout
|
||
counts={}
|
||
for line in apps.splitlines():
|
||
gpu,_=line.split(',',1);counts[gpu.strip()]=counts.get(gpu.strip(),0)+1
|
||
return [dict(uuid=u.strip(),name=n.strip(),total_mib=float(t),free_mib=float(f),processes=counts.get(u.strip(),0)) for u,n,t,f in (line.split(',') for line in rows.splitlines())]
|
||
|
||
def select_gpus(devices,model_size,encoder_size,vae_size):
|
||
needed=(model_size+vae_size+3*GIB)/1024**2
|
||
images=[g for g in devices if 'RTX 5080' in g['name'] and g['processes']==0 and g['free_mib']>needed]
|
||
encoders=[g for g in devices if 'RTX 3060' in g['name'] and g['processes']==0 and g['free_mib']>(encoder_size+2*GIB)/1024**2]
|
||
if not encoders:raise ValueError('Textencoder benötigt eine freie RTX 3060 mit ausreichend VRAM. Kein automatischer CPU-Fallback.')
|
||
if not images:raise ValueError('Bildmodell benötigt eine freie RTX 5080 mit ausreichend VRAM. Deck stoppt keine anderen Modelle.')
|
||
return max(images,key=lambda g:g['free_mib']),encoders[0]
|
||
|
||
def workflow(prompt,params,seed):
|
||
return {
|
||
'1':{'class_type':'UnetLoaderGGUF','inputs':{'unet_name':'model.gguf'}},
|
||
'2':{'class_type':'DeckTextEncoderLoader','inputs':{'clip_name':'encoder.safetensors'}},
|
||
'3':{'class_type':'VAELoader','inputs':{'vae_name':'vae.safetensors'}},
|
||
'4':{'class_type':'TextEncodeQwenImage21','inputs':{'clip':['2',0],'prompt':prompt,'negative_prompt':'','resolution':max(params['width'],params['height'])}},
|
||
'5':{'class_type':'EmptyLatentImage','inputs':{'width':params['width'],'height':params['height'],'batch_size':1}},
|
||
'6':{'class_type':'KSampler','inputs':{'model':['1',0],'positive':['4',0],'negative':['4',1],'latent_image':['5',0],'seed':seed,'steps':params['steps'],'cfg':params['guidance'],'sampler_name':'euler','scheduler':'simple','denoise':1}},
|
||
'7':{'class_type':'VAEDecode','inputs':{'samples':['6',0],'vae':['3',0]}},
|
||
'8':{'class_type':'SaveImage','inputs':{'filename_prefix':'result','images':['7',0]}}
|
||
}
|
||
|
||
class ImageTests:
|
||
def __init__(self,root,profiles):
|
||
self.root=Path(root);self.profiles=profiles;self.runtime=None;self.lock=threading.RLock();self.process=None;self.cancel=threading.Event();self.job=None;self.acquire=lambda wait=False:lambda:None
|
||
if (self.root/'status.json').exists():
|
||
self.job=json.loads((self.root/'status.json').read_text())
|
||
if self.job.get('state')=='running':self.job.update(state='interrupted',phase='Deck wurde neu gestartet; Test erneut starten.')
|
||
def _save(self):
|
||
self.root.mkdir(parents=True,exist_ok=True,mode=0o700)
|
||
p=self.root/'status.tmp';p.write_text(json.dumps(self.job));p.replace(self.root/'status.json')
|
||
def _phase(self,text):
|
||
with self.lock:
|
||
self.job['phase']=text;self._save()
|
||
def status(self):
|
||
with self.lock:
|
||
python,comfy=self.runtime.paths() if self.runtime else (PYTHON,COMFY)
|
||
return dict(job=dict(self.job) if self.job else None,runtime_installed=python.is_file() and (comfy/'main.py').is_file())
|
||
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)
|
||
try:process.wait(timeout=3)
|
||
except subprocess.TimeoutExpired:os.killpg(process.pid,signal.SIGKILL);process.wait()
|
||
except ProcessLookupError:pass
|
||
return {'cancellation_requested':True}
|
||
def model_path(self,item):
|
||
return (self.profiles.catalog.root/item['id']/('model'+Path(item['file']).suffix)).resolve()
|
||
def start(self,profile_id,prompt,wait=False,reserved=False):
|
||
if not isinstance(prompt,str) or not 1<=len(prompt.strip())<=4000:raise ValueError('Bitte einen Prompt mit 1–4000 Zeichen eingeben.')
|
||
release=(lambda:None) if reserved else self.acquire(wait)
|
||
try:return self._start(profile_id,prompt,release)
|
||
except Exception:release();raise
|
||
def _start(self,profile_id,prompt,release):
|
||
if not isinstance(prompt,str) or not 1<=len(prompt.strip())<=4000:raise ValueError('Bitte einen Prompt mit 1–4000 Zeichen eingeben.')
|
||
with self.lock:
|
||
if self.job and self.job['state']=='running':raise ValueError('Ein Bildtest läuft bereits.')
|
||
if self.runtime and (self.runtime.status()['job'] or {}).get('state')=='running':raise ValueError('Die Bildlaufzeit wird gerade installiert.')
|
||
if not self.status()['runtime_installed']:raise ValueError('Die eigene ComfyUI-Bildlaufzeit ist noch nicht installiert.')
|
||
profile=next((p for p in self.profiles.status()['profiles'] if p['id']==profile_id),None)
|
||
if not profile or profile['kind']!='image' or not profile['model']:raise ValueError('Bildprofil nicht verfügbar.')
|
||
model=profile['model']
|
||
if model['repo']!=QWEN_REPO or not model['file'].endswith('.gguf'):raise ValueError('Der Bildtest unterstützt zunächst Qwen-Image-2.1 GGUF aus dem hinterlegten Rezept.')
|
||
params=profile['parameters']
|
||
if params['width']>1024 or params['height']>1024:raise ValueError('Der isolierte Test unterstützt maximal 1024 × 1024 Pixel.')
|
||
encoder=self.profiles._component(model,'text_encoder',profile.get('components',{}).get('text_encoder'))
|
||
vae=self.profiles._component(model,'vae',profile.get('components',{}).get('vae'))
|
||
mem={line.split(':')[0]:int(line.split()[1])*1024 for line in Path('/proc/meminfo').read_text().splitlines() if line.startswith(('MemAvailable:','MemTotal:'))}
|
||
required_ram=max(20*GIB,encoder['size']*2+4*GIB)
|
||
headroom=cgroup_headroom()
|
||
if headroom is not None and headroom<required_ram:raise ValueError('Deck-RAM-Limit reicht für Textencoder und Arbeitsdaten nicht aus.')
|
||
if mem.get('MemAvailable',0)<required_ram+4*GIB:raise ValueError('Aktuell zu wenig freier System-RAM; produktive Dienste bleiben unverändert.')
|
||
gpu,encoder_gpu=select_gpus(probe(),model['size'],encoder['size'],vae['size']);self.root.mkdir(parents=True,exist_ok=True,mode=0o700)
|
||
if shutil.disk_usage(self.root).free<10*GIB:raise ValueError('Weniger als 10 GiB freier Plattenspeicher.')
|
||
job_id=uuid.uuid4().hex;seed=params['seed'] if params['seed']>=0 else secrets.randbelow(2147483648)
|
||
self.cancel.clear();self.job=dict(id=job_id,state='running',phase='Bildlaufzeit startet',profile_id=profile_id,profile_name=profile['name'],started_at=time.time(),gpu=gpu['name'],encoder_gpu=encoder_gpu['name'],seed=seed,parameters=params)
|
||
self._save();threading.Thread(target=self._run,args=(job_id,prompt,params,seed,gpu,encoder_gpu,model,encoder,vae,release),daemon=True).start()
|
||
return dict(self.job)
|
||
def _run(self,job_id,prompt,params,seed,gpu,encoder_gpu,model,encoder,vae,release=lambda:None):
|
||
directory=self.root/job_id;process=None
|
||
try:
|
||
directory.mkdir(mode=0o700)
|
||
for role,item,filename in [('unet',model,'model.gguf'),('clip',encoder,'encoder.safetensors'),('vae',vae,'vae.safetensors')]:
|
||
dest=directory/'models'/role;dest.mkdir(parents=True);(dest/filename).symlink_to(self.model_path(item))
|
||
for folder in ('output','temp','user','input'):(directory/folder).mkdir()
|
||
custom=directory/'custom_nodes'/'deck_encoder';custom.mkdir(parents=True)
|
||
shutil.copyfile(Path(__file__).parent/'image_encoder_node.py',custom/'__init__.py')
|
||
config={'deck_nodes':{'base_path':str(directory),'custom_nodes':'custom_nodes'},'deck':{'base_path':str(directory/'models'),'unet':'unet','clip':'clip','vae':'vae'}}
|
||
(directory/'paths.json').write_text(json.dumps(config)) # JSON is valid YAML.
|
||
with socket.socket() as sock:sock.bind(('127.0.0.1',0));port=sock.getsockname()[1]
|
||
env=dict(os.environ,CUDA_VISIBLE_DEVICES=gpu['uuid']+','+encoder_gpu['uuid'],OMP_NUM_THREADS='2',MKL_NUM_THREADS='2',HOME=str(directory),HF_HUB_OFFLINE='1',TRANSFORMERS_OFFLINE='1',PYTHONDONTWRITEBYTECODE='1')
|
||
args=[str(self.runtime.paths()[0] if self.runtime else PYTHON),str((self.runtime.paths()[1] if self.runtime else COMFY)/'main.py'),'--listen','127.0.0.1','--port',str(port),'--disable-auto-launch','--disable-metadata','--lowvram','--reserve-vram','1.5','--extra-model-paths-config',str(directory/'paths.json'),'--output-directory',str(directory/'output'),'--temp-directory',str(directory/'temp'),'--user-directory',str(directory/'user'),'--input-directory',str(directory/'input')]
|
||
with self.lock:
|
||
if self.cancel.is_set():raise InterruptedError()
|
||
process=subprocess.Popen(args,cwd=self.runtime.paths()[1] if self.runtime else COMFY,env=env,stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL,start_new_session=True);self.process=process
|
||
def request(path,data=None):
|
||
req=urllib.request.Request(f'http://127.0.0.1:{port}'+path,data=json.dumps(data).encode() if data is not None else None,headers={'Content-Type':'application/json'})
|
||
with urllib.request.urlopen(req,timeout=5) as r:
|
||
raw=r.read(4*1024*1024+1)
|
||
if len(raw)>4*1024*1024:raise ValueError('Bildlaufzeit-Antwort zu groß.')
|
||
return json.loads(raw)
|
||
def guard():
|
||
if self.cancel.is_set():raise InterruptedError()
|
||
if process.poll() is not None:raise ValueError('Bildlaufzeit wurde beendet (Speicherlimit oder Startfehler).')
|
||
current={g['uuid']:g for g in probe()}
|
||
for reserved in (gpu,encoder_gpu):
|
||
row=current.get(reserved['uuid'])
|
||
if not row or row['processes']>1:raise ValueError('Reservierte GPU wird inzwischen von einem weiteren Prozess verwendet. Deck-Test beendet.')
|
||
deadline=time.monotonic()+180
|
||
while True:
|
||
guard()
|
||
try:request('/system_stats');break
|
||
except (OSError,ValueError):
|
||
if time.monotonic()>deadline:raise ValueError('Bildlaufzeit wurde nicht rechtzeitig bereit.')
|
||
time.sleep(1)
|
||
nodes=request('/object_info')
|
||
required={'UnetLoaderGGUF','DeckTextEncoderLoader','VAELoader','TextEncodeQwenImage21','EmptyLatentImage','KSampler','VAEDecode','SaveImage'}
|
||
if not required.issubset(nodes):raise ValueError('Der installierten Bildlaufzeit fehlen erforderliche Qwen/GGUF-Nodes.')
|
||
self._phase('Auftrag wird verarbeitet · Textencoder auf RTX 3060 · Bildberechnung auf RTX 5080')
|
||
response=request('/prompt',{'prompt':workflow(prompt,params,seed),'client_id':'deck-'+job_id});prompt_id=response['prompt_id']
|
||
deadline=time.monotonic()+1800
|
||
while time.monotonic()<deadline:
|
||
guard();history=request('/history/'+prompt_id)
|
||
if prompt_id in history:
|
||
result=history[prompt_id]
|
||
if result.get('status',{}).get('status_str')=='error':
|
||
errors=[m[1] for m in result.get('status',{}).get('messages',[]) if m[0]=='execution_error']
|
||
error=errors[0] if errors else {}
|
||
# Do not include prompts, tensor dumps or upstream exception messages.
|
||
raise ValueError('Generierung fehlgeschlagen: '+str(error.get('node_type','Worker'))+' / '+str(error.get('exception_type','Fehler')))
|
||
images=[i for value in result.get('outputs',{}).values() for i in value.get('images',[])]
|
||
if images:
|
||
image=images[0];path=(directory/'output'/image.get('subfolder','')/image['filename']).resolve()
|
||
if not path.is_relative_to((directory/'output').resolve()) or path.stat().st_size>32*1024**2:raise ValueError('Ungültiges Ergebnisbild.')
|
||
with path.open('rb') as f:
|
||
if f.read(8)!=b'\x89PNG\r\n\x1a\n':raise ValueError('Ergebnis ist kein PNG.')
|
||
shutil.copyfile(path,directory/'result.png');break
|
||
time.sleep(2)
|
||
else:raise ValueError('Zeitlimit der Bildgenerierung erreicht.')
|
||
final_state='complete';phase='Bild fertig · Modell wurde entladen'
|
||
except InterruptedError:final_state='cancelled';phase='Bildtest abgebrochen'
|
||
except Exception as exc:final_state='failed';phase=str(exc) if isinstance(exc,ValueError) else 'Bildlaufzeit nicht erreichbar oder nicht bereit. Komponenten und Installation prüfen.'
|
||
finally:
|
||
try:
|
||
if process and process.poll() is None:
|
||
try:os.killpg(process.pid,signal.SIGTERM);process.wait(timeout=10)
|
||
except subprocess.TimeoutExpired:os.killpg(process.pid,signal.SIGKILL);process.wait()
|
||
except ProcessLookupError:pass
|
||
with self.lock:
|
||
self.process=None;self.job.update(state=final_state,phase=phase,finished_at=time.time());self._save()
|
||
finally:release()
|
||
def image(self,job_id):
|
||
with self.lock:
|
||
if not self.job or self.job['id']!=job_id or self.job['state']!='complete':raise ValueError('Ergebnisbild nicht verfügbar.')
|
||
return (self.root/job_id/'result.png').read_bytes()
|