471 lines
15 KiB
TypeScript
471 lines
15 KiB
TypeScript
import { NextRequest, NextResponse } from "next/server";
|
||
import { auth } from "@/app/(auth)/auth";
|
||
import { db } from "@/lib/db";
|
||
import { systemSettings, generatedImage } from "@/lib/db/schema";
|
||
import { eq } from "drizzle-orm";
|
||
import { uploadFile } from "@/lib/storage/service";
|
||
import {
|
||
createMediaAssetRelation,
|
||
getPointsSystemConfig,
|
||
getMediaAssetsByUserAndUrls,
|
||
upsertMediaAsset,
|
||
} from "@/lib/db/queries";
|
||
import { prepareFeaturePointsCharge } from "@/lib/points/feature-charge";
|
||
import { getImageToolCreditCost } from "@/lib/points/config";
|
||
|
||
// PicWish API 需要公网可访问的 HTTP/HTTPS URL
|
||
const WATERMARK_SYNC_WAIT_TIMEOUT_MS = 20_000;
|
||
const WATERMARK_POLL_INTERVAL_MS = 1200;
|
||
|
||
// 获取去水印API配置
|
||
async function getWatermarkRemoveConfig() {
|
||
try {
|
||
const setting = await db
|
||
.select()
|
||
.from(systemSettings)
|
||
.where(eq(systemSettings.key, "watermark_remove_config"))
|
||
.limit(1);
|
||
|
||
if (setting.length > 0 && setting[0].value) {
|
||
return setting[0].value as { apiKey: string };
|
||
}
|
||
return null;
|
||
} catch (error) {
|
||
console.error("[WatermarkRemove] 获取配置失败:", error);
|
||
return null;
|
||
}
|
||
}
|
||
|
||
// 轮询获取处理结果
|
||
async function pollForResult(
|
||
taskId: string,
|
||
apiKey: string,
|
||
timeoutMs = WATERMARK_SYNC_WAIT_TIMEOUT_MS
|
||
): Promise<{ success: boolean; imageUrl?: string; error?: string; pending?: boolean }> {
|
||
const maxAttempts = Math.max(1, Math.ceil(timeoutMs / WATERMARK_POLL_INTERVAL_MS));
|
||
const pollInterval = WATERMARK_POLL_INTERVAL_MS;
|
||
|
||
for (let attempt = 0; attempt < maxAttempts; attempt++) {
|
||
try {
|
||
const response = await fetch(
|
||
`https://techsz.aoscdn.com/api/tasks/visual/external/watermark-remove/${taskId}`,
|
||
{
|
||
method: "GET",
|
||
headers: {
|
||
"X-API-KEY": apiKey,
|
||
},
|
||
}
|
||
);
|
||
|
||
if (!response.ok) {
|
||
const errorText = await response.text();
|
||
console.error(`[WatermarkRemove] 轮询失败 (attempt ${attempt + 1}):`, errorText);
|
||
// 继续轮询
|
||
await new Promise(resolve => setTimeout(resolve, pollInterval));
|
||
continue;
|
||
}
|
||
|
||
const result = await response.json();
|
||
console.log(`[WatermarkRemove] 轮询结果 (attempt ${attempt + 1}):`, JSON.stringify(result).substring(0, 200));
|
||
|
||
const state = result.data?.state;
|
||
const stateDetail = result.data?.state_detail;
|
||
|
||
// 检查任务状态
|
||
if (state === 1) {
|
||
// 任务完成 - PicWish API 返回 "file" 字段
|
||
const imageUrl = result.data?.file || result.data?.image;
|
||
if (imageUrl) {
|
||
return { success: true, imageUrl };
|
||
}
|
||
return { success: false, error: "处理完成但未返回图片URL" };
|
||
} else if (typeof state === "number" && state < 0) {
|
||
// 任务失败
|
||
return { success: false, error: result.data?.error || "去水印处理失败" };
|
||
} else if (state === 4 || state === 0) {
|
||
// 明确处理中状态，继续轮询
|
||
console.log(`[WatermarkRemove] 任务处理中: state=${state}, detail=${stateDetail || "N/A"}`);
|
||
}
|
||
// 其他状态默认继续轮询
|
||
|
||
await new Promise(resolve => setTimeout(resolve, pollInterval));
|
||
} catch (error) {
|
||
console.error(`[WatermarkRemove] 轮询错误 (attempt ${attempt + 1}):`, error);
|
||
await new Promise(resolve => setTimeout(resolve, pollInterval));
|
||
}
|
||
}
|
||
|
||
return {
|
||
success: false,
|
||
pending: true,
|
||
error: "去水印任务仍在处理中",
|
||
};
|
||
}
|
||
|
||
function buildWatermarkTaskPendingResponse(taskId: string) {
|
||
return NextResponse.json(
|
||
{
|
||
success: true,
|
||
processing: true,
|
||
status: "running",
|
||
taskId,
|
||
pollAfterMs: WATERMARK_POLL_INTERVAL_MS,
|
||
},
|
||
{ status: 202 }
|
||
);
|
||
}
|
||
|
||
async function finalizeWatermarkTaskResult(params: {
|
||
taskId: string;
|
||
userId: string;
|
||
startedAt?: number;
|
||
apiKey: string;
|
||
sourceImageUrls?: string[];
|
||
}): Promise<{
|
||
success: boolean;
|
||
original_url: string;
|
||
thumbnail_url: string;
|
||
processing_time?: number;
|
||
}> {
|
||
const pollResult = await pollForResult(params.taskId, params.apiKey, WATERMARK_SYNC_WAIT_TIMEOUT_MS);
|
||
if (pollResult.pending) {
|
||
throw new Error("WATERMARK_TASK_PENDING");
|
||
}
|
||
if (!pollResult.success || !pollResult.imageUrl) {
|
||
throw new Error(pollResult.error || "去水印处理失败");
|
||
}
|
||
|
||
const processedImageUrl = pollResult.imageUrl;
|
||
try {
|
||
const imageResponse = await fetch(processedImageUrl);
|
||
if (!imageResponse.ok) {
|
||
throw new Error("下载处理后的图片失败");
|
||
}
|
||
const imageBuffer = await imageResponse.arrayBuffer();
|
||
const filename = `watermark-removed-${Date.now()}.png`;
|
||
const uploadResult = await uploadFile(filename, imageBuffer);
|
||
|
||
let saved = false;
|
||
let saveError: unknown = null;
|
||
let savedImage:
|
||
| {
|
||
id: string;
|
||
}
|
||
| null = null;
|
||
for (let attempt = 1; attempt <= 3; attempt += 1) {
|
||
try {
|
||
([savedImage] = await db.insert(generatedImage).values({
|
||
userId: params.userId,
|
||
thumbnailUrl: uploadResult.url,
|
||
originalUrl: uploadResult.url,
|
||
imageSize: "2K",
|
||
prompt: "去水印处理",
|
||
modelId: "watermark-remove",
|
||
}).returning());
|
||
saved = true;
|
||
break;
|
||
} catch (error) {
|
||
saveError = error;
|
||
if (attempt < 3) {
|
||
await new Promise((resolve) => setTimeout(resolve, 120 * attempt));
|
||
}
|
||
}
|
||
}
|
||
if (!saved) {
|
||
console.error("[WatermarkRemove] 保存到相册失败:", saveError);
|
||
throw new Error("去水印成功但保存到相册失败，请稍后重试");
|
||
}
|
||
|
||
if (savedImage) {
|
||
try {
|
||
const sourceAssets =
|
||
Array.isArray(params.sourceImageUrls) && params.sourceImageUrls.length > 0
|
||
? await getMediaAssetsByUserAndUrls({
|
||
userId: params.userId,
|
||
urls: params.sourceImageUrls,
|
||
})
|
||
: [];
|
||
const media = await upsertMediaAsset({
|
||
userId: params.userId,
|
||
mediaType: "image",
|
||
sourceType: "edited",
|
||
status: "ready",
|
||
originalUrl: uploadResult.url,
|
||
previewUrl: uploadResult.url,
|
||
thumbnailUrl: uploadResult.url,
|
||
displayName: "去水印处理",
|
||
prompt: "去水印处理",
|
||
modelId: "watermark-remove",
|
||
sourceTable: "GeneratedImage",
|
||
sourceRowId: savedImage.id,
|
||
primaryParentAssetId: sourceAssets[0]?.id ?? null,
|
||
metadata: {
|
||
imageSize: "2K",
|
||
},
|
||
});
|
||
if (media) {
|
||
for (const sourceAsset of sourceAssets) {
|
||
await createMediaAssetRelation({
|
||
userId: params.userId,
|
||
fromAssetId: sourceAsset.id,
|
||
toAssetId: media.id,
|
||
relationType: "edited_from",
|
||
metadata: {
|
||
prompt: "去水印处理",
|
||
},
|
||
});
|
||
}
|
||
}
|
||
} catch (assetError) {
|
||
console.error("[WatermarkRemove] ⚠️ MediaAsset 注册失败:", assetError);
|
||
}
|
||
}
|
||
|
||
return {
|
||
success: true,
|
||
original_url: uploadResult.url,
|
||
thumbnail_url: uploadResult.url,
|
||
processing_time:
|
||
typeof params.startedAt === "number" ? (Date.now() - params.startedAt) / 1000 : undefined,
|
||
};
|
||
} catch (uploadError) {
|
||
console.error("[WatermarkRemove] 上传结果失败，回退原始URL:", uploadError);
|
||
return {
|
||
success: true,
|
||
original_url: processedImageUrl,
|
||
thumbnail_url: processedImageUrl,
|
||
processing_time:
|
||
typeof params.startedAt === "number" ? (Date.now() - params.startedAt) / 1000 : undefined,
|
||
};
|
||
}
|
||
}
|
||
|
||
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 config = await getWatermarkRemoveConfig();
|
||
if (!config?.apiKey) {
|
||
return NextResponse.json(
|
||
{ error: "去水印服务未配置，请联系管理员在系统设置中配置API密钥" },
|
||
{ status: 500 }
|
||
);
|
||
}
|
||
|
||
try {
|
||
const result = await finalizeWatermarkTaskResult({
|
||
taskId,
|
||
userId: session.user.id,
|
||
apiKey: config.apiKey,
|
||
});
|
||
return NextResponse.json(result);
|
||
} catch (error) {
|
||
if (error instanceof Error && error.message === "WATERMARK_TASK_PENDING") {
|
||
return buildWatermarkTaskPendingResponse(taskId);
|
||
}
|
||
throw error;
|
||
}
|
||
} catch (error) {
|
||
console.error("[WatermarkRemove][TaskStatus] 处理错误:", error);
|
||
return NextResponse.json(
|
||
{ error: 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 startTime = Date.now();
|
||
|
||
// 获取请求参数
|
||
const body = await request.json();
|
||
const { image_url } = body;
|
||
|
||
if (!image_url) {
|
||
return NextResponse.json(
|
||
{ error: "请提供图片URL" },
|
||
{ status: 400 }
|
||
);
|
||
}
|
||
|
||
// 获取API配置
|
||
const config = await getWatermarkRemoveConfig();
|
||
if (!config?.apiKey) {
|
||
return NextResponse.json(
|
||
{ error: "去水印服务未配置，请联系管理员在系统设置中配置API密钥" },
|
||
{ status: 500 }
|
||
);
|
||
}
|
||
|
||
console.log("[WatermarkRemove] 使用原图URL:", image_url.substring(0, 100) + "...");
|
||
|
||
// 处理图片URL，确保获取公网可访问的URL
|
||
let publicImageUrl: string;
|
||
|
||
// 如果是 base64 格式的 data URL，先上传到 S3
|
||
if (image_url.startsWith("data:")) {
|
||
console.log("[WatermarkRemove] 检测到base64格式，上传到S3...");
|
||
try {
|
||
// 避免对超长 base64 使用正则，直接按 data URL 头部拆分
|
||
const commaIndex = image_url.indexOf(",");
|
||
if (commaIndex <= 0) {
|
||
return NextResponse.json(
|
||
{ error: "无效的图片格式" },
|
||
{ status: 400 }
|
||
);
|
||
}
|
||
const header = image_url.slice(0, commaIndex);
|
||
const base64Data = image_url.slice(commaIndex + 1);
|
||
const headerMatch = header.match(/^data:image\/([a-zA-Z0-9.+-]+);base64$/i);
|
||
if (!headerMatch || !base64Data) {
|
||
return NextResponse.json(
|
||
{ error: "无效的图片格式" },
|
||
{ status: 400 }
|
||
);
|
||
}
|
||
const ext = headerMatch[1].toLowerCase() === "jpeg" ? "jpg" : headerMatch[1].toLowerCase();
|
||
const buffer = Buffer.from(base64Data, "base64");
|
||
|
||
// 上传到 S3
|
||
const filename = `watermark-input-${Date.now()}.${ext}`;
|
||
const uploadResult = await uploadFile(filename, buffer.buffer.slice(buffer.byteOffset, buffer.byteOffset + buffer.byteLength));
|
||
publicImageUrl = uploadResult.url;
|
||
console.log("[WatermarkRemove] base64图片已上传到S3:", publicImageUrl);
|
||
} catch (uploadError) {
|
||
console.error("[WatermarkRemove] base64上传失败:", uploadError);
|
||
return NextResponse.json(
|
||
{ error: "图片上传失败" },
|
||
{ status: 500 }
|
||
);
|
||
}
|
||
}
|
||
// 如果是 HTTP/HTTPS URL，下载后重新上传确保公网可访问
|
||
else if (image_url.startsWith("http://") || image_url.startsWith("https://")) {
|
||
try {
|
||
console.log("[WatermarkRemove] 下载原图并上传到S3...");
|
||
const imageResponse = await fetch(image_url);
|
||
if (!imageResponse.ok) {
|
||
throw new Error("下载原图失败");
|
||
}
|
||
const imageBuffer = await imageResponse.arrayBuffer();
|
||
const filename = `watermark-input-${Date.now()}.jpg`;
|
||
const uploadResult = await uploadFile(filename, imageBuffer);
|
||
publicImageUrl = uploadResult.url;
|
||
console.log("[WatermarkRemove] 图片已上传到S3:", publicImageUrl);
|
||
} catch (downloadError) {
|
||
console.error("[WatermarkRemove] 下载/上传原图失败:", downloadError);
|
||
// 如果下载失败，尝试使用原始URL
|
||
console.log("[WatermarkRemove] 尝试使用原始URL");
|
||
publicImageUrl = image_url;
|
||
}
|
||
} else {
|
||
return NextResponse.json(
|
||
{ error: "图片URL格式无效" },
|
||
{ status: 400 }
|
||
);
|
||
}
|
||
|
||
console.log("[WatermarkRemove] 开始处理图片:", publicImageUrl);
|
||
|
||
const pointsConfig = await getPointsSystemConfig();
|
||
const pointsCharge = await prepareFeaturePointsCharge({
|
||
userId: session.user.id,
|
||
unitCost: getImageToolCreditCost(pointsConfig, "watermark-remove"),
|
||
reason: "image_watermark_remove",
|
||
modelId: "watermark-remove",
|
||
note: "去水印消耗积分",
|
||
pointsConfig,
|
||
});
|
||
if (!pointsCharge.ok) {
|
||
return NextResponse.json(pointsCharge.payload, { status: pointsCharge.status });
|
||
}
|
||
|
||
// 创建去水印任务
|
||
const formData = new FormData();
|
||
formData.append("sync", "0");
|
||
formData.append("url", publicImageUrl);
|
||
|
||
const createResponse = await fetch(
|
||
"https://techsz.aoscdn.com/api/tasks/visual/external/watermark-remove",
|
||
{
|
||
method: "POST",
|
||
headers: {
|
||
"X-API-KEY": config.apiKey,
|
||
},
|
||
body: formData,
|
||
}
|
||
);
|
||
|
||
if (!createResponse.ok) {
|
||
const errorText = await createResponse.text();
|
||
console.error("[WatermarkRemove] 创建任务失败:", errorText);
|
||
return NextResponse.json(
|
||
{ error: "创建去水印任务失败" },
|
||
{ status: 500 }
|
||
);
|
||
}
|
||
|
||
const createResult = await createResponse.json();
|
||
console.log("[WatermarkRemove] 创建任务结果:", JSON.stringify(createResult).substring(0, 200));
|
||
|
||
const taskId = createResult.data?.task_id;
|
||
if (!taskId) {
|
||
return NextResponse.json(
|
||
{ error: "未获取到任务ID" },
|
||
{ status: 500 }
|
||
);
|
||
}
|
||
|
||
if (pointsCharge.chargeable) {
|
||
try {
|
||
await pointsCharge.consume({
|
||
note: `去水印任务消耗积分（${taskId}）`,
|
||
});
|
||
} catch (pointsError) {
|
||
console.error("[Points] Failed to consume watermark remove points:", pointsError);
|
||
}
|
||
}
|
||
|
||
try {
|
||
const result = await finalizeWatermarkTaskResult({
|
||
taskId,
|
||
userId: session.user.id,
|
||
startedAt: startTime,
|
||
apiKey: config.apiKey,
|
||
sourceImageUrls:
|
||
typeof image_url === "string" && /^https?:\/\//i.test(image_url)
|
||
? [image_url]
|
||
: undefined,
|
||
});
|
||
return NextResponse.json(result);
|
||
} catch (error) {
|
||
if (error instanceof Error && error.message === "WATERMARK_TASK_PENDING") {
|
||
return buildWatermarkTaskPendingResponse(taskId);
|
||
}
|
||
throw error;
|
||
}
|
||
} catch (error) {
|
||
console.error("[WatermarkRemove] 处理错误:", error);
|
||
return NextResponse.json(
|
||
{ error: error instanceof Error ? error.message : "去水印处理失败" },
|
||
{ status: 500 }
|
||
);
|
||
}
|
||
}
|