ai-chatbot/lib/ai/rmbg/service.ts
2026-07-04 12:02:29 +08:00

742 lines
22 KiB
TypeScript
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import {
createMediaAssetRelation,
getMediaAssetsByUserAndUrls,
getSystemSetting,
upsertMediaAsset,
} from "@/lib/db/queries";
import { db } from "@/lib/db";
import { generatedImage } from "@/lib/db/schema";
import { createLogger } from "@/lib/observability/logger";
import { uploadFile } from "@/lib/storage/service";
const logger = createLogger("rmbg-service");
export type RmbgProvider = "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 = {
userId: string;
imageBuffer: Buffer;
imageBase64?: string;
prompt?: string;
chatId?: string | null;
messageId?: string | null;
skipGallerySave?: boolean;
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(message);
this.name = "RmbgServiceError";
this.status = status;
this.detail = detail;
}
}
const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));
const AOS_MAX_RESOLUTION = 4096;
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 = 60_000;
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);
}
};
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 }> {
try {
const sharp = await import("sharp");
const image = sharp.default(imageBuffer);
const metadata = await image.metadata();
const width = metadata.width;
const height = metadata.height;
const format = (metadata.format || "png").toLowerCase();
const ext = format === "jpeg" ? "jpg" : format;
const filenameExt =
ext === "jpg" || ext === "png" || ext === "webp" || ext === "tiff" ? ext : "png";
const mimeType =
ext === "jpg"
? "image/jpeg"
: ext === "webp"
? "image/webp"
: ext === "tiff"
? "image/tiff"
: "image/png";
if (!width || !height || (width <= AOS_MAX_RESOLUTION && height <= AOS_MAX_RESOLUTION)) {
return {
buffer: imageBuffer,
mimeType,
filename: `image.${filenameExt}`,
};
}
const resizedBuffer = await image
.resize(AOS_MAX_RESOLUTION, AOS_MAX_RESOLUTION, {
fit: "inside",
withoutEnlargement: true,
})
.toBuffer();
return {
buffer: resizedBuffer,
mimeType,
filename: `image.${filenameExt}`,
};
} catch (error) {
logger.warn("aos_image_preprocess_failed", error);
return { buffer: imageBuffer, mimeType: "image/png", filename: "image.png" };
}
}
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 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";
}
): 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 };
try {
const normalizedImage = await prepareImageForAos(imageBuffer);
const createForm = new FormData();
const imageBlob = new Blob([new Uint8Array(normalizedImage.buffer)], {
type: normalizedImage.mimeType,
});
createForm.append("image_file", imageBlob, normalizedImage.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 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任务创建失败: ${createPayload?.message || createResponse.statusText}`,
};
}
const taskId = createPayload.data.task_id as string;
const taskUrl = `${createUrl}/${taskId}`;
const maxPollCount = 60;
for (let pollCount = 0; pollCount < maxPollCount; pollCount += 1) {
await sleep(1000);
const pollResponse = await withTimeout(
(signal) =>
fetch(taskUrl, {
method: "GET",
headers,
signal,
}),
AOS_POLL_TIMEOUT_MS
);
const pollPayload = await pollResponse.json().catch(() => null);
if (!pollResponse.ok || !pollPayload?.data) {
logger.warn("aos_task_poll_failed", {
hasPayload: pollPayload != null,
status: pollResponse.status,
});
return {
success: false,
error: `AOS任务查询失败: ${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,
});
const imageResponse = await withTimeout(
(signal) =>
fetch(imageUrl, {
method: "GET",
signal,
}),
AOS_DOWNLOAD_TIMEOUT_MS
);
if (!imageResponse.ok) {
logger.warn("aos_result_download_failed", {
status: imageResponse.status,
taskId,
});
return {
success: false,
error: `AOS结果下载失败: ${imageResponse.status} ${imageResponse.statusText}`,
};
}
const resultBuffer = Buffer.from(await imageResponse.arrayBuffer());
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任务失败: ${errMessage || stateDetail || pollPayload?.message || "状态异常"}`,
};
}
}
logger.warn("aos_task_timeout", { maxPollCount });
return { success: false, error: "AOS任务超时(最多60秒)" };
} catch (error) {
if (error instanceof Error && error.name === "AbortError") {
return { success: false, error: "AOS请求超时,请稍后重试" };
}
logger.warn("aos_provider_request_error", error);
return { success: false, error: error instanceof Error ? error.message : "未知错误" };
}
}
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 === "photoroom"
? "photoroom-hd"
: provider === "aoscdn"
? "aos-segmentation"
: "rmbg-2.0";
}
export async function fetchImageBufferFromUrl(imageUrl: string): Promise<Buffer> {
try {
return await withTimeout(async (signal) => {
const imageResponse = await fetch(imageUrl, { 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;
}
logger.warn("source_image_fetch_failed", error);
throw new RmbgServiceError("无法获取图片,请检查图片链接是否有效", 400);
}
}
export async function executeRmbgJob(
params: ExecuteRmbgJobParams
): Promise<ExecuteRmbgJobResult> {
const {
userId,
imageBuffer,
imageBase64 = imageBuffer.toString("base64"),
prompt = "抠图处理",
chatId = null,
messageId = null,
skipGallerySave = false,
sourceImageUrl = null,
} = params;
const { provider, config } = await getValidatedRmbgConfig();
logger.info("job_started", {
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 === "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: config.aosForegroundType || "",
crop: config.aosCrop === 1 ? 1 : 0,
format: config.aosFormat === "jpg" ? "jpg" : "png",
});
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", {
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", {
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,
};
}