Add RTX 5080 FLUX hot-swap worker
This commit is contained in:
Mikei386
+85
-20
@@ -119,6 +119,8 @@ 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_WORKER_URL = os.environ.get("IMAGE_WORKER_URL", "").rstrip("/")
|
||||
IMAGE_WORKER_TOKEN = os.environ.get("IMAGE_WORKER_TOKEN", "").strip()
|
||||
IMAGE_DIR = os.environ.get(
|
||||
"IMAGE_DIR", "/opt/mike-ai/ai-profile-router/images")
|
||||
IMAGE_WORKER_LOG = os.environ.get(
|
||||
@@ -147,10 +149,10 @@ IMAGE_SIZES = {
|
||||
"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}
|
||||
# FLUX.2 Klein Distilled ist fest auf vier Schritte und Guidance 1.0
|
||||
# destilliert. Qualitätsstufen bleiben aus OpenAI-Kompatibilitätsgründen
|
||||
# akzeptiert, ändern aber bewusst nicht die offiziellen Sampling-Werte.
|
||||
IMAGE_QUALITY = {"standard": 4, "high": 4}
|
||||
IMAGE_DEFAULT_QUALITY = "standard"
|
||||
IMAGE_MAX_N = 4
|
||||
|
||||
@@ -693,11 +695,66 @@ def _worker() -> _Worker:
|
||||
if not img.worker or not img.worker.alive():
|
||||
if img.worker:
|
||||
img.worker.stop()
|
||||
img.worker = _Worker()
|
||||
img.worker = _RemoteWorker() if IMAGE_WORKER_URL else _Worker()
|
||||
img.worker.start()
|
||||
return img.worker
|
||||
|
||||
|
||||
class _RemoteWorker:
|
||||
"""Docker-Worker, dessen Lebenszyklus nur der Controller steuert."""
|
||||
|
||||
model_loaded = False
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.running = False
|
||||
|
||||
def alive(self) -> bool:
|
||||
return self.running
|
||||
|
||||
def _request(self, method: str, path: str, payload: dict | None = None,
|
||||
timeout: float = 120) -> dict:
|
||||
body = None if payload is None else json.dumps(payload).encode()
|
||||
headers = {"Authorization": f"Bearer {IMAGE_WORKER_TOKEN}"}
|
||||
if body is not None:
|
||||
headers["Content-Type"] = "application/json"
|
||||
req = urllib.request.Request(IMAGE_WORKER_URL + path, data=body,
|
||||
method=method, headers=headers)
|
||||
with urllib.request.urlopen(req, timeout=timeout) as response:
|
||||
return json.load(response)
|
||||
|
||||
def start(self) -> None:
|
||||
if not IMAGE_WORKER_TOKEN or len(IMAGE_WORKER_TOKEN) < 32:
|
||||
raise RuntimeError("Bild-Worker-Token fehlt oder ist zu kurz")
|
||||
_profile_controller_request("POST", "/workers/image/start")
|
||||
deadline = time.monotonic() + IMAGE_START_TIMEOUT
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
self._request("GET", "/health", timeout=3)
|
||||
self.running = True
|
||||
return
|
||||
except (OSError, urllib.error.URLError, TimeoutError):
|
||||
time.sleep(1)
|
||||
self.stop()
|
||||
raise RuntimeError("Bild-Worker hat nicht gestartet")
|
||||
|
||||
def request(self, payload: dict, timeout: float) -> dict:
|
||||
if payload.get("cmd") != "generate":
|
||||
raise RuntimeError("Remote-Bild-Worker erlaubt nur generate")
|
||||
clean = dict(payload)
|
||||
clean.pop("cmd", None)
|
||||
output = clean.pop("output", "")
|
||||
clean["filename"] = os.path.basename(output)
|
||||
return self._request("POST", "/generate", clean, timeout)
|
||||
|
||||
def stop(self) -> None:
|
||||
try:
|
||||
_profile_controller_request("POST", "/workers/image/stop")
|
||||
finally:
|
||||
self.running = False
|
||||
self.model_loaded = False
|
||||
RUNTIME.clear_worker("image")
|
||||
|
||||
|
||||
def _wait_upstream_down(deadline: float) -> None:
|
||||
"""Wartet, bis llama.cpp den Port freigegeben hat (VRAM frei)."""
|
||||
while time.monotonic() < deadline:
|
||||
@@ -753,6 +810,11 @@ def _wait_vram_free(threshold_mib: int = 1000,
|
||||
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)
|
||||
if PROFILE_CONTROL_URL:
|
||||
_profile_controller_request("POST", f"/profiles/{profile}/activate")
|
||||
_wait_ready(profile, time.monotonic() + SWITCH_TIMEOUT)
|
||||
RUNTIME.save(last_profile=profile, phase="idle")
|
||||
return
|
||||
try:
|
||||
proc = subprocess.run([SYSTEMCTL_BIN, "start", LLAMA_SERVICE],
|
||||
stdin=subprocess.DEVNULL,
|
||||
@@ -799,15 +861,18 @@ def generate_image(prompt: str, width: int, height: int, steps: int,
|
||||
# 1) Qwen stoppen (VRAM freigeben).
|
||||
img.phase = "stopping-qwen"
|
||||
RUNTIME.save(last_profile=profile, phase=img.phase)
|
||||
proc = subprocess.run([SYSTEMCTL_BIN, "stop", LLAMA_SERVICE],
|
||||
stdin=subprocess.DEVNULL,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT, timeout=120)
|
||||
if proc.returncode != 0:
|
||||
out = proc.stdout.decode(errors="replace").strip()
|
||||
raise RuntimeError(
|
||||
f"systemctl stop {LLAMA_SERVICE} fehlgeschlagen "
|
||||
f"(Exit {proc.returncode}): {out[-500:]}")
|
||||
if PROFILE_CONTROL_URL:
|
||||
_profile_controller_request("POST", "/inference/stop")
|
||||
else:
|
||||
proc = subprocess.run([SYSTEMCTL_BIN, "stop", LLAMA_SERVICE],
|
||||
stdin=subprocess.DEVNULL,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT, timeout=120)
|
||||
if proc.returncode != 0:
|
||||
out = proc.stdout.decode(errors="replace").strip()
|
||||
raise RuntimeError(
|
||||
f"systemctl stop {LLAMA_SERVICE} fehlgeschlagen "
|
||||
f"(Exit {proc.returncode}): {out[-500:]}")
|
||||
_wait_upstream_down(time.monotonic() + 60)
|
||||
|
||||
# 2) Worker starten (Modell wird beim ersten generate geladen).
|
||||
@@ -848,7 +913,7 @@ def generate_image(prompt: str, width: int, height: int, steps: int,
|
||||
"guidance": guidance,
|
||||
"quality": quality,
|
||||
"seconds": resp.get("seconds"),
|
||||
"model": "FLUX.2-klein-base-4B",
|
||||
"model": "FLUX.2-klein-4B",
|
||||
"created": time.strftime("%Y-%m-%dT%H:%M:%S"),
|
||||
}
|
||||
meta_path = os.path.join(IMAGE_DIR, filename[:-4] + ".json")
|
||||
@@ -1363,19 +1428,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 not 4 <= steps <= 150:
|
||||
self._send_error(400, "'steps' muss eine Ganzzahl 4..150 sein",
|
||||
if not isinstance(steps, int) or isinstance(steps, bool) or steps != 4:
|
||||
self._send_error(400, "FLUX.2 Klein Distilled erfordert 'steps'=4",
|
||||
"invalid_request_error", "invalid_steps")
|
||||
return
|
||||
guidance = data.get("guidance", 4.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 not 1.0 <= guidance <= 10.0:
|
||||
self._send_error(400, "'guidance' muss zwischen 1.0 und 10.0 sein",
|
||||
if guidance != 1.0:
|
||||
self._send_error(400, "FLUX.2 Klein Distilled erfordert 'guidance'=1.0",
|
||||
"invalid_request_error", "invalid_guidance")
|
||||
return
|
||||
|
||||
|
||||
Reference in New Issue
Block a user