Frank Schwenk
83d4468b69
Downscale to 2000px JPEG before rembg, random blend opacity, timeout retries, per-variant remix prefill, lightbox, and longer SFTP idle. Co-authored-by: Cursor <cursoragent@cursor.com>
550 lines
20 KiB
Python
550 lines
20 KiB
Python
"""Image compose pipeline — ported from `make_random.py`.
|
|
|
|
Core building blocks (gmic filter application, blending, filter-retry with
|
|
timeout) are kept as close to the reference script as possible. On top of
|
|
that this module adds job-directory bookkeeping, kept intermediates, and a
|
|
manifest format the web app can read/append to (for remix).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import random
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
import uuid
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any, Iterator
|
|
|
|
from . import config
|
|
|
|
logger = logging.getLogger("livef12.pipeline")
|
|
|
|
_ANSI_RE = re.compile(r"\x1b\[[0-9;]*m")
|
|
|
|
# Optional override for run_gmic timeout (used by the long-retry pass).
|
|
_filter_timeout_override: int | None = None
|
|
|
|
|
|
class PipelineError(RuntimeError):
|
|
"""Raised when a job (or a single variant) cannot be completed."""
|
|
|
|
|
|
class FilterNotFoundError(PipelineError):
|
|
"""Raised when no working filter could be picked after retries."""
|
|
|
|
|
|
class FilterTimeoutError(PipelineError):
|
|
"""Raised when a gmic step hits the active filter timeout."""
|
|
|
|
|
|
@dataclass
|
|
class FilterAssets:
|
|
background_names: list[str]
|
|
foreground_names: list[str]
|
|
blend_modes: list[str]
|
|
commands: dict[str, str]
|
|
|
|
|
|
# --- low-level helpers, ported from make_random.py ------------------------
|
|
|
|
|
|
def strip_ansi(text: str) -> str:
|
|
return _ANSI_RE.sub("", text)
|
|
|
|
|
|
def _nice(cmd: list[str]) -> list[str]:
|
|
if config.NICE_LEVEL <= 0:
|
|
return cmd
|
|
return ["nice", "-n", str(config.NICE_LEVEL), *cmd]
|
|
|
|
|
|
def _active_filter_timeout() -> int:
|
|
return _filter_timeout_override if _filter_timeout_override is not None else config.FILTER_TIMEOUT
|
|
|
|
|
|
@contextmanager
|
|
def filter_timeout(seconds: int) -> Iterator[None]:
|
|
"""Temporarily override FILTER_TIMEOUT for run_gmic (long-retry pass)."""
|
|
global _filter_timeout_override
|
|
previous = _filter_timeout_override
|
|
_filter_timeout_override = seconds
|
|
try:
|
|
yield
|
|
finally:
|
|
_filter_timeout_override = previous
|
|
|
|
|
|
def _magick_bin() -> str:
|
|
if config.MAGICK_BIN:
|
|
return config.MAGICK_BIN
|
|
for name in ("magick", "convert"):
|
|
if shutil.which(name):
|
|
return name
|
|
raise PipelineError("ImageMagick nicht installiert (magick/convert fehlt)")
|
|
|
|
|
|
def load_lines(path: Path) -> list[str]:
|
|
if not path.exists():
|
|
raise FileNotFoundError(f"Datei nicht gefunden: {path}")
|
|
lines = [line.strip() for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
|
|
if not lines:
|
|
raise ValueError(f"Datei ist leer: {path}")
|
|
return lines
|
|
|
|
|
|
def load_filter_commands() -> dict[str, str]:
|
|
data = json.loads(config.FILTERS_JSON.read_text(encoding="utf-8"))
|
|
return {item["plain_name"]: item["full_command"] for item in data}
|
|
|
|
|
|
def load_assets() -> FilterAssets:
|
|
"""Load background/foreground/blend-mode lists and cross-check against
|
|
filters.json, dropping any names without a known command (mirrors the
|
|
warning+filter behaviour in make_random.py's main())."""
|
|
bg_names = load_lines(config.BACKGROUND_FILE)
|
|
fg_names = load_lines(config.FOREGROUND_FILE)
|
|
blend_modes = load_lines(config.BLEND_MODES_FILE)
|
|
commands = load_filter_commands()
|
|
|
|
missing_bg = [name for name in bg_names if name not in commands]
|
|
missing_fg = [name for name in fg_names if name not in commands]
|
|
if missing_bg:
|
|
logger.warning("%d background filter names have no command, dropping them", len(missing_bg))
|
|
bg_names = [name for name in bg_names if name in commands]
|
|
if missing_fg:
|
|
logger.warning("%d foreground filter names have no command, dropping them", len(missing_fg))
|
|
fg_names = [name for name in fg_names if name in commands]
|
|
if not bg_names or not fg_names:
|
|
raise PipelineError("Keine gueltigen Filter in background/foreground Listen")
|
|
|
|
return FilterAssets(background_names=bg_names, foreground_names=fg_names, blend_modes=blend_modes, commands=commands)
|
|
|
|
|
|
def run_gmic(args: list[str], output_image: Path) -> tuple[bool, str, float]:
|
|
output_image.parent.mkdir(parents=True, exist_ok=True)
|
|
timeout = _active_filter_timeout()
|
|
cmd = _nice([config.GMIC_BIN, *args, "-o", str(output_image)])
|
|
start = time.monotonic()
|
|
try:
|
|
proc = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
|
|
except subprocess.TimeoutExpired:
|
|
return False, f"Timeout nach {timeout}s", time.monotonic() - start
|
|
elapsed = time.monotonic() - start
|
|
if proc.returncode != 0:
|
|
err = strip_ansi((proc.stderr or proc.stdout or "").strip())
|
|
err = err.splitlines()[-1] if err else f"Exit code {proc.returncode}"
|
|
return False, err[:500], elapsed
|
|
if not output_image.exists() or output_image.stat().st_size == 0:
|
|
return False, "Keine Ausgabedatei erzeugt", elapsed
|
|
return True, "", elapsed
|
|
|
|
|
|
def _raise_on_gmic_failure(ok: bool, err: str, label: str) -> None:
|
|
if ok:
|
|
return
|
|
if err.startswith("Timeout nach"):
|
|
raise FilterTimeoutError(f"{label}: {err}")
|
|
raise PipelineError(f"{label}: {err}")
|
|
|
|
|
|
def gmic_filter_args(image: Path, full_command: str) -> list[str]:
|
|
if " " in full_command:
|
|
name, args = full_command.split(" ", 1)
|
|
return [str(image), name, args]
|
|
return [str(image), full_command]
|
|
|
|
|
|
def apply_filter(image: Path, full_command: str, output_image: Path) -> tuple[bool, str]:
|
|
ok, err, _ = run_gmic(gmic_filter_args(image, full_command), output_image)
|
|
return ok, err
|
|
|
|
|
|
def apply_filter_chain(image: Path, commands: tuple[str, ...], output_image: Path, tmp_dir: Path, prefix: str) -> tuple[bool, str]:
|
|
current = image
|
|
for i, command in enumerate(commands):
|
|
target = output_image if i == len(commands) - 1 else tmp_dir / f"{prefix}_post_{i}.png"
|
|
ok, err = apply_filter(current, command, target)
|
|
if not ok:
|
|
return False, err
|
|
current = target
|
|
return True, ""
|
|
|
|
|
|
def blend_layers(base: Path, overlay: Path, mode: str, output_image: Path) -> tuple[bool, str]:
|
|
ok, err, _ = run_gmic([str(base), str(overlay), "blend", f"{mode},{config.BLEND_OPACITY}"], output_image)
|
|
return ok, err
|
|
|
|
|
|
def blend_layers_opacity(base: Path, overlay: Path, mode: str, opacity: str, output_image: Path) -> tuple[bool, str]:
|
|
ok, err, _ = run_gmic([str(base), str(overlay), "blend", f"{mode},{opacity}"], output_image)
|
|
return ok, err
|
|
|
|
|
|
def alpha_composite(base: Path, overlay: Path, output_image: Path) -> tuple[bool, str]:
|
|
ok, err, _ = run_gmic([str(base), str(overlay), "blend", "alpha"], output_image)
|
|
return ok, err
|
|
|
|
|
|
def run_rembg(input_image: Path, output_image: Path) -> tuple[bool, str]:
|
|
output_image.parent.mkdir(parents=True, exist_ok=True)
|
|
cmd = _nice([sys.executable, "-m", "app.rembg_cli", str(input_image), str(output_image)])
|
|
try:
|
|
proc = subprocess.run(cmd, capture_output=True, text=True, timeout=config.REMBG_TIMEOUT)
|
|
except subprocess.TimeoutExpired:
|
|
return False, f"rembg Timeout nach {config.REMBG_TIMEOUT}s"
|
|
if proc.returncode != 0 or not output_image.exists() or output_image.stat().st_size == 0:
|
|
lines = strip_ansi((proc.stderr or proc.stdout or "").strip()).splitlines()
|
|
return False, (lines[-1] if lines else f"Exit code {proc.returncode}")
|
|
return True, ""
|
|
|
|
|
|
def preprocess_original(paths: JobPaths) -> JobPaths:
|
|
"""Downscale to MAX_EDGE_PX and normalize to sRGB JPEG as `original.jpg`.
|
|
|
|
Replaces any non-jpg original in the job dir. Side-effect free for the
|
|
inbox copy — only mutates jobs/<id>/.
|
|
"""
|
|
src = paths.original
|
|
if not src.exists():
|
|
raise PipelineError(f"Original fehlt: {src}")
|
|
|
|
dest = paths.root / "original.jpg"
|
|
resize = f"{config.MAX_EDGE_PX}x{config.MAX_EDGE_PX}>"
|
|
# Write to a temp name first so we can replace an existing original.jpg
|
|
# in-place without reading/writing the same path.
|
|
tmp = paths.root / f".original_pre_{uuid.uuid4().hex[:8]}.jpg"
|
|
cmd = _nice(
|
|
[
|
|
_magick_bin(),
|
|
str(src),
|
|
"-auto-orient",
|
|
"-resize",
|
|
resize,
|
|
"-colorspace",
|
|
"sRGB",
|
|
"-quality",
|
|
"92",
|
|
str(tmp),
|
|
]
|
|
)
|
|
try:
|
|
proc = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
|
|
except subprocess.TimeoutExpired as exc:
|
|
tmp.unlink(missing_ok=True)
|
|
raise PipelineError("Preprocess Timeout nach 120s") from exc
|
|
if proc.returncode != 0 or not tmp.exists() or tmp.stat().st_size == 0:
|
|
tmp.unlink(missing_ok=True)
|
|
err = strip_ansi((proc.stderr or proc.stdout or "").strip())
|
|
err = err.splitlines()[-1] if err else f"Exit code {proc.returncode}"
|
|
raise PipelineError(f"Preprocess fehlgeschlagen: {err}")
|
|
|
|
tmp.replace(dest)
|
|
if src.resolve() != dest.resolve():
|
|
src.unlink(missing_ok=True)
|
|
|
|
logger.info(
|
|
"preprocessed %s -> %s (max_edge=%d)",
|
|
src.name,
|
|
dest.name,
|
|
config.MAX_EDGE_PX,
|
|
)
|
|
return JobPaths(
|
|
root=paths.root,
|
|
original=dest,
|
|
rembg=paths.rembg,
|
|
intermediates=paths.intermediates,
|
|
variants=paths.variants,
|
|
manifest=paths.manifest,
|
|
status=paths.status,
|
|
)
|
|
|
|
|
|
def pick_working_filter(
|
|
names: list[str],
|
|
commands: dict[str, str],
|
|
image: Path,
|
|
tmp_dir: Path,
|
|
label: str,
|
|
rng: random.Random,
|
|
) -> tuple[str, str]:
|
|
tried: set[str] = set()
|
|
for _ in range(config.MAX_FILTER_ATTEMPTS):
|
|
candidates = [n for n in names if n not in tried]
|
|
if not candidates:
|
|
break
|
|
name = rng.choice(candidates)
|
|
tried.add(name)
|
|
command = commands.get(name)
|
|
if not command:
|
|
continue
|
|
probe = tmp_dir / f"probe_{label}_{uuid.uuid4().hex[:8]}.png"
|
|
ok, err = apply_filter(image, command, probe)
|
|
probe.unlink(missing_ok=True)
|
|
if ok:
|
|
return name, command
|
|
logger.info("skip %s filter %r: %s", label, name, err)
|
|
raise FilterNotFoundError(f"Kein funktionierender {label}-Filter nach {config.MAX_FILTER_ATTEMPTS} Versuchen")
|
|
|
|
|
|
# --- job-level orchestration ------------------------------------------------
|
|
|
|
|
|
def now_iso() -> str:
|
|
return datetime.now(timezone.utc).isoformat(timespec="seconds")
|
|
|
|
|
|
def new_job_id() -> str:
|
|
stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
|
|
return f"{stamp}-{uuid.uuid4().hex[:6]}"
|
|
|
|
|
|
@dataclass
|
|
class JobPaths:
|
|
root: Path
|
|
original: Path
|
|
rembg: Path
|
|
intermediates: Path
|
|
variants: Path
|
|
manifest: Path
|
|
status: Path
|
|
|
|
|
|
def job_paths(job_id: str, original_suffix: str = ".jpg") -> JobPaths:
|
|
root = config.JOBS_DIR / job_id
|
|
return JobPaths(
|
|
root=root,
|
|
original=root / f"original{original_suffix}",
|
|
rembg=root / "rembg.png",
|
|
intermediates=root / "intermediates",
|
|
variants=root / "variants",
|
|
manifest=root / "manifest.json",
|
|
status=root / "status.json",
|
|
)
|
|
|
|
|
|
def read_status(job_id: str) -> dict[str, Any] | None:
|
|
paths = job_paths(job_id)
|
|
if not paths.status.exists():
|
|
return None
|
|
return json.loads(paths.status.read_text(encoding="utf-8"))
|
|
|
|
|
|
def write_status(paths: JobPaths, status: str, **extra: Any) -> None:
|
|
data = {}
|
|
if paths.status.exists():
|
|
try:
|
|
data = json.loads(paths.status.read_text(encoding="utf-8"))
|
|
except (OSError, ValueError):
|
|
data = {}
|
|
data["status"] = status
|
|
data["updated_at"] = now_iso()
|
|
data.setdefault("created_at", data["updated_at"])
|
|
data.update(extra)
|
|
paths.status.write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")
|
|
|
|
|
|
def read_manifest(job_id: str) -> dict[str, Any] | None:
|
|
paths = job_paths(job_id)
|
|
if not paths.manifest.exists():
|
|
return None
|
|
return json.loads(paths.manifest.read_text(encoding="utf-8"))
|
|
|
|
|
|
def write_manifest(paths: JobPaths, manifest: dict[str, Any]) -> None:
|
|
manifest["updated_at"] = now_iso()
|
|
paths.manifest.write_text(json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8")
|
|
|
|
|
|
def next_variant_id(manifest: dict[str, Any]) -> str:
|
|
existing = {v["id"] for v in manifest.get("variants", [])}
|
|
i = len(manifest.get("variants", [])) + 1
|
|
while f"v{i}" in existing:
|
|
i += 1
|
|
return f"v{i}"
|
|
|
|
|
|
def compose_variant(
|
|
paths: JobPaths,
|
|
assets: FilterAssets,
|
|
*,
|
|
variant_id: str,
|
|
source: str,
|
|
bg_name: str | None = None,
|
|
bg_mode: str | None = None,
|
|
fg_name: str | None = None,
|
|
fg_mode: str | None = None,
|
|
opacity: str | None = None,
|
|
rng: random.Random | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Compose one variant image from the job's cached original + rembg.
|
|
|
|
If bg_name/fg_name/bg_mode/fg_mode are given (remix path) they are used
|
|
directly. Otherwise a random working filter is picked with retries,
|
|
exactly like make_random.py's compose_one().
|
|
"""
|
|
rng = rng or random.Random()
|
|
tmp_dir = paths.intermediates
|
|
tmp_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
bg_mode = bg_mode or rng.choice(assets.blend_modes)
|
|
fg_mode = fg_mode or rng.choice(assets.blend_modes)
|
|
if opacity is None:
|
|
lo = min(config.BLEND_OPACITY_MIN, config.BLEND_OPACITY_MAX)
|
|
hi = max(config.BLEND_OPACITY_MIN, config.BLEND_OPACITY_MAX)
|
|
opacity = f"{rng.randint(lo, hi)}%"
|
|
|
|
if bg_name:
|
|
bg_command = assets.commands.get(bg_name)
|
|
if not bg_command:
|
|
raise PipelineError(f"Unbekannter Background-Filter: {bg_name}")
|
|
else:
|
|
bg_name, bg_command = pick_working_filter(assets.background_names, assets.commands, paths.original, tmp_dir, "bg", rng)
|
|
|
|
if fg_name:
|
|
fg_command = assets.commands.get(fg_name)
|
|
if not fg_command:
|
|
raise PipelineError(f"Unbekannter Foreground-Filter: {fg_name}")
|
|
else:
|
|
fg_name, fg_command = pick_working_filter(assets.foreground_names, assets.commands, paths.rembg, tmp_dir, "fg", rng)
|
|
|
|
p = {
|
|
"bg_filtered": tmp_dir / f"{variant_id}_bg_filtered.png",
|
|
"step1": tmp_dir / f"{variant_id}_bg_blend.png",
|
|
"step2": tmp_dir / f"{variant_id}_rembg_alpha.png",
|
|
"fg_filtered": tmp_dir / f"{variant_id}_fg_filtered.png",
|
|
"composed": tmp_dir / f"{variant_id}_composed.png",
|
|
"final": paths.variants / f"{variant_id}.png",
|
|
}
|
|
paths.variants.mkdir(parents=True, exist_ok=True)
|
|
|
|
ok, err = apply_filter(paths.original, bg_command, p["bg_filtered"])
|
|
_raise_on_gmic_failure(ok, err, "Background-Filter fehlgeschlagen")
|
|
|
|
ok, err = blend_layers_opacity(paths.original, p["bg_filtered"], bg_mode, opacity, p["step1"])
|
|
_raise_on_gmic_failure(ok, err, "Background-Blend fehlgeschlagen")
|
|
|
|
ok, err = alpha_composite(p["step1"], paths.rembg, p["step2"])
|
|
_raise_on_gmic_failure(ok, err, "Rembg-Alpha fehlgeschlagen")
|
|
|
|
ok, err = apply_filter(paths.rembg, fg_command, p["fg_filtered"])
|
|
_raise_on_gmic_failure(ok, err, "Foreground-Filter fehlgeschlagen")
|
|
|
|
ok, err = blend_layers_opacity(p["step2"], p["fg_filtered"], fg_mode, opacity, p["composed"])
|
|
_raise_on_gmic_failure(ok, err, "Foreground-Blend fehlgeschlagen")
|
|
|
|
ok, err = apply_filter_chain(p["composed"], config.POST_FILTERS, p["final"], tmp_dir, variant_id)
|
|
_raise_on_gmic_failure(ok, err, "Post-Processing fehlgeschlagen")
|
|
|
|
rel = lambda path: str(path.relative_to(paths.root))
|
|
return {
|
|
"id": variant_id,
|
|
"source": source,
|
|
"file": rel(p["final"]),
|
|
"background_filter": bg_name,
|
|
"background_command": bg_command,
|
|
"background_blend": bg_mode,
|
|
"foreground_filter": fg_name,
|
|
"foreground_command": fg_command,
|
|
"foreground_blend": fg_mode,
|
|
"blend_opacity": opacity,
|
|
"post_filters": list(config.POST_FILTERS),
|
|
"created_at": now_iso(),
|
|
"intermediates": {
|
|
"bg_filtered": rel(p["bg_filtered"]),
|
|
"bg_blend": rel(p["step1"]),
|
|
"rembg_alpha": rel(p["step2"]),
|
|
"fg_filtered": rel(p["fg_filtered"]),
|
|
"composed": rel(p["composed"]),
|
|
},
|
|
}
|
|
|
|
|
|
def process_job(job_id: str, source_path: Path, original_suffix: str) -> None:
|
|
"""Full pipeline for a freshly ingested inbox file: preprocess, rembg
|
|
once, generate OUTPUT_COUNT variants, write manifest + status.
|
|
|
|
`source_path` must already be a private copy (jobs/<id>/original.*) —
|
|
callers (worker.py) are responsible for copying out of inbox first, so
|
|
the inbox file itself is never touched here.
|
|
|
|
Variants that hit FILTER_TIMEOUT during the normal pass are retried
|
|
once at the end with FILTER_TIMEOUT_LONG. Filter probes in
|
|
pick_working_filter stay on the short timeout (skip, don't escalate).
|
|
"""
|
|
paths = job_paths(job_id, original_suffix)
|
|
paths.root.mkdir(parents=True, exist_ok=True)
|
|
paths.intermediates.mkdir(parents=True, exist_ok=True)
|
|
paths.variants.mkdir(parents=True, exist_ok=True)
|
|
|
|
write_status(paths, "processing", source_file=str(source_path.name))
|
|
|
|
try:
|
|
paths = preprocess_original(paths)
|
|
|
|
manifest: dict[str, Any] = {
|
|
"job_id": job_id,
|
|
"original_file": paths.original.name,
|
|
"rembg_file": paths.rembg.name,
|
|
"created_at": now_iso(),
|
|
"variants": [],
|
|
}
|
|
|
|
ok, err = run_rembg(paths.original, paths.rembg)
|
|
if not ok:
|
|
raise PipelineError(f"rembg fehlgeschlagen: {err}")
|
|
|
|
assets = load_assets()
|
|
rng = random.Random()
|
|
variant_errors: list[str] = []
|
|
timeout_retries: list[str] = []
|
|
|
|
for i in range(1, config.OUTPUT_COUNT + 1):
|
|
variant_id = f"v{i}"
|
|
try:
|
|
entry = compose_variant(paths, assets, variant_id=variant_id, source="auto", rng=rng)
|
|
manifest["variants"].append(entry)
|
|
write_manifest(paths, manifest)
|
|
except FilterTimeoutError as exc:
|
|
logger.warning("[%s] variant %s timed out (will retry long): %s", job_id, variant_id, exc)
|
|
timeout_retries.append(variant_id)
|
|
except PipelineError as exc:
|
|
logger.error("[%s] variant %s failed: %s", job_id, variant_id, exc)
|
|
variant_errors.append(f"{variant_id}: {exc}")
|
|
|
|
if timeout_retries:
|
|
logger.info(
|
|
"[%s] long-retry %d variant(s) with timeout=%ds: %s",
|
|
job_id,
|
|
len(timeout_retries),
|
|
config.FILTER_TIMEOUT_LONG,
|
|
", ".join(timeout_retries),
|
|
)
|
|
with filter_timeout(config.FILTER_TIMEOUT_LONG):
|
|
for variant_id in timeout_retries:
|
|
try:
|
|
entry = compose_variant(paths, assets, variant_id=variant_id, source="auto", rng=rng)
|
|
manifest["variants"].append(entry)
|
|
write_manifest(paths, manifest)
|
|
except PipelineError as exc:
|
|
logger.error("[%s] variant %s long-retry failed: %s", job_id, variant_id, exc)
|
|
variant_errors.append(f"{variant_id}: {exc}")
|
|
|
|
if not manifest["variants"]:
|
|
raise PipelineError("Keine Variante erfolgreich erzeugt: " + "; ".join(variant_errors))
|
|
|
|
write_status(paths, "done", variant_errors=variant_errors)
|
|
except PipelineError as exc:
|
|
logger.error("[%s] job failed: %s", job_id, exc)
|
|
write_status(paths, "error", error=str(exc))
|
|
except Exception as exc: # noqa: BLE001 - keep the worker loop alive
|
|
logger.exception("[%s] unexpected error", job_id)
|
|
write_status(paths, "error", error=f"Unerwarteter Fehler: {exc}")
|