354 lines
10 KiB
TypeScript
354 lines
10 KiB
TypeScript
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";
|
||
|
||
type ExtendConfig = {
|
||
apiKey?: string;
|
||
xScale?: number;
|
||
yScale?: number;
|
||
};
|
||
|
||
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 : "未知错误",
|
||
);
|
||
}
|
||
}
|
||
|
||
async function extendImage(
|
||
imageUrl: string,
|
||
apiKey: string,
|
||
xScale: number,
|
||
yScale: number,
|
||
): Promise<string> {
|
||
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 ${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 ${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.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.apiKey!, 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: "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: "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,
|
||
};
|
||
}
|