Your Name f1ffeed9f2
Some checks failed
Lint / build (20) (push) Has been cancelled
Playwright Tests / test (push) Has been cancelled
feat: add canvas layer split workflow
2026-03-24 02:07:45 +08:00

660 lines
20 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 "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 { uploadFile } from "@/lib/storage/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;
};
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<LayerSplitConfig> {
const raw =
value && typeof value === "object" && !Array.isArray(value)
? (value as Record<string, unknown>)
: {};
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<string, unknown>;
}) {
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<ReplicatePredictionResponse> {
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<string, unknown>)) {
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<Array<{ buffer: Buffer; filename: string; mimeType: string }>> {
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<Array<{ buffer: Buffer; filename: string; mimeType: string }>> {
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<Buffer> {
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) {
console.error("[LayerSplit] Failed to create thumbnail:", error);
return buffer;
}
}
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),
]);
return {
thumbnail: thumbnailResult.url,
original: originalResult.url,
};
}
async function runReplicateLayerSplit(params: {
apiKey: string;
modelId: string;
imageUrl: string;
config: Required<LayerSplitConfig>;
prompt?: string;
}): Promise<ReplicatePredictionResponse> {
const baseInput: Record<string, unknown> = {
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;
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)) {
continue;
}
}
throw new LayerSplitServiceError(
"图层拆分请求失败",
502,
lastError || toReplicateErrorMessage(lastPayload || {}, "Replicate 请求失败")
);
}
export async function executeLayerSplitJob(params: LayerSplitParams): Promise<ExecuteLayerSplitJobResult> {
const startedAt = Date.now();
const config = await getValidatedLayerSplitConfig();
const payload = await runReplicateLayerSplit({
apiKey: config.replicateApiToken,
modelId: config.replicateModel,
imageUrl: params.imageUrl,
config,
prompt: params.prompt,
});
const layerBuffers = await resolveLayerArtifacts(payload);
const totalLayers = layerBuffers.length;
const sourceAssets = await getMediaAssetsByUserAndUrls({
userId: params.userId,
urls: [params.imageUrl],
}).catch(() => []);
const results: LayerSplitResultItem[] = [];
for (let index = 0; index < layerBuffers.length; index += 1) {
const layer = layerBuffers[index];
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: params.imageUrl,
sourceFilename: layer.filename,
},
});
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,
});
}
return {
layers: results,
provider: "replicate",
model: config.replicateModel,
processingTime: (Date.now() - startedAt) / 1000,
};
}