Add local FLUX image editing

router/image_worker.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-"""FLUX.2 [klein] 4B Base – Bild-Worker.
+"""FLUX.2 [klein] 4B – Bild-Worker with reference-image editing.
 
 Protokoll: zeilenbasiertes JSON über stdin/stdout.
 
@@ -32,7 +32,7 @@ os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
 os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
 
 MODEL_DIR = os.environ.get(
-    "FLUX_MODEL_DIR", "/opt/mike-ai/models/FLUX.2-klein-base-4B")
+    "FLUX_MODEL_DIR", "/opt/mike-ai/models/FLUX.2-klein-4B")
 
 _pipe = None          # geladene Pipeline (None = entladen)
 _load_seconds = 0.0   # Dauer des letzten Ladens
@@ -83,12 +83,13 @@ def _unload() -> None:
 
 def _generate(req: dict) -> dict:
     import torch
+    from PIL import Image
 
     prompt = req["prompt"]
     width = int(req.get("width", 1024))
     height = int(req.get("height", 1024))
-    steps = int(req.get("steps", 50))
-    guidance = float(req.get("guidance", 4.0))
+    steps = int(req.get("steps", 4))
+    guidance = float(req.get("guidance", 1.0))
     seed = req.get("seed")
     output = req["output"]
 
@@ -98,14 +99,27 @@ def _generate(req: dict) -> dict:
     generator = None
     if seed is not None:
         generator = torch.Generator(device="cuda").manual_seed(int(seed))
-    image = _pipe(
+    kwargs = dict(
         prompt=prompt,
         height=height,
         width=width,
         guidance_scale=guidance,
         num_inference_steps=steps,
         generator=generator,
-    ).images[0]
+    )
+    source_files = req.get("source_files") or []
+    if not isinstance(source_files, list) or len(source_files) > 4:
+        raise ValueError("invalid source image list")
+    sources = []
+    for source in source_files:
+        if not isinstance(source, str):
+            raise ValueError("invalid source image filename")
+        path = os.path.join(os.path.dirname(output), source)
+        with Image.open(path) as opened:
+            sources.append(opened.convert("RGB"))
+    if sources:
+        kwargs["image"] = sources[0] if len(sources) == 1 else sources
+    image = _pipe(**kwargs).images[0]
 
     os.makedirs(os.path.dirname(output) or ".", exist_ok=True)
     image.save(output)