Split image count into separate generation jobs

2026-05-13 02:10:32 +00:00
parent 7b3235b218
commit 8bc18c6cd8
3 changed files with 128 additions and 42 deletions
--- a/src/app/api/generate/image/route.ts
+++ b/src/app/api/generate/image/route.ts
@@ -127,9 +127,15 @@ function normalizeNewApiImageCount(modelName: string | undefined, count: number)
  return Math.min(10, Math.max(1, Math.floor(count)));
 }

-function isNewApiConfig(config: Pick<CustomApiConfig, 'provider'> | undefined): boolean {
+function isNewApiConfig(config: Pick<CustomApiConfig, 'provider' | 'apiUrl' | 'modelName'> | undefined): boolean {
  const provider = (config?.provider || '').trim().toLowerCase();
-  return provider === 'newapi' || provider === 'new api';
+  const apiUrl = (config?.apiUrl || '').trim().toLowerCase();
+  return provider === 'newapi'
+    || provider === 'new api'
+    || isGptImageModel(config?.modelName)
+    || isDallE3Model(config?.modelName)
+    || isDallE2Model(config?.modelName)
+    || /\/v1\/images\/(generations|edits)\b/i.test(apiUrl);
 }

 function normalizeImageOutputFormat(value: unknown): 'png' | 'jpeg' | 'webp' {
@@ -356,7 +362,7 @@ async function requestQualifiedCustomImages(
 ): Promise<{ images: string[]; rejected: string[]; upstreamError?: { status: number; text: string } }> {
  const accepted: string[] = [];
  const rejected: string[] = [];
-  const maxAttempts = Math.max(targetCount * 3, 3);
+  const maxAttempts = 1;

  for (let attempt = 1; attempt <= maxAttempts && accepted.length < targetCount; attempt += 1) {
    const remaining = targetCount - accepted.length;
--- a/src/app/api/generation-jobs/route.ts
+++ b/src/app/api/generation-jobs/route.ts
@@ -50,13 +50,41 @@ export async function POST(request: NextRequest) {
      etaSource = estimate.source;
      etaSampleCount = estimate.sampleCount;
      etaWindowDays = estimate.windowDays;
+      const payloadJson = JSON.stringify(payload);
+      const existing = await client.query(
+        `SELECT id, status, progress
+         FROM generation_jobs
+         WHERE user_id = $1
+           AND type = $2
+           AND status IN ('queued', 'running')
+           AND payload = $3::jsonb
+         ORDER BY created_at DESC
+         LIMIT 1`,
+        [userId, type, payloadJson],
+      );
+      if (existing.rows.length > 0) {
+        const row = existing.rows[0];
+        return NextResponse.json({
+          jobId: row.id,
+          status: row.status,
+          estimateSeconds,
+          progress: row.progress || {},
+          eta: {
+            estimateSeconds,
+            source: etaSource,
+            sampleCount: etaSampleCount,
+            windowDays: etaWindowDays,
+          },
+          deduplicated: true,
+        }, { status: 202 });
+      }
      const result = await client.query(
        `INSERT INTO generation_jobs (type, status, payload, user_id, provider, model_name, api_url, progress)
         VALUES ($1, 'queued', $2::jsonb, $3, $4, $5, $6, $7::jsonb)
         RETURNING id`,
        [
          type,
-          JSON.stringify(payload),
+          payloadJson,
          userId,
          identity.provider,
          identity.modelName,
--- a/src/components/create/text-to-image.tsx
+++ b/src/components/create/text-to-image.tsx
@@ -71,6 +71,7 @@ export function TextToImagePanel() {
  const [results, setResults] = useState<string[]>([]);
  const [generationError, setGenerationError] = useState<GenerationErrorState | null>(null);
  const [optimizing, setOptimizing] = useState(false);
+  const activeSubmissionSignaturesRef = useRef(new Set<string>());
  const generating = activeTasks.length > 0;

  // History state
@@ -197,7 +198,9 @@ export function TextToImagePanel() {
    [stylePresets, selectedStyleLabel],
  );
  const creditCount = count === 'auto' ? (inferredImageParams.count ?? 1) : (Number(count) || 1);
-  const credits = calcImageCredits(selectedModel, resolution, aspectRatio, creditCount);
+  const creditResolution = resolution === 'auto' ? inferredImageParams.resolution : resolution;
+  const creditAspectRatio = aspectRatio === 'auto' ? inferredImageParams.aspectRatio : aspectRatio;
+  const credits = calcImageCredits(selectedModel, creditResolution, creditAspectRatio, creditCount);

  const resolveGenerationParams = useCallback((): { aspectRatio: string; resolution: string; count: number } | null => {
    const resolvedAspectRatio = aspectRatio === 'auto' ? inferredImageParams.aspectRatio : aspectRatio;
@@ -230,35 +233,26 @@ export function TextToImagePanel() {
    if (!user) { toast.error('请先登录'); return; }

    setGenerationError(null);
-    const taskId = `text2img-${Date.now()}-${Math.random().toString(36).slice(2, 7)}`;
+    let submissionSignature: string | null = null;
    try {
      const resolvedParams = resolveGenerationParams();
      if (!resolvedParams) return;
-      setActiveTasks(prev => [
-        ...prev,
-        {
-          id: taskId,
-          title: '正在生成图片',
-          startedAt: Date.now(),
-          estimateSeconds: 90,
-          jobStatus: null,
-          finalCountdownSeconds: null,
-        },
-      ]);
+      const taskCount = resolvedParams.count;
+
      // Keep custom/system API size aligned with the selected resolution.
      const useCustomApiSize = isCustomModel(selectedModel) || isSystemModel(selectedModel);
      const resolvedSize = useCustomApiSize
        ? resolveCustomApiImageSize(resolvedParams.aspectRatio, resolvedParams.resolution)
        : resolveImageSize(resolvedParams.aspectRatio, resolvedParams.resolution);

-      let requestBody: Record<string, unknown> = {
+      let requestBodyBase: Record<string, unknown> = {
        prompt: prompt.trim(),
        negativePrompt: negativePrompt.trim() || undefined,
        model: selectedModel,
        aspectRatio: resolvedParams.aspectRatio,
        resolution: resolvedParams.resolution,
        size: resolvedSize,
-        count: resolvedParams.count,
+        count: 1,
        outputFormat,
        imageQuality,
        styleLabel: selectedStylePreset?.label,
@@ -269,26 +263,72 @@ export function TextToImagePanel() {
      if (isCustomModel(selectedModel)) {
        const key = imageKeys.find(k => k.id === getCustomKeyId(selectedModel));
        if (key) {
-          requestBody = { ...requestBody, model: key.modelName, customApiConfig: { customApiKeyId: key.id, modelName: key.modelName } };
+          requestBodyBase = { ...requestBodyBase, model: key.modelName, customApiConfig: { customApiKeyId: key.id, modelName: key.modelName } };
        }
      } else if (isSystemModel(selectedModel)) {
        const api = systemImageApis.find(a => a.id === getSystemApiId(selectedModel));
        if (api) {
-          requestBody = { ...requestBody, model: api.modelName, customApiConfig: { systemApiId: api.id, modelName: api.modelName } };
+          requestBodyBase = { ...requestBodyBase, model: api.modelName, customApiConfig: { systemApiId: api.id, modelName: api.modelName } };
        }
      }

-      // Fetch with 180s timeout for image generation
-      const data = await runGenerationJob<{ images?: string[]; error?: string }>(
-        'image',
-        requestBody,
-        { timeoutMs: 900_000, onStatus: (status: GenerationJobStatus) => updateActiveTask(taskId, { jobStatus: status }) },
-      );
-      await runGenerationFinalCountdown((seconds) => updateActiveTask(taskId, { finalCountdownSeconds: seconds }), 3);
-      if (data.images && data.images.length > 0) {
-        setResults(prev => [...data.images!, ...prev]);
+      submissionSignature = JSON.stringify({
+        prompt: prompt.trim(),
+        negativePrompt: negativePrompt.trim(),
+        model: selectedModel,
+        aspectRatio: resolvedParams.aspectRatio,
+        resolution: resolvedParams.resolution,
+        count: taskCount,
+        outputFormat,
+        imageQuality,
+        styleLabel: selectedStylePreset?.label || '',
+        guidanceScale,
+      });
+      if (activeSubmissionSignaturesRef.current.has(submissionSignature)) {
+        toast.info('相同任务正在生成中，请勿重复提交');
+        return;
+      }
+      activeSubmissionSignaturesRef.current.add(submissionSignature);
+
+      const batchId = `text2img-${Date.now()}-${Math.random().toString(36).slice(2, 7)}`;
+      const taskIds = Array.from({ length: taskCount }, (_, index) => `${batchId}-${index + 1}`);
+      setActiveTasks(prev => [
+        ...prev,
+        ...taskIds.map(taskId => ({
+          id: taskId,
+          title: '正在生成图片',
+          startedAt: Date.now(),
+          estimateSeconds: 90,
+          jobStatus: null,
+          finalCountdownSeconds: null,
+        })),
+      ]);
+
+      const runSingleTask = async (taskId: string, index: number) => {
+        try {
+          const data = await runGenerationJob<{ images?: string[]; error?: string }>(
+            'image',
+            { ...requestBodyBase, count: 1, clientRequestId: `${batchId}-${index + 1}` },
+            { timeoutMs: 900_000, onStatus: (status: GenerationJobStatus) => updateActiveTask(taskId, { jobStatus: status }) },
+          );
+          await runGenerationFinalCountdown((seconds) => updateActiveTask(taskId, { finalCountdownSeconds: seconds }), 3);
+          if (!data.images || data.images.length === 0) {
+            throw new Error(data.error || '图片生成失败');
+          }
+          return data.images;
+        } finally {
+          removeActiveTask(taskId);
+        }
+      };
+
+      const settled = await Promise.allSettled(taskIds.map((taskId, index) => runSingleTask(taskId, index)));
+      const generatedImages = settled.flatMap(result => result.status === 'fulfilled' ? result.value : []);
+      const failedResults = settled.filter((result): result is PromiseRejectedResult => result.status === 'rejected');
+
+      if (generatedImages.length > 0) {
+        setResults(prev => [...generatedImages, ...prev]);
        setGenerationError(null);
-        for (const url of data.images) {
+        for (const url of generatedImages) {
          addRecord({
            type: 'image', url, prompt: prompt.trim(),
            negativePrompt: negativePrompt.trim() || undefined,
@@ -299,7 +339,8 @@ export function TextToImagePanel() {
              creationMode: 'text2img',
              aspectRatio: resolvedParams.aspectRatio,
              resolution: resolvedParams.resolution,
-              count: resolvedParams.count,
+              count: 1,
+              batchCount: taskCount,
              outputFormat,
              imageQuality,
              styleLabel: selectedStylePreset?.label,
@@ -307,24 +348,33 @@ export function TextToImagePanel() {
            },
          });
        }
-        toast.success(`生成 ${data.images.length} 张图片`);
-        // Record credit consumption (custom/system models cost 0 credits)
-        if (membershipEnabled && credits > 0 && user) {
+        toast.success(`生成 ${generatedImages.length} 张图片`);
+
+        const chargedCredits = calcImageCredits(selectedModel, resolvedParams.resolution, resolvedParams.aspectRatio, generatedImages.length);
+        if (membershipEnabled && chargedCredits > 0 && user) {
          const currentCredits = typeof user.creditsBalance === 'number' ? user.creditsBalance : 0;
          addCreditRecord({
            type: 'consume',
-            amount: -credits,
-            balanceAfter: Math.max(0, currentCredits - credits),
+            amount: -chargedCredits,
+            balanceAfter: Math.max(0, currentCredits - chargedCredits),
            description: `文生图 - ${getCurrentModelLabel()}`,
          });
        }
-      } else {
-        setGenerationError(createGenerationError(data.error || '图片生成失败'));
+      }
+
+      if (failedResults.length > 0) {
+        const stillRunning = failedResults.some(result => result.reason instanceof GenerationJobStillRunningError);
+        if (generatedImages.length === 0) {
+          const firstError = failedResults[0]?.reason;
+          setGenerationError(createGenerationError(firstError instanceof Error ? firstError.message : '图片生成失败'));
+        } else {
+          toast.error(`${failedResults.length} 个生成任务失败`);
+        }
+        if (stillRunning) toast.info('部分生成任务仍在执行，可稍后在创作历史中查看');
      }
    } catch (err: unknown) {
      if (err instanceof GenerationJobStillRunningError) {
        setGenerationError(null);
-        removeActiveTask(taskId);
        toast.info('生成任务仍在执行，可稍后在创作历史中查看');
      } else if (err instanceof DOMException && err.name === 'AbortError') {
        setGenerationError(createGenerationError('请求超时，请尝试减少生成数量或降低分辨率'));
@@ -332,8 +382,10 @@ export function TextToImagePanel() {
        setGenerationError(createGenerationError(err instanceof Error ? err.message : '网络错误，请重试'));
      }
    }
-    finally { removeActiveTask(taskId); }
-  }, [prompt, negativePrompt, selectedModel, outputFormat, imageQuality, selectedStylePreset, guidanceScale, user, imageKeys, systemImageApis, getCurrentModelLabel, addRecord, credits, membershipEnabled, resolveGenerationParams, removeActiveTask, updateActiveTask]);
+    finally {
+      if (submissionSignature) activeSubmissionSignaturesRef.current.delete(submissionSignature);
+    }
+  }, [prompt, negativePrompt, selectedModel, outputFormat, imageQuality, selectedStylePreset, guidanceScale, user, imageKeys, systemImageApis, getCurrentModelLabel, addRecord, membershipEnabled, resolveGenerationParams, removeActiveTask, updateActiveTask]);

  // Download
  const handleDownload = useCallback(async (url: string, index: number) => {