263 lines
19 KiB
Python
263 lines
19 KiB
Python
"""Owned, serial ComfyUI test jobs. Never talks to the production image router."""
|
||
import json
|
||
import base64
|
||
import binascii
|
||
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 image_adapter, image_recipe, image_capabilities
|
||
|
||
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,offload=False):
|
||
needed=(min(model_size,8*GIB)+vae_size+3*GIB if offload else 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,family='qwen',references=()):
|
||
if family not in ('qwen','flux'):raise ValueError('Für diese Modellfamilie fehlt eine Bildlaufzeit-Anbindung.')
|
||
if len(references)>4:raise ValueError('Maximal vier Referenzbilder pro Auftrag.')
|
||
if family=='flux':
|
||
graph={
|
||
'1':{'class_type':'UNETLoader','inputs':{'unet_name':'model.safetensors','weight_dtype':'default'}},
|
||
'2':{'class_type':'DeckTextEncoderLoader','inputs':{'clip_name':'encoder.safetensors','family':'flux2'}},
|
||
'3':{'class_type':'VAELoader','inputs':{'vae_name':'vae.safetensors'}},
|
||
'4':{'class_type':'CLIPTextEncode','inputs':{'clip':['2',0],'text':prompt}},
|
||
'5':{'class_type':'EmptyFlux2LatentImage','inputs':{'width':params['width'],'height':params['height'],'batch_size':1}},
|
||
'6':{'class_type':'SamplerCustomAdvanced','inputs':{'noise':['9',0],'guider':['10',0],'sampler':['11',0],'sigmas':['12',0],'latent_image':['5',0]}},
|
||
'7':{'class_type':'VAEDecode','inputs':{'samples':['6',0],'vae':['3',0]}},
|
||
'8':{'class_type':'SaveImage','inputs':{'filename_prefix':'result','images':['7',0]}},
|
||
'9':{'class_type':'RandomNoise','inputs':{'noise_seed':seed}},
|
||
'10':{'class_type':'CFGGuider','inputs':{'model':['1',0],'positive':['4',0],'negative':['13',0],'cfg':params['guidance']}},
|
||
'11':{'class_type':'KSamplerSelect','inputs':{'sampler_name':'euler'}},
|
||
'12':{'class_type':'Flux2Scheduler','inputs':{'steps':params['steps'],'width':params['width'],'height':params['height']}},
|
||
'13':{'class_type':'ConditioningZeroOut','inputs':{'conditioning':['4',0]}}
|
||
}
|
||
|
||
# FLUX.2 editing uses reference latents on both conditioning branches,
|
||
# with fresh output noise; never silently discard uploaded references.
|
||
positive=['4',0];negative=['13',0]
|
||
for index,name in enumerate(references):
|
||
load,scale,encode,pos,neg=map(str,range(20+index*5,25+index*5))
|
||
graph[load]={'class_type':'LoadImage','inputs':{'image':name}}
|
||
graph[scale]={'class_type':'ImageScaleToTotalPixels','inputs':{'image':[load,0],'upscale_method':'lanczos','megapixels':1.0,'resolution_steps':1}}
|
||
graph[encode]={'class_type':'VAEEncode','inputs':{'pixels':[scale,0],'vae':['3',0]}}
|
||
graph[pos]={'class_type':'ReferenceLatent','inputs':{'conditioning':positive,'latent':[encode,0]}}
|
||
graph[neg]={'class_type':'ReferenceLatent','inputs':{'conditioning':negative,'latent':[encode,0]}}
|
||
positive=[pos,0];negative=[neg,0]
|
||
graph['10']['inputs'].update(positive=positive,negative=negative)
|
||
return graph
|
||
|
||
graph={
|
||
'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]}}
|
||
}
|
||
if references:
|
||
graph['4']['inputs']['vae']=['3',0]
|
||
for index,name in enumerate(references,1):
|
||
graph[str(20+index)]={'class_type':'LoadImage','inputs':{'image':name}}
|
||
graph['4']['inputs'][f'images.image_{index}']=[str(20+index),0]
|
||
graph['6']['inputs']['latent_image']=['4',2]
|
||
return graph
|
||
|
||
def decode_references(values,limit):
|
||
if not isinstance(values,list) or len(values)>limit:raise ValueError(f'Dieses Bildprofil erlaubt höchstens {limit} Referenzbilder.')
|
||
result=[]
|
||
for value in values:
|
||
if not isinstance(value,str) or not value.startswith(('data:image/png;base64,','data:image/jpeg;base64,','data:image/webp;base64,')) or len(value)>14*1024*1024:raise ValueError('Referenzbild muss PNG, JPEG oder WebP bis 10 MiB sein.')
|
||
try:raw=base64.b64decode(value.split(',',1)[1],validate=True)
|
||
except (ValueError,binascii.Error):raise ValueError('Ungültige Bildkodierung.') from None
|
||
if not 0<len(raw)<=10*1024*1024:raise ValueError('Referenzbild muss 1 bis 10 MiB groß sein.')
|
||
kind=value.split(';',1)[0].split('/')[-1]
|
||
signatures={'png':b'\x89PNG\r\n\x1a\n','jpeg':b'\xff\xd8\xff','webp':b'RIFF'}
|
||
if not raw.startswith(signatures[kind]) or kind=='webp' and raw[8:12]!=b'WEBP':raise ValueError('Dateiinhalt entspricht nicht dem Bildformat.')
|
||
result.append((raw,'jpg' if kind=='jpeg' else kind))
|
||
return result
|
||
|
||
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,reference_images=None,skip_enhancer=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,reference_images or [],skip_enhancer)
|
||
except Exception:release();raise
|
||
def _start(self,profile_id,prompt,release,reference_images,skip_enhancer=False):
|
||
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 not image_recipe(model):raise ValueError('Für dieses Bildmodell fehlt ein Komponentenrezept.')
|
||
references=decode_references(reference_images,image_capabilities(model)['reference_images'])
|
||
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,model['size']+8*GIB if image_adapter(model)['offload'] else 0)
|
||
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'],offload=image_adapter(model)['offload']);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,reference_count=len(references))
|
||
self._save();threading.Thread(target=self._run,args=(job_id,prompt,params,seed,gpu,encoder_gpu,model,encoder,vae,release,references,profile,skip_enhancer),daemon=True).start()
|
||
return dict(self.job)
|
||
def _run(self,job_id,prompt,params,seed,gpu,encoder_gpu,model,encoder,vae,release=lambda:None,references=(),profile=None,skip_enhancer=False):
|
||
directory=self.root/job_id;process=None
|
||
try:
|
||
if not skip_enhancer and profile and getattr(self,'enhancer',None) and (profile.get('prompt_enhancer') or {}).get('i2i' if references else 't2i'):
|
||
self._phase('Prompt-Aufwerter lädt; Bildmodell wartet')
|
||
result=self.enhancer.rewrite(profile,prompt,references,self.cancel)
|
||
prompt=result['prompt']
|
||
with self.lock:self.job['prompt_enhanced']=True;self.job['prompt_enhancer']=result['source'];self._save()
|
||
if self.cancel.is_set():raise InterruptedError()
|
||
self._phase('Prompt-Aufwerter entladen; Bildlaufzeit startet')
|
||
directory.mkdir(mode=0o700)
|
||
for role,item,filename in [('unet',model,image_adapter(model)['model_filename']),('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()
|
||
names=[]
|
||
for index,(raw,extension) in enumerate(references,1):
|
||
name=f'reference-{index}.{extension}';(directory/'input'/name).write_bytes(raw);names.append(name)
|
||
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')
|
||
graph=workflow(prompt,params,seed,image_adapter(model)['workflow'],names)
|
||
required={node['class_type'] for node in graph.values()}
|
||
if not required.issubset(nodes):raise ValueError('Der installierten Bildlaufzeit fehlen erforderliche Nodes für dieses Bildmodell.')
|
||
self._phase('Auftrag wird verarbeitet · Textencoder auf RTX 3060 · Bildberechnung auf RTX 5080')
|
||
response=request('/prompt',{'prompt':graph,'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()
|
||
for name in (f'reference-{index}.{extension}' for index,(_,extension) in enumerate(references,1)):
|
||
(directory/'input'/name).unlink(missing_ok=True)
|
||
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()
|