mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
fix(ai): advance the progress bar during upscale and background removal (#608)
RealESRGAN's enhance() and rembg's remove() run in one opaque call, so the progress bar froze at 30% for the whole inference. Add a time-based heartbeat that advances the bar in a background thread while the model runs and stops when it returns, so the bar moves instead of freezing. Verified end to end on a GPU box (forced CPU): a 55s upscale emitted 26 steady ticks then completed; background removal too. Fixes #591
This commit is contained in:
@@ -0,0 +1,54 @@
|
||||
import { spawnSync } from "node:child_process";
|
||||
import { dirname, resolve } from "node:path";
|
||||
import { fileURLToPath } from "node:url";
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { hasPython, pythonBin } from "../../helpers/python-gate.js";
|
||||
|
||||
const here = dirname(fileURLToPath(import.meta.url));
|
||||
const pyDir = resolve(here, "../../../packages/ai/python");
|
||||
|
||||
// The heartbeat is pure stdlib (threading), so this runs anywhere python3 is
|
||||
// present, no AI bundle required. It guards the contract the upscale / remove-bg
|
||||
// scripts rely on: advance the bar during an opaque model call, stop the moment
|
||||
// it returns, pass the return value through, and propagate exceptions (#591).
|
||||
const SCRIPT = `
|
||||
import sys, time
|
||||
sys.path.insert(0, ${JSON.stringify(pyDir)})
|
||||
from progress_heartbeat import run_with_heartbeat
|
||||
|
||||
emitted = []
|
||||
def emit(pct, stage): emitted.append(pct)
|
||||
|
||||
def slow(): time.sleep(1.6); return "RESULT"
|
||||
value = run_with_heartbeat(slow, emit, 30, 80, "Working", interval=0.5)
|
||||
assert value == "RESULT", "must return the wrapped call's value"
|
||||
assert emitted, "must emit at least one tick during a slow call"
|
||||
assert emitted[0] == 31, ("first tick is start+1", emitted)
|
||||
assert emitted == sorted(emitted), ("monotonic", emitted)
|
||||
assert max(emitted) <= 79, ("never reaches end", emitted)
|
||||
|
||||
before = len(emitted)
|
||||
time.sleep(1.0)
|
||||
assert len(emitted) == before, "heartbeat must stop once the call returns"
|
||||
|
||||
def boom(): raise ValueError("boom")
|
||||
try:
|
||||
run_with_heartbeat(boom, emit, 30, 80, "Working", interval=0.2)
|
||||
raise SystemExit("exception was swallowed")
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
print("OK")
|
||||
`;
|
||||
|
||||
describe.skipIf(!hasPython)("progress heartbeat", () => {
|
||||
it("advances during a slow call, stops after, returns value, propagates errors", () => {
|
||||
const res = spawnSync(pythonBin as string, ["-c", SCRIPT], {
|
||||
encoding: "utf8",
|
||||
timeout: 30_000,
|
||||
});
|
||||
expect(res.stderr).toBe("");
|
||||
expect(res.status).toBe(0);
|
||||
expect(res.stdout).toContain("OK");
|
||||
}, 40_000);
|
||||
});
|
||||
Reference in New Issue
Block a user