902 lines
29 KiB
TypeScript
902 lines
29 KiB
TypeScript
import {
|
||
createMediaAssetRelation,
|
||
getMediaAssetsByUserAndUrls,
|
||
getSystemSetting,
|
||
upsertMediaAsset,
|
||
} from "@/lib/db/queries";
|
||
import { executeConfiguredVip98kTask } from "@/lib/ai/vip98k-runtime";
|
||
import { db } from "@/lib/db";
|
||
import { generatedImage } from "@/lib/db/schema";
|
||
import { createLogger } from "@/lib/observability/logger";
|
||
import { trackRuntimeWork } from "@/lib/observability/active-work";
|
||
import { uploadFile } from "@/lib/storage/service";
|
||
|
||
const logger = createLogger("rmbg-service");
|
||
|
||
export type RmbgProvider = "vip98k" | "rmbg" | "photoroom" | "aoscdn";
|
||
|
||
type RmbgConfig = {
|
||
provider?: RmbgProvider;
|
||
apiUrl?: string;
|
||
apiKey?: string;
|
||
photoroomApiKey?: string;
|
||
aosApiKey?: string;
|
||
aosApiBaseUrl?: string;
|
||
aosForegroundType?: "" | "person" | "object" | "stamp";
|
||
aosCrop?: 0 | 1;
|
||
aosFormat?: "png" | "jpg";
|
||
};
|
||
|
||
export type ExecuteRmbgJobParams = {
|
||
requestId?: string;
|
||
userId: string;
|
||
imageBuffer: Buffer;
|
||
imageBase64?: string;
|
||
foregroundType?: "" | "person" | "object" | "stamp";
|
||
prompt?: string;
|
||
chatId?: string | null;
|
||
messageId?: string | null;
|
||
skipGallerySave?: boolean;
|
||
// 仅当 imageBuffer 与 URL 对应的输入文件一致时提供。
|
||
inputImageUrl?: string | null;
|
||
sourceImageUrl?: string | null;
|
||
};
|
||
|
||
export type ExecuteRmbgJobResult = {
|
||
image_base64: string;
|
||
thumbnail_url: string;
|
||
original_url: string;
|
||
processing_time: number;
|
||
provider: RmbgProvider;
|
||
mediaAssetId?: string | null;
|
||
generatedImageId?: string | null;
|
||
};
|
||
|
||
export class RmbgServiceError extends Error {
|
||
status: number;
|
||
detail?: string;
|
||
|
||
constructor(message: string, status = 500, detail?: string) {
|
||
super(detail && detail !== message ? `${message}：${detail}` : message);
|
||
this.name = "RmbgServiceError";
|
||
this.status = status;
|
||
this.detail = detail;
|
||
}
|
||
}
|
||
|
||
const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));
|
||
// https://picwish.cn/background-removal-api-doc
|
||
const AOS_MAX_RESOLUTION = 4096;
|
||
const AOS_MAX_INPUT_BYTES = 50 * 1024 * 1024;
|
||
// The shared AOS state -5 documents a separate limit for image_url downloads.
|
||
const AOS_MAX_URL_INPUT_BYTES = 30 * 1024 * 1024;
|
||
const AOS_CREATE_TIMEOUT_MS = 120_000;
|
||
const AOS_POLL_TIMEOUT_MS = 60_000;
|
||
const AOS_DOWNLOAD_TIMEOUT_MS = 60_000;
|
||
const IMAGE_URL_FETCH_TIMEOUT_MS = 90_000;
|
||
const IMAGE_URL_FETCH_RETRIES = 2;
|
||
|
||
async function createThumbnail(buffer: Buffer, mimeType: string): Promise<Buffer> {
|
||
try {
|
||
const sharp = await import("sharp");
|
||
|
||
return await sharp.default(buffer)
|
||
.resize(500, null, {
|
||
fit: "inside",
|
||
withoutEnlargement: true,
|
||
})
|
||
.png({ quality: 80 })
|
||
.toBuffer();
|
||
} catch (error) {
|
||
logger.warn("thumbnail_creation_failed", error);
|
||
return buffer;
|
||
}
|
||
}
|
||
|
||
async function uploadBufferToS3(
|
||
buffer: Buffer,
|
||
mimeType: string = "image/png"
|
||
): Promise<{ thumbnail: string; original: string }> {
|
||
try {
|
||
const ext = mimeType.split("/")[1] || "png";
|
||
const timestamp = Date.now();
|
||
const randomStr = Math.random().toString(36).slice(2, 8);
|
||
const thumbnailFilename = `rmbg-thumb-${timestamp}-${randomStr}.${ext}`;
|
||
const originalFilename = `rmbg-orig-${timestamp}-${randomStr}.${ext}`;
|
||
|
||
const thumbnailBuffer = await createThumbnail(buffer, mimeType);
|
||
|
||
const thumbnailArrayBuffer = thumbnailBuffer.buffer.slice(
|
||
thumbnailBuffer.byteOffset,
|
||
thumbnailBuffer.byteOffset + thumbnailBuffer.byteLength
|
||
) as ArrayBuffer;
|
||
const originalArrayBuffer = buffer.buffer.slice(
|
||
buffer.byteOffset,
|
||
buffer.byteOffset + buffer.byteLength
|
||
) as ArrayBuffer;
|
||
|
||
const [thumbnailResult, originalResult] = await Promise.all([
|
||
uploadFile(thumbnailFilename, thumbnailArrayBuffer),
|
||
uploadFile(originalFilename, originalArrayBuffer),
|
||
]);
|
||
|
||
logger.info("storage_upload_completed", {
|
||
mimeType,
|
||
originalBytes: buffer.length,
|
||
originalUrl: originalResult.url,
|
||
thumbnailBytes: thumbnailBuffer.length,
|
||
thumbnailUrl: thumbnailResult.url,
|
||
});
|
||
|
||
return {
|
||
thumbnail: thumbnailResult.url,
|
||
original: originalResult.url,
|
||
};
|
||
} catch (error) {
|
||
logger.error("storage_upload_failed", error);
|
||
throw new RmbgServiceError(
|
||
"图片上传到存储桶失败",
|
||
500,
|
||
error instanceof Error ? error.message : "未知错误"
|
||
);
|
||
}
|
||
}
|
||
|
||
const withTimeout = async <T,>(
|
||
fn: (signal: AbortSignal) => Promise<T>,
|
||
timeoutMs: number
|
||
): Promise<T> => {
|
||
const controller = new AbortController();
|
||
const timeoutId = setTimeout(() => controller.abort(), timeoutMs);
|
||
try {
|
||
return await fn(controller.signal);
|
||
} finally {
|
||
clearTimeout(timeoutId);
|
||
}
|
||
};
|
||
|
||
async function readAosResult<T>(read: (signal: AbortSignal) => Promise<T>, timeoutMs: number): Promise<T> {
|
||
for (let attempt = 0; ; attempt++) {
|
||
try {
|
||
return await withTimeout(read, timeoutMs);
|
||
} catch (error) {
|
||
const name = error instanceof Error ? error.name : "";
|
||
if (attempt >= 2 || (name !== "TypeError" && name !== "AbortError")) throw error;
|
||
await sleep(1000 * (attempt + 1));
|
||
}
|
||
}
|
||
}
|
||
|
||
function getAosFriendlyErrorMessage(detail: string): string | null {
|
||
const resolutionMatch = detail.match(/resolution excess,\s*width=(\d+),\s*height=(\d+)/i);
|
||
if (resolutionMatch) {
|
||
const [, width, height] = resolutionMatch;
|
||
return `图片分辨率超限（${width}x${height}），AOS要求宽高都不超过 ${AOS_MAX_RESOLUTION}px。请缩小后重试。`;
|
||
}
|
||
|
||
if (/invalid file/i.test(detail)) {
|
||
return "图片文件无法识别，请尝试转换为 JPG/PNG 后重试。";
|
||
}
|
||
|
||
return null;
|
||
}
|
||
|
||
async function prepareImageForAos(
|
||
imageBuffer: Buffer
|
||
): Promise<{ buffer: Buffer; mimeType: string; filename: string }> {
|
||
if (imageBuffer.length > AOS_MAX_INPUT_BYTES) {
|
||
throw new RmbgServiceError(
|
||
`抠图参考图文件大小超限（${(imageBuffer.length / 1024 / 1024).toFixed(2)} MB），当前 AOS 上游最大支持 50 MB。`,
|
||
400
|
||
);
|
||
}
|
||
|
||
let metadata: import("sharp").Metadata;
|
||
try {
|
||
const sharp = await import("sharp");
|
||
metadata = await sharp.default(imageBuffer, { limitInputPixels: false }).metadata();
|
||
} catch (error) {
|
||
// AOS also accepts BMP/RAW/RGB formats that Sharp cannot identify.
|
||
logger.warn("aos_image_metadata_unavailable", error);
|
||
return { buffer: imageBuffer, mimeType: "application/octet-stream", filename: "image" };
|
||
}
|
||
|
||
const { width, height } = metadata;
|
||
if (width && height && (width > AOS_MAX_RESOLUTION || height > AOS_MAX_RESOLUTION)) {
|
||
throw new RmbgServiceError(
|
||
`抠图参考图分辨率超限（${width}×${height} 像素），当前 AOS 上游要求宽、高均不超过 ${AOS_MAX_RESOLUTION} 像素。`,
|
||
400
|
||
);
|
||
}
|
||
|
||
const format = metadata.format || "";
|
||
if (["gif", "svg", "heif", "avif", "jp2", "jxl"].includes(format)) {
|
||
const name = format === "heif" ? "HEIF/AVIF" : format.toUpperCase();
|
||
throw new RmbgServiceError(`抠图参考图格式不支持（${name}），当前 AOS 上游不支持此格式。`, 400);
|
||
}
|
||
const ext = format === "jpeg" ? "jpg" : format;
|
||
return {
|
||
buffer: imageBuffer,
|
||
mimeType: format ? `image/${format}` : "application/octet-stream",
|
||
filename: ext ? `image.${ext}` : "image",
|
||
};
|
||
}
|
||
|
||
async function processWithRMBG(
|
||
imageBase64: string,
|
||
config: { apiUrl: string; apiKey?: string }
|
||
): Promise<{ success: boolean; image_base64?: string; processing_time?: number; error?: string }> {
|
||
const apiUrl = config.apiUrl.endsWith("/")
|
||
? `${config.apiUrl}rmbg/base64`
|
||
: `${config.apiUrl}/rmbg/base64`;
|
||
|
||
const headers: HeadersInit = {
|
||
"Content-Type": "application/json",
|
||
};
|
||
|
||
if (config.apiKey) {
|
||
headers["X-API-Key"] = config.apiKey;
|
||
}
|
||
|
||
const startTime = Date.now();
|
||
const controller = new AbortController();
|
||
const timeoutId = setTimeout(() => controller.abort(), 30000);
|
||
|
||
try {
|
||
const response = await fetch(apiUrl, {
|
||
method: "POST",
|
||
headers,
|
||
body: JSON.stringify({ image_base64: imageBase64 }),
|
||
signal: controller.signal,
|
||
});
|
||
|
||
clearTimeout(timeoutId);
|
||
const processingTime = (Date.now() - startTime) / 1000;
|
||
|
||
if (!response.ok) {
|
||
const errorText = await response.text();
|
||
logger.warn("rmbg_provider_request_failed", {
|
||
errorLength: errorText.length,
|
||
status: response.status,
|
||
});
|
||
return { success: false, error: errorText || response.statusText };
|
||
}
|
||
|
||
const result = await response.json();
|
||
if (!result.success || !result.image_base64) {
|
||
return { success: false, error: result.detail || "未知错误" };
|
||
}
|
||
|
||
return {
|
||
success: true,
|
||
image_base64: result.image_base64,
|
||
processing_time: result.processing_time || processingTime,
|
||
};
|
||
} catch (error) {
|
||
clearTimeout(timeoutId);
|
||
if (error instanceof Error && error.name === "AbortError") {
|
||
return { success: false, error: "请求超时（30秒）" };
|
||
}
|
||
logger.warn("rmbg_provider_request_error", error);
|
||
return { success: false, error: error instanceof Error ? error.message : "未知错误" };
|
||
}
|
||
}
|
||
|
||
async function processWithVip98k(
|
||
imageBuffer: Buffer,
|
||
): Promise<{ success: boolean; imageBuffer?: Buffer; processing_time?: number; error?: string }> {
|
||
const startedAt = Date.now();
|
||
try {
|
||
const result = await executeConfiguredVip98kTask({
|
||
modelId: "remove-background",
|
||
kind: "image-tool",
|
||
prompt: "Remove the image background and return a transparent PNG.",
|
||
imageUrls: [`data:image/png;base64,${imageBuffer.toString("base64")}`],
|
||
});
|
||
const outputUrl = result.outputUrls[0];
|
||
if (!outputUrl) {
|
||
return { success: false, error: "VIP98K 抠图未返回图片" };
|
||
}
|
||
const response = await fetch(outputUrl, {
|
||
cache: "no-store",
|
||
signal: AbortSignal.timeout(120_000),
|
||
});
|
||
if (!response.ok) {
|
||
return {
|
||
success: false,
|
||
error: `VIP98K 抠图结果下载失败（HTTP ${response.status}）`,
|
||
};
|
||
}
|
||
return {
|
||
success: true,
|
||
imageBuffer: Buffer.from(await response.arrayBuffer()),
|
||
processing_time: (Date.now() - startedAt) / 1000,
|
||
};
|
||
} catch (error) {
|
||
return {
|
||
success: false,
|
||
error: error instanceof Error ? error.message : "VIP98K 抠图失败",
|
||
};
|
||
}
|
||
}
|
||
|
||
async function processWithPhotoRoom(
|
||
imageBuffer: Buffer,
|
||
apiKey: string
|
||
): Promise<{ success: boolean; imageBuffer?: Buffer; processing_time?: number; error?: string }> {
|
||
const startTime = Date.now();
|
||
|
||
try {
|
||
const formData = new FormData();
|
||
const uint8Array = new Uint8Array(imageBuffer);
|
||
const blob = new Blob([uint8Array], { type: "image/png" });
|
||
formData.append("imageFile", blob, "image.png");
|
||
formData.append("referenceBox", "originalImage");
|
||
|
||
const controller = new AbortController();
|
||
const timeoutId = setTimeout(() => controller.abort(), 60000);
|
||
|
||
const response = await fetch("https://image-api.photoroom.com/v2/edit", {
|
||
method: "POST",
|
||
headers: {
|
||
"x-api-key": apiKey,
|
||
"pr-hd-background-removal": "auto",
|
||
},
|
||
body: formData,
|
||
signal: controller.signal,
|
||
});
|
||
|
||
clearTimeout(timeoutId);
|
||
const processingTime = (Date.now() - startTime) / 1000;
|
||
|
||
if (!response.ok) {
|
||
const errorText = await response.text();
|
||
logger.warn("photoroom_provider_request_failed", {
|
||
errorLength: errorText.length,
|
||
status: response.status,
|
||
});
|
||
return { success: false, error: `PhotoRoom API错误: ${response.status} - ${errorText}` };
|
||
}
|
||
|
||
const resultBuffer = Buffer.from(await response.arrayBuffer());
|
||
return {
|
||
success: true,
|
||
imageBuffer: resultBuffer,
|
||
processing_time: processingTime,
|
||
};
|
||
} catch (error) {
|
||
if (error instanceof Error && error.name === "AbortError") {
|
||
return { success: false, error: "PhotoRoom请求超时（60秒）" };
|
||
}
|
||
logger.warn("photoroom_provider_request_error", error);
|
||
return { success: false, error: error instanceof Error ? error.message : "未知错误" };
|
||
}
|
||
}
|
||
|
||
async function processWithAosSegmentation(
|
||
imageBuffer: Buffer,
|
||
config: {
|
||
apiKey: string;
|
||
apiBaseUrl?: string;
|
||
foregroundType?: "" | "person" | "object" | "stamp";
|
||
crop?: 0 | 1;
|
||
format?: "png" | "jpg";
|
||
inputImageUrl?: string | null;
|
||
requestId?: string;
|
||
}
|
||
): Promise<{ success: boolean; imageBuffer?: Buffer; processing_time?: number; error?: string }> {
|
||
const startTime = Date.now();
|
||
const baseUrl = (config.apiBaseUrl || "https://techsz.aoscdn.com").replace(/\/+$/, "");
|
||
const createUrl = `${baseUrl}/api/tasks/visual/segmentation`;
|
||
const headers = { "X-API-KEY": config.apiKey };
|
||
let stage = "create";
|
||
let providerTaskId: string | undefined;
|
||
const sourceImage = await prepareImageForAos(imageBuffer);
|
||
|
||
try {
|
||
const createForm = new FormData();
|
||
let publicSourceUrl: string | null = null;
|
||
if (config.inputImageUrl && imageBuffer.length <= AOS_MAX_URL_INPUT_BYTES) {
|
||
try {
|
||
const source = new URL(config.inputImageUrl);
|
||
const cdn = new URL(process.env.NEXT_PUBLIC_CANVAS_CDN_ORIGIN || "");
|
||
if (["http:", "https:"].includes(source.protocol) && source.hostname === cdn.hostname) {
|
||
publicSourceUrl = source.toString();
|
||
}
|
||
} catch { /* Uploaded files and other sources use the file API. */ }
|
||
}
|
||
if (publicSourceUrl) {
|
||
createForm.append("image_url", publicSourceUrl);
|
||
} else {
|
||
const imageBlob = new Blob([new Uint8Array(sourceImage.buffer)], {
|
||
type: sourceImage.mimeType,
|
||
});
|
||
createForm.append("image_file", imageBlob, sourceImage.filename);
|
||
}
|
||
createForm.append("sync", "0");
|
||
createForm.append("return_type", "1");
|
||
createForm.append("output_type", "2");
|
||
createForm.append("format", config.format || "png");
|
||
createForm.append("crop", config.crop === 1 ? "1" : "0");
|
||
if (config.foregroundType) {
|
||
createForm.append("type", config.foregroundType);
|
||
}
|
||
|
||
const createStartedAt = Date.now();
|
||
logger.info("aos_task_create_started", { requestId: config.requestId, sourceMode: publicSourceUrl ? "url" : "file", inputBytes: imageBuffer.length, uploadBytes: publicSourceUrl ? 0 : sourceImage.buffer.length });
|
||
|
||
const createResponse = await withTimeout(
|
||
(signal) =>
|
||
fetch(createUrl, {
|
||
method: "POST",
|
||
headers,
|
||
body: createForm,
|
||
signal,
|
||
}),
|
||
AOS_CREATE_TIMEOUT_MS
|
||
);
|
||
|
||
const createPayload = await createResponse.json().catch(() => null);
|
||
if (!createResponse.ok || !createPayload?.data?.task_id) {
|
||
logger.warn("aos_task_create_failed", {
|
||
hasPayload: createPayload != null,
|
||
status: createResponse.status,
|
||
});
|
||
return {
|
||
success: false,
|
||
error: `AOS任务创建失败（HTTP ${createResponse.status}）: ${createPayload?.message || createResponse.statusText}`,
|
||
};
|
||
}
|
||
|
||
const taskId = createPayload.data.task_id as string;
|
||
providerTaskId = taskId;
|
||
logger.info("aos_task_created", { requestId: config.requestId, taskId, createMs: Date.now() - createStartedAt });
|
||
const taskUrl = `${createUrl}/${taskId}`;
|
||
const maxPollCount = 60;
|
||
|
||
for (let pollCount = 0; pollCount < maxPollCount; pollCount += 1) {
|
||
await sleep(1000);
|
||
stage = "poll";
|
||
|
||
const pollResponse = await readAosResult(
|
||
async (signal) => {
|
||
const response = await fetch(taskUrl, {
|
||
method: "GET",
|
||
headers,
|
||
signal,
|
||
});
|
||
const payload = await response.json().catch((error) => {
|
||
if (error instanceof Error && error.name !== "SyntaxError") throw error;
|
||
return null;
|
||
});
|
||
return { ok: response.ok, status: response.status, statusText: response.statusText, payload };
|
||
},
|
||
AOS_POLL_TIMEOUT_MS
|
||
);
|
||
|
||
const pollPayload = pollResponse.payload;
|
||
if (!pollResponse.ok || !pollPayload?.data) {
|
||
logger.warn("aos_task_poll_failed", {
|
||
hasPayload: pollPayload != null,
|
||
status: pollResponse.status,
|
||
});
|
||
return {
|
||
success: false,
|
||
error: `AOS任务查询失败（HTTP ${pollResponse.status}，任务 ${taskId}）: ${pollPayload?.message || pollResponse.statusText}`,
|
||
};
|
||
}
|
||
|
||
const state = Number(pollPayload.data.state);
|
||
const imageUrl = pollPayload.data.image as string | undefined;
|
||
const stateDetail = String(pollPayload.data.state_detail || "");
|
||
const errMessage = String(pollPayload.data.err_message || pollPayload.data.err_info || "");
|
||
|
||
if (state === 1 && imageUrl) {
|
||
logger.debug("aos_task_completed", {
|
||
pollCount,
|
||
taskId,
|
||
});
|
||
stage = "download";
|
||
const imageResponse = await readAosResult(
|
||
async (signal) => {
|
||
const response = await fetch(imageUrl, {
|
||
method: "GET",
|
||
signal,
|
||
});
|
||
return { ok: response.ok, status: response.status, statusText: response.statusText, buffer: response.ok ? Buffer.from(await response.arrayBuffer()) : null };
|
||
},
|
||
AOS_DOWNLOAD_TIMEOUT_MS
|
||
);
|
||
|
||
if (!imageResponse.ok) {
|
||
logger.warn("aos_result_download_failed", {
|
||
status: imageResponse.status,
|
||
taskId,
|
||
});
|
||
return {
|
||
success: false,
|
||
error: `AOS结果下载失败（HTTP ${imageResponse.status}，任务 ${taskId}）: ${imageResponse.statusText}`,
|
||
};
|
||
}
|
||
|
||
const resultBuffer = imageResponse.buffer!;
|
||
return {
|
||
success: true,
|
||
imageBuffer: resultBuffer,
|
||
processing_time: (Date.now() - startTime) / 1000,
|
||
};
|
||
}
|
||
|
||
if (state < 0) {
|
||
logger.warn("aos_task_failed", {
|
||
hasErrorMessage: Boolean(errMessage),
|
||
state,
|
||
taskId,
|
||
});
|
||
return {
|
||
success: false,
|
||
error: `AOS任务失败（状态 ${state}，任务 ${taskId}）: ${errMessage || stateDetail || pollPayload?.message || "未提供详情"}`,
|
||
};
|
||
}
|
||
}
|
||
|
||
logger.warn("aos_task_timeout", { requestId: config.requestId, taskId, maxPollCount });
|
||
return { success: false, error: `AOS查询任务超时，结果尚未确认（任务 ${taskId}）` };
|
||
} catch (error) {
|
||
if (error instanceof Error && error.name === "AbortError") {
|
||
logger.warn("aos_request_timeout", { requestId: config.requestId, stage, taskId: providerTaskId, inputBytes: imageBuffer.length });
|
||
const message = stage === "create"
|
||
? "AOS创建任务超时（120秒），结果尚未确认"
|
||
: stage === "poll"
|
||
? `AOS查询任务超时，结果尚未确认（任务 ${providerTaskId}）`
|
||
: `AOS结果下载超时，任务已完成（任务 ${providerTaskId}）`;
|
||
return { success: false, error: message };
|
||
}
|
||
logger.warn("aos_provider_request_error", error, { requestId: config.requestId, stage, taskId: providerTaskId });
|
||
const detail = error instanceof Error ? error.message : "未知错误";
|
||
return { success: false, error: stage === "create"
|
||
? `AOS创建任务连接失败，结果尚未确认：${detail}`
|
||
: stage === "poll"
|
||
? `AOS查询任务连接失败，结果尚未确认（任务 ${providerTaskId}）：${detail}`
|
||
: `AOS结果下载连接失败，任务已完成（任务 ${providerTaskId}）：${detail}` };
|
||
}
|
||
}
|
||
|
||
async function getValidatedRmbgConfig(): Promise<{ provider: RmbgProvider; config: RmbgConfig }> {
|
||
const rmbgConfig = await getSystemSetting({ key: "rmbg_api_config" });
|
||
|
||
if (!rmbgConfig || !rmbgConfig.value) {
|
||
throw new RmbgServiceError("抠图API配置未设置，请在系统设置中配置", 400);
|
||
}
|
||
|
||
const config = rmbgConfig.value as RmbgConfig;
|
||
const provider: RmbgProvider = config.provider || "rmbg";
|
||
|
||
if (provider === "rmbg" && !config.apiUrl) {
|
||
throw new RmbgServiceError("RMBG API地址未配置", 400);
|
||
}
|
||
|
||
if (provider === "photoroom" && !config.photoroomApiKey) {
|
||
throw new RmbgServiceError("PhotoRoom API密钥未配置", 400);
|
||
}
|
||
|
||
if (provider === "aoscdn" && !config.aosApiKey) {
|
||
throw new RmbgServiceError("AOS API密钥未配置", 400);
|
||
}
|
||
|
||
return { provider, config };
|
||
}
|
||
|
||
function getModelId(provider: RmbgProvider) {
|
||
return provider === "vip98k"
|
||
? "remove-background"
|
||
: provider === "photoroom"
|
||
? "photoroom-hd"
|
||
: provider === "aoscdn"
|
||
? "aos-segmentation"
|
||
: "rmbg-2.0";
|
||
}
|
||
|
||
function getImageUrlLogMeta(imageUrl: string) {
|
||
try {
|
||
const parsed = new URL(imageUrl);
|
||
return {
|
||
host: parsed.host,
|
||
pathname: parsed.pathname.slice(0, 96),
|
||
};
|
||
} catch {
|
||
return { host: "", pathname: "" };
|
||
}
|
||
}
|
||
|
||
export async function fetchImageBufferFromUrl(imageUrl: string): Promise<Buffer> {
|
||
let lastError: unknown = null;
|
||
|
||
for (let attempt = 1; attempt <= IMAGE_URL_FETCH_RETRIES + 1; attempt += 1) {
|
||
try {
|
||
return await withTimeout(async (signal) => {
|
||
const imageResponse = await fetch(imageUrl, {
|
||
headers: {
|
||
Accept: "image/avif,image/webp,image/apng,image/svg+xml,image/*,*/*;q=0.8",
|
||
"User-Agent": "ai-chatbot-image-fetch/1.0",
|
||
},
|
||
signal,
|
||
});
|
||
if (!imageResponse.ok) {
|
||
throw new RmbgServiceError(
|
||
`无法获取图片: ${imageResponse.status} ${imageResponse.statusText}`,
|
||
400
|
||
);
|
||
}
|
||
return Buffer.from(await imageResponse.arrayBuffer());
|
||
}, IMAGE_URL_FETCH_TIMEOUT_MS);
|
||
} catch (error) {
|
||
if (error instanceof RmbgServiceError) {
|
||
throw error;
|
||
}
|
||
lastError = error;
|
||
logger.warn("source_image_fetch_failed", error, {
|
||
attempt,
|
||
maxAttempts: IMAGE_URL_FETCH_RETRIES + 1,
|
||
...getImageUrlLogMeta(imageUrl),
|
||
});
|
||
if (attempt <= IMAGE_URL_FETCH_RETRIES) {
|
||
await sleep(attempt * 1500);
|
||
}
|
||
}
|
||
}
|
||
|
||
throw new RmbgServiceError(
|
||
"无法获取图片，请检查图片链接是否有效",
|
||
400,
|
||
lastError instanceof Error ? lastError.message : String(lastError || "unknown")
|
||
);
|
||
}
|
||
|
||
export async function executeRmbgJob(
|
||
params: ExecuteRmbgJobParams
|
||
): Promise<ExecuteRmbgJobResult> {
|
||
return trackRuntimeWork(() => runRmbgJob(params));
|
||
}
|
||
|
||
async function runRmbgJob(
|
||
params: ExecuteRmbgJobParams
|
||
): Promise<ExecuteRmbgJobResult> {
|
||
const {
|
||
requestId = crypto.randomUUID(),
|
||
userId,
|
||
imageBuffer,
|
||
imageBase64 = imageBuffer.toString("base64"),
|
||
foregroundType,
|
||
prompt = "抠图处理",
|
||
chatId = null,
|
||
messageId = null,
|
||
skipGallerySave = false,
|
||
inputImageUrl = null,
|
||
sourceImageUrl = null,
|
||
} = params;
|
||
|
||
const { provider, config } = await getValidatedRmbgConfig();
|
||
logger.info("job_started", {
|
||
requestId,
|
||
hasChatId: Boolean(chatId),
|
||
hasMessageId: Boolean(messageId),
|
||
hasSourceImageUrl: Boolean(sourceImageUrl),
|
||
inputBytes: imageBuffer.length,
|
||
provider,
|
||
skipGallerySave,
|
||
});
|
||
|
||
let resultBuffer: Buffer;
|
||
let resultBase64: string;
|
||
let processingTime: number;
|
||
|
||
if (provider === "vip98k") {
|
||
const result = await processWithVip98k(imageBuffer);
|
||
if (!result.success || !result.imageBuffer) {
|
||
logger.warn("job_provider_failed", {
|
||
provider,
|
||
reasonLength: result.error?.length ?? 0,
|
||
});
|
||
throw new RmbgServiceError("抠图处理失败", 500, result.error);
|
||
}
|
||
resultBuffer = result.imageBuffer;
|
||
resultBase64 = resultBuffer.toString("base64");
|
||
processingTime = result.processing_time || 0;
|
||
} else if (provider === "photoroom") {
|
||
const result = await processWithPhotoRoom(imageBuffer, config.photoroomApiKey!);
|
||
if (!result.success || !result.imageBuffer) {
|
||
logger.warn("job_provider_failed", {
|
||
provider,
|
||
reasonLength: result.error?.length ?? 0,
|
||
});
|
||
throw new RmbgServiceError("抠图处理失败", 500, result.error);
|
||
}
|
||
resultBuffer = result.imageBuffer;
|
||
resultBase64 = resultBuffer.toString("base64");
|
||
processingTime = result.processing_time || 0;
|
||
} else if (provider === "aoscdn") {
|
||
const result = await processWithAosSegmentation(imageBuffer, {
|
||
apiKey: config.aosApiKey!,
|
||
apiBaseUrl: config.aosApiBaseUrl,
|
||
foregroundType: foregroundType ?? config.aosForegroundType ?? "",
|
||
crop: config.aosCrop === 1 ? 1 : 0,
|
||
format: config.aosFormat === "jpg" ? "jpg" : "png",
|
||
inputImageUrl,
|
||
requestId,
|
||
});
|
||
|
||
if (!result.success || !result.imageBuffer) {
|
||
const detail = result.error || "AOS处理失败";
|
||
const friendlyMessage = getAosFriendlyErrorMessage(detail);
|
||
logger.warn("job_provider_failed", {
|
||
friendly: Boolean(friendlyMessage),
|
||
provider,
|
||
reasonLength: detail.length,
|
||
});
|
||
throw new RmbgServiceError(
|
||
"抠图处理失败",
|
||
friendlyMessage ? 400 : 500,
|
||
friendlyMessage || detail
|
||
);
|
||
}
|
||
|
||
resultBuffer = result.imageBuffer;
|
||
resultBase64 = resultBuffer.toString("base64");
|
||
processingTime = result.processing_time || 0;
|
||
} else {
|
||
const result = await processWithRMBG(imageBase64, {
|
||
apiUrl: config.apiUrl!,
|
||
apiKey: config.apiKey,
|
||
});
|
||
|
||
if (!result.success || !result.image_base64) {
|
||
logger.warn("job_provider_failed", {
|
||
provider,
|
||
reasonLength: result.error?.length ?? 0,
|
||
});
|
||
throw new RmbgServiceError("抠图处理失败", 500, result.error);
|
||
}
|
||
|
||
resultBase64 = result.image_base64;
|
||
resultBuffer = Buffer.from(resultBase64, "base64");
|
||
processingTime = result.processing_time || 0;
|
||
}
|
||
|
||
const uploadResult = await uploadBufferToS3(resultBuffer, "image/png");
|
||
const thumbnailUrl = uploadResult.thumbnail;
|
||
const originalUrl = uploadResult.original;
|
||
|
||
if (!skipGallerySave) {
|
||
let saved = false;
|
||
let saveError: unknown = null;
|
||
let mediaAssetId: string | null = null;
|
||
let savedImage:
|
||
| {
|
||
id: string;
|
||
}
|
||
| null = null;
|
||
|
||
for (let attempt = 1; attempt <= 3; attempt += 1) {
|
||
try {
|
||
[savedImage] = await db
|
||
.insert(generatedImage)
|
||
.values({
|
||
userId,
|
||
chatId,
|
||
messageId,
|
||
thumbnailUrl,
|
||
originalUrl,
|
||
imageSize: "2K",
|
||
prompt,
|
||
modelId: getModelId(provider),
|
||
})
|
||
.returning();
|
||
saved = true;
|
||
break;
|
||
} catch (dbError) {
|
||
saveError = dbError;
|
||
if (attempt < 3) {
|
||
await sleep(120 * attempt);
|
||
}
|
||
}
|
||
}
|
||
|
||
if (!saved) {
|
||
logger.error("gallery_save_failed", saveError);
|
||
throw new RmbgServiceError("图片处理成功但保存到相册失败，请稍后重试", 500);
|
||
}
|
||
|
||
if (savedImage) {
|
||
try {
|
||
const sourceAssets = sourceImageUrl
|
||
? await getMediaAssetsByUserAndUrls({
|
||
userId,
|
||
urls: [sourceImageUrl],
|
||
})
|
||
: [];
|
||
|
||
const media = await upsertMediaAsset({
|
||
userId,
|
||
chatId,
|
||
mediaType: "image",
|
||
sourceType: "edited",
|
||
status: "ready",
|
||
originalUrl,
|
||
previewUrl: originalUrl,
|
||
thumbnailUrl,
|
||
displayName: prompt.slice(0, 80),
|
||
prompt,
|
||
modelId: getModelId(provider),
|
||
sourceMessageId: messageId,
|
||
sourceTable: "GeneratedImage",
|
||
sourceRowId: savedImage.id,
|
||
primaryParentAssetId: sourceAssets[0]?.id ?? null,
|
||
metadata: {
|
||
provider,
|
||
imageSize: "2K",
|
||
},
|
||
});
|
||
mediaAssetId = media?.id || null;
|
||
|
||
if (media) {
|
||
for (const sourceAsset of sourceAssets) {
|
||
await createMediaAssetRelation({
|
||
userId,
|
||
chatId,
|
||
fromAssetId: sourceAsset.id,
|
||
toAssetId: media.id,
|
||
relationType: "edited_from",
|
||
sourceMessageId: messageId,
|
||
metadata: {
|
||
prompt,
|
||
provider,
|
||
},
|
||
});
|
||
}
|
||
}
|
||
} catch (assetError) {
|
||
logger.warn("media_asset_registration_failed", assetError);
|
||
}
|
||
|
||
logger.info("job_completed", {
|
||
requestId,
|
||
generatedImageId: savedImage.id,
|
||
mediaAssetId,
|
||
processingSeconds: Number(processingTime.toFixed(2)),
|
||
provider,
|
||
savedToGallery: true,
|
||
});
|
||
|
||
return {
|
||
image_base64: resultBase64,
|
||
thumbnail_url: thumbnailUrl,
|
||
original_url: originalUrl,
|
||
processing_time: processingTime,
|
||
provider,
|
||
mediaAssetId,
|
||
generatedImageId: savedImage.id,
|
||
};
|
||
}
|
||
} else {
|
||
logger.info("gallery_save_skipped", { provider });
|
||
}
|
||
|
||
logger.info("job_completed", {
|
||
requestId,
|
||
processingSeconds: Number(processingTime.toFixed(2)),
|
||
provider,
|
||
savedToGallery: false,
|
||
});
|
||
|
||
return {
|
||
image_base64: resultBase64,
|
||
thumbnail_url: thumbnailUrl,
|
||
original_url: originalUrl,
|
||
processing_time: processingTime,
|
||
provider,
|
||
mediaAssetId: null,
|
||
generatedImageId: null,
|
||
};
|
||
}
|