diff --git a/packages/ai/src/background-removal.ts b/packages/ai/src/background-removal.ts index 387df72d..2a82b97f 100644 --- a/packages/ai/src/background-removal.ts +++ b/packages/ai/src/background-removal.ts @@ -1,4 +1,4 @@ -import { runPythonScript } from "./bridge.js"; +import { runPythonWithProgress, type ProgressCallback } from "./bridge.js"; import { writeFile, readFile, unlink } from "node:fs/promises"; import { join } from "node:path"; import { tmpdir } from "node:os"; @@ -13,6 +13,7 @@ export async function removeBackground( inputBuffer: Buffer, outputDir: string, options: RemoveBackgroundOptions = {}, + onProgress?: ProgressCallback, ): Promise { const id = randomUUID(); const inputPath = join(tmpdir(), `rembg_in_${id}.png`); @@ -22,11 +23,11 @@ export async function removeBackground( try { // BiRefNet models need longer timeout (up to 10 min for first load) const timeout = options.model?.startsWith("birefnet") ? 600000 : 300000; - const { stdout } = await runPythonScript("remove_bg.py", [ + const { stdout } = await runPythonWithProgress("remove_bg.py", [ inputPath, outputPath, JSON.stringify(options), - ], timeout); + ], { onProgress, timeout }); const result = JSON.parse(stdout); if (!result.success) { diff --git a/packages/ai/src/face-detection.ts b/packages/ai/src/face-detection.ts index 3535c171..05695592 100644 --- a/packages/ai/src/face-detection.ts +++ b/packages/ai/src/face-detection.ts @@ -1,4 +1,4 @@ -import { runPythonScript } from "./bridge.js"; +import { runPythonWithProgress, type ProgressCallback } from "./bridge.js"; import { writeFile, readFile } from "node:fs/promises"; import { join } from "node:path"; @@ -24,16 +24,17 @@ export async function blurFaces( inputBuffer: Buffer, outputDir: string, options: BlurFacesOptions = {}, + onProgress?: ProgressCallback, ): Promise { const inputPath = join(outputDir, "input_faces.png"); const outputPath = join(outputDir, "output_faces.png"); await writeFile(inputPath, inputBuffer); - const { stdout } = await runPythonScript("detect_faces.py", [ + const { stdout } = await runPythonWithProgress("detect_faces.py", [ inputPath, outputPath, JSON.stringify(options), - ]); + ], { onProgress }); const result = JSON.parse(stdout); if (!result.success) { diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 6a7d54a3..e6d1736d 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -1,4 +1,6 @@ export { runPythonScript } from "./bridge.js"; +export { runPythonWithProgress } from "./bridge.js"; +export type { ProgressCallback } from "./bridge.js"; export { removeBackground } from "./background-removal.js"; export type { RemoveBackgroundOptions } from "./background-removal.js"; export { upscale } from "./upscaling.js"; diff --git a/packages/ai/src/inpainting.ts b/packages/ai/src/inpainting.ts index b4810d53..44b8884a 100644 --- a/packages/ai/src/inpainting.ts +++ b/packages/ai/src/inpainting.ts @@ -1,4 +1,4 @@ -import { runPythonScript } from "./bridge.js"; +import { runPythonWithProgress, type ProgressCallback } from "./bridge.js"; import { writeFile, readFile } from "node:fs/promises"; import { join } from "node:path"; @@ -6,6 +6,7 @@ export async function inpaint( inputBuffer: Buffer, maskBuffer: Buffer, outputDir: string, + onProgress?: ProgressCallback, ): Promise { const inputPath = join(outputDir, "input_inpaint.png"); const maskPath = join(outputDir, "mask_inpaint.png"); @@ -14,11 +15,11 @@ export async function inpaint( await writeFile(inputPath, inputBuffer); await writeFile(maskPath, maskBuffer); - const { stdout } = await runPythonScript("inpaint.py", [ + const { stdout } = await runPythonWithProgress("inpaint.py", [ inputPath, maskPath, outputPath, - ]); + ], { onProgress }); const result = JSON.parse(stdout); if (!result.success) { diff --git a/packages/ai/src/ocr.ts b/packages/ai/src/ocr.ts index da950cb4..f490f383 100644 --- a/packages/ai/src/ocr.ts +++ b/packages/ai/src/ocr.ts @@ -1,4 +1,4 @@ -import { runPythonScript } from "./bridge.js"; +import { runPythonWithProgress, type ProgressCallback } from "./bridge.js"; import { writeFile } from "node:fs/promises"; import { join } from "node:path"; @@ -16,14 +16,15 @@ export async function extractText( inputBuffer: Buffer, outputDir: string, options: OcrOptions = {}, + onProgress?: ProgressCallback, ): Promise { const inputPath = join(outputDir, "input_ocr.png"); await writeFile(inputPath, inputBuffer); - const { stdout } = await runPythonScript("ocr.py", [ + const { stdout } = await runPythonWithProgress("ocr.py", [ inputPath, JSON.stringify(options), - ]); + ], { onProgress }); const result = JSON.parse(stdout); if (!result.success) { diff --git a/packages/ai/src/upscaling.ts b/packages/ai/src/upscaling.ts index f082e4e1..5b9b176b 100644 --- a/packages/ai/src/upscaling.ts +++ b/packages/ai/src/upscaling.ts @@ -1,4 +1,4 @@ -import { runPythonScript } from "./bridge.js"; +import { runPythonWithProgress, type ProgressCallback } from "./bridge.js"; import { writeFile, readFile } from "node:fs/promises"; import { join } from "node:path"; @@ -17,16 +17,17 @@ export async function upscale( inputBuffer: Buffer, outputDir: string, options: UpscaleOptions = {}, + onProgress?: ProgressCallback, ): Promise { const inputPath = join(outputDir, "input_upscale.png"); const outputPath = join(outputDir, "output_upscale.png"); await writeFile(inputPath, inputBuffer); - const { stdout } = await runPythonScript("upscale.py", [ + const { stdout } = await runPythonWithProgress("upscale.py", [ inputPath, outputPath, JSON.stringify(options), - ]); + ], { onProgress }); const result = JSON.parse(stdout); if (!result.success) {