ai-chatbot/lib/ai/extend/service.ts

477 lines
14 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 { executeConfiguredVip98kTask } from "@/lib/ai/vip98k-runtime";
import { db } from "@/lib/db";
import { generatedImage } from "@/lib/db/schema";
import { uploadFile } from "@/lib/storage/service";
type ExtendConfig = {
provider?: "vip98k" | "dashscope";
apiKey?: string;
xScale?: number;
yScale?: number;
};
const VIP98K_EXPANSION_FALLBACK_MODEL = "google/nano-banana-pro/edit";
const VIP98K_SOURCE_DOWNLOAD_TIMEOUT_MS = 120_000;
export interface ExecuteExtendJobParams {
userId: string;
imageUrl?: string | null;
imageBase64?: string | null;
prompt?: string;
chatId?: string | null;
messageId?: string | null;
xScale?: number;
yScale?: number;
includeBase64?: boolean;
skipGallerySave?: boolean;
sourceImageUrl?: string | null;
requestId?: string | null;
}
export interface ExecuteExtendJobResult {
image_base64?: string;
thumbnail_url: string;
original_url: string;
processing_time: number;
mediaAssetId?: string | null;
generatedImageId?: string | null;
}
export class ExtendServiceError extends Error {
status: number;
detail?: string;
constructor(message: string, status = 500, detail?: string) {
super(message);
this.name = "ExtendServiceError";
this.status = status;
this.detail = detail;
}
}
async function createThumbnail(buffer: Buffer): Promise<Buffer> {
try {
const sharp = await import("sharp");
return await sharp.default(buffer)
.resize(500, null, {
fit: "inside",
withoutEnlargement: true,
})
.jpeg({ quality: 80 })
.toBuffer();
} catch (error) {
console.error("[Extend Service] Failed to create thumbnail:", error);
return buffer;
}
}
async function uploadBufferPairToS3(buffer: Buffer): Promise<{ thumbnail: string; original: string }> {
try {
const timestamp = Date.now();
const randomStr = Math.random().toString(36).slice(2, 8);
const thumbnailFilename = `extend-thumb-${timestamp}-${randomStr}.jpg`;
const originalFilename = `extend-orig-${timestamp}-${randomStr}.jpg`;
const thumbnailBuffer = await createThumbnail(buffer);
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),
]);
return {
thumbnail: thumbnailResult.url,
original: originalResult.url,
};
} catch (error) {
console.error("[Extend Service] Failed to upload to S3:", error);
throw new ExtendServiceError(
"图片上传到存储桶失败",
500,
error instanceof Error ? error.message : "未知错误",
);
}
}
async function uploadImageToS3ForAPI(imageBase64: string): Promise<string> {
try {
const buffer = Buffer.from(imageBase64, "base64");
const timestamp = Date.now();
const randomStr = Math.random().toString(36).slice(2, 8);
const filename = `extend-input-${timestamp}-${randomStr}.jpg`;
const arrayBuffer = buffer.buffer.slice(
buffer.byteOffset,
buffer.byteOffset + buffer.byteLength,
) as ArrayBuffer;
const result = await uploadFile(filename, arrayBuffer);
return result.url;
} catch (error) {
console.error("[Extend Service] Failed to upload input to S3:", error);
throw new ExtendServiceError(
"图片上传到存储桶失败",
500,
error instanceof Error ? error.message : "未知错误",
);
}
}
function isTemporaryVip98kExpansionError(error: unknown) {
const message = error instanceof Error ? error.message : String(error || "");
return /HTTP\s+(?:429|5\d\d)|temporar(?:y|ily)|timeout|timed out|unavailable/i.test(
message,
);
}
async function createVip98kExpansionCanvas(
imageUrl: string,
xScale: number,
yScale: number,
) {
const response = await fetch(imageUrl, {
cache: "no-store",
redirect: "follow",
signal: AbortSignal.timeout(VIP98K_SOURCE_DOWNLOAD_TIMEOUT_MS),
});
if (!response.ok) {
throw new ExtendServiceError(
"下载待拓展图片失败",
502,
`${response.status} ${response.statusText}`,
);
}
const sourceBuffer = Buffer.from(await response.arrayBuffer());
const sharp = await import("sharp");
const image = sharp.default(sourceBuffer, { failOn: "none" });
const metadata = await image.metadata();
const width = Number(metadata.width || 0);
const height = Number(metadata.height || 0);
if (!width || !height) {
throw new ExtendServiceError("无法识别待拓展图片尺寸", 422);
}
const targetWidth = Math.max(width, Math.round(width * xScale));
const targetHeight = Math.max(height, Math.round(height * yScale));
const left = Math.floor((targetWidth - width) / 2);
const top = Math.floor((targetHeight - height) / 2);
const canvas = await image
.ensureAlpha()
.extend({
left,
right: targetWidth - width - left,
top,
bottom: targetHeight - height - top,
background: { r: 0, g: 0, b: 0, alpha: 0 },
})
.png()
.toBuffer();
const arrayBuffer = canvas.buffer.slice(
canvas.byteOffset,
canvas.byteOffset + canvas.byteLength,
) as ArrayBuffer;
const uploaded = await uploadFile(
`extend-vip98k-canvas-${Date.now()}-${Math.random().toString(36).slice(2, 8)}.png`,
arrayBuffer,
);
return {
imageUrl: uploaded.url,
width,
height,
targetWidth,
targetHeight,
};
}
async function extendImageWithVip98kEditFallback(
imageUrl: string,
xScale: number,
yScale: number,
) {
const canvas = await createVip98kExpansionCanvas(imageUrl, xScale, yScale);
const result = await executeConfiguredVip98kTask({
modelId: VIP98K_EXPANSION_FALLBACK_MODEL,
kind: "image",
prompt: [
"Outpaint only the transparent regions of this image.",
"Keep every original pixel, subject, scale, position, perspective, lighting and style unchanged.",
`Complete the ${canvas.targetWidth}x${canvas.targetHeight} canvas naturally from the centered ${canvas.width}x${canvas.height} source.`,
"Do not crop, zoom, move, replace or redraw the original content.",
].join(" "),
imageUrls: [canvas.imageUrl],
});
const outputUrl = result.outputUrls[0];
if (!outputUrl) {
throw new ExtendServiceError("VIP98K 生成式扩图未返回图片", 502);
}
return outputUrl;
}
async function extendImage(
imageUrl: string,
config: ExtendConfig,
xScale: number,
yScale: number,
): Promise<string> {
if (config.provider === "vip98k") {
try {
const result = await executeConfiguredVip98kTask({
modelId: "image-expansion",
kind: "image-tool",
prompt: `Expand the image canvas naturally to ${xScale}x width and ${yScale}x height. Preserve the original subject, perspective, lighting and style.`,
imageUrls: [imageUrl],
imageParams: { xScale, yScale },
});
const outputUrl = result.outputUrls[0];
if (!outputUrl) {
throw new ExtendServiceError("VIP98K 图像拓展未返回图片", 502);
}
return outputUrl;
} catch (error) {
if (!isTemporaryVip98kExpansionError(error)) throw error;
console.warn(
"[Extend Service] VIP98K expansion tool is temporarily unavailable; using VIP98K generative outpainting fallback",
error instanceof Error ? error.message : String(error),
);
return extendImageWithVip98kEditFallback(imageUrl, xScale, yScale);
}
}
if (!config.apiKey) {
throw new ExtendServiceError("图像拓展API密钥未配置", 400);
}
const requestBody = {
model: "image-out-painting",
input: {
image_url: imageUrl,
},
parameters: {
x_scale: xScale,
y_scale: yScale,
best_quality: false,
limit_image_size: true,
add_watermark: false,
},
};
const response = await fetch(
"https://dashscope.aliyuncs.com/api/v1/services/aigc/image2image/out-painting",
{
method: "POST",
headers: {
"Content-Type": "application/json",
Authorization: `Bearer ${config.apiKey}`,
"X-DashScope-Async": "enable",
},
body: JSON.stringify(requestBody),
},
);
const initialResult = await response.json();
if (initialResult.code) {
throw new ExtendServiceError(
"图像拓展API错误",
500,
initialResult.message || JSON.stringify(initialResult),
);
}
const taskId = initialResult.output.task_id;
const taskStatusUrl = `https://dashscope.aliyuncs.com/api/v1/tasks/${taskId}`;
await new Promise((resolve) => setTimeout(resolve, 3000));
for (let attempt = 0; attempt < 60; attempt += 1) {
const statusResponse = await fetch(taskStatusUrl, {
headers: {
Authorization: `Bearer ${config.apiKey}`,
},
});
const result = await statusResponse.json();
if (result.output?.task_status === "SUCCEEDED") {
return result.output.output_image_url;
}
if (result.output?.task_status === "FAILED") {
throw new ExtendServiceError(
"图像拓展任务失败",
500,
result.output.message || result.output.code || "未知错误",
);
}
await new Promise((resolve) => setTimeout(resolve, 2000));
}
throw new ExtendServiceError("图像拓展任务超时", 500);
}
export function inferExtendOptionsFromText(text: string): Pick<ExecuteExtendJobParams, "xScale" | "yScale"> {
const normalizedText = String(text || "").toLowerCase();
const scaleMatch = normalizedText.match(/([1-4](?:\.\d+)?)\s*倍/);
const scale = scaleMatch ? Number(scaleMatch[1]) : undefined;
return {
xScale: scale,
yScale: scale,
};
}
async function getValidatedExtendConfig(): Promise<ExtendConfig> {
const extendConfig = await getSystemSetting({ key: "extend_api_config" });
if (!extendConfig || !extendConfig.value) {
throw new ExtendServiceError("图像拓展API配置未设置,请在系统设置中配置", 400);
}
const config = extendConfig.value as ExtendConfig;
if (config.provider !== "vip98k" && !config.apiKey) {
throw new ExtendServiceError("图像拓展API密钥未配置", 400);
}
return config;
}
export async function executeExtendJob(
params: ExecuteExtendJobParams,
): Promise<ExecuteExtendJobResult> {
const {
userId,
imageUrl = null,
imageBase64 = null,
prompt = "图像拓展处理",
chatId = null,
messageId = null,
includeBase64 = false,
skipGallerySave = false,
sourceImageUrl = imageUrl || null,
requestId = null,
} = params;
if (!imageUrl && !imageBase64) {
throw new ExtendServiceError("缺少 image_url 或 image_base64 参数", 400);
}
const config = await getValidatedExtendConfig();
const xScale = Math.max(1.1, Math.min(Number(params.xScale || config.xScale || 1.5), 4));
const yScale = Math.max(1.1, Math.min(Number(params.yScale || config.yScale || 1.5), 4));
const inputImageUrl = imageUrl || (imageBase64 ? await uploadImageToS3ForAPI(imageBase64) : "");
if (!inputImageUrl) {
throw new ExtendServiceError("无法获取输入图片URL", 400);
}
const startTime = Date.now();
const extendedImageUrl = await extendImage(inputImageUrl, config, xScale, yScale);
const imageResponse = await fetch(extendedImageUrl);
if (!imageResponse.ok) {
throw new ExtendServiceError(
"下载拓展图片失败",
500,
`${imageResponse.status} ${imageResponse.statusText}`,
);
}
const buffer = Buffer.from(await imageResponse.arrayBuffer());
const imageBase64Result = includeBase64 ? buffer.toString("base64") : undefined;
const uploadResult = await uploadBufferPairToS3(buffer);
console.log("[Extend Service] Result uploaded:", {
requestId,
includeBase64,
inputImageUrl: inputImageUrl.slice(0, 120),
extendedImageUrl: extendedImageUrl.slice(0, 120),
outputBytes: buffer.length,
});
let mediaAssetId: string | null = null;
let generatedImageId: string | null = null;
if (!skipGallerySave) {
try {
const [savedImage] = await db
.insert(generatedImage)
.values({
userId,
chatId,
messageId,
thumbnailUrl: uploadResult.thumbnail,
originalUrl: uploadResult.original,
prompt,
modelId: config.provider === "vip98k" ? "image-expansion" : "extend-1.0",
imageSize: "2K",
})
.returning();
generatedImageId = savedImage?.id || null;
const sourceAssets = sourceImageUrl
? await getMediaAssetsByUserAndUrls({
userId,
urls: [sourceImageUrl],
})
: [];
const media = await upsertMediaAsset({
userId,
chatId,
mediaType: "image",
sourceType: "edited",
status: "ready",
originalUrl: uploadResult.original,
previewUrl: uploadResult.original,
thumbnailUrl: uploadResult.thumbnail,
displayName: prompt.slice(0, 80),
prompt,
modelId: config.provider === "vip98k" ? "image-expansion" : "extend-1.0",
sourceMessageId: messageId,
sourceTable: "GeneratedImage",
sourceRowId: generatedImageId,
primaryParentAssetId: sourceAssets[0]?.id ?? null,
metadata: {
xScale,
yScale,
imageSize: "2K",
},
});
mediaAssetId = media?.id || null;
if (media?.id && sourceAssets.length > 0) {
for (const sourceAsset of sourceAssets) {
await createMediaAssetRelation({
userId,
chatId,
fromAssetId: sourceAsset.id,
toAssetId: media.id,
relationType: "edited_from",
sourceMessageId: messageId,
metadata: {
prompt,
tool: "extend",
xScale,
yScale,
},
});
}
}
} catch (error) {
console.error("[Extend Service] Failed to save result:", error);
throw new ExtendServiceError(
"图像拓展成功但保存结果失败",
500,
error instanceof Error ? error.message : "未知错误",
);
}
}
return {
image_base64: imageBase64Result,
thumbnail_url: uploadResult.thumbnail,
original_url: uploadResult.original,
processing_time: Number(((Date.now() - startTime) / 1000).toFixed(2)),
mediaAssetId,
generatedImageId,
};
}