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
@@ -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