新增AOS抠图服务商配置并接入异步任务调用

2026-02-15 22:56:34 +08:00 · 2026-02-15 22:56:34 +08:00 · 2987bd793e
commit 2987bd793e
parent 67c18dd820
2 changed files with 307 additions and 10 deletions
--- a/app/(admin)/settings/admin-settings-client.tsx
+++ b/app/(admin)/settings/admin-settings-client.tsx
@ -541,11 +541,17 @@ export function AdminSettingsClient() {

  // 抠图API配置状态
  const [rmbgConfig, setRmbgConfig] = useState({
-    provider: "rmbg" as "rmbg" | "photoroom", // 抠图服务提供商
+    provider: "rmbg" as "rmbg" | "photoroom" | "aoscdn", // 抠图服务提供商
    apiUrl: "",
    apiKey: "",
    // PhotoRoom 配置
    photoroomApiKey: "",
+    // AOS 配置
+    aosApiKey: "",
+    aosApiBaseUrl: "https://techsz.aoscdn.com",
+    aosForegroundType: "auto" as "auto" | "person" | "object" | "stamp",
+    aosCrop: false,
+    aosFormat: "png" as "png" | "jpg",
  });

  // 变清晰API配置状态
@ -598,16 +604,31 @@ export function AdminSettingsClient() {
  useEffect(() => {
    if (rmbgConfigData?.setting?.value) {
      const config = rmbgConfigData.setting.value as { 
-        provider?: "rmbg" | "photoroom";
+        provider?: "rmbg" | "photoroom" | "aoscdn";
        apiUrl?: string; 
        apiKey?: string;
        photoroomApiKey?: string;
+        aosApiKey?: string;
+        aosApiBaseUrl?: string;
+        aosForegroundType?: "" | "person" | "object" | "stamp";
+        aosCrop?: 0 | 1;
+        aosFormat?: "png" | "jpg";
      };
      setRmbgConfig({
        provider: config.provider || "rmbg",
        apiUrl: config.apiUrl || "",
        apiKey: config.apiKey || "",
        photoroomApiKey: config.photoroomApiKey || "",
+        aosApiKey: config.aosApiKey || "",
+        aosApiBaseUrl: config.aosApiBaseUrl || "https://techsz.aoscdn.com",
+        aosForegroundType:
+          config.aosForegroundType === "person" ||
+          config.aosForegroundType === "object" ||
+          config.aosForegroundType === "stamp"
+            ? config.aosForegroundType
+            : "auto",
+        aosCrop: config.aosCrop === 1,
+        aosFormat: config.aosFormat === "jpg" ? "jpg" : "png",
      });
    }
  }, [rmbgConfigData]);
@ -1289,6 +1310,10 @@ export function AdminSettingsClient() {
        toast.error("请填写PhotoRoom API密钥");
        return;
      }
+      if (rmbgConfig.provider === "aoscdn" && !rmbgConfig.aosApiKey.trim()) {
+        toast.error("请填写AOS API密钥");
+        return;
+      }

      const response = await fetch("/api/admin/settings", {
          method: "POST",
@ -1300,6 +1325,11 @@ export function AdminSettingsClient() {
            apiUrl: rmbgConfig.apiUrl.trim(),
            apiKey: rmbgConfig.apiKey.trim() || undefined,
            photoroomApiKey: rmbgConfig.photoroomApiKey.trim() || undefined,
+            aosApiKey: rmbgConfig.aosApiKey.trim() || undefined,
+            aosApiBaseUrl: rmbgConfig.aosApiBaseUrl.trim() || "https://techsz.aoscdn.com",
+            aosForegroundType: rmbgConfig.aosForegroundType === "auto" ? "" : rmbgConfig.aosForegroundType,
+            aosCrop: rmbgConfig.aosCrop ? 1 : 0,
+            aosFormat: rmbgConfig.aosFormat,
          },
        }),
      });
