From c27636ac34a15d1166932a5d61d108eea77c6429 Mon Sep 17 00:00:00 2001 From: Mikei386 <44135113+Mikei386@users.noreply.github.com> Date: Sun, 30 Aug 2026 10:07:40 +0200 Subject: [PATCH] Add local FLUX image editing --- .env.example | 2 +- ATHENA.md | 4 +- README.md | 2 +- compose.yaml | 6 +- config/install.env.example | 2 +- deploy/install.sh | 9 +- dev/test_local.sh | 27 +++++ docs/ARCHITECTURE.md | 4 +- install.sh | 14 +-- platform/docker/image-worker/image_worker.py | 60 +++++++--- platform/models/manifest.example.yaml | 5 +- router/ai_profile_router.py | 112 ++++++++++++++++--- router/image_worker.py | 26 ++++- 13 files changed, 209 insertions(+), 64 deletions(-) diff --git a/.env.example b/.env.example index 12993a9..6664b35 100644 --- a/.env.example +++ b/.env.example @@ -3,7 +3,7 @@ AI_BIND_ADDRESS=10.77.0.2 MODEL_DIR=/data/models ROUTER_API_KEY=GENERATED_BY_INSTALLER CONTROLLER_TOKEN=GENERATED_BY_INSTALLER -Z_IMAGE_MODEL_DIR=/data/models/Z-Image-Turbo +FLUX_MODEL_DIR=/data/models/FLUX.2-klein-4B IMAGE_GPU_DEVICES=1 PIPER_TTS_VERSION=1.6.0 PIPER_VOICE=de_DE-thorsten-high diff --git a/ATHENA.md b/ATHENA.md index 9a17f6e..44e0ace 100644 --- a/ATHENA.md +++ b/ATHENA.md @@ -10,7 +10,7 @@ Sie betreibt: - llama.cpp mit genau einem aktiven Qwen-Profil, - den OpenAI-kompatiblen Profile Router, -- Z-Image-Turbo für Bilder, +- FLUX.2-klein-4B für Textbilder und Referenzbild-Bearbeitung, - XTTS und Piper für Sprache, - das Athena-Dashboard, - WireGuard-Gateway und Datenbackup, @@ -48,7 +48,7 @@ Qwen-Profil wird vom Profile Controller verwaltet. - Fast: kurze, interaktive Aufgaben - Medium/Large/Ultra: steigende Kontextgrößen desselben lokalen Qwen-Modells - Uncensored: separates lokales Profil -- Z-Image-Turbo: Bildgenerierung; Qwen wird dafür kurz entladen und danach +- FLUX.2-klein-4B: Bildgenerierung und Editing; Qwen wird dafür kurz entladen und danach automatisch wiederhergestellt - XTTS: RTX 3060; Piper bleibt CPU-Fallback diff --git a/README.md b/README.md index ede56f1..ba4e68a 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ Bild- und Sprachausgabe. **Hermes und die Fach-MCPs laufen auf Unraid.** - genau ein aktives llama.cpp-Profil: Fast, Medium, Large, Ultra oder Uncensored - Profile Router auf Port 8081 -- Z-Image-Turbo als exklusiver Bild-Worker auf der RTX 5080 +- FLUX.2-klein-4B für Textbilder und Referenzbild-Bearbeitung auf der RTX 5080 - XTTS auf der RTX 3060 mit Piper als CPU-Fallback - Live-Dashboard mit 21 Tagen Detailhistorie auf Port 8099 - WireGuard-Gateway, Datenbackup und Athena-Operator diff --git a/compose.yaml b/compose.yaml index ab49ff6..c4308df 100644 --- a/compose.yaml +++ b/compose.yaml @@ -573,7 +573,7 @@ services: IMAGE_DIR: /data/images IMAGE_WORKER_URL: http://image-worker:8086 IMAGE_WORKER_TOKEN: "${CONTROLLER_TOKEN:?CONTROLLER_TOKEN is required}" - IMAGE_MODEL_NAME: Z-Image-Turbo + IMAGE_MODEL_NAME: FLUX.2-klein-4B CHAT_IMAGE_ALLOW_REMOTE_URLS: "false" ENABLE_IMAGE_GENERATION: "true" ENABLE_TTS: "true" @@ -628,13 +628,13 @@ services: read_only: true tmpfs: ["/tmp:size=1g,mode=1777"] volumes: - - "${Z_IMAGE_MODEL_DIR:-/data/models/Z-Image-Turbo}:/models/Z-Image-Turbo:ro" + - "${FLUX_MODEL_DIR:-/data/models/FLUX.2-klein-4B}:/models/FLUX.2-klein-4B:ro" - router-images:/data/images environment: NVIDIA_VISIBLE_DEVICES: ${IMAGE_GPU_DEVICES:-1} NVIDIA_DRIVER_CAPABILITIES: compute,utility WORKER_TOKEN: "${CONTROLLER_TOKEN:?CONTROLLER_TOKEN is required}" - Z_IMAGE_MODEL_DIR: /models/Z-Image-Turbo + FLUX_MODEL_DIR: /models/FLUX.2-klein-4B IMAGE_DIR: /data/images networks: [inference] security_opt: ["no-new-privileges:true"] diff --git a/config/install.env.example b/config/install.env.example index 6ad7d4d..5422a90 100644 --- a/config/install.env.example +++ b/config/install.env.example @@ -16,7 +16,7 @@ NVIDIA_MIN_DRIVER_MAJOR=570 TEXT_GPU_DEVICES=0 SECONDARY_GPU_DEVICES=1 IMAGE_GPU_DEVICES=1 -Z_IMAGE_MODEL_DIR=/data/models/Z-Image-Turbo +FLUX_MODEL_DIR=/data/models/FLUX.2-klein-4B # Headless remote reachability. Firmware power-loss recovery is configured # separately once at the physical machine. diff --git a/deploy/install.sh b/deploy/install.sh index f2d727c..1e64e57 100755 --- a/deploy/install.sh +++ b/deploy/install.sh @@ -15,7 +15,7 @@ OLD_SERVICE=mike-ai-local-llm-router.service OLD_DIR=/opt/mike-ai/local-llm-router BACKUP_DIR=/opt/mike-ai/.backup-ai-profile-router-$(date +%Y%m%d-%H%M%S) VENV="$INSTALL_DIR/venv" -MODEL_DIR=/opt/mike-ai/models/FLUX.2-klein-base-4B +MODEL_DIR=/opt/mike-ai/models/FLUX.2-klein-4B IMAGE_DIR="$INSTALL_DIR/images" XTTS_VENV=/opt/mike-ai/xtts/venv XTTS_CACHE=/opt/mike-ai/xtts/.cache @@ -72,13 +72,14 @@ echo "-- Installiere festgeschriebene Bild-Abhängigkeiten" if [ -f "$MODEL_DIR/model_index.json" ]; then echo "-- FLUX-Modell vorhanden: $MODEL_DIR" else - echo "-- Lade FLUX.2-klein-base-4B nach $MODEL_DIR (kann dauern)" + echo "-- Lade FLUX.2-klein-4B nach $MODEL_DIR (kann dauern)" "$VENV/bin/python" - <<'PY' import os from huggingface_hub import snapshot_download snapshot_download( - repo_id="black-forest-labs/FLUX.2-klein-base-4B", - local_dir="/opt/mike-ai/models/FLUX.2-klein-base-4B", + repo_id="black-forest-labs/FLUX.2-klein-4B", + revision="e7b7dc27f91deacad38e78976d1f2b499d76a294", + local_dir="/opt/mike-ai/models/FLUX.2-klein-4B", local_dir_use_symlinks=False, ) print("Modell-Download abgeschlossen") diff --git a/dev/test_local.sh b/dev/test_local.sh index a6d0789..8d5b810 100755 --- a/dev/test_local.sh +++ b/dev/test_local.sh @@ -343,6 +343,33 @@ CTYPE=$(grep -i content-type /tmp/hdr12.txt | tr -d "\r") [ "$CODE" = "200" ] && [ -s /tmp/test_dl.png ] && echo "$CTYPE" | grep -qi "image/png" \ && ok "PNG-Download (200, $CTYPE)" || bad "PNG-Download (Code $CODE, $CTYPE)" +# --- 12b. Referenzbild-Bearbeitung ------------------------------------------------------ +echo "== Test 12b: POST /v1/images/edits mit lokalem Referenzbild" +python3 - <<'PY' >/tmp/edit-request.json +import base64, json +png = open('/tmp/test_dl.png', 'rb').read() +print(json.dumps({ + 'prompt': 'Behalte die Person bei und ändere nur den Hintergrund', + 'size': '1024x1024', + 'steps': 4, + 'guidance': 1.0, + 'response_format': 'b64_json', + 'image_b64': base64.b64encode(png).decode(), +})) +PY +rm -f /tmp/test_worker_requests.jsonl +RESP=$(curl -sf "$BASE/v1/images/edits" -H "Content-Type: application/json" \ + --data-binary @/tmp/edit-request.json) +echo "$RESP" | python3 -c ' +import base64,json,sys +d=json.load(sys.stdin) +assert base64.b64decode(d["data"][0]["b64_json"])[:4] == b"\x89PNG" +' || bad "Bildbearbeitung liefert kein PNG" +REFS=$(tail -1 /tmp/test_worker_requests.jsonl | python3 -c \ + 'import json,sys; print(len(json.load(sys.stdin).get("source_files", [])))') +[ "$REFS" = "1" ] && ! find /tmp/test-images -name '.edit-*.ref' | grep -q . \ + && ok "Referenzbild übergeben und danach gelöscht" || bad "Referenzbild-Pfad/Cleanup" + # --- 13. Bild-Liste --------------------------------------------------------------------- echo "== Test 13: GET /images" RESP=$(curl -sf "$BASE/images") diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index a919c82..b809bde 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -6,7 +6,7 @@ flowchart LR H -->|OpenAI API| R[Profile Router
Athena :8081] R --> P[Profile Controller] P --> Q[genau ein llama.cpp-Profil
Qwen Fast / Medium / Large / Ultra / Uncensored] - R --> I[Z-Image-Turbo
RTX 5080, bei Bedarf] + R --> I[FLUX.2-klein-4B
RTX 5080, Text + Editing] R --> T[XTTS RTX 3060
Piper CPU-Fallback] H --> U[MUA / Unraid MCP] @@ -49,4 +49,4 @@ Kontextgröße: Die visuelle Fassung liegt als `athena-architecture-map.png` neben dieser Datei. Eine zweite Detailkarte, `athena-gpu-allocation-map.png`, zeigt die profilabhängige Layer-Verteilung auf RTX 5080 und RTX 3060 sowie die festen -GPU-Zuordnungen von Z-Image, Vision-Projektor und XTTS. +GPU-Zuordnungen von FLUX.2, Vision-Projektor und XTTS. diff --git a/install.sh b/install.sh index 9860403..1363f8c 100755 --- a/install.sh +++ b/install.sh @@ -341,7 +341,7 @@ UNCENSORED_TENSOR_SPLIT=${UNCENSORED_TENSOR_SPLIT:-90,10} UNCENSORED_MTP_MAX=${UNCENSORED_MTP_MAX:-2} EXPERIMENTAL_GPU_DEVICES=${TEXT_GPU_DEVICES:-0} IMAGE_GPU_DEVICES=${IMAGE_GPU_DEVICES:-${TEXT_GPU_DEVICES:-0}} -Z_IMAGE_MODEL_DIR=${Z_IMAGE_MODEL_DIR:-/data/models/Z-Image-Turbo} +FLUX_MODEL_DIR=${FLUX_MODEL_DIR:-/data/models/FLUX.2-klein-4B} LLAMA_THREADS=${LLAMA_THREADS:-6} LLAMA_THREADS_BATCH=${LLAMA_THREADS_BATCH:-6} EOF @@ -460,14 +460,14 @@ build_and_start() { docker build --progress=plain --build-arg LLAMA_CPP_COMMIT="$commit" \ -f platform/docker/llama-cpp/Dockerfile -t mike-ai/llama.cpp:local . docker compose --env-file "$SECRETS_DIR/stack.env" --profile image build image-worker - if [[ ! -s ${Z_IMAGE_MODEL_DIR:-/data/models/Z-Image-Turbo}/model_index.json ]]; then - log "Z-Image-Turbo laden" - install -d -m 0755 "${Z_IMAGE_MODEL_DIR:-/data/models/Z-Image-Turbo}" + if [[ ! -s ${FLUX_MODEL_DIR:-/data/models/FLUX.2-klein-4B}/model_index.json ]]; then + log "FLUX.2-klein-4B laden" + install -d -m 0755 "${FLUX_MODEL_DIR:-/data/models/FLUX.2-klein-4B}" docker run --rm --entrypoint python \ - -v "${Z_IMAGE_MODEL_DIR:-/data/models/Z-Image-Turbo}:/download" \ + -v "${FLUX_MODEL_DIR:-/data/models/FLUX.2-klein-4B}:/download" \ mike-ai/image-worker:local -c \ - "from huggingface_hub import snapshot_download; snapshot_download('Tongyi-MAI/Z-Image-Turbo', revision='f332072aa78be7aecdf3ee76d5c247082da564a6', local_dir='/download')" - chmod -R a-w "${Z_IMAGE_MODEL_DIR:-/data/models/Z-Image-Turbo}" + "from huggingface_hub import snapshot_download; snapshot_download('black-forest-labs/FLUX.2-klein-4B', revision='e7b7dc27f91deacad38e78976d1f2b499d76a294', local_dir='/download')" + chmod -R a-w "${FLUX_MODEL_DIR:-/data/models/FLUX.2-klein-4B}" fi # Creates the tools network and deploys the only host-bound MCP: Operator. # Portable MCPs and Hermes live on Unraid and are restored through Appdata. diff --git a/platform/docker/image-worker/image_worker.py b/platform/docker/image-worker/image_worker.py index 60ede3a..ab43729 100644 --- a/platform/docker/image-worker/image_worker.py +++ b/platform/docker/image-worker/image_worker.py @@ -1,5 +1,10 @@ #!/usr/bin/env python3 -"""Private Z-Image-Turbo worker used only during a GPU hot swap.""" +"""Private FLUX.2 Klein 4B worker used only during a GPU hot swap. + +The same pipeline handles text-to-image and local reference-image editing. +Reference images are exchanged with the router through the shared image +volume; request bodies therefore never contain private image bytes here. +""" from __future__ import annotations @@ -14,7 +19,7 @@ from pathlib import Path HOST = os.environ.get("WORKER_HOST", "0.0.0.0") PORT = int(os.environ.get("WORKER_PORT", "8086")) TOKEN = os.environ.get("WORKER_TOKEN", "").strip() -MODEL_DIR = os.environ.get("Z_IMAGE_MODEL_DIR", "/models/Z-Image-Turbo") +MODEL_DIR = os.environ.get("FLUX_MODEL_DIR", "/models/FLUX.2-klein-4B") OUTPUT_DIR = Path(os.environ.get("IMAGE_DIR", "/data/images")).resolve() PIPE = None LOAD_SECONDS = 0.0 @@ -33,15 +38,13 @@ def load_pipeline() -> None: if PIPE is not None: return import torch - from diffusers import ZImagePipeline + from diffusers import Flux2KleinPipeline started = time.monotonic() - PIPE = ZImagePipeline.from_pretrained( + PIPE = Flux2KleinPipeline.from_pretrained( MODEL_DIR, torch_dtype=torch.bfloat16, low_cpu_mem_usage=False) - # The Qwen text encoder and the DiT do not fit together in the usable - # 16 GiB of the RTX 5080. Sequential offload keeps only the active - # submodule on CUDA. This is slower than a fully resident pipeline, but - # deterministic and leaves the RTX 3060 available for XTTS. - PIPE.enable_sequential_cpu_offload() + # Officially supported low-VRAM path. It keeps the complete pipeline + # within the usable 16 GiB of the RTX 5080 and leaves the RTX 3060 alone. + PIPE.enable_model_cpu_offload() if hasattr(PIPE, "enable_vae_slicing"): PIPE.enable_vae_slicing() if hasattr(PIPE, "enable_vae_tiling"): @@ -51,6 +54,7 @@ def load_pipeline() -> None: def generate(data: dict) -> dict: import torch + from PIL import Image prompt = data.get("prompt") filename = data.get("filename") if not isinstance(prompt, str) or not prompt.strip() or len(prompt) > 8000: @@ -62,17 +66,37 @@ def generate(data: dict) -> dict: if (width, height) not in {(1024, 1024), (1536, 1024), (1024, 1536), (1920, 1088), (1088, 1920)}: raise ValueError("unsupported image size") - steps = int(data.get("steps", 9)) - guidance = float(data.get("guidance", 0.0)) - if steps != 9 or guidance != 0.0: - raise ValueError("Z-Image-Turbo requires steps=9 and guidance=0.0") + steps = int(data.get("steps", 4)) + guidance = float(data.get("guidance", 1.0)) + if steps != 4 or guidance != 1.0: + raise ValueError("FLUX.2-klein-4B requires steps=4 and guidance=1.0") + source_files = data.get("source_files") or [] + if not isinstance(source_files, list) or len(source_files) > 4: + raise ValueError("invalid source image list") + source_images = [] + for source_name in source_files: + if not isinstance(source_name, str) or Path(source_name).name != source_name: + raise ValueError("invalid source image filename") + source = (OUTPUT_DIR / source_name).resolve() + if source.parent != OUTPUT_DIR or not source.is_file(): + raise ValueError("source image not found") + with Image.open(source) as opened: + source_images.append(opened.convert("RGB")) seed = data.get("seed") generator = None if seed is None else torch.Generator(device="cuda").manual_seed(int(seed)) load_pipeline() started = time.monotonic() - image = PIPE(prompt=prompt, height=height, width=width, - num_inference_steps=9, guidance_scale=0.0, - generator=generator).images[0] + kwargs = { + "prompt": prompt, + "height": height, + "width": width, + "num_inference_steps": 4, + "guidance_scale": 1.0, + "generator": generator, + } + if source_images: + kwargs["image"] = source_images[0] if len(source_images) == 1 else source_images + image = PIPE(**kwargs).images[0] OUTPUT_DIR.mkdir(parents=True, exist_ok=True) output = OUTPUT_DIR / filename image.save(output) @@ -84,7 +108,7 @@ def generate(data: dict) -> dict: class Handler(BaseHTTPRequestHandler): def log_message(self, fmt: str, *args: object) -> None: # Never log request bodies/prompts. - print(f"[z-image-worker] {self.client_address[0]} {fmt % args}", flush=True) + print(f"[flux-image-worker] {self.client_address[0]} {fmt % args}", flush=True) def reply(self, status: int, payload: dict) -> None: body = json.dumps(payload, separators=(",", ":")).encode() @@ -113,7 +137,7 @@ class Handler(BaseHTTPRequestHandler): raise ValueError("invalid request size") self.reply(200, generate(json.loads(self.rfile.read(length)))) except Exception as exc: - print(f"[z-image-worker] generation failed: " + print(f"[flux-image-worker] generation failed: " f"{type(exc).__name__}: {str(exc)[:1000]}", flush=True) self.reply(400, {"status": "error", "message": str(exc)}) diff --git a/platform/models/manifest.example.yaml b/platform/models/manifest.example.yaml index 8bd7368..0b62217 100644 --- a/platform/models/manifest.example.yaml +++ b/platform/models/manifest.example.yaml @@ -28,8 +28,9 @@ models: sha256: "REPLACE_AFTER_VERIFICATION" image: role: image-generation - source: Tongyi-MAI/Z-Image-Turbo - target: /data/models/Z-Image-Turbo + source: black-forest-labs/FLUX.2-klein-4B + revision: e7b7dc27f91deacad38e78976d1f2b499d76a294 + target: /data/models/FLUX.2-klein-4B revision: "f332072aa78be7aecdf3ee76d5c247082da564a6" xtts: role: text-to-speech diff --git a/router/ai_profile_router.py b/router/ai_profile_router.py index 5f71941..ae2681a 100755 --- a/router/ai_profile_router.py +++ b/router/ai_profile_router.py @@ -18,8 +18,9 @@ Virtuelle Modelle: qwen-fast, qwen-medium, qwen-large, qwen-ultra, Kommandos: POST /fast, /medium, /large, /ultra, /uncensored GET /status (Zustand) -Bildgenerierung (Z-Image-Turbo): +Bildgenerierung und Editing (FLUX.2-klein-4B): POST /v1/images/generations (OpenAI-kompatibel) + POST /v1/images/edits (lokal, Referenzbilder) GET /images (Liste) GET /images/ (PNG-Download) @@ -37,7 +38,7 @@ Der Router leitet /v1/audio/speech und /v1/audio/transcriptions per HTTP an die Worker weiter. Der Router agiert als Modell-Orchestrator: vor der Generierung wird -llama.cpp gestoppt, der Bild-Worker lädt Z-Image, generiert und entlädt +llama.cpp gestoppt, der Bild-Worker lädt FLUX.2, generiert/bearbeitet und entlädt das Modell wieder; danach wird das vorherige Qwen-Profil wiederher- gestellt und erst dann geantwortet (try/finally – Qwen wird auch bei Fehlgeschlagener Generierung wiederhergestellt). @@ -116,7 +117,7 @@ CONNECT_TIMEOUT = float(os.environ.get("CONNECT_TIMEOUT", "10")) # s, Connect POLL_INTERVAL = float(os.environ.get("POLL_INTERVAL", "2")) # s, Polling-Intervall MAX_GENERATION_TOKENS = int(os.environ.get("MAX_GENERATION_TOKENS", "8192")) -# --- Bildgenerierung (Z-Image-Turbo) --- +# --- Bildgenerierung und Referenzbild-Bearbeitung (FLUX.2 Klein 4B) --- LLAMA_SERVICE = os.environ.get("LLAMA_SERVICE", "mike-ai-llama-ui.service") SYSTEMCTL_BIN = os.environ.get("SYSTEMCTL_BIN", "systemctl") IMAGE_WORKER = os.environ.get( @@ -125,7 +126,7 @@ IMAGE_PYTHON = os.environ.get( "IMAGE_PYTHON", "/opt/mike-ai/ai-profile-router/venv/bin/python") IMAGE_WORKER_URL = os.environ.get("IMAGE_WORKER_URL", "").rstrip("/") IMAGE_WORKER_TOKEN = os.environ.get("IMAGE_WORKER_TOKEN", "").strip() -IMAGE_MODEL_NAME = os.environ.get("IMAGE_MODEL_NAME", "Z-Image-Turbo") +IMAGE_MODEL_NAME = os.environ.get("IMAGE_MODEL_NAME", "FLUX.2-klein-4B") IMAGE_DIR = os.environ.get( "IMAGE_DIR", "/opt/mike-ai/ai-profile-router/images") IMAGE_WORKER_LOG = os.environ.get( @@ -153,8 +154,8 @@ IMAGE_SIZES = { "1920x1088": (1920, 1088), "1088x1920": (1088, 1920), } -# Z-Image-Turbo nutzt neun Scheduler-Schritte (acht DiT-Forwards) ohne CFG. -IMAGE_QUALITY = {"standard": 9, "high": 9} +# Das destillierte FLUX.2-klein-4B ist auf vier Schritte ausgelegt. +IMAGE_QUALITY = {"standard": 4, "high": 4} IMAGE_DEFAULT_QUALITY = "standard" IMAGE_MAX_N = 4 @@ -742,7 +743,7 @@ def switch_profile(profile: str, implicit: bool = False) -> None: # --------------------------------------------------------------------------- -# Bildgenerierung (Z-Image-Turbo) +# Bildgenerierung und Editing (FLUX.2-klein-4B) # --------------------------------------------------------------------------- class _Worker: @@ -986,7 +987,8 @@ def _restore_qwen(profile: str) -> None: def generate_image(prompt: str, width: int, height: int, steps: int, guidance: float, seed: int | None, n: int, - quality: str = "standard" + quality: str = "standard", + source_files: list[str] | None = None, ) -> tuple[list[str], str | None]: """Orchestriert die Bildgenerierung inkl. Qwen-Hotswap. @@ -1047,6 +1049,7 @@ def generate_image(prompt: str, width: int, height: int, steps: int, "guidance": guidance, "seed": seed, "output": output, + "source_files": source_files or [], }, timeout=IMAGE_GEN_TIMEOUT) if resp.get("status") != "ok": raise RuntimeError( @@ -1065,6 +1068,8 @@ def generate_image(prompt: str, width: int, height: int, steps: int, "steps": steps, "guidance": guidance, "quality": quality, + "mode": "image-edit" if source_files else "text-to-image", + "reference_images": len(source_files or []), "seconds": resp.get("seconds"), "model": IMAGE_MODEL_NAME, "created": time.strftime("%Y-%m-%dT%H:%M:%S"), @@ -1403,6 +1408,12 @@ class Handler(BaseHTTPRequestHandler): else: self._send_error(503, "Bildgenerierung ist nicht installiert", "server_error", "feature_disabled") + elif path == "/v1/images/edits" and self.command == "POST": + if ENABLE_IMAGE_GENERATION: + self._image_edit() + else: + self._send_error(503, "Bildbearbeitung ist nicht installiert", + "server_error", "feature_disabled") elif path == "/v1/audio/speech" and self.command == "POST": if ENABLE_TTS: self._speech() @@ -1627,22 +1638,89 @@ class Handler(BaseHTTPRequestHandler): # ---------- Bildgenerierung ---------- def _image_generate(self) -> None: + data = self._read_image_request() + if data is not None: + self._image_request(data, []) + + def _image_edit(self) -> None: + """Edit with local image bytes supplied by the private Hermes plugin.""" + data = self._read_image_request() + if data is None: + return + encoded: list[str] = [] + primary = data.pop("image_b64", None) + if isinstance(primary, str) and primary: + encoded.append(primary) + references = data.pop("reference_images_b64", []) + if references is None: + references = [] + if not isinstance(references, list) or any( + not isinstance(item, str) for item in references): + self._send_error(400, "'reference_images_b64' muss eine Liste sein", + "invalid_request_error", "invalid_references") + return + encoded.extend(references) + if not encoded: + self._send_error(400, "Referenzbild fehlt", + "invalid_request_error", "missing_image") + return + if len(encoded) > 4: + self._send_error(400, "höchstens vier Referenzbilder erlaubt", + "invalid_request_error", "too_many_images") + return + + source_files: list[str] = [] + try: + for item in encoded: + if item.startswith("data:"): + header, separator, item = item.partition(",") + if not separator or not header.lower().startswith("data:image/"): + raise ValueError("ungültige Bild-Data-URI") + try: + raw = base64.b64decode(item, validate=True) + except Exception as exc: + raise ValueError("ungültige Base64-Bilddaten") from exc + if not raw or len(raw) > CHAT_IMAGE_MAX_BYTES: + raise ValueError( + f"Referenzbild muss 1..{CHAT_IMAGE_MAX_BYTES} Bytes groß sein") + name = f".edit-{os.urandom(12).hex()}.ref" + os.makedirs(IMAGE_DIR, exist_ok=True) + with open(os.path.join(IMAGE_DIR, name), "xb") as output: + output.write(raw) + source_files.append(name) + self._image_request(data, source_files) + except ValueError as exc: + self._send_error(400, str(exc), + "invalid_request_error", "invalid_image") + finally: + for name in source_files: + try: + os.unlink(os.path.join(IMAGE_DIR, name)) + except FileNotFoundError: + pass + except OSError as exc: + log.warning("temporäres Referenzbild nicht gelöscht: %s", exc) + + def _read_image_request(self) -> dict | None: try: body = self._read_body() except ValueError as e: self._send_error(400, str(e), "invalid_request_error", "invalid_body") - return + return None try: data = json.loads(body) except ValueError: self._send_error(400, "ungültiges JSON", "invalid_request_error", "invalid_json") - return + return None if not isinstance(data, dict): self._send_error(400, "Request muss ein JSON-Objekt sein", "invalid_request_error", "invalid_request") - return + return None + return data + + def _image_request(self, data: dict, source_files: list[str]) -> None: prompt = data.get("prompt") if not isinstance(prompt, str) or not prompt.strip(): @@ -1679,19 +1757,19 @@ class Handler(BaseHTTPRequestHandler): "invalid_request_error", "invalid_quality") return steps = data.get("steps", IMAGE_QUALITY[quality]) - if not isinstance(steps, int) or isinstance(steps, bool) or steps != 9: - self._send_error(400, "Z-Image-Turbo erfordert 'steps'=9", + if not isinstance(steps, int) or isinstance(steps, bool) or steps != 4: + self._send_error(400, "FLUX.2-klein-4B erfordert 'steps'=4", "invalid_request_error", "invalid_steps") return - guidance = data.get("guidance", 0.0) + guidance = data.get("guidance", 1.0) try: guidance = float(guidance) except (TypeError, ValueError): self._send_error(400, "'guidance' muss eine Zahl sein", "invalid_request_error", "invalid_guidance") return - if guidance != 0.0: - self._send_error(400, "Z-Image-Turbo erfordert 'guidance'=0.0", + if guidance != 1.0: + self._send_error(400, "FLUX.2-klein-4B erfordert 'guidance'=1.0", "invalid_request_error", "invalid_guidance") return @@ -1719,7 +1797,7 @@ class Handler(BaseHTTPRequestHandler): try: results, warning = generate_image( prompt.strip(), width, height, steps, guidance, seed, n, - quality) + quality, source_files) except (ValueError, RuntimeError) as e: self._send_error(503, str(e), "server_error", "image_generation_failed") return diff --git a/router/image_worker.py b/router/image_worker.py index 1e625a2..a92b328 100644 --- a/router/image_worker.py +++ b/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)