1089 lines
33 KiB
TypeScript
1089 lines
33 KiB
TypeScript
import { createHash } from "node:crypto";
|
||
import { and, eq, or } from "drizzle-orm";
|
||
import { NextRequest, NextResponse } from "next/server";
|
||
import { auth } from "@/app/(auth)/auth";
|
||
import { db } from "@/lib/db";
|
||
import {
|
||
consumeUserPoints,
|
||
createMediaAssetRelation,
|
||
getPointsSystemConfig,
|
||
getMediaAssetsByUserAndUrls,
|
||
getSystemSetting,
|
||
getUserById,
|
||
upsertMediaAsset,
|
||
} from "@/lib/db/queries";
|
||
import {
|
||
enhanceProTask,
|
||
generatedImage,
|
||
pointsLedger,
|
||
} from "@/lib/db/schema";
|
||
import {
|
||
ENHANCE_PRO_KREA_SETTING_KEY,
|
||
normalizeEnhanceProKreaConfig,
|
||
type EnhanceProKreaConfig,
|
||
type EnhanceProKreaModel,
|
||
type EnhanceProOutputFormat,
|
||
type EnhanceProSubjectDetection,
|
||
} from "@/lib/enhance-pro/config";
|
||
import { calculateEnhanceProPointsEstimate } from "@/lib/points/config";
|
||
import { uploadFile } from "@/lib/storage/service";
|
||
|
||
const ENHANCE_PRO_CREATE_TIMEOUT_MS = 120_000;
|
||
const ENHANCE_PRO_TASK_POLL_INTERVAL_MS = 1_500;
|
||
const ENHANCE_PRO_MODEL_ID = "krea-topaz-generative";
|
||
const ENHANCE_PRO_MAX_RESULT_IMAGES = 4;
|
||
let enhanceProKeyPoolCursor = 0;
|
||
|
||
type EnhanceProRequestBody = {
|
||
imageUrl?: string;
|
||
image_base64?: string;
|
||
mimeType?: string;
|
||
prompt?: string;
|
||
chatId?: string | null;
|
||
messageId?: string | null;
|
||
clientRequestId?: string | null;
|
||
width?: number;
|
||
height?: number;
|
||
batchSize?: number;
|
||
seed?: number;
|
||
output_format?: EnhanceProOutputFormat;
|
||
subject_detection?: EnhanceProSubjectDetection;
|
||
face_enhancement?: boolean;
|
||
face_enhancement_creativity?: number;
|
||
face_enhancement_strength?: number;
|
||
crop_to_fill?: boolean;
|
||
model?: EnhanceProKreaModel;
|
||
upscaling_activated?: boolean;
|
||
image_scaling_factor?: number;
|
||
creativity?: number;
|
||
texture?: number;
|
||
sharpen?: number;
|
||
denoise?: number;
|
||
detail?: number;
|
||
};
|
||
|
||
type FinalizedImageResult = {
|
||
thumbnailUrl: string;
|
||
originalUrl: string;
|
||
width?: number;
|
||
height?: number;
|
||
};
|
||
|
||
type EnhanceProStoredRequest = {
|
||
prompt: string;
|
||
sourceImageUrl: string;
|
||
width: number;
|
||
height: number;
|
||
batchSize: number;
|
||
model: EnhanceProKreaModel;
|
||
upscalingActivated: boolean;
|
||
imageScalingFactor: number;
|
||
chatId: string | null;
|
||
messageId: string | null;
|
||
};
|
||
|
||
type EnhanceProSuccessPayload = Awaited<
|
||
ReturnType<typeof finalizeEnhanceProTaskResult>
|
||
> & {
|
||
chargedPoints: number;
|
||
};
|
||
|
||
type EnhanceProTaskRecord = typeof enhanceProTask.$inferSelect;
|
||
|
||
function clampNumber(value: unknown, min: number, max: number, fallback: number) {
|
||
const parsed = Number(value);
|
||
if (!Number.isFinite(parsed)) return fallback;
|
||
return Math.min(max, Math.max(min, parsed));
|
||
}
|
||
|
||
function clampInteger(value: unknown, min: number, max: number, fallback: number) {
|
||
return Math.round(clampNumber(value, min, max, fallback));
|
||
}
|
||
|
||
function normalizeBoolean(value: unknown, fallback = false) {
|
||
if (typeof value === "boolean") return value;
|
||
if (value === "true") return true;
|
||
if (value === "false") return false;
|
||
return fallback;
|
||
}
|
||
|
||
function normalizeJobStatus(value: unknown) {
|
||
return String(value || "").trim().toLowerCase();
|
||
}
|
||
|
||
function normalizeClientRequestId(value: unknown) {
|
||
return String(value || "").trim().slice(0, 64);
|
||
}
|
||
|
||
function getEnhanceProKeyFingerprint(apiKey: string) {
|
||
return createHash("sha256")
|
||
.update(String(apiKey || "").trim())
|
||
.digest("hex")
|
||
.slice(0, 16);
|
||
}
|
||
|
||
function getEnhanceProEffectiveApiKeys(config: EnhanceProKreaConfig) {
|
||
const pooledKeys = Array.isArray(config.apiKeys) ? config.apiKeys : [];
|
||
if (pooledKeys.length > 0) return pooledKeys;
|
||
return config.apiKey ? [config.apiKey] : [];
|
||
}
|
||
|
||
function resolveEnhanceProApiKey(
|
||
config: EnhanceProKreaConfig,
|
||
preferredKeyId?: string | null,
|
||
) {
|
||
const keys = getEnhanceProEffectiveApiKeys(config);
|
||
if (keys.length === 0) return null;
|
||
|
||
const normalizedPreferredKeyId = String(preferredKeyId || "").trim();
|
||
if (normalizedPreferredKeyId) {
|
||
const matchedKey = keys.find(
|
||
(apiKey) => getEnhanceProKeyFingerprint(apiKey) === normalizedPreferredKeyId,
|
||
);
|
||
if (matchedKey) {
|
||
return {
|
||
apiKey: matchedKey,
|
||
keyId: normalizedPreferredKeyId,
|
||
};
|
||
}
|
||
}
|
||
|
||
const selectedKey = keys[enhanceProKeyPoolCursor % keys.length];
|
||
enhanceProKeyPoolCursor = (enhanceProKeyPoolCursor + 1) % keys.length;
|
||
return {
|
||
apiKey: selectedKey,
|
||
keyId: getEnhanceProKeyFingerprint(selectedKey),
|
||
};
|
||
}
|
||
|
||
function buildEnhanceProTaskPendingResponse(taskId: string) {
|
||
return NextResponse.json(
|
||
{
|
||
success: true,
|
||
processing: true,
|
||
status: "running",
|
||
taskId,
|
||
pollAfterMs: ENHANCE_PRO_TASK_POLL_INTERVAL_MS,
|
||
},
|
||
{ status: 202 }
|
||
);
|
||
}
|
||
|
||
function normalizeBaseUrl(baseUrl: string) {
|
||
return String(baseUrl || "").trim().replace(/\/+$/, "");
|
||
}
|
||
|
||
function sanitizeDimension(value: unknown, fallback: number) {
|
||
return clampInteger(value, 1, 32000, fallback);
|
||
}
|
||
|
||
function clampDimensions(width: number, height: number, maxEdge: number) {
|
||
const safeWidth = sanitizeDimension(width, 1024);
|
||
const safeHeight = sanitizeDimension(height, 1024);
|
||
const longestEdge = Math.max(safeWidth, safeHeight);
|
||
|
||
if (longestEdge <= maxEdge) {
|
||
return { width: safeWidth, height: safeHeight };
|
||
}
|
||
|
||
const scale = maxEdge / longestEdge;
|
||
return {
|
||
width: Math.max(1, Math.round(safeWidth * scale)),
|
||
height: Math.max(1, Math.round(safeHeight * scale)),
|
||
};
|
||
}
|
||
|
||
async function createThumbnail(buffer: Buffer): Promise<Buffer> {
|
||
try {
|
||
const sharp = await import("sharp");
|
||
return await sharp.default(buffer)
|
||
.resize(500, 500, {
|
||
fit: "inside",
|
||
withoutEnlargement: true,
|
||
})
|
||
.jpeg({ quality: 82 })
|
||
.toBuffer();
|
||
} catch (error) {
|
||
console.error("[Enhance Pro] 生成缩略图失败，回退使用原图:", error);
|
||
return buffer;
|
||
}
|
||
}
|
||
|
||
async function uploadBufferPairToStorage(
|
||
buffer: Buffer,
|
||
mimeType = "image/jpeg",
|
||
prefix = "enhance-pro",
|
||
): Promise<{ thumbnail: string; original: string }> {
|
||
const ext = mimeType.split("/")[1] || "jpg";
|
||
const timestamp = Date.now();
|
||
const randomStr = Math.random().toString(36).slice(2, 8);
|
||
const originalFilename = `${prefix}-orig-${timestamp}-${randomStr}.${ext}`;
|
||
const thumbnailFilename = `${prefix}-thumb-${timestamp}-${randomStr}.jpg`;
|
||
const thumbnailBuffer = await createThumbnail(buffer);
|
||
|
||
const [thumbnailResult, originalResult] = await Promise.all([
|
||
uploadFile(
|
||
thumbnailFilename,
|
||
thumbnailBuffer.buffer.slice(
|
||
thumbnailBuffer.byteOffset,
|
||
thumbnailBuffer.byteOffset + thumbnailBuffer.byteLength,
|
||
) as ArrayBuffer,
|
||
),
|
||
uploadFile(
|
||
originalFilename,
|
||
buffer.buffer.slice(buffer.byteOffset, buffer.byteOffset + buffer.byteLength) as ArrayBuffer,
|
||
),
|
||
]);
|
||
|
||
return {
|
||
thumbnail: thumbnailResult.url,
|
||
original: originalResult.url,
|
||
};
|
||
}
|
||
|
||
async function getImageDimensionsFromBuffer(buffer: Buffer) {
|
||
try {
|
||
const sharp = await import("sharp");
|
||
const metadata = await sharp.default(buffer).metadata();
|
||
return {
|
||
width: metadata.width || 1024,
|
||
height: metadata.height || 1024,
|
||
};
|
||
} catch (error) {
|
||
console.error("[Enhance Pro] 获取图片尺寸失败，回退默认值:", error);
|
||
return {
|
||
width: 1024,
|
||
height: 1024,
|
||
};
|
||
}
|
||
}
|
||
|
||
async function fetchBufferFromUrl(url: string) {
|
||
const response = await fetch(url);
|
||
if (!response.ok) {
|
||
throw new Error(`无法下载图片资源: ${response.status} ${response.statusText}`);
|
||
}
|
||
const mimeType = response.headers.get("content-type") || "image/jpeg";
|
||
const arrayBuffer = await response.arrayBuffer();
|
||
return {
|
||
buffer: Buffer.from(arrayBuffer),
|
||
mimeType,
|
||
};
|
||
}
|
||
|
||
async function resolveSourceInput(body: EnhanceProRequestBody) {
|
||
const imageUrl = String(body.imageUrl || "").trim();
|
||
const imageBase64 = String(body.image_base64 || "").trim();
|
||
const requestedMimeType = String(body.mimeType || "").trim() || "image/jpeg";
|
||
|
||
if (!imageUrl && !imageBase64) {
|
||
throw new Error("缺少 imageUrl 或 image_base64 参数");
|
||
}
|
||
|
||
if (imageUrl) {
|
||
const hasRequestedDimensions =
|
||
Number.isFinite(Number(body.width)) && Number.isFinite(Number(body.height));
|
||
if (hasRequestedDimensions) {
|
||
return {
|
||
sourceImageUrl: imageUrl,
|
||
width: sanitizeDimension(body.width, 1024),
|
||
height: sanitizeDimension(body.height, 1024),
|
||
};
|
||
}
|
||
|
||
const downloaded = await fetchBufferFromUrl(imageUrl);
|
||
const dimensions = await getImageDimensionsFromBuffer(downloaded.buffer);
|
||
return {
|
||
sourceImageUrl: imageUrl,
|
||
width: dimensions.width,
|
||
height: dimensions.height,
|
||
};
|
||
}
|
||
|
||
const sourceBuffer = Buffer.from(imageBase64, "base64");
|
||
const stored = await uploadBufferPairToStorage(sourceBuffer, requestedMimeType, "enhance-pro-source");
|
||
const dimensions = await getImageDimensionsFromBuffer(sourceBuffer);
|
||
return {
|
||
sourceImageUrl: stored.original,
|
||
width: dimensions.width,
|
||
height: dimensions.height,
|
||
};
|
||
}
|
||
|
||
async function createEnhanceProJob(params: {
|
||
config: EnhanceProKreaConfig;
|
||
requestBody: EnhanceProRequestBody;
|
||
sourceImageUrl: string;
|
||
width: number;
|
||
height: number;
|
||
}) {
|
||
const apiUrl = normalizeBaseUrl(params.config.baseUrl);
|
||
const requestPayload = {
|
||
width: params.width,
|
||
height: params.height,
|
||
image_url: params.sourceImageUrl,
|
||
batchSize: params.requestBody.batchSize,
|
||
seed: params.requestBody.seed,
|
||
prompt: String(params.requestBody.prompt || "").trim(),
|
||
output_format: params.requestBody.output_format || params.config.defaultOutputFormat,
|
||
subject_detection: params.requestBody.subject_detection || "All",
|
||
face_enhancement: normalizeBoolean(params.requestBody.face_enhancement, false),
|
||
face_enhancement_creativity: clampNumber(
|
||
params.requestBody.face_enhancement_creativity,
|
||
0,
|
||
1,
|
||
0.5,
|
||
),
|
||
face_enhancement_strength: clampNumber(
|
||
params.requestBody.face_enhancement_strength,
|
||
0,
|
||
1,
|
||
0.5,
|
||
),
|
||
crop_to_fill: normalizeBoolean(params.requestBody.crop_to_fill, false),
|
||
model: params.requestBody.model || params.config.defaultModel,
|
||
upscaling_activated: normalizeBoolean(params.requestBody.upscaling_activated, false),
|
||
image_scaling_factor: clampNumber(
|
||
params.requestBody.image_scaling_factor,
|
||
1,
|
||
params.config.maxUpscaleFactor,
|
||
1,
|
||
),
|
||
creativity: clampInteger(params.requestBody.creativity, 1, 6, 3),
|
||
texture: clampInteger(params.requestBody.texture, 1, 5, 3),
|
||
sharpen: clampNumber(params.requestBody.sharpen, 0, 1, 0.5),
|
||
denoise: clampNumber(params.requestBody.denoise, 0, 1, 0.5),
|
||
detail: clampNumber(params.requestBody.detail, 0, 1, 0.5),
|
||
};
|
||
|
||
const controller = new AbortController();
|
||
const timeoutId = setTimeout(() => controller.abort(), ENHANCE_PRO_CREATE_TIMEOUT_MS);
|
||
|
||
try {
|
||
const response = await fetch(`${apiUrl}/generate/enhance/topaz/generative-enhance`, {
|
||
method: "POST",
|
||
headers: {
|
||
"Content-Type": "application/json",
|
||
Authorization: `Bearer ${params.config.apiKey}`,
|
||
},
|
||
body: JSON.stringify(requestPayload),
|
||
signal: controller.signal,
|
||
});
|
||
|
||
if (!response.ok) {
|
||
const detail = await response.text();
|
||
throw new Error(`创建任务失败: ${response.status} ${detail || response.statusText}`);
|
||
}
|
||
|
||
return await response.json();
|
||
} finally {
|
||
clearTimeout(timeoutId);
|
||
}
|
||
}
|
||
|
||
async function fetchEnhanceProJobStatus(config: EnhanceProKreaConfig, jobId: string) {
|
||
const apiUrl = normalizeBaseUrl(config.baseUrl);
|
||
const response = await fetch(`${apiUrl}/jobs/${jobId}`, {
|
||
method: "GET",
|
||
headers: {
|
||
Authorization: `Bearer ${config.apiKey}`,
|
||
},
|
||
cache: "no-store",
|
||
});
|
||
|
||
if (!response.ok) {
|
||
const detail = await response.text();
|
||
throw new Error(`查询任务状态失败: ${response.status} ${detail || response.statusText}`);
|
||
}
|
||
|
||
return await response.json();
|
||
}
|
||
|
||
async function saveGeneratedImageRecord(params: {
|
||
userId: string;
|
||
chatId?: string | null;
|
||
messageId?: string | null;
|
||
prompt: string;
|
||
originalUrl: string;
|
||
thumbnailUrl: string;
|
||
imageSize: string;
|
||
}) {
|
||
const [savedImage] = await db
|
||
.insert(generatedImage)
|
||
.values({
|
||
userId: params.userId,
|
||
chatId: params.chatId || null,
|
||
messageId: params.messageId || null,
|
||
thumbnailUrl: params.thumbnailUrl,
|
||
originalUrl: params.originalUrl,
|
||
imageSize: params.imageSize,
|
||
prompt: params.prompt,
|
||
modelId: ENHANCE_PRO_MODEL_ID,
|
||
})
|
||
.returning();
|
||
|
||
return savedImage;
|
||
}
|
||
|
||
async function finalizeEnhanceProTaskResult(params: {
|
||
userId: string;
|
||
taskResult: any;
|
||
prompt?: string;
|
||
chatId?: string | null;
|
||
messageId?: string | null;
|
||
startedAt?: number;
|
||
sourceImageUrls?: string[];
|
||
width: number;
|
||
height: number;
|
||
}) {
|
||
const resultUrls = Array.isArray(params.taskResult?.result?.urls)
|
||
? params.taskResult.result.urls.filter((url: unknown) => typeof url === "string" && url)
|
||
: [];
|
||
|
||
if (resultUrls.length === 0) {
|
||
throw new Error("Krea 任务未返回结果图片");
|
||
}
|
||
|
||
const limitedUrls = resultUrls.slice(0, ENHANCE_PRO_MAX_RESULT_IMAGES);
|
||
const finalPrompt = String(params.prompt || "").trim() || "图像增强 Pro";
|
||
const imageSize = `${params.width}x${params.height}`;
|
||
const sourceAssets =
|
||
Array.isArray(params.sourceImageUrls) && params.sourceImageUrls.length > 0
|
||
? await getMediaAssetsByUserAndUrls({
|
||
userId: params.userId,
|
||
urls: params.sourceImageUrls,
|
||
})
|
||
: [];
|
||
|
||
const finalizedImages: FinalizedImageResult[] = [];
|
||
|
||
for (const url of limitedUrls) {
|
||
const downloaded = await fetchBufferFromUrl(url);
|
||
const uploaded = await uploadBufferPairToStorage(downloaded.buffer, downloaded.mimeType, "enhance-pro");
|
||
const savedImage = await saveGeneratedImageRecord({
|
||
userId: params.userId,
|
||
chatId: params.chatId,
|
||
messageId: params.messageId,
|
||
prompt: finalPrompt,
|
||
originalUrl: uploaded.original,
|
||
thumbnailUrl: uploaded.thumbnail,
|
||
imageSize,
|
||
});
|
||
|
||
try {
|
||
const media = await upsertMediaAsset({
|
||
userId: params.userId,
|
||
chatId: params.chatId || null,
|
||
mediaType: "image",
|
||
sourceType: "edited",
|
||
status: "ready",
|
||
originalUrl: uploaded.original,
|
||
previewUrl: uploaded.original,
|
||
thumbnailUrl: uploaded.thumbnail,
|
||
displayName: finalPrompt.slice(0, 80),
|
||
prompt: finalPrompt,
|
||
modelId: ENHANCE_PRO_MODEL_ID,
|
||
sourceMessageId: params.messageId || null,
|
||
sourceTable: "GeneratedImage",
|
||
sourceRowId: savedImage.id,
|
||
primaryParentAssetId: sourceAssets[0]?.id ?? null,
|
||
metadata: {
|
||
imageSize,
|
||
provider: "krea-topaz-generative",
|
||
},
|
||
});
|
||
|
||
if (media) {
|
||
for (const sourceAsset of sourceAssets) {
|
||
await createMediaAssetRelation({
|
||
userId: params.userId,
|
||
chatId: params.chatId || null,
|
||
fromAssetId: sourceAsset.id,
|
||
toAssetId: media.id,
|
||
relationType: "edited_from",
|
||
sourceMessageId: params.messageId || null,
|
||
metadata: {
|
||
prompt: finalPrompt,
|
||
},
|
||
});
|
||
}
|
||
}
|
||
} catch (assetError) {
|
||
console.error("[Enhance Pro] MediaAsset 注册失败:", assetError);
|
||
}
|
||
|
||
finalizedImages.push({
|
||
thumbnailUrl: uploaded.thumbnail,
|
||
originalUrl: uploaded.original,
|
||
width: params.width,
|
||
height: params.height,
|
||
});
|
||
}
|
||
|
||
const processingTime =
|
||
typeof params.startedAt === "number" ? (Date.now() - params.startedAt) / 1000 : undefined;
|
||
|
||
return {
|
||
success: true,
|
||
images: finalizedImages.map((image, index) => ({
|
||
index,
|
||
thumbnail_url: image.thumbnailUrl,
|
||
original_url: image.originalUrl,
|
||
width: image.width,
|
||
height: image.height,
|
||
})),
|
||
thumbnail_url: finalizedImages[0]?.thumbnailUrl || null,
|
||
original_url: finalizedImages[0]?.originalUrl || null,
|
||
processing_time: processingTime,
|
||
model_id: ENHANCE_PRO_MODEL_ID,
|
||
};
|
||
}
|
||
|
||
function parseEnhanceProStoredRequest(task: EnhanceProTaskRecord): EnhanceProStoredRequest {
|
||
const request =
|
||
task.request && typeof task.request === "object"
|
||
? (task.request as Partial<EnhanceProStoredRequest>)
|
||
: {};
|
||
|
||
const width = sanitizeDimension(request.width, 1024);
|
||
const height = sanitizeDimension(request.height, 1024);
|
||
|
||
return {
|
||
prompt: String(request.prompt || task.prompt || "").trim() || "图像增强 Pro",
|
||
sourceImageUrl:
|
||
String(request.sourceImageUrl || task.sourceImageUrl || "").trim(),
|
||
width,
|
||
height,
|
||
batchSize: clampInteger(request.batchSize, 1, 8, 1),
|
||
model:
|
||
request.model === "Recovery" ||
|
||
request.model === "Recovery V2" ||
|
||
request.model === "Reimagine" ||
|
||
request.model === "Redefine"
|
||
? request.model
|
||
: "Recovery V2",
|
||
upscalingActivated: normalizeBoolean(request.upscalingActivated, false),
|
||
imageScalingFactor: clampNumber(request.imageScalingFactor, 1, 8, 1),
|
||
chatId: request.chatId ? String(request.chatId) : task.chatId || null,
|
||
messageId: request.messageId ? String(request.messageId) : task.messageId || null,
|
||
};
|
||
}
|
||
|
||
async function getEnhanceProTaskForUser(params: { taskId: string; userId: string }) {
|
||
const rows = await db
|
||
.select()
|
||
.from(enhanceProTask)
|
||
.where(and(eq(enhanceProTask.id, params.taskId), eq(enhanceProTask.userId, params.userId)))
|
||
.limit(1);
|
||
return rows[0] ?? null;
|
||
}
|
||
|
||
async function getEnhanceProTaskByClientRequestId(params: {
|
||
userId: string;
|
||
clientRequestId: string;
|
||
}) {
|
||
const rows = await db
|
||
.select()
|
||
.from(enhanceProTask)
|
||
.where(
|
||
and(
|
||
eq(enhanceProTask.userId, params.userId),
|
||
eq(enhanceProTask.clientRequestId, params.clientRequestId),
|
||
),
|
||
)
|
||
.limit(1);
|
||
return rows[0] ?? null;
|
||
}
|
||
|
||
async function markEnhanceProTaskRunning(taskId: string) {
|
||
await db
|
||
.update(enhanceProTask)
|
||
.set({
|
||
status: "running",
|
||
updatedAt: new Date(),
|
||
error: null,
|
||
})
|
||
.where(eq(enhanceProTask.id, taskId));
|
||
}
|
||
|
||
async function claimEnhanceProTaskFinalization(taskId: string) {
|
||
const rows = await db
|
||
.update(enhanceProTask)
|
||
.set({
|
||
status: "finalizing",
|
||
updatedAt: new Date(),
|
||
error: null,
|
||
})
|
||
.where(
|
||
and(
|
||
eq(enhanceProTask.id, taskId),
|
||
or(eq(enhanceProTask.status, "pending"), eq(enhanceProTask.status, "running")),
|
||
),
|
||
)
|
||
.returning();
|
||
return rows[0] ?? null;
|
||
}
|
||
|
||
async function markEnhanceProTaskFailed(taskId: string, errorMessage: string) {
|
||
await db
|
||
.update(enhanceProTask)
|
||
.set({
|
||
status: "failed",
|
||
error: errorMessage,
|
||
completedAt: new Date(),
|
||
updatedAt: new Date(),
|
||
})
|
||
.where(eq(enhanceProTask.id, taskId));
|
||
}
|
||
|
||
async function persistEnhanceProTaskSuccess(params: {
|
||
taskId: string;
|
||
payload: EnhanceProSuccessPayload;
|
||
pointsCharged: boolean;
|
||
}) {
|
||
await db
|
||
.update(enhanceProTask)
|
||
.set({
|
||
status: "succeeded",
|
||
result: params.payload as Record<string, unknown>,
|
||
error: null,
|
||
pointsCharged: params.pointsCharged,
|
||
chargedAt: params.pointsCharged ? new Date() : null,
|
||
completedAt: new Date(),
|
||
updatedAt: new Date(),
|
||
})
|
||
.where(eq(enhanceProTask.id, params.taskId));
|
||
}
|
||
|
||
async function consumeEnhanceProPointsOnce(
|
||
task: EnhanceProTaskRecord,
|
||
request: EnhanceProStoredRequest,
|
||
) {
|
||
if (task.pointsCost <= 0) {
|
||
return { ok: true, skipped: true };
|
||
}
|
||
|
||
const note = `图像增强 Pro 消耗积分（任务 ${task.id} / ${task.providerTaskId}）`;
|
||
const existingLedger = await db
|
||
.select({ id: pointsLedger.id })
|
||
.from(pointsLedger)
|
||
.where(
|
||
and(
|
||
eq(pointsLedger.userId, task.userId),
|
||
eq(pointsLedger.reason, "enhance_pro_generation"),
|
||
eq(pointsLedger.note, note),
|
||
),
|
||
)
|
||
.limit(1);
|
||
|
||
if (existingLedger[0]) {
|
||
return { ok: true, skipped: true };
|
||
}
|
||
|
||
return await consumeUserPoints({
|
||
userId: task.userId,
|
||
amount: task.pointsCost,
|
||
reason: "enhance_pro_generation",
|
||
modelId: `${ENHANCE_PRO_MODEL_ID}:${request.model}`,
|
||
chatId: request.chatId,
|
||
messageId: request.messageId,
|
||
note,
|
||
});
|
||
}
|
||
|
||
function getStoredEnhanceProSuccessPayload(task: EnhanceProTaskRecord) {
|
||
if (!task.result || typeof task.result !== "object") {
|
||
return null;
|
||
}
|
||
return task.result as EnhanceProSuccessPayload;
|
||
}
|
||
|
||
async function finalizeEnhanceProTask(
|
||
task: EnhanceProTaskRecord,
|
||
taskResult: any,
|
||
startedAt?: number,
|
||
) {
|
||
const existingPayload = getStoredEnhanceProSuccessPayload(task);
|
||
if (task.status === "succeeded" && existingPayload) {
|
||
return existingPayload;
|
||
}
|
||
|
||
const claimedTask =
|
||
task.status === "finalizing" ? task : await claimEnhanceProTaskFinalization(task.id);
|
||
if (!claimedTask) {
|
||
return null;
|
||
}
|
||
|
||
const request = parseEnhanceProStoredRequest(claimedTask);
|
||
|
||
try {
|
||
const responsePayload = await finalizeEnhanceProTaskResult({
|
||
userId: claimedTask.userId,
|
||
taskResult,
|
||
prompt: request.prompt,
|
||
chatId: request.chatId,
|
||
messageId: request.messageId,
|
||
startedAt,
|
||
sourceImageUrls: request.sourceImageUrl ? [request.sourceImageUrl] : undefined,
|
||
width: request.width,
|
||
height: request.height,
|
||
});
|
||
|
||
let pointsCharged = claimedTask.pointsCharged;
|
||
if (!pointsCharged && claimedTask.pointsCost > 0) {
|
||
try {
|
||
const chargeResult = await consumeEnhanceProPointsOnce(claimedTask, request);
|
||
pointsCharged = chargeResult.ok || chargeResult.skipped;
|
||
} catch (pointsError) {
|
||
console.error("[Points] Failed to consume Enhance Pro points:", pointsError);
|
||
}
|
||
}
|
||
|
||
const payload: EnhanceProSuccessPayload = {
|
||
...responsePayload,
|
||
chargedPoints: pointsCharged ? claimedTask.pointsCost : 0,
|
||
};
|
||
|
||
await persistEnhanceProTaskSuccess({
|
||
taskId: claimedTask.id,
|
||
payload,
|
||
pointsCharged,
|
||
});
|
||
|
||
return payload;
|
||
} catch (error) {
|
||
const errorMessage = error instanceof Error ? error.message : "未知错误";
|
||
await markEnhanceProTaskFailed(claimedTask.id, errorMessage);
|
||
throw error;
|
||
}
|
||
}
|
||
|
||
function buildEnhanceProTaskFailedResponse(task: EnhanceProTaskRecord) {
|
||
return NextResponse.json(
|
||
{
|
||
error: "图像增强 Pro 处理失败",
|
||
detail: String(task.error || "任务失败"),
|
||
},
|
||
{ status: 500 },
|
||
);
|
||
}
|
||
|
||
async function loadEnhanceProConfig() {
|
||
const setting = await getSystemSetting({ key: ENHANCE_PRO_KREA_SETTING_KEY });
|
||
const existingValue =
|
||
setting?.value && typeof setting.value === "object"
|
||
? (setting.value as Record<string, unknown>)
|
||
: {};
|
||
return normalizeEnhanceProKreaConfig(setting?.value, String(existingValue.apiKey || ""));
|
||
}
|
||
|
||
export const maxDuration = 180;
|
||
|
||
export async function GET(request: NextRequest) {
|
||
try {
|
||
const session = await auth();
|
||
if (!session?.user?.id) {
|
||
return NextResponse.json({ error: "未授权，请先登录" }, { status: 401 });
|
||
}
|
||
|
||
const { searchParams } = new URL(request.url);
|
||
const taskId = String(searchParams.get("taskId") || "").trim();
|
||
if (!taskId) {
|
||
return NextResponse.json({ error: "缺少 taskId 参数" }, { status: 400 });
|
||
}
|
||
|
||
const task = await getEnhanceProTaskForUser({
|
||
taskId,
|
||
userId: session.user.id,
|
||
});
|
||
if (!task) {
|
||
return NextResponse.json({ error: "任务不存在" }, { status: 404 });
|
||
}
|
||
|
||
const storedPayload = getStoredEnhanceProSuccessPayload(task);
|
||
if (task.status === "succeeded" && storedPayload) {
|
||
return NextResponse.json(storedPayload);
|
||
}
|
||
if (task.status === "failed") {
|
||
return buildEnhanceProTaskFailedResponse(task);
|
||
}
|
||
if (task.status === "finalizing") {
|
||
return buildEnhanceProTaskPendingResponse(task.id);
|
||
}
|
||
|
||
const config = await loadEnhanceProConfig();
|
||
if (!config.enabled) {
|
||
return NextResponse.json({ error: "图像增强 Pro 功能未开启" }, { status: 403 });
|
||
}
|
||
const keySelection = resolveEnhanceProApiKey(config, task.taskKeyId);
|
||
if (!keySelection || !config.baseUrl) {
|
||
return NextResponse.json({ error: "图像增强 Pro API 尚未配置" }, { status: 400 });
|
||
}
|
||
|
||
const taskResult = await fetchEnhanceProJobStatus(
|
||
{ ...config, apiKey: keySelection.apiKey },
|
||
task.providerTaskId,
|
||
);
|
||
const providerStatus = normalizeJobStatus(taskResult?.status);
|
||
|
||
if (providerStatus === "completed") {
|
||
const completedPayload = await finalizeEnhanceProTask(task, taskResult);
|
||
if (completedPayload) {
|
||
return NextResponse.json(completedPayload);
|
||
}
|
||
|
||
const refreshedTask = await getEnhanceProTaskForUser({
|
||
taskId,
|
||
userId: session.user.id,
|
||
});
|
||
const refreshedPayload = refreshedTask
|
||
? getStoredEnhanceProSuccessPayload(refreshedTask)
|
||
: null;
|
||
if (refreshedTask?.status === "succeeded" && refreshedPayload) {
|
||
return NextResponse.json(refreshedPayload);
|
||
}
|
||
if (refreshedTask?.status === "failed") {
|
||
return buildEnhanceProTaskFailedResponse(refreshedTask);
|
||
}
|
||
return buildEnhanceProTaskPendingResponse(task.id);
|
||
}
|
||
|
||
if (providerStatus === "failed" || providerStatus === "cancelled") {
|
||
const detail = `Krea 任务${providerStatus === "failed" ? "失败" : "已取消"}`;
|
||
await markEnhanceProTaskFailed(task.id, detail);
|
||
return NextResponse.json(
|
||
{
|
||
error: "图像增强 Pro 处理失败",
|
||
detail,
|
||
},
|
||
{ status: 500 },
|
||
);
|
||
}
|
||
|
||
await markEnhanceProTaskRunning(task.id);
|
||
return buildEnhanceProTaskPendingResponse(task.id);
|
||
} catch (error) {
|
||
console.error("[Enhance Pro API][TaskStatus] 错误:", error);
|
||
return NextResponse.json(
|
||
{ error: "服务器内部错误", detail: error instanceof Error ? error.message : "未知错误" },
|
||
{ status: 500 },
|
||
);
|
||
}
|
||
}
|
||
|
||
export async function POST(request: NextRequest) {
|
||
try {
|
||
const session = await auth();
|
||
if (!session?.user?.id) {
|
||
return NextResponse.json({ error: "未授权，请先登录" }, { status: 401 });
|
||
}
|
||
|
||
const config = await loadEnhanceProConfig();
|
||
if (!config.enabled) {
|
||
return NextResponse.json({ error: "图像增强 Pro 功能未开启" }, { status: 403 });
|
||
}
|
||
const keySelection = resolveEnhanceProApiKey(config);
|
||
if (!keySelection || !config.baseUrl) {
|
||
return NextResponse.json({ error: "图像增强 Pro API 尚未配置" }, { status: 400 });
|
||
}
|
||
const requestConfig = { ...config, apiKey: keySelection.apiKey };
|
||
|
||
const body = (await request.json()) as EnhanceProRequestBody;
|
||
const clientRequestId = normalizeClientRequestId(body.clientRequestId);
|
||
if (clientRequestId) {
|
||
const existingTask = await getEnhanceProTaskByClientRequestId({
|
||
userId: session.user.id,
|
||
clientRequestId,
|
||
});
|
||
if (existingTask) {
|
||
const existingPayload = getStoredEnhanceProSuccessPayload(existingTask);
|
||
if (existingTask.status === "succeeded" && existingPayload) {
|
||
return NextResponse.json(existingPayload);
|
||
}
|
||
if (existingTask.status === "failed") {
|
||
return buildEnhanceProTaskFailedResponse(existingTask);
|
||
}
|
||
return buildEnhanceProTaskPendingResponse(existingTask.id);
|
||
}
|
||
}
|
||
|
||
const prompt = String(body.prompt || "").trim();
|
||
const chatId = body.chatId || null;
|
||
const messageId = body.messageId || null;
|
||
const sourceInput = await resolveSourceInput(body);
|
||
|
||
const requestedWidth = sanitizeDimension(body.width, sourceInput.width);
|
||
const requestedHeight = sanitizeDimension(body.height, sourceInput.height);
|
||
const clampedDimensions = clampDimensions(
|
||
requestedWidth,
|
||
requestedHeight,
|
||
config.maxOutputEdge,
|
||
);
|
||
|
||
const batchSize = config.allowBatchSize
|
||
? clampInteger(body.batchSize, 1, config.maxBatchSize, 1)
|
||
: 1;
|
||
const selectedModel: EnhanceProKreaModel =
|
||
body.model === "Recovery" ||
|
||
body.model === "Recovery V2" ||
|
||
body.model === "Reimagine" ||
|
||
body.model === "Redefine"
|
||
? body.model
|
||
: config.defaultModel;
|
||
const upscalingActivated = body.upscaling_activated === true;
|
||
const imageScalingFactor = clampNumber(
|
||
body.image_scaling_factor,
|
||
1,
|
||
config.maxUpscaleFactor,
|
||
1,
|
||
);
|
||
const pointsConfig = await getPointsSystemConfig();
|
||
const pointsEstimate = calculateEnhanceProPointsEstimate(pointsConfig, {
|
||
width: clampedDimensions.width,
|
||
height: clampedDimensions.height,
|
||
batchSize,
|
||
model: selectedModel,
|
||
upscalingActivated,
|
||
scaleFactor: imageScalingFactor,
|
||
});
|
||
|
||
if (pointsEstimate.chargeable && pointsConfig.blockOnInsufficient !== false) {
|
||
const userData = await getUserById({ id: session.user.id });
|
||
if (!userData?.isHighVip) {
|
||
const currentBalance = Math.max(0, Number(userData?.pointsBalance || 0));
|
||
if (currentBalance < pointsEstimate.cost) {
|
||
const message = `积分不足，当前余额 ${currentBalance}，本次预计消耗 ${pointsEstimate.cost}。`;
|
||
return NextResponse.json(
|
||
{
|
||
error: message,
|
||
code: "insufficient_points",
|
||
message,
|
||
balance: currentBalance,
|
||
expectedCost: pointsEstimate.cost,
|
||
unitCost: pointsEstimate.cost,
|
||
},
|
||
{ status: 402 },
|
||
);
|
||
}
|
||
}
|
||
}
|
||
|
||
const startTime = Date.now();
|
||
const createdJob = await createEnhanceProJob({
|
||
config: requestConfig,
|
||
requestBody: {
|
||
...body,
|
||
batchSize,
|
||
model: selectedModel,
|
||
upscaling_activated: upscalingActivated,
|
||
image_scaling_factor: imageScalingFactor,
|
||
},
|
||
sourceImageUrl: sourceInput.sourceImageUrl,
|
||
width: clampedDimensions.width,
|
||
height: clampedDimensions.height,
|
||
});
|
||
|
||
const createdStatus = normalizeJobStatus(createdJob?.status);
|
||
const providerTaskId = String(createdJob?.job_id || "").trim();
|
||
if (!providerTaskId) {
|
||
return NextResponse.json(
|
||
{ error: "创建图像增强 Pro 任务失败", detail: "Krea 未返回 job_id" },
|
||
{ status: 500 },
|
||
);
|
||
}
|
||
|
||
const storedRequest: EnhanceProStoredRequest = {
|
||
prompt: prompt || "图像增强 Pro",
|
||
sourceImageUrl: sourceInput.sourceImageUrl,
|
||
width: clampedDimensions.width,
|
||
height: clampedDimensions.height,
|
||
batchSize,
|
||
model: selectedModel,
|
||
upscalingActivated,
|
||
imageScalingFactor,
|
||
chatId,
|
||
messageId,
|
||
};
|
||
|
||
let createdTask: EnhanceProTaskRecord | null = null;
|
||
try {
|
||
const inserted = await db
|
||
.insert(enhanceProTask)
|
||
.values({
|
||
userId: session.user.id,
|
||
chatId,
|
||
messageId,
|
||
clientRequestId: clientRequestId || null,
|
||
providerTaskId,
|
||
taskKeyId: keySelection.keyId,
|
||
status:
|
||
createdStatus === "completed"
|
||
? "running"
|
||
: createdStatus === "failed" || createdStatus === "cancelled"
|
||
? "failed"
|
||
: "pending",
|
||
prompt: storedRequest.prompt,
|
||
sourceImageUrl: storedRequest.sourceImageUrl,
|
||
request: storedRequest as unknown as Record<string, unknown>,
|
||
error:
|
||
createdStatus === "failed" || createdStatus === "cancelled"
|
||
? `Krea 任务${createdStatus === "failed" ? "失败" : "已取消"}`
|
||
: null,
|
||
pointsCost: pointsEstimate.chargeable ? pointsEstimate.cost : 0,
|
||
pointsCharged: false,
|
||
completedAt:
|
||
createdStatus === "failed" || createdStatus === "cancelled" ? new Date() : null,
|
||
createdAt: new Date(),
|
||
updatedAt: new Date(),
|
||
})
|
||
.returning();
|
||
createdTask = inserted[0] ?? null;
|
||
} catch (insertError) {
|
||
if (clientRequestId) {
|
||
const existingTask = await getEnhanceProTaskByClientRequestId({
|
||
userId: session.user.id,
|
||
clientRequestId,
|
||
});
|
||
if (existingTask) {
|
||
const existingPayload = getStoredEnhanceProSuccessPayload(existingTask);
|
||
if (existingTask.status === "succeeded" && existingPayload) {
|
||
return NextResponse.json(existingPayload);
|
||
}
|
||
if (existingTask.status === "failed") {
|
||
return buildEnhanceProTaskFailedResponse(existingTask);
|
||
}
|
||
return buildEnhanceProTaskPendingResponse(existingTask.id);
|
||
}
|
||
}
|
||
throw insertError;
|
||
}
|
||
|
||
if (!createdTask) {
|
||
return NextResponse.json(
|
||
{ error: "创建图像增强 Pro 任务失败", detail: "任务入库失败" },
|
||
{ status: 500 },
|
||
);
|
||
}
|
||
|
||
if (createdStatus === "completed") {
|
||
const responsePayload = await finalizeEnhanceProTask(createdTask, createdJob, startTime);
|
||
if (responsePayload) {
|
||
return NextResponse.json(responsePayload);
|
||
}
|
||
return buildEnhanceProTaskPendingResponse(createdTask.id);
|
||
}
|
||
|
||
if (createdStatus === "failed" || createdStatus === "cancelled") {
|
||
return buildEnhanceProTaskFailedResponse(createdTask);
|
||
}
|
||
|
||
return buildEnhanceProTaskPendingResponse(createdTask.id);
|
||
} catch (error) {
|
||
console.error("[Enhance Pro API] 错误:", error);
|
||
return NextResponse.json(
|
||
{ error: "服务器内部错误", detail: error instanceof Error ? error.message : "未知错误" },
|
||
{ status: 500 },
|
||
);
|
||
}
|
||
}
|