From 40a09869682510f54c0efd7f26aa1a044def89ec Mon Sep 17 00:00:00 2001 From: viborc Date: Thu, 18 Jun 2026 18:46:38 +0200 Subject: [PATCH] feat(setup): improve demon-setup terminal UX with downstream-safe color Add a tiny stdlib styler (acestep/ui.py) and use it to colorize the human-facing output of demon-setup and the TensorRT build matrix: green status ([ok]), yellow/red warn/fail, light-orange phase banners and titles, green summary labels, and a blue/yellow launch block. Plain color over the existing text - no structural or wording changes. Color is emitted ONLY when the target stream is an interactive TTY. The gate is pure sys.stdout.isatty(), with NO_COLOR / TERM=dumb honored and intentionally no FORCE_COLOR override. On any pipe, file, or captured subprocess - i.e. every test and every downstream consumer - the gate is off and style() returns its input unchanged, so emitted bytes are byte-identical to before. SGR codes are ASCII-only (no encoding risk); no text, labels, or phase numbers are altered. No new dependency (stdlib os/sys only). obs.py structured logs and build_report.csv are untouched. Verified: the three setup/build unit tests pass unchanged, and a piped byte-diff against pre-change output is identical (apart from the live 'disk: NNN GB free' number). --- acestep/engine/trt/build.py | 16 +++++++---- acestep/setup.py | 32 +++++++++++---------- acestep/ui.py | 55 +++++++++++++++++++++++++++++++++++++ 3 files changed, 82 insertions(+), 21 deletions(-) create mode 100644 acestep/ui.py diff --git a/acestep/engine/trt/build.py b/acestep/engine/trt/build.py index 33ccc944..383f1ce4 100644 --- a/acestep/engine/trt/build.py +++ b/acestep/engine/trt/build.py @@ -1056,6 +1056,7 @@ def _print_matrix(durations, build_vae, build_decoder, output_dir, batch_max, # Defer to TRTBuildConfig.engine_filename so this preview can't # drift from what _build_decoder_engine writes. from .export import TRTBuildConfig + from acestep.ui import style def _decoder_dir_name(dur: int) -> str: cfg = TRTBuildConfig( @@ -1094,10 +1095,10 @@ def _decoder_dir_name(dur: int) -> str: engine_file = os.path.join(output_dir, dir_name, f"{dir_name}.engine") if os.path.exists(engine_file): size_mb = os.path.getsize(engine_file) / 1e6 - lines.append(f" [exists] {label} ({size_mb:.0f} MB)") + lines.append(f" {style('[exists]', 'dim')} {label} ({size_mb:.0f} MB)") to_skip += 1 else: - lines.append(f" [build] {label}") + lines.append(f" {style('[build]', 'green')} {label}") to_build += 1 print(f"\nBuild matrix: {to_build} to build, {to_skip} existing (skipped)") @@ -1109,11 +1110,14 @@ def _decoder_dir_name(dur: int) -> str: def _print_summary(results, output_dir): """Print build summary and list engines on disk.""" - print(f"\n{'=' * 60}") - print("BUILD SUMMARY") - print(f"{'=' * 60}") + from acestep.ui import style + print(f"\n{style('=' * 60, 'orange')}") + print(style("BUILD SUMMARY", "orange_b")) + print(style('=' * 60, "orange")) + _verb = {"OK": "ok", "FAILED": "fail", "SKIPPED": "dim"} for label, path, elapsed, status in results: - print(f" {status:7s} {elapsed:6.0f}s {label}") + _s = style(f"{status:7s}", _verb.get(status.strip(), "dim")) + print(f" {_s} {elapsed:6.0f}s {label}") failures = sum(1 for _, _, _, s in results if s == "FAILED") if failures: diff --git a/acestep/setup.py b/acestep/setup.py index f52a1c98..7c37bfec 100644 --- a/acestep/setup.py +++ b/acestep/setup.py @@ -34,6 +34,8 @@ import subprocess import sys +from acestep.ui import style + # Reused by the demo server's preflight banner so the "how do I fix # this" hint can't drift from what this script actually does. SETUP_COMMAND = "uv run demon-setup" @@ -84,22 +86,22 @@ def _env_skip_loras() -> bool: def _ok(msg: str) -> None: - print(f" [ok] {msg}") + print(f" {style('[ok]', 'ok')} {msg}") def _warn(msg: str) -> None: - print(f" [warn] {msg}") + print(f" {style('[warn]', 'warn')} {msg}") def _fail(msg: str) -> None: - print(f" [FAIL] {msg}") + print(f" {style('[FAIL]', 'fail')} {msg}") def _header(title: str) -> None: print() - print("=" * 64) - print(f" {title}") - print("=" * 64) + print(style("=" * 64, "orange")) + print(f" {style(title, 'orange_b')}") + print(style("=" * 64, "orange")) def _doctor() -> bool: @@ -269,9 +271,9 @@ def _build_engines(extra_args: list[str]) -> bool: from acestep.paths import trt_engines_dir _header("4/4 TensorRT engines (minimal preset)") - print(f" destination: {trt_engines_dir()}") - print(" set: 60s decoder + 60s VAE encode/decode + fixed 1s windowed " - "VAE decode") + print(style(f" destination: {trt_engines_dir()}", "green")) + print(style(" set: 60s decoder + 60s VAE encode/decode + fixed 1s windowed " + "VAE decode", "green")) print(" ONNX is fetched prebuilt from huggingface.co/daydreamlive/" "demon-onnx;") print(" expect a few minutes on a recent GPU (under 2 minutes of TRT") @@ -298,7 +300,7 @@ def _summary(*, engines_skipped: bool) -> None: from acestep.paths import models_dir, trt_engines_dir _header("Setup complete") - print(f" models dir: {models_dir()}") + print(f" {style('models dir', 'green')}: {models_dir()}") trt_dir = trt_engines_dir() if trt_dir.is_dir(): engines = sorted( @@ -307,13 +309,13 @@ def _summary(*, engines_skipped: bool) -> None: and (d / f"{d.name}.engine").exists() ) if engines: - print(" engines:") + print(f" {style('engines', 'green')}:") for name in engines: print(f" {name}") from acestep.paths import discover_loras n_loras = len(discover_loras()) if n_loras: - print(f" loras: {n_loras} in the library") + print(f" {style('loras', 'green')}: {n_loras} in the library") if engines_skipped: print("\n Engines were skipped (--skip-engines). The demo's " "default TRT mode") @@ -322,9 +324,9 @@ def _summary(*, engines_skipped: bool) -> None: print(f" or launch with `-- --accel compile` (slow warmup, no " f"engines needed).") print() - print(" Launch the web demo:") - print(f" {DEMO_COMMAND}") - print(" then open http://localhost:6660") + print(f" {style('Launch the web demo:', 'blue')}") + print(f" {style(DEMO_COMMAND, 'yellow')}") + print(f" {style('then open', 'blue')} {style('http://localhost:6660', 'yellow')}") print() diff --git a/acestep/ui.py b/acestep/ui.py new file mode 100644 index 00000000..94c053cd --- /dev/null +++ b/acestep/ui.py @@ -0,0 +1,55 @@ +"""Minimal terminal styling for human-facing CLI output (demon-setup and +the TensorRT build matrix). + +Color is added ONLY when the target stream is an interactive terminal. +On any pipe, file, or captured subprocess — i.e. every test and every +downstream consumer — the gate is off and :func:`style` returns its input +unchanged, so the emitted bytes are byte-identical to the uncolored +output. There is intentionally no ``FORCE_COLOR`` override (it would +inject ANSI into captured pipes and break that guarantee); ``NO_COLOR`` +and ``TERM=dumb`` hard-disable per no-color.org. + +SGR codes are pure ASCII, so they can never raise ``UnicodeEncodeError`` +and are only ever written to a real terminal anyway. Stdlib only. +""" + +from __future__ import annotations + +import os +import sys + +_SGR = { + "ok": "\x1b[32m", # green + "green": "\x1b[32m", # green (alias) + "warn": "\x1b[33m", # yellow + "fail": "\x1b[1;31m", # bold red + "yellow": "\x1b[33m", # yellow + "blue": "\x1b[34m", # blue + "magenta": "\x1b[35m", # magenta / purple + "orange": "\x1b[38;2;255;184;108m", # light orange (truecolor #FFB86C) + "orange_b": "\x1b[1;38;2;255;184;108m", # bold light orange + "header": "\x1b[1;35m",# bold magenta (phase banners + titles) + "accent": "\x1b[1m", # bold + "rule": "\x1b[2m", # dim + "dim": "\x1b[2m", # dim secondary text +} +_RESET = "\x1b[0m" + + +def should_color(stream=None) -> bool: + """True only on an interactive TTY; ``NO_COLOR`` / ``TERM=dumb`` win.""" + stream = sys.stdout if stream is None else stream + if os.environ.get("NO_COLOR"): + return False + if os.environ.get("TERM") == "dumb": + return False + return bool(getattr(stream, "isatty", lambda: False)()) + + +def style(text: str, token: str, stream=None) -> str: + """Wrap ``text`` in the SGR code for ``token`` on a TTY; return ``text`` + unchanged otherwise. Off-TTY this is a pure no-op, so output stays + byte-identical to the uncolored form.""" + if not should_color(stream): + return text + return f"{_SGR.get(token, '')}{text}{_RESET}"