@ -5940,7 +5970,7 @@ export function AdminSettingsClient() {
            <CardHeader>
              <CardTitle>抠图API配置</CardTitle>
              <CardDescription>
-                配置背景移除API，用于画图标注功能中的抠图处理。支持自部署RMBG和PhotoRoom云服务。
+                配置背景移除API，用于画图标注功能中的抠图处理。支持自部署RMBG、PhotoRoom和AOS云服务。
              </CardDescription>
            </CardHeader>
            <CardContent className="space-y-6 pt-6">
@ -5970,6 +6000,17 @@ export function AdminSettingsClient() {
                    />
                    <span className="text-sm">PhotoRoom (云服务)</span>
                  </label>
+                  <label className="flex items-center gap-2 cursor-pointer">
+                    <input
+                      type="radio"
+                      name="rmbgProvider"
+                      value="aoscdn"
+                      checked={rmbgConfig.provider === "aoscdn"}
+                      onChange={() => setRmbgConfig({ ...rmbgConfig, provider: "aoscdn" })}
+                      className="w-4 h-4"
+                    />
+                    <span className="text-sm">AOS Segmentation (云服务)</span>
+                  </label>
                </div>
              </div>

@ -6041,6 +6082,93 @@ export function AdminSettingsClient() {
                </>
              )}

+              {/* AOS 配置 */}
+              {rmbgConfig.provider === "aoscdn" && (
+                <>
+                  <div className="text-sm font-medium text-muted-foreground">AOS Segmentation 配置</div>
+                  <div className="space-y-2">
+                    <Label htmlFor="aosApiKey">API密钥 *</Label>
+                    <Input
+                      id="aosApiKey"
+                      type="password"
+                      value={rmbgConfig.aosApiKey}
+                      onChange={(e) => setRmbgConfig({ ...rmbgConfig, aosApiKey: e.target.value })}
+                      placeholder="输入 AOS X-API-KEY"
+                    />
+                  </div>
+                  <div className="space-y-2">
+                    <Label htmlFor="aosApiBaseUrl">API地址</Label>
+                    <Input
+                      id="aosApiBaseUrl"
+                      value={rmbgConfig.aosApiBaseUrl}
+                      onChange={(e) => setRmbgConfig({ ...rmbgConfig, aosApiBaseUrl: e.target.value })}
+                      placeholder="https://techsz.aoscdn.com"
+                    />
+                    <div className="text-xs text-muted-foreground">
+                      默认使用官方地址，可按需替换。
+                    </div>
+                  </div>
+                  <div className="grid gap-4 md:grid-cols-2">
+                    <div className="space-y-2">
+                      <Label>前景类型</Label>
+                      <Select
+                        value={rmbgConfig.aosForegroundType}
+                        onValueChange={(value: "auto" | "person" | "object" | "stamp") =>
+                          setRmbgConfig({ ...rmbgConfig, aosForegroundType: value })
+                        }
+                      >
+                        <SelectTrigger>
+                          <SelectValue />
+                        </SelectTrigger>
+                        <SelectContent>
+                          <SelectItem value="auto">自动识别</SelectItem>
+                          <SelectItem value="person">人物</SelectItem>
+                          <SelectItem value="object">物体</SelectItem>
+                          <SelectItem value="stamp">图章</SelectItem>
+                        </SelectContent>
+                      </Select>
+                    </div>
+                    <div className="space-y-2">
+                      <Label>输出格式</Label>
+                      <Select
+                        value={rmbgConfig.aosFormat}
+                        onValueChange={(value: "png" | "jpg") =>
+                          setRmbgConfig({ ...rmbgConfig, aosFormat: value })
+                        }
+                      >
+                        <SelectTrigger>
+                          <SelectValue />
+                        </SelectTrigger>
+                        <SelectContent>
+                          <SelectItem value="png">PNG（透明背景）</SelectItem>
+                          <SelectItem value="jpg">JPG（白底）</SelectItem>
+                        </SelectContent>
+                      </Select>
+                    </div>
+                  </div>
+                  <div className="flex items-center justify-between rounded-lg border p-3">
+                    <div className="space-y-0.5">
+                      <Label className="text-sm">裁剪到主体边缘</Label>
+                      <div className="text-xs text-muted-foreground">
+                        关闭时保持原图尺寸，开启时按主体区域裁剪。
+                      </div>
+                    </div>
+                    <Switch
+                      checked={rmbgConfig.aosCrop}
+                      onCheckedChange={(checked) => setRmbgConfig({ ...rmbgConfig, aosCrop: checked })}
+                    />
+                  </div>
+                  <div className="rounded-lg border border-indigo-200 bg-indigo-50 p-4 space-y-2">
+                    <div className="text-sm font-medium text-indigo-900">AOS 说明</div>
+                    <div className="text-xs text-indigo-800 space-y-1">
+                      <p>• 使用异步任务模式（1秒轮询，最多30秒）</p>
+                      <p>• 接口返回链接有效期1小时，系统会自动转存到你的存储桶</p>
+                      <p>• 官方限制：图片最大20MB，分辨率最大4096x4096</p>
+                    </div>
+                  </div>
+                </>
+              )}
+
              <Button onClick={handleSaveRmbgConfig}>
                <Save className="mr-2 h-4 w-4" /> 保存配置
              </Button>
--- a/app/api/rmbg/route.ts
+++ b/app/api/rmbg/route.ts
@ -5,6 +5,8 @@ import { uploadFile } from "@/lib/storage/service";
 import { db } from "@/lib/db";
 import { generatedImage } from "@/lib/db/schema";

+type RmbgProvider = "rmbg" | "photoroom" | "aoscdn";
+
 /**
 * 压缩图片生成缩略图（使用 sharp）
 */
@ -86,6 +88,21 @@ async function uploadBase64ToS3(base64Data: string, mimeType: string = "image/pn
  return uploadBufferToS3(buffer, mimeType);
 }

+const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));
+
+const withTimeout = async <T,>(
+  fn: (signal: AbortSignal) => Promise<T>,
+  timeoutMs: number
+): Promise<T> => {
+  const controller = new AbortController();
+  const timeoutId = setTimeout(() => controller.abort(), timeoutMs);
+  try {
+    return await fn(controller.signal);
+  } finally {
+    clearTimeout(timeoutId);
+  }
+};
+
 /**
 * 使用 RMBG-2.0 API 进行抠图
 */
