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 { 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 = { userId: string; imageBuffer: Buffer; imageBase64?: string; foregroundType?: "" | "person" | "object" | "stamp"; 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 = 90_000; const IMAGE_URL_FETCH_RETRIES = 2; async function createThumbnail(buffer: Buffer, mimeType: string): Promise { 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 ( fn: (signal: AbortSignal) => Promise, timeoutMs: number ): Promise => { 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 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"; } ): 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 === "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 { 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 { const { userId, imageBuffer, imageBase64 = imageBuffer.toString("base64"), foregroundType, 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 === "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", }); 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, }; }