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
2026-08-19 08:51:37 +02:00
parent c5d92acd93
commit 7c5bbe2ffb
15 changed files with 1796 additions and 68 deletions
+95
View File
@@ -0,0 +1,95 @@
#!/usr/bin/env python3
"""FLUX.2 [klein] 4B Base – GPU-Resident-Benchmark (OHNE CPU-Offload).
Ziel: Prüfen, ob das Modell vollständig auf der RTX 5080 (16 GB) läuft.
Messen:
- from_pretrained-Zeit
- .to("cuda")-Zeit
- VRAM (nvidia-smi + torch.cuda.memory_allocated / max_memory_allocated)
- Generierungszeit, Peak-VRAM pro Auflösung
Auflösungen: 512x512 (10 steps) → 1024x1024 (50 steps) → 1920x1088 (50 steps)
Bei OOM wird abgebrochen (CUDA-Kontext danach nicht mehr verlässlich).
"""
import os
import subprocess
import time
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
os.environ["TRANSFORMERS_VERBOSITY"] = "error"
os.environ["TOKENIZERS_PARALLELISM"] = "false"
MODEL = "/opt/mike-ai/models/FLUX.2-klein-base-4B"
PROMPT = ("A detailed photograph of a red cube on a white marble table, "
"soft studio lighting, shallow depth of field")
# (breite, hoehe, steps, seed, name)
CASES = [
(512, 512, 10, 0, "512x512-10s"),
(1024, 1024, 50, 0, "1024x1024-50s"),
(1920, 1088, 50, 0, "1920x1088-50s"),
]
def nvidia_vram() -> int:
out = subprocess.check_output(
["nvidia-smi", "--query-gpu=memory.used",
"--format=csv,noheader,nounits"]).decode().strip()
return int(out.split()[0])
def main() -> None:
import torch
print(f"torch {torch.__version__} | cuda {torch.cuda.is_available()} "
f"| {torch.cuda.get_device_name(0)}", flush=True)
from diffusers import Flux2KleinPipeline
# --- Laden (zuerst auf CPU, dann vollständig auf GPU) ---
t0 = time.monotonic()
pipe = Flux2KleinPipeline.from_pretrained(MODEL, torch_dtype=torch.bfloat16)
t_load = time.monotonic() - t0
print(f"[load] from_pretrained: {t_load:.1f} s", flush=True)
t1 = time.monotonic()
pipe.to("cuda")
torch.cuda.synchronize()
t_to = time.monotonic() - t1
print(f"[load] .to(cuda): {t_to:.1f} s", flush=True)
print(f"[load] VRAM nvidia-smi: {nvidia_vram()} MiB | "
f"torch allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB",
flush=True)
# --- Generierung ---
for width, height, steps, seed, name in CASES:
out = f"/tmp/flux-bench-{name}.png"
torch.cuda.synchronize()
torch.cuda.reset_peak_memory_stats()
t = time.monotonic()
try:
img = pipe(
prompt=PROMPT,
height=height,
width=width,
guidance_scale=4.0,
num_inference_steps=steps,
generator=torch.Generator(device="cuda").manual_seed(seed),
).images[0]
dt = time.monotonic() - t
img.save(out)
peak = torch.cuda.max_memory_allocated() / 1e9
print(f"[gen] {name}: {dt:.1f} s | peak torch {peak:.2f} GB | "
f"nvidia-smi {nvidia_vram()} MiB | {out}", flush=True)
except Exception as e: # noqa: BLE001
dt = time.monotonic() - t
print(f"[gen] {name}: FEHLER nach {dt:.1f} s: {e!r}", flush=True)
if "out of memory" in str(e).lower():
print("[gen] OOM – Abbruch, größere Auflösungen nicht getestet",
flush=True)
break
print("DONE", flush=True)
if __name__ == "__main__":
main()