import "server-only"; import JSZip from "jszip"; 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 { ensurePublicImageUrl } from "@/lib/realtime/image-input"; import { uploadFile } from "@/lib/storage/service"; const logger = createLogger("layer-split-service"); const LAYER_SPLIT_CONFIG_KEY = "layer_split_config"; const DEFAULT_REPLICATE_BASE_URL = "https://api.replicate.com/v1"; const DEFAULT_LAYER_SPLIT_MODEL = "qwen/qwen-image-layered"; const DEFAULT_LAYERS = 4; const DEFAULT_RESOLUTION = 640; const DEFAULT_TRUE_CFG_SCALE = 4; const DEFAULT_INFERENCE_STEPS = 30; const DEFAULT_TIMEOUT_MS = 5 * 60 * 1000; type LayerSplitConfig = { enabled?: boolean; provider?: "replicate"; replicateApiToken?: string; replicateModel?: string; layers?: number; resolution?: number; trueCfgScale?: number; numInferenceSteps?: number; cfgNormalize?: boolean; useEnPrompt?: boolean; negativePrompt?: string; timeoutMs?: number; }; type LayerSplitParams = { userId: string; imageUrl: string; prompt?: string; chatId?: string | null; messageId?: string | null; sourceNodeId?: string | null; skipGallerySave?: boolean; }; type ReplicatePredictionResponse = { id?: string; status?: string; output?: unknown; error?: string; detail?: string; message?: string; urls?: { get?: string; }; }; export type LayerSplitResultItem = { originalUrl: string; thumbnailUrl: string; generatedImageId?: string | null; mediaAssetId?: string | null; title: string; layerIndex: number; totalLayers: number; bounds?: { x: number; y: number; width: number; height: number; originalWidth: number; originalHeight: number; coverage: number; } | null; }; export type ExecuteLayerSplitJobResult = { layers: LayerSplitResultItem[]; provider: "replicate"; model: string; processingTime: number; }; export class LayerSplitServiceError extends Error { status: number; detail?: string; constructor(message: string, status = 500, detail?: string) { super(message); this.name = "LayerSplitServiceError"; this.status = status; this.detail = detail; } } function trimString(value: unknown) { return String(value || "").trim(); } function splitReplicateModelIdentifier(modelId: string) { const [owner, ...rest] = modelId.split("/").filter(Boolean); const name = rest.join("/"); if (!owner || !name) { throw new LayerSplitServiceError(`无效的 Replicate 模型标识:${modelId}`, 400); } return { owner, name }; } function normalizeLayerSplitConfig(value: unknown): Required { const raw = value && typeof value === "object" && !Array.isArray(value) ? (value as Record) : {}; return { enabled: raw.enabled === true, provider: "replicate", replicateApiToken: trimString(raw.replicateApiToken), replicateModel: trimString(raw.replicateModel) || DEFAULT_LAYER_SPLIT_MODEL, layers: Math.max(2, Math.min(12, Math.round(Number(raw.layers) || DEFAULT_LAYERS))), resolution: Math.max(256, Math.min(1536, Math.round(Number(raw.resolution) || DEFAULT_RESOLUTION))), trueCfgScale: Math.max(1, Math.min(12, Number(raw.trueCfgScale) || DEFAULT_TRUE_CFG_SCALE)), numInferenceSteps: Math.max( 1, Math.min(80, Math.round(Number(raw.numInferenceSteps) || DEFAULT_INFERENCE_STEPS)) ), cfgNormalize: raw.cfgNormalize !== false, useEnPrompt: raw.useEnPrompt !== false, negativePrompt: trimString(raw.negativePrompt), timeoutMs: Math.max(30_000, Math.min(15 * 60_000, Math.round(Number(raw.timeoutMs) || DEFAULT_TIMEOUT_MS))), }; } async function getValidatedLayerSplitConfig() { const setting = await getSystemSetting({ key: LAYER_SPLIT_CONFIG_KEY }); const config = normalizeLayerSplitConfig(setting?.value); if (!config.enabled) { throw new LayerSplitServiceError("图层拆分功能未启用,请先在本地AI设置中开启", 400); } if (!config.replicateApiToken) { throw new LayerSplitServiceError("未配置 Replicate API Token,请先在本地AI设置中保存", 400); } return config; } function toReplicateErrorMessage(payload: any, fallback: string) { const message = trimString(payload?.detail || payload?.error || payload?.message); return message || fallback; } function shouldRetryWithAlternativeInput(errorMessage: string) { const normalized = errorMessage.toLowerCase(); return ( normalized.includes("validation") || normalized.includes("schema") || normalized.includes("image") || normalized.includes("input") || normalized.includes("field") ); } async function createPrediction(params: { apiKey: string; modelId: string; input: Record; }) { const { owner, name } = splitReplicateModelIdentifier(params.modelId); const endpoint = `${DEFAULT_REPLICATE_BASE_URL}/models/${encodeURIComponent(owner)}/${encodeURIComponent(name)}/predictions`; const response = await fetch(endpoint, { method: "POST", headers: { Authorization: `Bearer ${params.apiKey}`, "Content-Type": "application/json", Prefer: "wait=5", }, body: JSON.stringify({ input: params.input }), cache: "no-store", }); const payload = (await response.json().catch(() => null)) as ReplicatePredictionResponse | null; return { response, payload }; } async function pollPrediction(params: { apiKey: string; taskId?: string; pollingUrl?: string; timeoutMs: number; }): Promise { const deadline = Date.now() + params.timeoutMs; let lastMessage = ""; while (Date.now() < deadline) { const taskId = trimString(params.taskId); const url = trimString(params.pollingUrl) || `${DEFAULT_REPLICATE_BASE_URL}/predictions/${encodeURIComponent(taskId)}`; const response = await fetch(url, { method: "GET", headers: { Authorization: `Bearer ${params.apiKey}`, Accept: "application/json", }, cache: "no-store", }); const payload = (await response.json().catch(() => null)) as ReplicatePredictionResponse | null; const status = trimString(payload?.status).toLowerCase(); if (status) { lastMessage = status; } if (response.ok && status === "succeeded") { return payload || {}; } if (status === "failed" || status === "canceled") { throw new LayerSplitServiceError( "图层拆分任务失败", 502, toReplicateErrorMessage(payload, "Replicate 任务执行失败") ); } if (!response.ok) { throw new LayerSplitServiceError( `图层拆分轮询失败(${response.status})`, 502, toReplicateErrorMessage(payload, `Replicate 轮询失败(${response.status})`) ); } await new Promise((resolve) => setTimeout(resolve, 600)); } throw new LayerSplitServiceError( "图层拆分任务超时", 504, `Replicate 轮询超时${lastMessage ? `(${lastMessage})` : ""}` ); } function collectArtifactStrings(value: unknown, bucket: string[] = []): string[] { if (typeof value === "string") { const next = trimString(value); if (next) bucket.push(next); return bucket; } if (Array.isArray(value)) { for (const item of value) { collectArtifactStrings(item, bucket); } return bucket; } if (value && typeof value === "object") { for (const [key, item] of Object.entries(value as Record)) { if ( key === "url" || key === "href" || key === "download" || key === "file" || key === "image" || key === "get" || key === "output" || key === "outputs" || key === "layers" || key === "files" || key === "result" ) { collectArtifactStrings(item, bucket); } } } return bucket; } function dedupeStrings(values: string[]) { return Array.from(new Set(values.map((value) => trimString(value)).filter(Boolean))); } function extFromFilename(filename: string): string { const clean = trimString(filename).split("?")[0]; const ext = clean.split(".").pop()?.toLowerCase() || ""; return ext; } function deriveFilenameFromUrl(url: string, fallback: string) { try { const parsed = new URL(url); const candidate = parsed.pathname.split("/").pop() || ""; return candidate || fallback; } catch { return fallback; } } function detectMimeType(buffer: Buffer, fallback?: string): string { if (buffer.length >= 8 && buffer.subarray(0, 8).equals(Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]))) { return "image/png"; } if (buffer.length >= 3 && buffer[0] === 0xff && buffer[1] === 0xd8 && buffer[2] === 0xff) { return "image/jpeg"; } if (buffer.length >= 12 && buffer.toString("ascii", 0, 4) === "RIFF" && buffer.toString("ascii", 8, 12) === "WEBP") { return "image/webp"; } if (buffer.length >= 4 && buffer.toString("ascii", 0, 4) === "PK\u0003\u0004") { return "application/zip"; } if (buffer.length >= 4 && buffer.toString("ascii", 0, 4) === "8BPS") { return "image/vnd.adobe.photoshop"; } const normalizedFallback = trimString(fallback).toLowerCase(); if (normalizedFallback.startsWith("image/")) return normalizedFallback; if (normalizedFallback === "application/zip") return normalizedFallback; return "application/octet-stream"; } function extFromMimeType(mimeType: string) { if (mimeType === "image/png") return "png"; if (mimeType === "image/jpeg") return "jpg"; if (mimeType === "image/webp") return "webp"; if (mimeType === "application/zip") return "zip"; if (mimeType === "image/vnd.adobe.photoshop") return "psd"; return "bin"; } async function fetchArtifact(url: string): Promise<{ buffer: Buffer; mimeType: string; filename: string }> { const response = await fetch(url, { headers: { "User-Agent": "AI-Chatbot-LayerSplit/1.0", }, cache: "no-store", }); if (!response.ok) { throw new LayerSplitServiceError( `下载图层拆分结果失败(${response.status})`, 502, `${response.status} ${response.statusText}` ); } const contentType = trimString(response.headers.get("content-type")); const buffer = Buffer.from(await response.arrayBuffer()); const filename = deriveFilenameFromUrl(url, `artifact.${extFromMimeType(contentType || "application/octet-stream")}`); const mimeType = detectMimeType(buffer, contentType); return { buffer, mimeType, filename }; } async function extractLayerBuffersFromZip(zipBuffer: Buffer): Promise> { const zip = await JSZip.loadAsync(zipBuffer); const entries: Array<{ buffer: Buffer; filename: string; mimeType: string }> = []; for (const [name, file] of Object.entries(zip.files)) { if (file.dir) continue; const ext = extFromFilename(name); if (!["png", "jpg", "jpeg", "webp"].includes(ext)) continue; const fileBuffer = await file.async("nodebuffer"); const mimeType = detectMimeType(fileBuffer, ext === "png" ? "image/png" : ext === "webp" ? "image/webp" : "image/jpeg"); if (!mimeType.startsWith("image/")) continue; entries.push({ buffer: fileBuffer, filename: name.split("/").pop() || `layer-${entries.length + 1}.${ext || "png"}`, mimeType, }); } return entries; } async function resolveLayerArtifacts(payload: ReplicatePredictionResponse): Promise> { const artifacts = dedupeStrings(collectArtifactStrings(payload?.output ?? payload)); if (artifacts.length === 0) { throw new LayerSplitServiceError("图层拆分未返回任何可用结果", 502); } const imageEntries: Array<{ buffer: Buffer; filename: string; mimeType: string }> = []; for (const artifact of artifacts) { if (artifact.startsWith("data:")) { const match = artifact.match(/^data:([^;,]+)?;base64,(.+)$/i); if (!match) continue; const mimeType = detectMimeType(Buffer.from(match[2], "base64"), match[1] || "image/png"); const buffer = Buffer.from(match[2], "base64"); if (mimeType.startsWith("image/")) { imageEntries.push({ buffer, mimeType, filename: `layer-${imageEntries.length + 1}.${extFromMimeType(mimeType)}`, }); continue; } if (mimeType === "application/zip") { const zipped = await extractLayerBuffersFromZip(buffer); imageEntries.push(...zipped); } continue; } if (!/^https?:\/\//i.test(artifact)) continue; const downloaded = await fetchArtifact(artifact); if (downloaded.mimeType.startsWith("image/")) { imageEntries.push(downloaded); continue; } if (downloaded.mimeType === "application/zip") { const zipped = await extractLayerBuffersFromZip(downloaded.buffer); imageEntries.push(...zipped); continue; } } if (imageEntries.length === 0) { throw new LayerSplitServiceError("图层拆分返回了结果文件,但未找到可回画布的图片图层", 502); } return imageEntries; } async function createThumbnail(buffer: Buffer, mimeType: string): Promise { try { const sharp = await import("sharp"); const pipeline = sharp.default(buffer).resize(500, null, { fit: "inside", withoutEnlargement: true, }); if (mimeType === "image/jpeg") { return await pipeline.jpeg({ quality: 82 }).toBuffer(); } if (mimeType === "image/webp") { return await pipeline.webp({ quality: 82 }).toBuffer(); } return await pipeline.png({ compressionLevel: 8 }).toBuffer(); } catch (error) { logger.warn("thumbnail_creation_failed", error); return buffer; } } async function detectVisibleBounds(buffer: Buffer) { try { const sharp = await import("sharp"); const { data, info } = await sharp.default(buffer) .ensureAlpha() .raw() .toBuffer({ resolveWithObject: true }); const width = Number(info.width || 0); const height = Number(info.height || 0); const channels = Number(info.channels || 4); if (width <= 0 || height <= 0 || channels < 4) { return null; } let minX = width; let minY = height; let maxX = -1; let maxY = -1; let visiblePixels = 0; for (let y = 0; y < height; y += 1) { for (let x = 0; x < width; x += 1) { const offset = (y * width + x) * channels; const alpha = data[offset + 3]; if (alpha <= 8) continue; visiblePixels += 1; if (x < minX) minX = x; if (y < minY) minY = y; if (x > maxX) maxX = x; if (y > maxY) maxY = y; } } if (maxX < minX || maxY < minY || visiblePixels <= 0) { return null; } return { x: minX, y: minY, width: maxX - minX + 1, height: maxY - minY + 1, originalWidth: width, originalHeight: height, coverage: visiblePixels / (width * height), }; } catch (error) { logger.warn("visible_bounds_detection_failed", error); return null; } } async function uploadBufferPair(buffer: Buffer, mimeType: string, prefix: string) { const ext = extFromMimeType(mimeType); const timestamp = Date.now(); const random = Math.random().toString(36).slice(2, 8); const thumbFilename = `${prefix}-thumb-${timestamp}-${random}.${ext}`; const origFilename = `${prefix}-orig-${timestamp}-${random}.${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(thumbFilename, thumbnailArrayBuffer), uploadFile(origFilename, originalArrayBuffer), ]); logger.info("layer_upload_completed", { mimeType, originalBytes: buffer.length, originalUrl: originalResult.url, prefix, thumbnailBytes: thumbnailBuffer.length, thumbnailUrl: thumbnailResult.url, }); return { thumbnail: thumbnailResult.url, original: originalResult.url, }; } async function runReplicateLayerSplit(params: { apiKey: string; modelId: string; imageUrl: string; config: Required; prompt?: string; }): Promise { const baseInput: Record = { max_num_layers: params.config.layers, resolution: params.config.resolution, true_cfg_scale: params.config.trueCfgScale, num_inference_steps: params.config.numInferenceSteps, cfg_normalize: params.config.cfgNormalize, use_en_prompt: params.config.useEnPrompt, }; const negativePrompt = trimString(params.config.negativePrompt); if (negativePrompt) { baseInput.negative_prompt = negativePrompt; } const prompt = trimString(params.prompt); if (prompt) { baseInput.prompt = prompt; } const candidateInputs = [ { ...baseInput, image: params.imageUrl }, { ...baseInput, input_image: params.imageUrl }, { ...baseInput, images: [params.imageUrl] }, ]; let lastPayload: ReplicatePredictionResponse | null = null; let lastError = ""; for (let index = 0; index < candidateInputs.length; index += 1) { const { response, payload } = await createPrediction({ apiKey: params.apiKey, modelId: params.modelId, input: candidateInputs[index], }); lastPayload = payload; logger.debug("replicate_prediction_attempt_completed", { attempt: index + 1, hasPayload: payload != null, status: response.status, }); if ( response.ok && (trimString(payload?.status) === "succeeded" || trimString(payload?.id) || trimString(payload?.urls?.get)) ) { if (trimString(payload?.status).toLowerCase() === "succeeded") { return payload || {}; } return await pollPrediction({ apiKey: params.apiKey, taskId: trimString(payload?.id), pollingUrl: trimString(payload?.urls?.get), timeoutMs: params.config.timeoutMs, }); } const errorMessage = toReplicateErrorMessage(payload, `Replicate 请求失败(${response.status})`); lastError = errorMessage; if (response.status === 401 || response.status === 403) { throw new LayerSplitServiceError("Replicate API Token 无效或无权限", 401, errorMessage); } if (response.status === 404) { throw new LayerSplitServiceError(`Replicate 模型不存在:${params.modelId}`, 400, errorMessage); } if (index < candidateInputs.length - 1 && shouldRetryWithAlternativeInput(errorMessage)) { logger.warn("replicate_retrying_with_alternative_input", { attempt: index + 1, reasonLength: errorMessage.length, }); continue; } } throw new LayerSplitServiceError( "图层拆分请求失败", 502, lastError || toReplicateErrorMessage(lastPayload || {}, "Replicate 请求失败") ); } export async function executeLayerSplitJob(params: LayerSplitParams): Promise { const startedAt = Date.now(); const config = await getValidatedLayerSplitConfig(); logger.info("job_started", { hasChatId: Boolean(params.chatId), hasMessageId: Boolean(params.messageId), hasPrompt: Boolean(params.prompt), hasSourceNodeId: Boolean(params.sourceNodeId), skipGallerySave: Boolean(params.skipGallerySave), }); const normalizedImageUrl = await ensurePublicImageUrl({ imageUrl: params.imageUrl, prefix: "layer-split-input", }); const payload = await runReplicateLayerSplit({ apiKey: config.replicateApiToken, modelId: config.replicateModel, imageUrl: normalizedImageUrl, config, prompt: params.prompt, }); const layerBuffers = await resolveLayerArtifacts(payload); const totalLayers = layerBuffers.length; logger.info("layer_artifacts_resolved", { model: config.replicateModel, totalLayers, }); const sourceAssets = await getMediaAssetsByUserAndUrls({ userId: params.userId, urls: normalizedImageUrl === params.imageUrl ? [params.imageUrl] : [params.imageUrl, normalizedImageUrl], }).catch(() => []); const results: LayerSplitResultItem[] = []; for (let index = 0; index < layerBuffers.length; index += 1) { const layer = layerBuffers[index]; const visibleBounds = await detectVisibleBounds(layer.buffer); const uploadResult = await uploadBufferPair(layer.buffer, layer.mimeType, `layer-split-${index + 1}`); let generatedImageId: string | null = null; let mediaAssetId: string | null = null; if (!params.skipGallerySave) { const [savedImage] = await db .insert(generatedImage) .values({ userId: params.userId, chatId: params.chatId || null, messageId: params.messageId || null, thumbnailUrl: uploadResult.thumbnail, originalUrl: uploadResult.original, prompt: params.prompt || `图层拆分 ${index + 1}/${totalLayers}`, modelId: config.replicateModel, createdAt: new Date(), }) .returning({ id: generatedImage.id }); generatedImageId = savedImage?.id || null; const media = await upsertMediaAsset({ userId: params.userId, chatId: params.chatId || null, mediaType: "image", sourceType: "edited", status: "ready", originalUrl: uploadResult.original, previewUrl: uploadResult.original, thumbnailUrl: uploadResult.thumbnail, displayName: `图层拆分 ${index + 1}/${totalLayers}`.slice(0, 80), prompt: params.prompt || `图层拆分 ${index + 1}/${totalLayers}`, modelId: config.replicateModel, sourceMessageId: params.messageId || null, sourceTable: "GeneratedImage", sourceRowId: generatedImageId, sourceNodeId: params.sourceNodeId || null, primaryParentAssetId: sourceAssets[0]?.id ?? null, metadata: { layerIndex: index + 1, totalLayers, provider: "replicate", model: config.replicateModel, sourceImageUrl: normalizedImageUrl, sourceFilename: layer.filename, bounds: visibleBounds, }, }); mediaAssetId = media?.id || null; if (media && sourceAssets.length > 0) { for (const sourceAsset of sourceAssets) { await createMediaAssetRelation({ userId: params.userId, fromAssetId: sourceAsset.id, toAssetId: media.id, relationType: "edited_from", chatId: params.chatId || null, sourceMessageId: params.messageId || null, sourceNodeId: params.sourceNodeId || null, metadata: { tool: "layer_split", layerIndex: index + 1, totalLayers, }, }).catch(() => null); } } } results.push({ originalUrl: uploadResult.original, thumbnailUrl: uploadResult.thumbnail, generatedImageId, mediaAssetId, title: `图层${index + 1}`, layerIndex: index + 1, totalLayers, bounds: visibleBounds, }); } const processingTime = (Date.now() - startedAt) / 1000; logger.info("job_completed", { layerCount: results.length, processingSeconds: Number(processingTime.toFixed(2)), savedToGallery: !params.skipGallerySave, }); return { layers: results, provider: "replicate", model: config.replicateModel, processingTime, }; }