278 lines
9.4 KiB
TypeScript
278 lines
9.4 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";
|
||
|
||
// PicWish API 需要公网可访问的 HTTP/HTTPS URL
|
||
|
||
// 获取去水印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): Promise<{ success: boolean; imageUrl?: string; error?: string }> {
|
||
const maxAttempts = 30; // 最多轮询30次（30秒）
|
||
const pollInterval = 1000; // 每次间隔1秒
|
||
|
||
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));
|
||
|
||
// 检查任务状态
|
||
if (result.data?.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 (result.data?.state === -1 || result.data?.state < 0) {
|
||
// 任务失败
|
||
return { success: false, error: result.data?.error || "去水印处理失败" };
|
||
}
|
||
// state === 0 表示处理中，继续轮询
|
||
|
||
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, error: "处理超时，请稍后重试" };
|
||
}
|
||
|
||
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 matches = image_url.match(/^data:image\/(\w+);base64,(.+)$/);
|
||
if (!matches) {
|
||
return NextResponse.json(
|
||
{ error: "无效的图片格式" },
|
||
{ status: 400 }
|
||
);
|
||
}
|
||
const ext = matches[1] === "jpeg" ? "jpg" : matches[1];
|
||
const base64Data = matches[2];
|
||
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 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 }
|
||
);
|
||
}
|
||
|
||
// 轮询获取结果
|
||
const pollResult = await pollForResult(taskId, config.apiKey);
|
||
|
||
if (!pollResult.success) {
|
||
return NextResponse.json(
|
||
{ error: pollResult.error || "去水印处理失败" },
|
||
{ status: 500 }
|
||
);
|
||
}
|
||
|
||
// 下载处理后的图片并上传到S3
|
||
const processedImageUrl = pollResult.imageUrl!;
|
||
console.log("[WatermarkRemove] 处理完成，结果URL:", processedImageUrl);
|
||
|
||
try {
|
||
// 下载处理后的图片
|
||
const imageResponse = await fetch(processedImageUrl);
|
||
if (!imageResponse.ok) {
|
||
throw new Error("下载处理后的图片失败");
|
||
}
|
||
const imageBuffer = await imageResponse.arrayBuffer();
|
||
|
||
// 上传到S3
|
||
const filename = `watermark-removed-${Date.now()}.png`;
|
||
const uploadResult = await uploadFile(filename, imageBuffer);
|
||
|
||
const processingTime = (Date.now() - startTime) / 1000;
|
||
console.log(`[WatermarkRemove] 处理完成，耗时: ${processingTime}s`);
|
||
|
||
// 保存到相册/历史出图记录
|
||
try {
|
||
await db.insert(generatedImage).values({
|
||
userId: session.user.id,
|
||
thumbnailUrl: uploadResult.url,
|
||
originalUrl: uploadResult.url,
|
||
imageSize: "2K",
|
||
prompt: "去水印处理",
|
||
modelId: "watermark-remove",
|
||
});
|
||
console.log(`[WatermarkRemove] 已保存到相册: ${uploadResult.url.substring(0, 80)}...`);
|
||
} catch (saveError) {
|
||
console.error("[WatermarkRemove] 保存到相册失败:", saveError);
|
||
}
|
||
|
||
return NextResponse.json({
|
||
success: true,
|
||
original_url: uploadResult.url,
|
||
thumbnail_url: uploadResult.url,
|
||
processing_time: processingTime,
|
||
});
|
||
} catch (uploadError) {
|
||
console.error("[WatermarkRemove] 上传结果失败:", uploadError);
|
||
// 如果上传失败，直接返回原始URL
|
||
const processingTime = (Date.now() - startTime) / 1000;
|
||
return NextResponse.json({
|
||
success: true,
|
||
original_url: processedImageUrl,
|
||
thumbnail_url: processedImageUrl,
|
||
processing_time: processingTime,
|
||
});
|
||
}
|
||
} catch (error) {
|
||
console.error("[WatermarkRemove] 处理错误:", error);
|
||
return NextResponse.json(
|
||
{ error: error instanceof Error ? error.message : "去水印处理失败" },
|
||
{ status: 500 }
|
||
);
|
||
}
|
||
}
|