router: Bildgenerierung mit FLUX.2 [klein] 4B Base (GPU-Hotswap)
- POST /v1/images/generations (OpenAI-kompatibel, prompt/size/n/seed/quality) - quality: standard=30 Steps (Default), high=50 Steps - Größen: 1024x1024, 1536x1024, 1024x1536, 1920x1088, 1088x1920 - GPU-Hotswap: Qwen stoppen -> FLUX laden -> Bild -> FLUX entladen -> Qwen wiederherstellen (exakt vorheriges Profil) - Zentrales GPU/Modell-Lock (Profilwechsel und Bild teilen sich das Lock) - Chat-Requests warten während Bild-Job (kein 502), Timeout CHAT_WAIT_TIMEOUT - Robuste Recovery: try/finally, Worker-Beendigung, VRAM-Check, Qwen-Readiness - /status: image.phase, image.worker, image.model_loaded, qwen.available, qwen.active_chats - GET /images, GET /images/<datei> (validiert, nur images/-Verzeichnis) - image_worker.py: FLUX-Worker (eigener Prozess, JSON-Protokoll, bf16 + enable_model_cpu_offload) - deploy: venv (torch/diffusers/transformers/accelerate), Modell-Download, Image-Dir, systemd-Unit mit Image-Umgebungsvariablen - dev: Mock-Worker, fake-systemctl, Benchmarks (GPU-Resident, Offload, Steps, Quality-Compare), 32 lokale Tests - README: Bildgenerierung, Hotswap, Recovery, Benchmarks (RTX 5080), Python-Pakete Benchmarks (RTX 5080, 16 GB, CPU-Offload): - 512x512 / 10 Steps: ~9.3 s - 1024x1024 / 30 Steps: ~31.3 s - 1024x1024 / 50 Steps: ~45.3 s - 1920x1088 / 50 Steps: ~91 s - Peak-VRAM: ~8.4-8.9 GB - Hotswap-Gesamtzeit: ~41-42 s (1024x1024 / 30 Steps)
This commit is contained in:
Mikei386
+606
-38
@@ -15,19 +15,28 @@ Virtuelle Modelle: qwen-fast, qwen-medium, qwen-long
|
||||
Kommandos: POST /fast, /medium, /long (Profilwechsel)
|
||||
GET /status (Zustand)
|
||||
|
||||
Ein Profilwechsel führt PROFILE_SCRIPT <profil> aus (ohne Shell, feste
|
||||
Argumente → keine Injection), wartet dann, bis llama.cpp wieder erreichbar
|
||||
ist, und erst dann wird eine erfolgreiche Antwort geliefert bzw. der
|
||||
Request weitergeleitet.
|
||||
Bildgenerierung (FLUX.2 [klein] 4B Base):
|
||||
POST /v1/images/generations (OpenAI-kompatibel)
|
||||
GET /images (Liste)
|
||||
GET /images/<datei> (PNG-Download)
|
||||
|
||||
Der Router agiert als Modell-Orchestrator: vor der Generierung wird
|
||||
llama.cpp gestoppt, der Bild-Worker lädt FLUX, generiert 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).
|
||||
|
||||
Nur Python-Standardbibliothek. Logging nach stdout (journald).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import queue
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
@@ -51,6 +60,42 @@ REQUEST_TIMEOUT = float(os.environ.get("REQUEST_TIMEOUT", "600")) # s, Read-Ti
|
||||
CONNECT_TIMEOUT = float(os.environ.get("CONNECT_TIMEOUT", "10")) # s, Connect-Timeout
|
||||
POLL_INTERVAL = float(os.environ.get("POLL_INTERVAL", "2")) # s, Polling-Intervall
|
||||
|
||||
# --- Bildgenerierung (FLUX.2 [klein] 4B Base) ---
|
||||
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(
|
||||
"IMAGE_WORKER", "/opt/mike-ai/ai-profile-router/image_worker.py")
|
||||
IMAGE_PYTHON = os.environ.get(
|
||||
"IMAGE_PYTHON", "/opt/mike-ai/ai-profile-router/venv/bin/python")
|
||||
IMAGE_DIR = os.environ.get(
|
||||
"IMAGE_DIR", "/opt/mike-ai/ai-profile-router/images")
|
||||
IMAGE_WORKER_LOG = os.environ.get(
|
||||
"IMAGE_WORKER_LOG", "/opt/mike-ai/ai-profile-router/image_worker.log")
|
||||
IMAGE_START_TIMEOUT = float(os.environ.get("IMAGE_START_TIMEOUT", "120")) # s, Worker-Start
|
||||
IMAGE_GEN_TIMEOUT = float(os.environ.get("IMAGE_GEN_TIMEOUT", "1800")) # s, pro Bild
|
||||
IMAGE_VRAM_FREE_TIMEOUT = float(os.environ.get("IMAGE_VRAM_FREE_TIMEOUT", "90")) # s, VRAM-Abgabe
|
||||
|
||||
# Erlaubte Auflösungen (Breite x Höhe). FLUX.2 klein ist für 1 MP
|
||||
# ausgelegt; 1920x1088 (≈2 MP) wird zusätzlich unterstützt.
|
||||
IMAGE_SIZES = {
|
||||
"1024x1024": (1024, 1024),
|
||||
"1536x1024": (1536, 1024),
|
||||
"1024x1536": (1024, 1536),
|
||||
"1920x1088": (1920, 1088),
|
||||
"1088x1920": (1088, 1920),
|
||||
}
|
||||
# Qualitätsstufen → Inference-Schritte (guidance bleibt offiziell 4.0).
|
||||
# Auf der RTX 5080 gemessen: 30 vs. 50 Steps liefern praktisch dieselbe
|
||||
# Qualität (1024x1024: 31,3 s vs. 45,3 s). Default ist daher "standard".
|
||||
IMAGE_QUALITY = {"standard": 30, "high": 50}
|
||||
IMAGE_DEFAULT_QUALITY = "standard"
|
||||
IMAGE_MAX_N = 4
|
||||
|
||||
# Chat-Waiting: Während eines Image-Jobs oder Profilwechsels ist Qwen
|
||||
# down. Chat-Requests warten (statt 502) bis Qwen wieder bereit ist.
|
||||
CHAT_WAIT_TIMEOUT = float(os.environ.get("CHAT_WAIT_TIMEOUT", "300")) # s, max. Warten
|
||||
CHAT_DRAIN_TIMEOUT = float(os.environ.get("CHAT_DRAIN_TIMEOUT", "60")) # s, max. Warten auf aktive Chats
|
||||
|
||||
PROFILES = {"fast": 73728, "medium": 94208, "long": 131072}
|
||||
VIRTUAL_MODELS = {f"qwen-{name}": name for name in PROFILES}
|
||||
|
||||
@@ -78,16 +123,68 @@ UPSTREAM_HOST, UPSTREAM_PORT = _parse_upstream(UPSTREAM_URL)
|
||||
# Zustand
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class _ImageState:
|
||||
"""Zustand der Bildgenerierung (nur für Status-Reporting)."""
|
||||
def __init__(self) -> None:
|
||||
self.phase = "idle" # siehe PHASES unten
|
||||
self.worker: "_Worker | None" = None
|
||||
self.last_error: str | None = None
|
||||
self.last_image: str | None = None
|
||||
self.last_seconds: float | None = None
|
||||
|
||||
|
||||
IMAGE_PHASES = (
|
||||
"idle", "stopping-qwen", "loading-image", "generating",
|
||||
"unloading-image", "restoring-qwen",
|
||||
)
|
||||
|
||||
|
||||
class _State:
|
||||
"""Gemeinsamer, thread-sicherer Zustand."""
|
||||
lock = threading.Lock() # serialisiert Profilwechsel
|
||||
"""Gemeinsamer, thread-sicherer Zustand.
|
||||
|
||||
lock : zentraler GPU-/Model-Lock. Wird von Profilwechsel UND
|
||||
Image-Generation gehalten → gegenseitiger Ausschluss,
|
||||
kein Race zwischen beiden.
|
||||
avail_lock : schützt qwen_unavailable + active_chats (Chat-Waiting).
|
||||
"""
|
||||
lock = threading.Lock() # GPU-/Model-Lock (Profilwechsel + Image)
|
||||
switching: str | None = None # Profil, das gerade gewechselt wird
|
||||
started = time.time()
|
||||
image = _ImageState()
|
||||
# Qwen-Verfügbarkeit für das Chat-Waiting:
|
||||
qwen_unavailable = False # True, wenn Qwen down/neu geladen wird
|
||||
active_chats = 0 # Anzahl laufender Chat-Requests
|
||||
avail_lock = threading.Lock() # schützt die beiden Felder oben
|
||||
|
||||
|
||||
STATE = _State()
|
||||
|
||||
|
||||
def _wait_chats_drained(timeout: float | None = None) -> None:
|
||||
"""Wartet, bis keine aktiven Chat-Requests mehr laufen.
|
||||
|
||||
Wird von Profilwechsel/Image-Job aufgerufen, BEVOR Qwen gestoppt wird.
|
||||
Verhindert, dass ein laufender Chat auf ein gestopptes Qwen trifft (502).
|
||||
"""
|
||||
timeout = CHAT_DRAIN_TIMEOUT if timeout is None else timeout
|
||||
deadline = time.monotonic() + timeout
|
||||
while True:
|
||||
with STATE.avail_lock:
|
||||
if STATE.active_chats == 0:
|
||||
return
|
||||
n = STATE.active_chats
|
||||
if time.monotonic() > deadline:
|
||||
log.warning("Chat-Drain-Timeout nach %.0f s (%d aktive Chats) – "
|
||||
"fahre trotzdem fort", timeout, n)
|
||||
return
|
||||
time.sleep(0.5)
|
||||
|
||||
|
||||
def _set_qwen_unavailable(unavailable: bool) -> None:
|
||||
with STATE.avail_lock:
|
||||
STATE.qwen_unavailable = unavailable
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Upstream (llama.cpp)
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -167,6 +264,9 @@ def switch_profile(profile: str, implicit: bool = False) -> None:
|
||||
if profile not in PROFILES:
|
||||
raise ValueError(f"unbekanntes Profil: {profile!r} "
|
||||
f"(erlaubt: {', '.join(PROFILES)})")
|
||||
# Kein Fast-Fail: Wenn ein Image-Job läuft (hält den GPU-Lock), wartet
|
||||
# der Profilwechsel auf den GPU-Lock (blockiert), bis der Image-Job
|
||||
# fertig ist. So bekommen Chat-Requests kein 502, sondern warten.
|
||||
with STATE.lock:
|
||||
STATE.switching = profile
|
||||
try:
|
||||
@@ -177,43 +277,299 @@ def switch_profile(profile: str, implicit: bool = False) -> None:
|
||||
if cur == profile and ready:
|
||||
log.info("Profil %s ist bereits aktiv", profile)
|
||||
return
|
||||
if cur == profile and up["reachable"] and not ready:
|
||||
# Modell wird gerade geladen (z.B. nach einem Wechsel)
|
||||
log.info("Warte, bis Profil %s geladen ist ...", profile)
|
||||
_wait_ready(profile, time.monotonic() + SWITCH_TIMEOUT)
|
||||
return
|
||||
if cur == profile and not up["reachable"] and implicit:
|
||||
raise RuntimeError(
|
||||
f"llama.cpp nicht erreichbar (Profil {profile} ist bereits "
|
||||
f"aktiv; Neustart über /{profile})")
|
||||
log.info("Profilwechsel: %s -> %s", cur, profile)
|
||||
# Qwen wird neu geladen/gewechselt → für Chats nicht verfügbar.
|
||||
_set_qwen_unavailable(True)
|
||||
try:
|
||||
proc = subprocess.run(
|
||||
[PROFILE_SCRIPT, profile],
|
||||
stdin=subprocess.DEVNULL,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
timeout=120,
|
||||
)
|
||||
out = proc.stdout.decode(errors="replace").strip()
|
||||
if out:
|
||||
log.info("llama-profile: %s", out[-500:])
|
||||
if proc.returncode != 0:
|
||||
# whiptail bricht das Skript ohne TTY ab – der Wechsel
|
||||
# selbst (cp + systemctl restart) ist dann aber erledigt.
|
||||
log.warning("llama-profile Exit-Code %d (ohne TTY erwartet)",
|
||||
proc.returncode)
|
||||
except subprocess.TimeoutExpired:
|
||||
log.error("llama-profile hat 120 s überschritten")
|
||||
if current_profile() != profile:
|
||||
raise RuntimeError(
|
||||
f"Profildatei wurde nicht gesetzt (erwartet: {profile})")
|
||||
log.info("Warte, bis llama.cpp das Profil geladen hat ...")
|
||||
_wait_ready(profile, time.monotonic() + SWITCH_TIMEOUT)
|
||||
_wait_chats_drained()
|
||||
if cur == profile and up["reachable"] and not ready:
|
||||
# Modell wird gerade geladen (z.B. nach einem Wechsel)
|
||||
log.info("Warte, bis Profil %s geladen ist ...", profile)
|
||||
_wait_ready(profile, time.monotonic() + SWITCH_TIMEOUT)
|
||||
return
|
||||
if cur == profile and not up["reachable"] and implicit:
|
||||
raise RuntimeError(
|
||||
f"llama.cpp nicht erreichbar (Profil {profile} ist "
|
||||
f"bereits aktiv; Neustart über /{profile})")
|
||||
log.info("Profilwechsel: %s -> %s", cur, profile)
|
||||
try:
|
||||
proc = subprocess.run(
|
||||
[PROFILE_SCRIPT, profile],
|
||||
stdin=subprocess.DEVNULL,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
timeout=120,
|
||||
)
|
||||
out = proc.stdout.decode(errors="replace").strip()
|
||||
if out:
|
||||
log.info("llama-profile: %s", out[-500:])
|
||||
if proc.returncode != 0:
|
||||
# whiptail bricht das Skript ohne TTY ab – der Wechsel
|
||||
# selbst (cp + systemctl restart) ist dann erledigt.
|
||||
log.warning("llama-profile Exit-Code %d (ohne TTY "
|
||||
"erwartet)", proc.returncode)
|
||||
except subprocess.TimeoutExpired:
|
||||
log.error("llama-profile hat 120 s überschritten")
|
||||
if current_profile() != profile:
|
||||
raise RuntimeError(
|
||||
f"Profildatei wurde nicht gesetzt (erwartet: {profile})")
|
||||
log.info("Warte, bis llama.cpp das Profil geladen hat ...")
|
||||
_wait_ready(profile, time.monotonic() + SWITCH_TIMEOUT)
|
||||
finally:
|
||||
_set_qwen_unavailable(False)
|
||||
finally:
|
||||
STATE.switching = None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Bildgenerierung (FLUX.2 [klein] 4B Base)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class _Worker:
|
||||
"""Verwaltet den Bild-Worker-Prozess (stdin/stdout-JSON-Protokoll)."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.proc: subprocess.Popen | None = None
|
||||
self.model_loaded = False
|
||||
self._queue: queue.Queue[dict] = queue.Queue()
|
||||
self._reader: threading.Thread | None = None
|
||||
|
||||
def alive(self) -> bool:
|
||||
return self.proc is not None and self.proc.poll() is None
|
||||
|
||||
def start(self) -> None:
|
||||
if self.alive():
|
||||
return
|
||||
log.info("starte Bild-Worker: %s %s", IMAGE_PYTHON, IMAGE_WORKER)
|
||||
logf = open(IMAGE_WORKER_LOG, "ab")
|
||||
self.proc = subprocess.Popen(
|
||||
[IMAGE_PYTHON, IMAGE_WORKER],
|
||||
stdin=subprocess.PIPE,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=logf,
|
||||
text=True,
|
||||
bufsize=1,
|
||||
)
|
||||
self._reader = threading.Thread(target=self._read_loop, daemon=True)
|
||||
self._reader.start()
|
||||
try:
|
||||
msg = self._queue.get(timeout=IMAGE_START_TIMEOUT)
|
||||
except queue.Empty:
|
||||
self.stop()
|
||||
raise RuntimeError("Bild-Worker hat nicht gestartet")
|
||||
if msg.get("status") != "ready":
|
||||
self.stop()
|
||||
raise RuntimeError(f"Bild-Worker-Startfehler: {msg}")
|
||||
log.info("Bild-Worker bereit")
|
||||
|
||||
def _read_loop(self) -> None:
|
||||
assert self.proc is not None and self.proc.stdout is not None
|
||||
for line in self.proc.stdout:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
self._queue.put(json.loads(line))
|
||||
except ValueError:
|
||||
log.warning("Worker-Zeile (kein JSON): %s", line[:200])
|
||||
|
||||
def request(self, payload: dict, timeout: float) -> dict:
|
||||
if not self.alive():
|
||||
raise RuntimeError("Bild-Worker ist nicht aktiv")
|
||||
assert self.proc is not None and self.proc.stdin is not None
|
||||
self.proc.stdin.write(json.dumps(payload) + "\n")
|
||||
self.proc.stdin.flush()
|
||||
try:
|
||||
return self._queue.get(timeout=timeout)
|
||||
except queue.Empty:
|
||||
raise RuntimeError(
|
||||
f"Bild-Worker hat nach {timeout:.0f} s nicht geantwortet "
|
||||
f"(cmd={payload.get('cmd')})")
|
||||
|
||||
def stop(self) -> None:
|
||||
if self.proc is not None and self.proc.poll() is None:
|
||||
self.proc.terminate()
|
||||
try:
|
||||
self.proc.wait(timeout=10)
|
||||
except subprocess.TimeoutExpired:
|
||||
self.proc.kill()
|
||||
self.proc = None
|
||||
self.model_loaded = False
|
||||
|
||||
|
||||
def _worker() -> _Worker:
|
||||
"""Worker-Instanz liefern (startet bei Bedarf)."""
|
||||
img = STATE.image
|
||||
if not img.worker or not img.worker.alive():
|
||||
if img.worker:
|
||||
img.worker.stop()
|
||||
img.worker = _Worker()
|
||||
img.worker.start()
|
||||
return img.worker
|
||||
|
||||
|
||||
def _wait_upstream_down(deadline: float) -> None:
|
||||
"""Wartet, bis llama.cpp den Port freigegeben hat (VRAM frei)."""
|
||||
while time.monotonic() < deadline:
|
||||
if not upstream_status()["reachable"]:
|
||||
return
|
||||
time.sleep(1)
|
||||
raise RuntimeError("llama.cpp gibt Port/VRAM nicht frei")
|
||||
|
||||
|
||||
def _vram_used_mib() -> int | None:
|
||||
"""Aktuelle VRAM-Belegung in MiB (via nvidia-smi), None bei Fehler."""
|
||||
try:
|
||||
out = subprocess.run(
|
||||
["nvidia-smi", "--query-gpu=memory.used",
|
||||
"--format=csv,noheader,nounits"],
|
||||
stdin=subprocess.DEVNULL, stdout=subprocess.PIPE,
|
||||
stderr=subprocess.DEVNULL, timeout=10,
|
||||
).stdout.decode().strip()
|
||||
return int(out.splitlines()[0].split()[0])
|
||||
except (OSError, ValueError, IndexError):
|
||||
return None
|
||||
|
||||
|
||||
def _wait_vram_free(threshold_mib: int = 1000,
|
||||
timeout: float | None = None) -> None:
|
||||
"""Wartet, bis der VRAM unter threshold_mib fällt (FLUX entladen).
|
||||
|
||||
Wird nach dem Beenden des Bild-Workers aufgerufen, um sicherzustellen,
|
||||
dass der VRAM (inkl. CUDA-Kontext) frei ist, bevor Qwen neu startet.
|
||||
Wenn nvidia-smi nicht verfügbar ist (z.B. lokale Tests), wird der
|
||||
Check übersprungen.
|
||||
"""
|
||||
timeout = IMAGE_VRAM_FREE_TIMEOUT if timeout is None else timeout
|
||||
deadline = time.monotonic() + timeout
|
||||
last = _vram_used_mib()
|
||||
if last is None:
|
||||
log.info("VRAM-Check übersprungen (nvidia-smi nicht verfügbar)")
|
||||
return
|
||||
while time.monotonic() < deadline:
|
||||
if last <= threshold_mib:
|
||||
log.info("VRAM frei: %d MiB", last)
|
||||
return
|
||||
time.sleep(1)
|
||||
last = _vram_used_mib()
|
||||
if last is None:
|
||||
log.info("VRAM-Check übersprungen (nvidia-smi nicht verfügbar)")
|
||||
return
|
||||
raise RuntimeError(
|
||||
f"VRAM nach {timeout:.0f} s nicht frei (letzte Messung: "
|
||||
f"{last} MiB, erwartet <= {threshold_mib} MiB)")
|
||||
|
||||
|
||||
def _restore_qwen(profile: str) -> None:
|
||||
"""Startet llama.cpp mit dem gemerkten Profil und wartet auf Readiness."""
|
||||
log.info("stelle Qwen-Profil %s wieder her ...", profile)
|
||||
try:
|
||||
subprocess.run([SYSTEMCTL_BIN, "start", LLAMA_SERVICE],
|
||||
stdin=subprocess.DEVNULL,
|
||||
stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
|
||||
timeout=120)
|
||||
except subprocess.TimeoutExpired:
|
||||
log.error("systemctl start hat 120 s überschritten")
|
||||
_wait_ready(profile, time.monotonic() + SWITCH_TIMEOUT)
|
||||
|
||||
|
||||
def generate_image(prompt: str, width: int, height: int, steps: int,
|
||||
guidance: float, seed: int | None, n: int
|
||||
) -> tuple[list[str], str | None]:
|
||||
"""Orchestriert die Bildgenerierung inkl. Qwen-Hotswap.
|
||||
|
||||
Hält den zentralen GPU-Lock (gegenseitiger Ausschluss mit Profilwechsel).
|
||||
Ablauf: Qwen stoppen → Worker laden → generieren → Worker beenden
|
||||
(VRAM + CUDA-Kontext frei) → Qwen wiederherstellen. Qwen wird auch bei
|
||||
Fehlern wiederhergestellt (try/finally).
|
||||
"""
|
||||
img = STATE.image
|
||||
with STATE.lock:
|
||||
if img.phase != "idle":
|
||||
raise RuntimeError(f"Bildgenerierung läuft ({img.phase})")
|
||||
profile = current_profile()
|
||||
if profile is None:
|
||||
raise RuntimeError("kein aktives Qwen-Profil (override.conf?)")
|
||||
os.makedirs(IMAGE_DIR, exist_ok=True)
|
||||
results: list[str] = []
|
||||
warning: str | None = None
|
||||
# Qwen wird gestoppt → für Chats nicht verfügbar (die warten).
|
||||
_set_qwen_unavailable(True)
|
||||
try:
|
||||
_wait_chats_drained()
|
||||
|
||||
# 1) Qwen stoppen (VRAM freigeben).
|
||||
img.phase = "stopping-qwen"
|
||||
subprocess.run([SYSTEMCTL_BIN, "stop", LLAMA_SERVICE],
|
||||
stdin=subprocess.DEVNULL,
|
||||
stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
|
||||
timeout=120)
|
||||
_wait_upstream_down(time.monotonic() + 60)
|
||||
|
||||
# 2) Worker starten (Modell wird beim ersten generate geladen).
|
||||
img.phase = "loading-image"
|
||||
worker = _worker()
|
||||
|
||||
# 3) Generieren.
|
||||
for i in range(n):
|
||||
img.phase = "generating"
|
||||
filename = time.strftime("%Y%m%d-%H%M%S") + \
|
||||
f"-{os.urandom(2).hex()}.png"
|
||||
output = os.path.join(IMAGE_DIR, filename)
|
||||
resp = worker.request({
|
||||
"cmd": "generate",
|
||||
"prompt": prompt,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"steps": steps,
|
||||
"guidance": guidance,
|
||||
"seed": seed,
|
||||
"output": output,
|
||||
}, timeout=IMAGE_GEN_TIMEOUT)
|
||||
if resp.get("status") != "ok":
|
||||
raise RuntimeError(
|
||||
resp.get("message", "Bildgenerierung fehlgeschlagen"))
|
||||
worker.model_loaded = True
|
||||
results.append(filename)
|
||||
img.last_image = filename
|
||||
img.last_seconds = resp.get("seconds")
|
||||
log.info("Bild %d/%d: %s (%.1f s)", i + 1, n, filename,
|
||||
resp.get("seconds", 0))
|
||||
|
||||
# 4) Worker vollständig beenden (VRAM + CUDA-Kontext freigeben).
|
||||
img.phase = "unloading-image"
|
||||
worker.stop()
|
||||
img.worker = None
|
||||
try:
|
||||
_wait_vram_free()
|
||||
except RuntimeError as e:
|
||||
log.warning("VRAM-Check: %s (fahre mit Qwen-Restore fort)", e)
|
||||
except Exception as e:
|
||||
img.last_error = str(e)
|
||||
log.error("Bildgenerierung fehlgeschlagen: %s", e)
|
||||
# Worker sicher beenden (falls noch aktiv), VRAM freigeben.
|
||||
if img.worker is not None:
|
||||
img.worker.stop()
|
||||
img.worker = None
|
||||
raise
|
||||
finally:
|
||||
# 5) Qwen immer wiederherstellen.
|
||||
img.phase = "restoring-qwen"
|
||||
try:
|
||||
_restore_qwen(profile)
|
||||
_set_qwen_unavailable(False)
|
||||
except Exception as e:
|
||||
warning = f"Qwen-Wiederherstellung fehlgeschlagen: {e}"
|
||||
img.last_error = warning
|
||||
log.error(warning)
|
||||
# Qwen ist down → qwen_unavailable bleibt True.
|
||||
img.phase = "idle"
|
||||
return results, warning
|
||||
|
||||
|
||||
def _image_filename_ok(name: str) -> bool:
|
||||
return bool(re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]*\.png", name))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# HTTP-Handler
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -238,6 +594,12 @@ class Handler(BaseHTTPRequestHandler):
|
||||
self._send_json(200, self._models_payload())
|
||||
elif path == "/status":
|
||||
self._send_json(200, self._status_payload())
|
||||
elif path == "/v1/images/generations" and self.command == "POST":
|
||||
self._image_generate()
|
||||
elif path == "/images" and self.command == "GET":
|
||||
self._images_list()
|
||||
elif path.startswith("/images/") and self.command == "GET":
|
||||
self._image_serve(path[len("/images/"):])
|
||||
elif path in ("/fast", "/medium", "/long"):
|
||||
self._switch(path[1:])
|
||||
elif (self.command == "POST" and path.startswith("/")
|
||||
@@ -277,6 +639,10 @@ class Handler(BaseHTTPRequestHandler):
|
||||
|
||||
def _status_payload(self) -> dict:
|
||||
up = upstream_status()
|
||||
img = STATE.image
|
||||
with STATE.avail_lock:
|
||||
qwen_unavailable = STATE.qwen_unavailable
|
||||
active_chats = STATE.active_chats
|
||||
return {
|
||||
"router": "ai-profile-router",
|
||||
"uptime_seconds": round(time.time() - STATE.started, 1),
|
||||
@@ -289,8 +655,174 @@ class Handler(BaseHTTPRequestHandler):
|
||||
"model": up.get("model"),
|
||||
"ctx": up.get("ctx"),
|
||||
},
|
||||
"qwen": {
|
||||
"available": not qwen_unavailable,
|
||||
"active_chats": active_chats,
|
||||
},
|
||||
"image": {
|
||||
"phase": img.phase,
|
||||
"worker": "running" if (img.worker and img.worker.alive())
|
||||
else "stopped",
|
||||
"model_loaded": bool(img.worker and img.worker.model_loaded),
|
||||
"last_image": img.last_image,
|
||||
"last_seconds": img.last_seconds,
|
||||
"last_error": img.last_error,
|
||||
},
|
||||
}
|
||||
|
||||
# ---------- Bildgenerierung ----------
|
||||
|
||||
def _image_generate(self) -> None:
|
||||
length = int(self.headers.get("Content-Length") or 0)
|
||||
try:
|
||||
data = json.loads(self.rfile.read(length))
|
||||
except ValueError:
|
||||
self._send_error(400, "ungültiges JSON",
|
||||
"invalid_request_error", "invalid_json")
|
||||
return
|
||||
if not isinstance(data, dict):
|
||||
self._send_error(400, "Request muss ein JSON-Objekt sein",
|
||||
"invalid_request_error", "invalid_request")
|
||||
return
|
||||
|
||||
prompt = data.get("prompt")
|
||||
if not isinstance(prompt, str) or not prompt.strip():
|
||||
self._send_error(400, "'prompt' fehlt oder ist leer",
|
||||
"invalid_request_error", "missing_prompt")
|
||||
return
|
||||
if len(prompt) > 8000:
|
||||
self._send_error(400, "'prompt' zu lang (max 8000 Zeichen)",
|
||||
"invalid_request_error", "prompt_too_long")
|
||||
return
|
||||
|
||||
# Größe
|
||||
size = data.get("size", "1024x1024")
|
||||
if size not in IMAGE_SIZES:
|
||||
self._send_error(
|
||||
400, f"ungültige Größe: {size!r} "
|
||||
f"(erlaubt: {', '.join(IMAGE_SIZES)})",
|
||||
"invalid_request_error", "invalid_size")
|
||||
return
|
||||
width, height = IMAGE_SIZES[size]
|
||||
|
||||
# Anzahl
|
||||
n = data.get("n", 1)
|
||||
if not isinstance(n, int) or isinstance(n, bool) or not 1 <= n <= IMAGE_MAX_N:
|
||||
self._send_error(400, f"'n' muss eine Ganzzahl 1..{IMAGE_MAX_N} sein",
|
||||
"invalid_request_error", "invalid_n")
|
||||
return
|
||||
|
||||
# Qualität / Schritte / Guidance
|
||||
quality = data.get("quality", IMAGE_DEFAULT_QUALITY)
|
||||
if quality not in IMAGE_QUALITY:
|
||||
self._send_error(400, f"ungültige Qualität: {quality!r} "
|
||||
f"(erlaubt: {', '.join(IMAGE_QUALITY)})",
|
||||
"invalid_request_error", "invalid_quality")
|
||||
return
|
||||
steps = data.get("steps", IMAGE_QUALITY[quality])
|
||||
if not isinstance(steps, int) or isinstance(steps, bool) or not 4 <= steps <= 150:
|
||||
self._send_error(400, "'steps' muss eine Ganzzahl 4..150 sein",
|
||||
"invalid_request_error", "invalid_steps")
|
||||
return
|
||||
guidance = data.get("guidance", 4.0)
|
||||
try:
|
||||
guidance = float(guidance)
|
||||
except (TypeError, ValueError):
|
||||
self._send_error(400, "'guidance' muss eine Zahl sein",
|
||||
"invalid_request_error", "invalid_guidance")
|
||||
return
|
||||
if not 1.0 <= guidance <= 10.0:
|
||||
self._send_error(400, "'guidance' muss zwischen 1.0 und 10.0 sein",
|
||||
"invalid_request_error", "invalid_guidance")
|
||||
return
|
||||
|
||||
seed = data.get("seed")
|
||||
if seed is not None:
|
||||
try:
|
||||
seed = int(seed)
|
||||
except (TypeError, ValueError):
|
||||
self._send_error(400, "'seed' muss eine Ganzzahl sein",
|
||||
"invalid_request_error", "invalid_seed")
|
||||
return
|
||||
if not 0 <= seed <= 2**32 - 1:
|
||||
self._send_error(400, "'seed' muss zwischen 0 und 4294967295 sein",
|
||||
"invalid_request_error", "invalid_seed")
|
||||
return
|
||||
|
||||
response_format = data.get("response_format", "url")
|
||||
if response_format not in ("url", "b64_json"):
|
||||
self._send_error(400, "'response_format' muss 'url' oder 'b64_json' sein",
|
||||
"invalid_request_error", "invalid_response_format")
|
||||
return
|
||||
|
||||
# Generierung (blockt mehrere Minuten – eigener Thread-Timeout).
|
||||
self.timeout = None
|
||||
try:
|
||||
results, warning = generate_image(
|
||||
prompt.strip(), width, height, steps, guidance, seed, n)
|
||||
except (ValueError, RuntimeError) as e:
|
||||
self._send_error(503, str(e), "server_error", "image_generation_failed")
|
||||
return
|
||||
|
||||
# Antwort bauen
|
||||
host = self.headers.get("Host") or f"{HOST}:{PORT}"
|
||||
if not host.startswith(("http://", "https://")):
|
||||
host = f"http://{host}"
|
||||
items = []
|
||||
for filename in results:
|
||||
path = os.path.join(IMAGE_DIR, filename)
|
||||
item: dict = {"url": f"{host}/images/{filename}", "b64_json": None}
|
||||
if response_format == "b64_json":
|
||||
with open(path, "rb") as f:
|
||||
item["b64_json"] = base64.b64encode(f.read()).decode()
|
||||
item["url"] = None
|
||||
items.append(item)
|
||||
payload: dict = {"created": int(time.time()), "data": items}
|
||||
if warning:
|
||||
payload["router_warning"] = warning
|
||||
self._send_json(200, payload)
|
||||
|
||||
def _images_list(self) -> None:
|
||||
if not os.path.isdir(IMAGE_DIR):
|
||||
self._send_json(200, {"images": []})
|
||||
return
|
||||
entries = []
|
||||
for name in sorted(os.listdir(IMAGE_DIR), reverse=True):
|
||||
if not _image_filename_ok(name):
|
||||
continue
|
||||
path = os.path.join(IMAGE_DIR, name)
|
||||
try:
|
||||
st = os.stat(path)
|
||||
except OSError:
|
||||
continue
|
||||
entries.append({
|
||||
"name": name,
|
||||
"url": f"/images/{name}",
|
||||
"bytes": st.st_size,
|
||||
"modified": int(st.st_mtime),
|
||||
})
|
||||
self._send_json(200, {"images": entries[:200]})
|
||||
|
||||
def _image_serve(self, name: str) -> None:
|
||||
if not _image_filename_ok(name):
|
||||
self._send_error(400, "ungültiger Dateiname",
|
||||
"invalid_request_error", "invalid_filename")
|
||||
return
|
||||
path = os.path.join(IMAGE_DIR, name)
|
||||
if not os.path.isfile(path):
|
||||
self._send_error(404, "Bild nicht gefunden",
|
||||
"invalid_request_error", "not_found")
|
||||
return
|
||||
data = open(path, "rb").read()
|
||||
self._last_code = 200
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "image/png")
|
||||
self.send_header("Content-Length", str(len(data)))
|
||||
self.send_header("Cache-Control", "public, max-age=86400")
|
||||
self.send_header("Connection", "close")
|
||||
self.end_headers()
|
||||
self.wfile.write(data)
|
||||
|
||||
def _switch(self, profile: str) -> None:
|
||||
if profile not in PROFILES:
|
||||
self._send_error(400, f"unbekanntes Profil: {profile}",
|
||||
@@ -343,6 +875,42 @@ class Handler(BaseHTTPRequestHandler):
|
||||
"invalid_request_error", "unknown_model")
|
||||
return
|
||||
|
||||
# An llama.cpp weiterleiten (mit Chat-Waiting, Streaming bleibt erhalten).
|
||||
self._proxy_with_wait(body)
|
||||
|
||||
def _proxy_with_wait(self, body: bytes | None) -> None:
|
||||
"""Leitet an llama.cpp weiter, wartet aber erst, bis Qwen verfügbar ist.
|
||||
|
||||
Während eines Image-Jobs oder Profilwechsels ist Qwen down. Statt
|
||||
502 zu liefern, wartet der Request (mit Timeout), bis Qwen wieder
|
||||
bereit ist. Mehrere Chats können parallel laufen (active_chats).
|
||||
|
||||
Race-frei: Der Check auf qwen_unavailable und das Inkrement von
|
||||
active_chats sind atomar (avail_lock). Ein Image-Job/Profilwechsel
|
||||
setzt qwen_unavailable=True und wartet auf active_chats==0, BEVOR
|
||||
er Qwen stoppt – ein laufender Chat wird daher nie unterbrochen.
|
||||
"""
|
||||
deadline = time.monotonic() + CHAT_WAIT_TIMEOUT
|
||||
while True:
|
||||
with STATE.avail_lock:
|
||||
if not STATE.qwen_unavailable:
|
||||
STATE.active_chats += 1
|
||||
break
|
||||
if time.monotonic() > deadline:
|
||||
self._send_error(
|
||||
503,
|
||||
"Qwen wird neu geladen (Image-Job oder Profilwechsel), "
|
||||
"bitte später erneut",
|
||||
"server_error", "qwen_reloading")
|
||||
return
|
||||
time.sleep(0.5)
|
||||
try:
|
||||
self._proxy(body)
|
||||
finally:
|
||||
with STATE.avail_lock:
|
||||
STATE.active_chats -= 1
|
||||
|
||||
def _proxy(self, body: bytes | None) -> None:
|
||||
# An llama.cpp weiterleiten (Streaming bleibt erhalten).
|
||||
try:
|
||||
conn = http.client.HTTPConnection(UPSTREAM_HOST, UPSTREAM_PORT,
|
||||
|
||||
@@ -0,0 +1,159 @@
|
||||
#!/usr/bin/env python3
|
||||
"""FLUX.2 [klein] 4B Base – Bild-Worker.
|
||||
|
||||
Protokoll: zeilenbasiertes JSON über stdin/stdout.
|
||||
|
||||
Start: Worker gibt {"status": "ready"} aus (Modell noch NICHT geladen).
|
||||
Request: {"cmd": "generate", "prompt": ..., "width": ..., "height": ...,
|
||||
"steps": ..., "guidance": ..., "seed": ..., "output": ...}
|
||||
Antwort: {"status": "ok", "path": ..., "seconds": ..., "load_seconds": ...}
|
||||
oder {"status": "error", "message": ...}
|
||||
Request: {"cmd": "unload"} -> {"status": "ok"}
|
||||
Request: {"cmd": "status"} -> {"status": "ok", "model_loaded": bool}
|
||||
|
||||
Das Modell wird beim ersten generate geladen (bf16, cpu_offload) und auf
|
||||
Anforderung wieder entladen (VRAM freigeben). Der Prozess bleibt danach
|
||||
laufen – ohne geladenes Modell belegt er kaum Ressourcen.
|
||||
|
||||
Alle torch-/diffusers-Logs gehen nach stderr, stdout ist reines Protokoll.
|
||||
"""
|
||||
|
||||
import gc
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
import sys
|
||||
import time
|
||||
|
||||
# stderr-Logs von torch & Co. unterdrücken, bevor importiert wird
|
||||
os.environ.setdefault("DIFFUSERS_VERBOSITY", "error")
|
||||
os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
|
||||
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")
|
||||
|
||||
_pipe = None # geladene Pipeline (None = entladen)
|
||||
_load_seconds = 0.0 # Dauer des letzten Ladens
|
||||
|
||||
|
||||
def _emit(payload: dict) -> None:
|
||||
sys.stdout.write(json.dumps(payload) + "\n")
|
||||
sys.stdout.flush()
|
||||
|
||||
|
||||
def _log(msg: str) -> None:
|
||||
print(f"[image-worker] {msg}", file=sys.stderr, flush=True)
|
||||
|
||||
|
||||
def _load() -> None:
|
||||
"""Pipeline laden (bf16, CPU-Offload)."""
|
||||
global _pipe, _load_seconds
|
||||
if _pipe is not None:
|
||||
return
|
||||
import torch
|
||||
from diffusers import Flux2KleinPipeline
|
||||
|
||||
t0 = time.monotonic()
|
||||
_log(f"lade Modell aus {MODEL_DIR} ...")
|
||||
_pipe = Flux2KleinPipeline.from_pretrained(
|
||||
MODEL_DIR, torch_dtype=torch.bfloat16)
|
||||
_pipe.enable_model_cpu_offload()
|
||||
_load_seconds = time.monotonic() - t0
|
||||
_log(f"Modell geladen in {_load_seconds:.1f} s")
|
||||
|
||||
|
||||
def _unload() -> None:
|
||||
"""Pipeline entladen und VRAM freigeben."""
|
||||
global _pipe
|
||||
if _pipe is None:
|
||||
return
|
||||
t0 = time.monotonic()
|
||||
del _pipe
|
||||
_pipe = None
|
||||
gc.collect()
|
||||
try:
|
||||
import torch
|
||||
torch.cuda.empty_cache()
|
||||
except Exception:
|
||||
pass
|
||||
_log(f"Modell entladen in {time.monotonic() - t0:.1f} s")
|
||||
|
||||
|
||||
def _generate(req: dict) -> dict:
|
||||
import torch
|
||||
|
||||
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))
|
||||
seed = req.get("seed")
|
||||
output = req["output"]
|
||||
|
||||
_load()
|
||||
|
||||
t0 = time.monotonic()
|
||||
generator = None
|
||||
if seed is not None:
|
||||
generator = torch.Generator(device="cuda").manual_seed(int(seed))
|
||||
image = _pipe(
|
||||
prompt=prompt,
|
||||
height=height,
|
||||
width=width,
|
||||
guidance_scale=guidance,
|
||||
num_inference_steps=steps,
|
||||
generator=generator,
|
||||
).images[0]
|
||||
|
||||
os.makedirs(os.path.dirname(output) or ".", exist_ok=True)
|
||||
image.save(output)
|
||||
seconds = time.monotonic() - t0
|
||||
_log(f"generiert {output} in {seconds:.1f} s "
|
||||
f"({width}x{height}, {steps} steps, seed={seed})")
|
||||
return {
|
||||
"status": "ok",
|
||||
"path": output,
|
||||
"seconds": round(seconds, 2),
|
||||
"load_seconds": round(_load_seconds, 2),
|
||||
}
|
||||
|
||||
|
||||
def _handle(line: str) -> None:
|
||||
try:
|
||||
req = json.loads(line)
|
||||
except ValueError:
|
||||
_emit({"status": "error", "message": "ungültiges JSON"})
|
||||
return
|
||||
|
||||
cmd = req.get("cmd")
|
||||
try:
|
||||
if cmd == "generate":
|
||||
_emit(_generate(req))
|
||||
elif cmd == "unload":
|
||||
_unload()
|
||||
_emit({"status": "ok"})
|
||||
elif cmd == "status":
|
||||
_emit({"status": "ok", "model_loaded": _pipe is not None})
|
||||
else:
|
||||
_emit({"status": "error", "message": f"unbekanntes Kommando: {cmd}"})
|
||||
except Exception as e: # noqa: BLE001 – Fehler ans Router-Protokoll
|
||||
_log(f"Fehler bei {cmd}: {e!r}")
|
||||
_emit({"status": "error", "message": str(e)})
|
||||
|
||||
|
||||
def main() -> None:
|
||||
signal.signal(signal.SIGTERM, lambda *_: sys.exit(0))
|
||||
_emit({"status": "ready"})
|
||||
for line in sys.stdin:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
_handle(line)
|
||||
if _pipe is None and line.startswith('{"cmd": "unload"'):
|
||||
pass # Worker bleibt laufen, Modell ist entladen
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user