diff --git a/apps/api/src/routes/features.ts b/apps/api/src/routes/features.ts index 4abfe22f..7bbfff29 100644 --- a/apps/api/src/routes/features.ts +++ b/apps/api/src/routes/features.ts @@ -202,9 +202,32 @@ export async function registerFeatureRoutes(app: FastifyInstance): Promise duration_ms: Date.now() - installStartTime, }); } else { - const errorDetail = - lastStderrLines.filter((l) => !l.startsWith("{")).join("\n") || stdoutBuffer.trim(); - const errorMsg = errorDetail || `Install failed with exit code ${code}`; + // Extract the structured error from Python's fail() function first. + // fail() writes {"error": "..."} to stderr — prefer this over raw lines. + let errorMsg: string | undefined; + for (let i = lastStderrLines.length - 1; i >= 0; i--) { + const line = lastStderrLines[i]; + if (line.startsWith("{")) { + try { + const parsed = JSON.parse(line) as Record; + if (typeof parsed.error === "string") { + errorMsg = parsed.error; + break; + } + } catch { + // Not valid JSON + } + } + } + if (!errorMsg) { + const meaningful = lastStderrLines.filter( + (l) => !l.startsWith("{") && !l.includes("pthread_setaffinity_np"), + ); + errorMsg = + meaningful.join("\n") || + stdoutBuffer.trim() || + `Install failed with exit code ${code}`; + } setInstallProgress(bundleId, null, errorMsg); updateSingleFileProgress({ jobId, phase: "failed", percent: 0, error: errorMsg }); } diff --git a/apps/web/src/stores/features-store.ts b/apps/web/src/stores/features-store.ts index 6cad7106..fb3f0bc9 100644 --- a/apps/web/src/stores/features-store.ts +++ b/apps/web/src/stores/features-store.ts @@ -191,7 +191,10 @@ export const useFeaturesStore = create((set, get) => { }); try { - const result = await apiPost<{ jobId: string }>(`/v1/admin/features/${bundleId}/install`); + const result = await apiPost<{ jobId: string }>( + `/v1/admin/features/${bundleId}/install`, + {}, + ); listenToProgress(bundleId, result.jobId); } catch (err) { const installing = { ...get().installing }; @@ -209,7 +212,7 @@ export const useFeaturesStore = create((set, get) => { uninstallBundle: async (bundleId: string) => { try { - await apiPost(`/v1/admin/features/${bundleId}/uninstall`); + await apiPost(`/v1/admin/features/${bundleId}/uninstall`, {}); await refreshBundles(); } catch (err) { set({ diff --git a/packages/ai/python/install_feature.py b/packages/ai/python/install_feature.py index c8228c43..f88975f7 100644 --- a/packages/ai/python/install_feature.py +++ b/packages/ai/python/install_feature.py @@ -277,6 +277,45 @@ def download_url_model(model: dict, models_dir: str) -> None: os.rename(tmp_path, final_path) +_matting_registered = False + + +def _register_birefnet_matting() -> None: + """Register the custom BiRefNet-matting ONNX session. + + This model is not built into rembg — it must be registered before + calling new_session("birefnet-matting"). The same registration is + done in remove_bg.py (runtime) and download_models.py (build-time). + """ + global _matting_registered + if _matting_registered: + return + _matting_registered = True + + import pooch + from rembg.sessions import sessions_class + from rembg.sessions.birefnet_general import BiRefNetSessionGeneral + + class BiRefNetMattingSession(BiRefNetSessionGeneral): + @classmethod + def download_models(cls, *args, **kwargs): + fname = f"{cls.name(*args, **kwargs)}.onnx" + pooch.retrieve( + "https://github.com/ZhengPeng7/BiRefNet/releases/download/v1/BiRefNet-matting-epoch_100.onnx", + None, + fname=fname, + path=cls.u2net_home(*args, **kwargs), + progressbar=True, + ) + return os.path.join(cls.u2net_home(*args, **kwargs), fname) + + @classmethod + def name(cls, *args, **kwargs): + return "birefnet-matting" + + sessions_class.append(BiRefNetMattingSession) + + def download_rembg_session(model: dict, models_dir: str) -> None: """Download a rembg model by initializing a session.""" args = model.get("args", []) @@ -291,6 +330,7 @@ def download_rembg_session(model: dict, models_dir: str) -> None: os.environ["U2NET_HOME"] = u2net_dir from rembg import new_session + _register_birefnet_matting() new_session(model_name) diff --git a/tests/unit/web/zustand-stores.test.ts b/tests/unit/web/zustand-stores.test.ts index a1e6f4d1..f79bd059 100644 --- a/tests/unit/web/zustand-stores.test.ts +++ b/tests/unit/web/zustand-stores.test.ts @@ -1966,7 +1966,7 @@ describe("useFeaturesStore", () => { await useFeaturesStore.getState().uninstallBundle("ai-rembg"); - expect(mockApiPost).toHaveBeenCalledWith("/v1/admin/features/ai-rembg/uninstall"); + expect(mockApiPost).toHaveBeenCalledWith("/v1/admin/features/ai-rembg/uninstall", {}); }); it("uninstallBundle sets error on failure", async () => { @@ -2207,8 +2207,8 @@ describe("useFeaturesStore", () => { await useFeaturesStore.getState().reinstallBundle("ai-rembg"); - expect(mockApiPost).toHaveBeenCalledWith("/v1/admin/features/ai-rembg/uninstall"); - expect(mockApiPost).toHaveBeenCalledWith("/v1/admin/features/ai-rembg/install"); + expect(mockApiPost).toHaveBeenCalledWith("/v1/admin/features/ai-rembg/uninstall", {}); + expect(mockApiPost).toHaveBeenCalledWith("/v1/admin/features/ai-rembg/install", {}); }); // -- EventSource progress handling ---------------------------------------- @@ -2452,7 +2452,7 @@ describe("useFeaturesStore", () => { expect(useFeaturesStore.getState().installAllActive).toBe(false); expect(useFeaturesStore.getState().queued).toEqual([]); // Only not-installed bundle should have been installed - expect(mockApiPost).toHaveBeenCalledWith("/v1/admin/features/ai-rembg/install"); + expect(mockApiPost).toHaveBeenCalledWith("/v1/admin/features/ai-rembg/install", {}); }); it("installAll skips bundles that are already installed", async () => {