448 lines
14 KiB
TypeScript
448 lines
14 KiB
TypeScript
import "server-only";
|
|
|
|
import sharp from "sharp";
|
|
import { constrainImageDimensions } from "@/lib/ai/image-size-constraints";
|
|
import { uploadFile } from "@/lib/storage/service";
|
|
|
|
interface StepfunImageConfig {
|
|
apiKey: string;
|
|
baseURL: string;
|
|
}
|
|
|
|
export interface StepfunImageParams {
|
|
aspectRatio?: string;
|
|
imageSize?: string;
|
|
imageCount?: number;
|
|
size?: string;
|
|
seed?: number;
|
|
negativePrompt?: string;
|
|
}
|
|
|
|
interface UploadedImage {
|
|
thumbnail: string;
|
|
original: string;
|
|
}
|
|
|
|
const DEFAULT_BASE_URL = "https://api.stepfun.com";
|
|
const DEFAULT_TIMEOUT_MS = 300000;
|
|
const DEFAULT_CFG_SCALE = 1.0;
|
|
const DEFAULT_STEPS = 8;
|
|
const DEFAULT_TEXT_MODE = true;
|
|
const MAX_PROMPT_LENGTH = 512;
|
|
|
|
function normalizeBaseURL(baseURL: string) {
|
|
const normalized = String(baseURL || "").trim().replace(/\/+$/, "");
|
|
const base = normalized || DEFAULT_BASE_URL;
|
|
if (/\/v1$/i.test(base)) return base;
|
|
return `${base}/v1`;
|
|
}
|
|
|
|
function joinEndpoint(baseURL: string, path: string) {
|
|
const normalizedPath = path.startsWith("/") ? path : `/${path}`;
|
|
return `${normalizeBaseURL(baseURL)}${normalizedPath}`;
|
|
}
|
|
|
|
function normalizePrompt(prompt: string) {
|
|
const normalized = String(prompt || "").trim();
|
|
if (normalized.length <= MAX_PROMPT_LENGTH) return normalized;
|
|
console.warn("[StepFun Image] Prompt exceeds 512 characters, truncating for step-image-edit-2");
|
|
return normalized.slice(0, MAX_PROMPT_LENGTH);
|
|
}
|
|
|
|
function parseAspectRatioValue(ratio?: string | null) {
|
|
const normalized = String(ratio || "").trim();
|
|
const [width, height] = normalized.split(":").map((item) => Number(item));
|
|
if (!Number.isFinite(width) || !Number.isFinite(height) || width <= 0 || height <= 0) {
|
|
return 1;
|
|
}
|
|
return width / height;
|
|
}
|
|
|
|
function resolveStepfunSize(aspectRatio?: string, explicitSize?: string) {
|
|
const normalizedExplicit = String(explicitSize || "").trim();
|
|
if (/^\d{3,4}x\d{3,4}$/i.test(normalizedExplicit)) {
|
|
const [rawWidth, rawHeight] = normalizedExplicit.split("x").map((item) => Number(item));
|
|
const { width, height } = constrainImageDimensions(
|
|
{ width: rawWidth, height: rawHeight },
|
|
{
|
|
maxEdge: 1360,
|
|
maxPixels: 1360 * 1360,
|
|
minEdge: 256,
|
|
multiple: 8,
|
|
}
|
|
);
|
|
return `${width}x${height}`;
|
|
}
|
|
|
|
const ratio = parseAspectRatioValue(aspectRatio);
|
|
const candidates = [
|
|
{ size: "1024x1024", ratio: 1 },
|
|
{ size: "768x1360", ratio: 1360 / 768 },
|
|
{ size: "896x1184", ratio: 1184 / 896 },
|
|
{ size: "1360x768", ratio: 768 / 1360 },
|
|
{ size: "1184x896", ratio: 896 / 1184 },
|
|
];
|
|
|
|
let best = candidates[0];
|
|
let minDiff = Number.POSITIVE_INFINITY;
|
|
for (const candidate of candidates) {
|
|
const diff = Math.abs(candidate.ratio - ratio);
|
|
if (diff < minDiff) {
|
|
best = candidate;
|
|
minDiff = diff;
|
|
}
|
|
}
|
|
return best.size;
|
|
}
|
|
|
|
function detectImageExtension(imageUrl: string, contentType?: string | null) {
|
|
const type = String(contentType || "").toLowerCase();
|
|
if (type.includes("jpeg") || type.includes("jpg")) return "jpg";
|
|
if (type.includes("webp")) return "webp";
|
|
if (type.includes("png")) return "png";
|
|
|
|
const cleanUrl = imageUrl.split("?")[0]?.toLowerCase() || "";
|
|
if (cleanUrl.endsWith(".jpg") || cleanUrl.endsWith(".jpeg")) return "jpg";
|
|
if (cleanUrl.endsWith(".webp")) return "webp";
|
|
if (cleanUrl.endsWith(".png")) return "png";
|
|
return "png";
|
|
}
|
|
|
|
async function createThumbnail(buffer: Buffer) {
|
|
try {
|
|
return await sharp(buffer)
|
|
.resize(400, 400, { fit: "inside", withoutEnlargement: true })
|
|
.jpeg({ quality: 80 })
|
|
.toBuffer();
|
|
} catch (error) {
|
|
console.error("[StepFun Image] Failed to create thumbnail:", error);
|
|
return buffer;
|
|
}
|
|
}
|
|
|
|
function generateFilename(prefix: string, extension: string) {
|
|
const timestamp = Date.now();
|
|
const random = Math.random().toString(36).slice(2, 8);
|
|
return `${prefix}-${timestamp}-${random}.${extension}`;
|
|
}
|
|
|
|
async function uploadGeneratedBuffer(
|
|
buffer: Buffer,
|
|
extension: string,
|
|
mimeType: string
|
|
): Promise<UploadedImage> {
|
|
const thumbnailBuffer = await createThumbnail(buffer);
|
|
|
|
const originalFilename = generateFilename("stepfun-image-orig", extension);
|
|
const originalArrayBuffer = buffer.buffer.slice(
|
|
buffer.byteOffset,
|
|
buffer.byteOffset + buffer.byteLength
|
|
) as ArrayBuffer;
|
|
const originalUploaded = await uploadFile(originalFilename, originalArrayBuffer);
|
|
|
|
const thumbFilename = generateFilename("stepfun-image-thumb", "jpg");
|
|
const thumbArrayBuffer = thumbnailBuffer.buffer.slice(
|
|
thumbnailBuffer.byteOffset,
|
|
thumbnailBuffer.byteOffset + thumbnailBuffer.byteLength
|
|
) as ArrayBuffer;
|
|
const thumbUploaded = await uploadFile(thumbFilename, thumbArrayBuffer);
|
|
|
|
console.log("[StepFun Image] Uploaded generated image:", {
|
|
original: originalUploaded.url,
|
|
thumbnail: thumbUploaded.url,
|
|
mimeType,
|
|
});
|
|
|
|
return {
|
|
thumbnail: thumbUploaded.url,
|
|
original: originalUploaded.url,
|
|
};
|
|
}
|
|
|
|
async function downloadAndUploadImage(imageUrl: string): Promise<UploadedImage> {
|
|
const response = await fetch(imageUrl, { cache: "no-store" });
|
|
if (!response.ok) {
|
|
throw new Error(`Failed to download StepFun generated image: ${response.status}`);
|
|
}
|
|
const buffer = Buffer.from(await response.arrayBuffer());
|
|
const extension = detectImageExtension(imageUrl, response.headers.get("content-type"));
|
|
const mimeType = response.headers.get("content-type") || `image/${extension === "jpg" ? "jpeg" : extension}`;
|
|
return uploadGeneratedBuffer(buffer, extension, mimeType);
|
|
}
|
|
|
|
async function uploadBase64ToS3(base64Data: string, mimeType: string) {
|
|
const buffer = Buffer.from(base64Data, "base64");
|
|
const extension = detectImageExtension("", mimeType);
|
|
return uploadGeneratedBuffer(
|
|
buffer,
|
|
extension,
|
|
mimeType || `image/${extension === "jpg" ? "jpeg" : extension}`
|
|
);
|
|
}
|
|
|
|
function extractPromptAndImages(messages: Array<{ role: string; content: any }>) {
|
|
const lastUserMessage = messages.filter((message) => message.role === "user").pop();
|
|
if (!lastUserMessage) {
|
|
throw new Error("No user message found");
|
|
}
|
|
|
|
let prompt = "";
|
|
const imageUrls: string[] = [];
|
|
|
|
const pushImageUrl = (value?: string) => {
|
|
if (!value) return;
|
|
const normalized = String(value).trim();
|
|
if (!normalized || imageUrls.includes(normalized)) return;
|
|
imageUrls.push(normalized);
|
|
};
|
|
|
|
if (typeof lastUserMessage.content === "string") {
|
|
prompt = lastUserMessage.content.trim();
|
|
} else if (Array.isArray(lastUserMessage.content)) {
|
|
for (const part of lastUserMessage.content) {
|
|
if (part?.type === "text" && typeof part.text === "string") {
|
|
prompt = `${prompt} ${part.text}`.trim();
|
|
} else if (part?.type === "image_url" && part.image_url?.url) {
|
|
pushImageUrl(part.image_url.url);
|
|
} else if (part?.type === "file") {
|
|
pushImageUrl(part.url || part.data);
|
|
} else if (part?.type === "image") {
|
|
pushImageUrl(part.image || part.url);
|
|
}
|
|
}
|
|
}
|
|
|
|
if (!prompt) {
|
|
throw new Error("No prompt found in user message");
|
|
}
|
|
|
|
return { prompt: normalizePrompt(prompt), imageUrls };
|
|
}
|
|
|
|
async function loadInputImage(source: string) {
|
|
const normalized = String(source || "").trim();
|
|
if (!normalized) {
|
|
throw new Error("输入图片地址为空");
|
|
}
|
|
|
|
if (normalized.startsWith("data:")) {
|
|
const match = normalized.match(/^data:([A-Za-z0-9.+/-]+);base64,(.+)$/);
|
|
if (!match) {
|
|
throw new Error("不支持的 Base64 图片格式");
|
|
}
|
|
const mimeType = match[1] || "image/png";
|
|
const buffer = Buffer.from(match[2], "base64");
|
|
const extension = detectImageExtension("", mimeType);
|
|
return {
|
|
buffer,
|
|
mimeType,
|
|
filename: generateFilename("stepfun-image-edit-input", extension),
|
|
};
|
|
}
|
|
|
|
const response = await fetch(normalized, { cache: "no-store" });
|
|
if (!response.ok) {
|
|
throw new Error(`Failed to fetch StepFun input image: ${response.status}`);
|
|
}
|
|
const buffer = Buffer.from(await response.arrayBuffer());
|
|
const mimeType = response.headers.get("content-type") || "image/png";
|
|
const extension = detectImageExtension(normalized, mimeType);
|
|
return {
|
|
buffer,
|
|
mimeType,
|
|
filename: generateFilename("stepfun-image-edit-input", extension),
|
|
};
|
|
}
|
|
|
|
function normalizeSeed(seed?: number) {
|
|
if (typeof seed !== "number" || !Number.isFinite(seed)) return undefined;
|
|
return Math.max(0, Math.floor(seed));
|
|
}
|
|
|
|
function extractErrorMessage(payload: any) {
|
|
if (!payload) return "StepFun Images 请求失败";
|
|
if (typeof payload === "string") return payload;
|
|
if (payload.error?.message) return String(payload.error.message);
|
|
if (payload.message) return String(payload.message);
|
|
try {
|
|
return JSON.stringify(payload);
|
|
} catch {
|
|
return "StepFun Images 请求失败";
|
|
}
|
|
}
|
|
|
|
async function parseJsonResponse(response: Response) {
|
|
const rawText = await response.text();
|
|
try {
|
|
return JSON.parse(rawText);
|
|
} catch {
|
|
return rawText;
|
|
}
|
|
}
|
|
|
|
async function requestStepfunImage(params: {
|
|
baseURL: string;
|
|
apiKey: string;
|
|
model: string;
|
|
prompt: string;
|
|
size?: string;
|
|
imageUrls: string[];
|
|
seed?: number;
|
|
}) {
|
|
const { baseURL, apiKey, model, prompt, size, imageUrls, seed } = params;
|
|
const controller = new AbortController();
|
|
const timeoutId = setTimeout(() => controller.abort(), DEFAULT_TIMEOUT_MS);
|
|
|
|
try {
|
|
if (imageUrls.length > 0) {
|
|
const endpoint = joinEndpoint(baseURL, "/images/edits");
|
|
const formData = new FormData();
|
|
formData.append("model", model);
|
|
formData.append("prompt", prompt);
|
|
formData.append("response_format", "b64_json");
|
|
formData.append("cfg_scale", String(DEFAULT_CFG_SCALE));
|
|
formData.append("steps", String(DEFAULT_STEPS));
|
|
formData.append("text_mode", String(DEFAULT_TEXT_MODE));
|
|
const normalizedSeed = normalizeSeed(seed);
|
|
if (normalizedSeed !== undefined) formData.append("seed", String(normalizedSeed));
|
|
|
|
for (const source of imageUrls.slice(0, 1)) {
|
|
const image = await loadInputImage(source);
|
|
formData.append("image", new Blob([image.buffer], { type: image.mimeType }), image.filename);
|
|
}
|
|
|
|
const response = await fetch(endpoint, {
|
|
method: "POST",
|
|
headers: {
|
|
Authorization: `Bearer ${apiKey}`,
|
|
},
|
|
body: formData,
|
|
signal: controller.signal,
|
|
});
|
|
const payload = await parseJsonResponse(response);
|
|
if (!response.ok) {
|
|
throw new Error(`${extractErrorMessage(payload)} (HTTP ${response.status})`);
|
|
}
|
|
return payload;
|
|
}
|
|
|
|
const endpoint = joinEndpoint(baseURL, "/images/generations");
|
|
const body: Record<string, unknown> = {
|
|
model,
|
|
prompt,
|
|
response_format: "b64_json",
|
|
cfg_scale: DEFAULT_CFG_SCALE,
|
|
steps: DEFAULT_STEPS,
|
|
text_mode: DEFAULT_TEXT_MODE,
|
|
...(size ? { size } : {}),
|
|
};
|
|
const normalizedSeed = normalizeSeed(seed);
|
|
if (normalizedSeed !== undefined) body.seed = normalizedSeed;
|
|
|
|
const response = await fetch(endpoint, {
|
|
method: "POST",
|
|
headers: {
|
|
Authorization: `Bearer ${apiKey}`,
|
|
"Content-Type": "application/json",
|
|
},
|
|
body: JSON.stringify(body),
|
|
signal: controller.signal,
|
|
});
|
|
const payload = await parseJsonResponse(response);
|
|
if (!response.ok) {
|
|
throw new Error(`${extractErrorMessage(payload)} (HTTP ${response.status})`);
|
|
}
|
|
return payload;
|
|
} finally {
|
|
clearTimeout(timeoutId);
|
|
}
|
|
}
|
|
|
|
async function extractUploadedImages(payload: any): Promise<UploadedImage[]> {
|
|
const data = Array.isArray(payload?.data) ? payload.data : [];
|
|
const uploadedImages: UploadedImage[] = [];
|
|
|
|
for (const item of data) {
|
|
if (item?.url && typeof item.url === "string") {
|
|
uploadedImages.push(await downloadAndUploadImage(item.url));
|
|
continue;
|
|
}
|
|
|
|
if (item?.b64_json && typeof item.b64_json === "string") {
|
|
uploadedImages.push(await uploadBase64ToS3(item.b64_json, "image/png"));
|
|
}
|
|
}
|
|
|
|
return uploadedImages;
|
|
}
|
|
|
|
function formatMarkdown(images: UploadedImage[], imageSize: string) {
|
|
return images
|
|
.map((image) => ``)
|
|
.join("\n\n");
|
|
}
|
|
|
|
export function createStepfunImage({ apiKey, baseURL }: StepfunImageConfig) {
|
|
return {
|
|
provider: "stepfun-image",
|
|
chat: (modelId: string) => ({
|
|
provider: "stepfun-image",
|
|
doGenerate: async (options: {
|
|
messages: Array<{ role: string; content: any }>;
|
|
imageParams?: StepfunImageParams;
|
|
overrideModelId?: string;
|
|
}) => {
|
|
const { prompt, imageUrls } = extractPromptAndImages(options.messages);
|
|
const activeModelId = options.overrideModelId || modelId;
|
|
const size =
|
|
imageUrls.length > 0
|
|
? undefined
|
|
: resolveStepfunSize(options.imageParams?.aspectRatio, options.imageParams?.size);
|
|
|
|
console.log("[StepFun Image] Submitting request", {
|
|
model: activeModelId,
|
|
size,
|
|
imageCount: imageUrls.length,
|
|
baseURL: normalizeBaseURL(baseURL),
|
|
});
|
|
|
|
const payload = await requestStepfunImage({
|
|
baseURL,
|
|
apiKey,
|
|
model: activeModelId,
|
|
prompt,
|
|
size,
|
|
imageUrls,
|
|
seed: options.imageParams?.seed,
|
|
});
|
|
|
|
const uploadedImages = await extractUploadedImages(payload);
|
|
if (uploadedImages.length === 0) {
|
|
throw new Error("StepFun Images 未返回可用图片");
|
|
}
|
|
|
|
const text = formatMarkdown(uploadedImages, options.imageParams?.imageSize || "1K");
|
|
return {
|
|
text,
|
|
finishReason: "stop" as const,
|
|
usage: {
|
|
promptTokens: 0,
|
|
completionTokens: text.length,
|
|
},
|
|
};
|
|
},
|
|
doStream: async function* (options: {
|
|
messages: Array<{ role: string; content: any }>;
|
|
imageParams?: StepfunImageParams;
|
|
overrideModelId?: string;
|
|
}) {
|
|
const result = await this.doGenerate(options);
|
|
yield {
|
|
type: "text-delta",
|
|
textDelta: result.text,
|
|
};
|
|
},
|
|
}),
|
|
};
|
|
}
|