michael/Athena-Deck
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Sort newest Hub models and run image encoder on second GPU

Browse Source
This commit is contained in:
Mikei386
2026-09-28 18:27:23 +02:00
parent fd797a6c8d
commit a0941c0549
15 changed files with 159 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_test.py
+23 -14
View File
@@ -25,10 +25,18 @@ def probe():
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':'CLIPLoader','inputs':{'clip_name':'encoder.safetensors','type':'qwen_image','device':'cpu'}},
'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}},
@@ -87,26 +95,25 @@ class ImageTests:
if raw!='max' and int(raw)-current<required_ram:raise ValueError('Deck-RAM-Limit reicht für Textencoder und Arbeitsdaten nicht aus.')
except FileNotFoundError:pass
if mem.get('MemAvailable',0)<required_ram+4*GIB:raise ValueError('Aktuell zu wenig freier System-RAM; produktive Dienste bleiben unverändert.')
needed=(model['size']+vae['size']+3*GIB)/1024**2
gpus=[g for g in probe() if g['processes']==0 and g['free_mib']>needed]
if not gpus:raise ValueError('Keine unbelegte GPU mit ausreichend freiem VRAM. Deck stoppt keine produktiven Modelle.')
gpu=max(gpus,key=lambda g:g['free_mib']);self.root.mkdir(parents=True,exist_ok=True,mode=0o700)
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'],seed=seed,parameters=params)
self._save();threading.Thread(target=self._run,args=(job_id,prompt,params,seed,gpu,model,encoder,vae),daemon=True).start()
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),daemon=True).start()
return dict(self.job)
def _run(self,job_id,prompt,params,seed,gpu,model,encoder,vae):
def _run(self,job_id,prompt,params,seed,gpu,encoder_gpu,model,encoder,vae):
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()
config={'deck':{'base_path':str(directory/'models'),'unet':'unet','clip':'clip','vae':'vae'}}
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'],OMP_NUM_THREADS='2',MKL_NUM_THREADS='2',HOME=str(directory),HF_HUB_OFFLINE='1',TRANSFORMERS_OFFLINE='1',PYTHONDONTWRITEBYTECODE='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()
@@ -120,8 +127,10 @@ class ImageTests:
def guard():
if self.cancel.is_set():raise InterruptedError()
if process.poll() is not None:raise ValueError('Bildlaufzeit wurde beendet (Speicherlimit oder Startfehler).')
current=next((g for g in probe() if g['uuid']==gpu['uuid']),None)
if not current or current['processes']>1:raise ValueError('GPU wird inzwischen von einem weiteren Prozess verwendet. Deck-Test beendet.')
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()
@@ -130,9 +139,9 @@ class ImageTests:
if time.monotonic()>deadline:raise ValueError('Bildlaufzeit wurde nicht rechtzeitig bereit.')
time.sleep(1)
nodes=request('/object_info')
required={'UnetLoaderGGUF','CLIPLoader','VAELoader','TextEncodeQwenImage21','EmptyLatentImage','KSampler','VAEDecode','SaveImage'}
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 CPU, Bildberechnung auf GPU')
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: