Add profile-aware image reference editing

image_test.py

@@ -1,5 +1,7 @@
 """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
@@ -11,7 +13,7 @@ import threading
 import time
 import urllib.request
 import uuid
-from profiles import QWEN_REPO, FLUX_REPO, image_recipe
+from profiles import QWEN_REPO, FLUX_REPO, image_recipe, image_capabilities
 
 PYTHON=Path('/opt/deck-image-python/bin/python')
 COMFY=Path('/opt/deck-comfy')
@@ -43,7 +45,8 @@ def select_gpus(devices,model_size,encoder_size,vae_size,offload=False):
     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'):
+def workflow(prompt,params,seed,family='qwen',references=()):
+    if references and family!='qwen':raise ValueError('Dieses Bildprofil unterstützt keine Referenzbilder.')
     if family=='flux':
         return {
           '1':{'class_type':'UNETLoader','inputs':{'unet_name':'model.safetensors','weight_dtype':'default'}},
@@ -61,7 +64,7 @@ def workflow(prompt,params,seed,family='qwen'):
           '13':{'class_type':'ConditioningZeroOut','inputs':{'conditioning':['4',0]}}
         }
 
-    return {
+    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'}},
@@ -71,6 +74,27 @@ def workflow(prompt,params,seed,family='qwen'):
       '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'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):
@@ -100,12 +124,12 @@ class ImageTests:
     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):
+ def start(self,profile_id,prompt,wait=False,reserved=False,reference_images=None):
     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)
+    try:return self._start(profile_id,prompt,release,reference_images or [])
     except Exception:release();raise
- def _start(self,profile_id,prompt,release):
+ def _start(self,profile_id,prompt,release,reference_images):
     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.')
@@ -115,6 +139,7 @@ class ImageTests:
         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'))
@@ -127,16 +152,19 @@ class ImageTests:
         gpu,encoder_gpu=select_gpus(probe(),model['size'],encoder['size'],vae['size'],offload=model['repo']==FLUX_REPO);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()
+        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),daemon=True).start()
         return dict(self.job)
- def _run(self,job_id,prompt,params,seed,gpu,encoder_gpu,model,encoder,vae,release=lambda:None):
+ def _run(self,job_id,prompt,params,seed,gpu,encoder_gpu,model,encoder,vae,release=lambda:None,references=()):
     directory=self.root/job_id;process=None
     try:
         directory.mkdir(mode=0o700)
         for role,item,filename in [('unet',model,'model.safetensors' if model['repo']==FLUX_REPO else '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()
+        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'}}
@@ -168,7 +196,7 @@ class ImageTests:
                 if time.monotonic()>deadline:raise ValueError('Bildlaufzeit wurde nicht rechtzeitig bereit.')
                 time.sleep(1)
         nodes=request('/object_info')
-        graph=workflow(prompt,params,seed,'flux' if model['repo']==FLUX_REPO else 'qwen')
+        graph=workflow(prompt,params,seed,'flux' if model['repo']==FLUX_REPO else 'qwen',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')
@@ -203,6 +231,8 @@ class ImageTests:
                 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: