418 lines
13 KiB
TypeScript
418 lines
13 KiB
TypeScript
import { NextResponse } from "next/server";
|
||
import { z } from "zod";
|
||
import { auth } from "@/app/(auth)/auth";
|
||
import { getModelGenerationCreditCost } from "@/lib/ai/image-model-pricing";
|
||
import { generateCanvasImage } from "@/lib/canvas/generate-image";
|
||
import { getSystemSetting } from "@/lib/db/queries";
|
||
import { prepareFeaturePointsCharge } from "@/lib/points/feature-charge";
|
||
import {
|
||
findStyleProImageModelByRef,
|
||
getStyleProImageModelEntries,
|
||
} from "@/lib/style-pro/model-options";
|
||
import {
|
||
STYLE_PRO_CONFIG_KEY,
|
||
type StyleProCategory,
|
||
type StyleProGenerationMode,
|
||
type StyleProOption,
|
||
normalizeStyleProConfig,
|
||
} from "@/lib/style-pro/config";
|
||
|
||
export const maxDuration = 120;
|
||
export const dynamic = "force-dynamic";
|
||
|
||
const generateRequestSchema = z.object({
|
||
canvasId: z.string().trim().min(1).max(160).optional(),
|
||
sourceImageUrl: z.string().trim().min(1).max(4000),
|
||
generationMode: z.enum(["style", "realistic"]).optional(),
|
||
modelId: z.string().trim().min(1).max(160).optional(),
|
||
selections: z.record(z.array(z.string().trim().min(1).max(120)).max(20)),
|
||
});
|
||
|
||
type SelectedOption = {
|
||
category: StyleProCategory;
|
||
option: StyleProOption;
|
||
};
|
||
|
||
function isDisallowedImageUrl(value: string) {
|
||
const trimmed = value.trim().toLowerCase();
|
||
return trimmed.startsWith("data:") || trimmed.startsWith("javascript:");
|
||
}
|
||
|
||
function buildSelectionPayload(selectedOptions: SelectedOption[]) {
|
||
return selectedOptions.map(({ category, option }) => ({
|
||
categoryId: category.id,
|
||
categoryTitle: category.title,
|
||
jsonKey: option.jsonKey || category.id,
|
||
optionId: option.id,
|
||
optionLabel: option.label,
|
||
prompt: option.prompt,
|
||
}));
|
||
}
|
||
|
||
function composePrompt({
|
||
defaultPrompt,
|
||
selectedOptions,
|
||
mode = "style-pro-final",
|
||
}: {
|
||
defaultPrompt: string;
|
||
selectedOptions: SelectedOption[];
|
||
mode?: string;
|
||
}) {
|
||
const selectionPayload = buildSelectionPayload(selectedOptions);
|
||
const groupedPrompt = selectionPayload.reduce<Record<string, string[]>>(
|
||
(acc, item) => {
|
||
const key = item.jsonKey || item.categoryId;
|
||
acc[key] = acc[key] || [];
|
||
acc[key].push(item.prompt);
|
||
return acc;
|
||
},
|
||
{}
|
||
);
|
||
|
||
return [
|
||
defaultPrompt,
|
||
"请严格依据以下风格化 PRO 选择执行最终图片生成。用户不可见配置 JSON：",
|
||
JSON.stringify(
|
||
{
|
||
mode,
|
||
selections: selectionPayload,
|
||
promptByKey: groupedPrompt,
|
||
},
|
||
null,
|
||
2
|
||
),
|
||
]
|
||
.filter(Boolean)
|
||
.join("\n\n");
|
||
}
|
||
|
||
function validateSelections(
|
||
categories: StyleProCategory[],
|
||
rawSelections: Record<string, string[]>
|
||
) {
|
||
const categoryById = new Map(
|
||
categories.map((category) => [category.id, category])
|
||
);
|
||
const selectedOptions: SelectedOption[] = [];
|
||
const normalizedSelections: Record<string, string[]> = {};
|
||
|
||
for (const [categoryId, optionIds] of Object.entries(rawSelections)) {
|
||
const category = categoryById.get(categoryId);
|
||
if (!category) {
|
||
return {
|
||
ok: false as const,
|
||
error: `无效分类：${categoryId}`,
|
||
};
|
||
}
|
||
|
||
const enabledOptions = new Map(
|
||
category.options
|
||
.filter((option) => option.enabled)
|
||
.map((option) => [option.id, option])
|
||
);
|
||
const uniqueOptionIds = Array.from(new Set(optionIds));
|
||
if (uniqueOptionIds.length > category.maxSelections) {
|
||
return {
|
||
ok: false as const,
|
||
error: `${category.title} 最多选择 ${category.maxSelections} 项`,
|
||
};
|
||
}
|
||
|
||
for (const optionId of uniqueOptionIds) {
|
||
const option = enabledOptions.get(optionId);
|
||
if (!option) {
|
||
return {
|
||
ok: false as const,
|
||
error: `无效选项：${category.title} / ${optionId}`,
|
||
};
|
||
}
|
||
selectedOptions.push({ category, option });
|
||
}
|
||
normalizedSelections[category.id] = uniqueOptionIds;
|
||
}
|
||
|
||
const missingCategory = categories.find((category) => {
|
||
if (!category.required && category.minSelections <= 0) return false;
|
||
const selectedCount = normalizedSelections[category.id]?.length || 0;
|
||
return (
|
||
selectedCount <
|
||
Math.max(category.required ? 1 : 0, category.minSelections)
|
||
);
|
||
});
|
||
if (missingCategory) {
|
||
return {
|
||
ok: false as const,
|
||
error: `请完成必选项：${missingCategory.title}`,
|
||
};
|
||
}
|
||
|
||
return {
|
||
ok: true as const,
|
||
selectedOptions,
|
||
normalizedSelections,
|
||
};
|
||
}
|
||
|
||
export async function POST(request: Request) {
|
||
try {
|
||
const session = await auth();
|
||
if (!session?.user?.id) {
|
||
return NextResponse.json(
|
||
{ success: false, error: "未登录" },
|
||
{ status: 401 }
|
||
);
|
||
}
|
||
|
||
const body = generateRequestSchema.parse(await request.json());
|
||
const generationMode: StyleProGenerationMode =
|
||
body.generationMode === "realistic" ? "realistic" : "style";
|
||
if (isDisallowedImageUrl(body.sourceImageUrl)) {
|
||
return NextResponse.json(
|
||
{ success: false, error: "请使用已上传的图片链接生成" },
|
||
{ status: 400 }
|
||
);
|
||
}
|
||
|
||
const setting = await getSystemSetting({ key: STYLE_PRO_CONFIG_KEY }).catch(
|
||
() => null
|
||
);
|
||
const config = normalizeStyleProConfig(setting?.value);
|
||
if (!config.enabled) {
|
||
return NextResponse.json(
|
||
{ success: false, error: "风格化 PRO 暂未启用" },
|
||
{ status: 403 }
|
||
);
|
||
}
|
||
|
||
const enabledCategories = config.categories
|
||
.filter((category) => category.enabled)
|
||
.map((category) => ({
|
||
...category,
|
||
options: category.options.filter((option) => option.enabled),
|
||
}))
|
||
.filter((category) => category.options.length > 0);
|
||
const selectionResult = validateSelections(
|
||
enabledCategories,
|
||
body.selections
|
||
);
|
||
if (!selectionResult.ok) {
|
||
return NextResponse.json(
|
||
{ success: false, error: selectionResult.error },
|
||
{ status: 400 }
|
||
);
|
||
}
|
||
|
||
const imageModelEntries = await getStyleProImageModelEntries({
|
||
includeHidden: generationMode === "realistic",
|
||
});
|
||
|
||
const selectedModels =
|
||
generationMode === "realistic"
|
||
? (() => {
|
||
const firstModel = findStyleProImageModelByRef(
|
||
imageModelEntries,
|
||
config.realisticFirstModelId
|
||
);
|
||
const secondModel = findStyleProImageModelByRef(
|
||
imageModelEntries,
|
||
config.realisticSecondModelId
|
||
);
|
||
return firstModel && secondModel
|
||
? [firstModel, secondModel]
|
||
: [];
|
||
})()
|
||
: (() => {
|
||
const configuredDefaultModel = findStyleProImageModelByRef(
|
||
imageModelEntries,
|
||
config.modelId
|
||
);
|
||
let selectedModel =
|
||
configuredDefaultModel ||
|
||
imageModelEntries.find((entry) => entry.model.isDefault) ||
|
||
imageModelEntries[0] ||
|
||
null;
|
||
|
||
if (body.modelId) {
|
||
selectedModel = findStyleProImageModelByRef(
|
||
imageModelEntries,
|
||
body.modelId
|
||
);
|
||
if (!selectedModel) {
|
||
return [];
|
||
}
|
||
}
|
||
return selectedModel ? [selectedModel] : [];
|
||
})();
|
||
|
||
if (generationMode === "realistic" && selectedModels.length < 2) {
|
||
return NextResponse.json(
|
||
{ success: false, error: "请先在后台配置实景化两段隐藏模型" },
|
||
{ status: 400 }
|
||
);
|
||
}
|
||
if (generationMode === "style" && body.modelId && selectedModels.length === 0) {
|
||
return NextResponse.json(
|
||
{ success: false, error: "所选生图模型不可用" },
|
||
{ status: 400 }
|
||
);
|
||
}
|
||
if (selectedModels.length === 0) {
|
||
return NextResponse.json(
|
||
{ success: false, error: "暂无可用的风格化 PRO 生图模型" },
|
||
{ status: 400 }
|
||
);
|
||
}
|
||
const unitCost = selectedModels.reduce(
|
||
(sum, entry) => sum + getModelGenerationCreditCost(entry.model, "image"),
|
||
0
|
||
);
|
||
const billingModelId =
|
||
generationMode === "realistic"
|
||
? selectedModels.map((entry) => entry.model.id).join("+")
|
||
: selectedModels[0].model.id;
|
||
const billingModelName =
|
||
generationMode === "realistic"
|
||
? selectedModels.map((entry) => entry.model.displayName).join(" + ")
|
||
: selectedModels[0].model.displayName;
|
||
const pointsCharge = await prepareFeaturePointsCharge({
|
||
userId: session.user.id,
|
||
unitCost,
|
||
reason:
|
||
generationMode === "realistic"
|
||
? "style_pro_realistic_generation"
|
||
: "style_pro_generation",
|
||
modelId: billingModelId,
|
||
note:
|
||
generationMode === "realistic"
|
||
? `风格化 PRO 实景化生图消耗积分（${billingModelName}）`
|
||
: `风格化 PRO 生图消耗积分（${billingModelName}）`,
|
||
});
|
||
if (!pointsCharge.ok) {
|
||
return NextResponse.json(pointsCharge.payload, {
|
||
status: pointsCharge.status,
|
||
});
|
||
}
|
||
|
||
const prompt = composePrompt({
|
||
defaultPrompt: config.defaultPrompt,
|
||
selectedOptions: selectionResult.selectedOptions,
|
||
});
|
||
const result =
|
||
generationMode === "realistic"
|
||
? await (async () => {
|
||
const firstModel = selectedModels[0];
|
||
const secondModel = selectedModels[1];
|
||
const firstResult = await generateCanvasImage({
|
||
prompt: config.realisticFirstPrompt,
|
||
negativePrompt: config.negativePrompt || undefined,
|
||
model: firstModel.model.id,
|
||
aspectRatio: config.aspectRatio || "auto",
|
||
imageSize: config.imageSize || "2K",
|
||
imageCount: 1,
|
||
referenceImage: body.sourceImageUrl,
|
||
ownerKey: session.user.id,
|
||
debugLabel: "style-pro-realistic-stage1",
|
||
debugMeta: {
|
||
canvasId: body.canvasId || null,
|
||
selectedOptionCount: selectionResult.selectedOptions.length,
|
||
},
|
||
});
|
||
const intermediateImage = firstResult.images[0]?.originalUrl;
|
||
if (!intermediateImage) {
|
||
throw new Error("实景化第一阶段生成失败");
|
||
}
|
||
return generateCanvasImage({
|
||
prompt: composePrompt({
|
||
defaultPrompt: config.realisticSecondPrompt,
|
||
selectedOptions: selectionResult.selectedOptions,
|
||
mode: "style-pro-realistic-final",
|
||
}),
|
||
negativePrompt: config.negativePrompt || undefined,
|
||
model: secondModel.model.id,
|
||
aspectRatio: config.aspectRatio || "auto",
|
||
imageSize: config.imageSize || "2K",
|
||
imageCount: 1,
|
||
referenceImage: intermediateImage,
|
||
ownerKey: session.user.id,
|
||
debugLabel: "style-pro-realistic-stage2",
|
||
debugMeta: {
|
||
canvasId: body.canvasId || null,
|
||
selectedOptionCount: selectionResult.selectedOptions.length,
|
||
firstStageModelRef: firstModel.model.id,
|
||
},
|
||
});
|
||
})()
|
||
: await generateCanvasImage({
|
||
prompt,
|
||
negativePrompt: config.negativePrompt || undefined,
|
||
model: selectedModels[0].model.id,
|
||
aspectRatio: config.aspectRatio || "auto",
|
||
imageSize: config.imageSize || "2K",
|
||
imageCount: 1,
|
||
referenceImage: body.sourceImageUrl,
|
||
ownerKey: session.user.id,
|
||
debugLabel: "style-pro",
|
||
debugMeta: {
|
||
canvasId: body.canvasId || null,
|
||
requestedStyleProModelId: body.modelId || null,
|
||
selectedOptionCount: selectionResult.selectedOptions.length,
|
||
},
|
||
});
|
||
const image = result.images[0];
|
||
if (!image?.originalUrl) {
|
||
return NextResponse.json(
|
||
{ success: false, error: "生成结果为空" },
|
||
{ status: 502 }
|
||
);
|
||
}
|
||
let chargedPoints = 0;
|
||
if (pointsCharge.chargeable) {
|
||
try {
|
||
const consumeResult = await pointsCharge.consume();
|
||
if (consumeResult?.ok && !consumeResult.skipped) {
|
||
chargedPoints = pointsCharge.unitCost;
|
||
} else if (!consumeResult?.ok) {
|
||
console.warn("[Style Pro Generate] Failed to consume points:", {
|
||
userId: session.user.id,
|
||
modelId: billingModelId,
|
||
unitCost: pointsCharge.unitCost,
|
||
});
|
||
}
|
||
} catch (pointsError) {
|
||
console.error("[Style Pro Generate] Points consume failed:", pointsError);
|
||
}
|
||
}
|
||
|
||
return NextResponse.json({
|
||
success: true,
|
||
data: {
|
||
imageUrl: image.originalUrl,
|
||
thumbnailUrl: image.thumbnailUrl || image.originalUrl,
|
||
model: result.model,
|
||
modelRef: result.modelRef,
|
||
generationMode,
|
||
chargedPoints,
|
||
selections: selectionResult.normalizedSelections,
|
||
optionLabels: selectionResult.selectedOptions.map(
|
||
({ category, option }) => `${category.title}：${option.label}`
|
||
),
|
||
},
|
||
});
|
||
} catch (error) {
|
||
console.error("[Style Pro Generate] Failed:", error);
|
||
if (error instanceof z.ZodError) {
|
||
return NextResponse.json(
|
||
{ success: false, error: "请求参数不完整" },
|
||
{ status: 400 }
|
||
);
|
||
}
|
||
return NextResponse.json(
|
||
{
|
||
success: false,
|
||
error: error instanceof Error ? error.message : "风格化 PRO 生成失败",
|
||
},
|
||
{ status: 500 }
|
||
);
|
||
}
|
||
}
|