@ -215,9 +232,129 @@ async function processWithPhotoRoom(
  }
 }

+/**
+ * 使用 AOS segmentation API 进行抠图（异步任务 + 轮询）
+ */
+async function processWithAosSegmentation(
+  imageBuffer: Buffer,
+  config: {
+    apiKey: string;
+    apiBaseUrl?: string;
+    foregroundType?: "" | "person" | "object" | "stamp";
+    crop?: 0 | 1;
+    format?: "png" | "jpg";
+  }
+): Promise<{ success: boolean; imageBuffer?: Buffer; processing_time?: number; error?: string }> {
+  const startTime = Date.now();
+  const baseUrl = (config.apiBaseUrl || "https://techsz.aoscdn.com").replace(/\/+$/, "");
+  const createUrl = `${baseUrl}/api/tasks/visual/segmentation`;
+  const headers = { "X-API-KEY": config.apiKey };
+
+  try {
+    const createForm = new FormData();
+    const imageBlob = new Blob([new Uint8Array(imageBuffer)], { type: "image/png" });
+    createForm.append("image_file", imageBlob, "image.png");
+    createForm.append("sync", "0");
+    createForm.append("return_type", "1");
+    createForm.append("output_type", "2");
+    createForm.append("format", config.format || "png");
+    createForm.append("crop", config.crop === 1 ? "1" : "0");
+    if (config.foregroundType) {
+      createForm.append("type", config.foregroundType);
+    }
+
+    const createResponse = await withTimeout(
+      (signal) =>
+        fetch(createUrl, {
+          method: "POST",
+          headers,
+          body: createForm,
+          signal,
+        }),
+      30000
+    );
+
+    const createPayload = await createResponse.json().catch(() => null);
+    if (!createResponse.ok || !createPayload?.data?.task_id) {
+      return {
+        success: false,
+        error: `AOS任务创建失败: ${createPayload?.message || createResponse.statusText}`,
+      };
+    }
+
+    const taskId = createPayload.data.task_id as string;
+    const taskUrl = `${createUrl}/${taskId}`;
+    const maxPollCount = 30;
+
+    for (let pollCount = 0; pollCount < maxPollCount; pollCount++) {
+      await sleep(1000);
+
+      const pollResponse = await withTimeout(
+        (signal) =>
+          fetch(taskUrl, {
+            method: "GET",
+            headers,
+            signal,
+          }),
+        30000
+      );
+
+      const pollPayload = await pollResponse.json().catch(() => null);
+      if (!pollResponse.ok || !pollPayload?.data) {
+        return {
+          success: false,
+          error: `AOS任务查询失败: ${pollPayload?.message || pollResponse.statusText}`,
+        };
+      }
+
+      const state = Number(pollPayload.data.state);
+      const imageUrl = pollPayload.data.image as string | undefined;
+
+      if (state === 1 && imageUrl) {
+        const imageResponse = await withTimeout(
+          (signal) =>
+            fetch(imageUrl, {
+              method: "GET",
+              signal,
+            }),
+          30000
+        );
+
+        if (!imageResponse.ok) {
+          return {
+            success: false,
+            error: `AOS结果下载失败: ${imageResponse.status} ${imageResponse.statusText}`,
+          };
+        }
+
+        const resultBuffer = Buffer.from(await imageResponse.arrayBuffer());
+        return {
+          success: true,
+          imageBuffer: resultBuffer,
+          processing_time: (Date.now() - startTime) / 1000,
+        };
+      }
+
+      if (state === 2 || state === 3 || state === -1) {
+        return {
+          success: false,
+          error: `AOS任务失败: ${pollPayload?.message || "状态异常"}`,
+        };
+      }
+    }
+
+    return { success: false, error: "AOS任务超时（30秒）" };
+  } catch (error) {
+    if (error instanceof Error && error.name === "AbortError") {
+      return { success: false, error: "AOS请求超时（30秒）" };
+    }
+    return { success: false, error: error instanceof Error ? error.message : "未知错误" };
+  }
+}
+
 /**
 * 抠图API路由
- * 从数据库读取API配置，支持 RMBG-2.0 和 PhotoRoom 两种服务
+ * 从数据库读取API配置，支持 RMBG-2.0、PhotoRoom、AOS 三种服务
 */
 export async function POST(request: NextRequest) {
  try {
@ -243,13 +380,18 @@ export async function POST(request: NextRequest) {
    }

    const config = rmbgConfig.value as {
-      provider?: "rmbg" | "photoroom";
+      provider?: RmbgProvider;
      apiUrl?: string;
      apiKey?: string;
      photoroomApiKey?: string;
+      aosApiKey?: string;
+      aosApiBaseUrl?: string;
+      aosForegroundType?: "" | "person" | "object" | "stamp";
+      aosCrop?: 0 | 1;
+      aosFormat?: "png" | "jpg";
    };

-    const provider = config.provider || "rmbg";
+    const provider: RmbgProvider = config.provider || "rmbg";
    
    // 验证配置
    if (provider === "rmbg" && !config.apiUrl) {
@ -266,6 +408,13 @@ export async function POST(request: NextRequest) {
      );
    }

+    if (provider === "aoscdn" && !config.aosApiKey) {
+      return NextResponse.json(
+        { error: "AOS API密钥未配置" },
+        { status: 400 }
+      );
+    }
+
    // 2. 获取请求体
    const contentType = request.headers.get("content-type") || "";
    
@ -280,11 +429,12 @@ export async function POST(request: NextRequest) {
      prompt = body.prompt || "抠图处理";
      chatId = body.chatId || null;
      messageId = body.messageId || null;
+      const imageUrl = body.imageUrl || body.image_url;
      
-      if (body.imageUrl) {
-        console.log(`[RMBG] 从 URL 获取图片: ${body.imageUrl}`);
+      if (imageUrl) {
+        console.log(`[RMBG] 从 URL 获取图片: ${imageUrl}`);
        try {
-          const imageResponse = await fetch(body.imageUrl);
+          const imageResponse = await fetch(imageUrl);
          if (!imageResponse.ok) {
            return NextResponse.json(
              { error: `无法获取图片: ${imageResponse.status} ${imageResponse.statusText}` },
@ -353,6 +503,25 @@ export async function POST(request: NextRequest) {
        );
      }
      
+      resultBuffer = result.imageBuffer;
+      resultBase64 = resultBuffer.toString("base64");
+      processingTime = result.processing_time || 0;
+    } else if (provider === "aoscdn") {
+      const result = await processWithAosSegmentation(imageBuffer, {
+        apiKey: config.aosApiKey!,
+        apiBaseUrl: config.aosApiBaseUrl,
+        foregroundType: config.aosForegroundType || "",
+        crop: config.aosCrop === 1 ? 1 : 0,
+        format: config.aosFormat === "jpg" ? "jpg" : "png",
+      });
+
+      if (!result.success || !result.imageBuffer) {
+        return NextResponse.json(
+          { error: "抠图处理失败", detail: result.error },
+          { status: 500 }
+        );
+      }
+
      resultBuffer = result.imageBuffer;
      resultBase64 = resultBuffer.toString("base64");
      processingTime = result.processing_time || 0;
@ -405,7 +574,7 @@ export async function POST(request: NextRequest) {
        originalUrl,
        imageSize: "2K",
        prompt: prompt || "抠图处理",
-        modelId: provider === "photoroom" ? "photoroom-hd" : "rmbg-2.0",
+        modelId: provider === "photoroom" ? "photoroom-hd" : provider === "aoscdn" ? "aos-segmentation" : "rmbg-2.0",
      }).returning();
      console.log(`[RMBG] ✅ 图片已保存到相册，ID: ${savedImage.id}`);
    } catch (dbError) {