460 lines
14 KiB
TypeScript
460 lines
14 KiB
TypeScript
import "server-only";
|
|
|
|
import sharp from "sharp";
|
|
import { uploadFile } from "@/lib/storage/service";
|
|
|
|
interface MidjourneyConfig {
|
|
apiKey: string;
|
|
baseURL: string;
|
|
}
|
|
|
|
export interface MidjourneyImageParams {
|
|
aspectRatio?: string;
|
|
imageSize?: string;
|
|
imageCount?: number;
|
|
}
|
|
|
|
interface UploadedImage {
|
|
thumbnail: string;
|
|
original: string;
|
|
}
|
|
|
|
interface MidjourneyTaskButton {
|
|
customId?: string;
|
|
label?: string;
|
|
emoji?: string;
|
|
style?: number;
|
|
type?: number;
|
|
}
|
|
|
|
interface MidjourneyTaskResult {
|
|
id?: string;
|
|
imageUrl?: string;
|
|
status?: string;
|
|
progress?: string;
|
|
failReason?: string;
|
|
description?: string;
|
|
prompt?: string;
|
|
buttons?: MidjourneyTaskButton[];
|
|
}
|
|
|
|
const DEFAULT_TIMEOUT_MS = 300000;
|
|
const POLL_INTERVAL_MS = 3500;
|
|
const DEFAULT_IMAGE_SIZE = "2K";
|
|
|
|
function normalizeBaseURL(baseURL: string) {
|
|
return String(baseURL || "").trim().replace(/\/+$/, "");
|
|
}
|
|
|
|
function resolveMode(baseURL: string, modelId: string) {
|
|
const normalizedBaseURL = normalizeBaseURL(baseURL).toLowerCase();
|
|
const normalizedModelId = String(modelId || "").trim().toLowerCase();
|
|
|
|
if (
|
|
normalizedBaseURL.endsWith("/relax") ||
|
|
normalizedBaseURL.endsWith("/mj-relax") ||
|
|
normalizedModelId.includes("relax")
|
|
) {
|
|
return "relax" as const;
|
|
}
|
|
|
|
return "fast" as const;
|
|
}
|
|
|
|
function buildEndpointCandidates(baseURL: string, mode: "fast" | "relax", suffix: string) {
|
|
const normalizedBaseURL = normalizeBaseURL(baseURL);
|
|
const normalizedSuffix = suffix.startsWith("/") ? suffix : `/${suffix}`;
|
|
const directPattern = /\/(fast|relax|mj-fast|mj-relax)$/i;
|
|
|
|
if (directPattern.test(normalizedBaseURL)) {
|
|
return [`${normalizedBaseURL}${normalizedSuffix}`];
|
|
}
|
|
|
|
const modePrefix = mode === "relax" ? "relax" : "fast";
|
|
const altPrefix = mode === "relax" ? "mj-relax" : "mj-fast";
|
|
|
|
return [
|
|
`${normalizedBaseURL}/${modePrefix}${normalizedSuffix}`,
|
|
`${normalizedBaseURL}/${altPrefix}${normalizedSuffix}`,
|
|
];
|
|
}
|
|
|
|
function normalizeImageSizeLabel(imageSize?: string): "1K" | "2K" | "4K" {
|
|
const normalized = String(imageSize || "").trim().toUpperCase();
|
|
if (normalized === "1K" || normalized === "2K" || normalized === "4K") return normalized;
|
|
return DEFAULT_IMAGE_SIZE;
|
|
}
|
|
|
|
function normalizeAspectRatio(value?: string) {
|
|
const normalized = String(value || "").trim();
|
|
if (!normalized || normalized.toLowerCase() === "auto") return undefined;
|
|
if (/^\d+:\d+$/.test(normalized)) return normalized;
|
|
return undefined;
|
|
}
|
|
|
|
function appendAspectRatio(prompt: string, aspectRatio?: string) {
|
|
const normalizedPrompt = String(prompt || "").trim();
|
|
const normalizedAspectRatio = normalizeAspectRatio(aspectRatio);
|
|
if (!normalizedAspectRatio) return normalizedPrompt;
|
|
if (/--ar\s+\d+:\d+/i.test(normalizedPrompt)) return normalizedPrompt;
|
|
return `${normalizedPrompt} --ar ${normalizedAspectRatio}`.trim();
|
|
}
|
|
|
|
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";
|
|
if (type.includes("gif")) return "gif";
|
|
|
|
const cleanUrl = imageUrl.split("?")[0]?.toLowerCase() || "";
|
|
if (cleanUrl.endsWith(".jpg") || cleanUrl.endsWith(".jpeg")) return "jpg";
|
|
if (cleanUrl.endsWith(".webp")) return "webp";
|
|
if (cleanUrl.endsWith(".gif")) return "gif";
|
|
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("[Midjourney] 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("midjourney-orig", extension);
|
|
const originalArrayBuffer = buffer.buffer.slice(
|
|
buffer.byteOffset,
|
|
buffer.byteOffset + buffer.byteLength
|
|
) as ArrayBuffer;
|
|
const originalUploaded = await uploadFile(originalFilename, originalArrayBuffer);
|
|
|
|
const thumbFilename = generateFilename("midjourney-thumb", "jpg");
|
|
const thumbArrayBuffer = thumbnailBuffer.buffer.slice(
|
|
thumbnailBuffer.byteOffset,
|
|
thumbnailBuffer.byteOffset + thumbnailBuffer.byteLength
|
|
) as ArrayBuffer;
|
|
const thumbUploaded = await uploadFile(thumbFilename, thumbArrayBuffer);
|
|
|
|
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 Midjourney 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);
|
|
}
|
|
|
|
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) return;
|
|
if (!imageUrls.includes(normalized)) {
|
|
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, imageUrls };
|
|
}
|
|
|
|
async function loadInputImageAsDataUrl(source: string) {
|
|
const normalized = String(source || "").trim();
|
|
if (!normalized) {
|
|
throw new Error("Midjourney 输入图片地址为空");
|
|
}
|
|
|
|
if (normalized.startsWith("data:")) {
|
|
return normalized;
|
|
}
|
|
|
|
const response = await fetch(normalized, { cache: "no-store" });
|
|
if (!response.ok) {
|
|
throw new Error(`Failed to fetch Midjourney input image: ${response.status}`);
|
|
}
|
|
|
|
const mimeType = response.headers.get("content-type") || "image/png";
|
|
const buffer = Buffer.from(await response.arrayBuffer());
|
|
return `data:${mimeType};base64,${buffer.toString("base64")}`;
|
|
}
|
|
|
|
async function parseJsonResponse(response: Response) {
|
|
const rawText = await response.text();
|
|
try {
|
|
return JSON.parse(rawText);
|
|
} catch {
|
|
return rawText;
|
|
}
|
|
}
|
|
|
|
function extractErrorMessage(payload: any) {
|
|
if (!payload) return "Midjourney 请求失败";
|
|
if (typeof payload === "string") return payload;
|
|
if (payload.failReason) return String(payload.failReason);
|
|
if (payload.description) return String(payload.description);
|
|
if (payload.message) return String(payload.message);
|
|
if (payload.error?.message) return String(payload.error.message);
|
|
try {
|
|
return JSON.stringify(payload);
|
|
} catch {
|
|
return "Midjourney 请求失败";
|
|
}
|
|
}
|
|
|
|
async function requestWithEndpointFallback(params: {
|
|
baseURL: string;
|
|
mode: "fast" | "relax";
|
|
suffix: string;
|
|
method: "GET" | "POST";
|
|
apiKey: string;
|
|
body?: string;
|
|
contentType?: string;
|
|
}) {
|
|
const { baseURL, mode, suffix, method, apiKey, body, contentType } = params;
|
|
const endpoints = buildEndpointCandidates(baseURL, mode, suffix);
|
|
let lastPayload: any = null;
|
|
let lastStatus: number | null = null;
|
|
let lastEndpoint = endpoints[0];
|
|
|
|
for (const endpoint of endpoints) {
|
|
lastEndpoint = endpoint;
|
|
const response = await fetch(endpoint, {
|
|
method,
|
|
headers: {
|
|
Authorization: `Bearer ${apiKey}`,
|
|
...(contentType ? { "Content-Type": contentType } : {}),
|
|
},
|
|
...(body ? { body } : {}),
|
|
signal: AbortSignal.timeout(45000),
|
|
});
|
|
const payload = await parseJsonResponse(response);
|
|
lastPayload = payload;
|
|
lastStatus = response.status;
|
|
|
|
if (response.ok || response.status !== 404) {
|
|
return { endpoint, response, payload };
|
|
}
|
|
}
|
|
|
|
throw new Error(
|
|
`Midjourney request failed: ${extractErrorMessage(lastPayload)} (${lastStatus || "no_status"}) @ ${lastEndpoint}`
|
|
);
|
|
}
|
|
|
|
async function submitImagineTask(params: {
|
|
baseURL: string;
|
|
apiKey: string;
|
|
mode: "fast" | "relax";
|
|
prompt: string;
|
|
imageUrls: string[];
|
|
}) {
|
|
const { baseURL, apiKey, mode, prompt, imageUrls } = params;
|
|
const base64Array = await Promise.all(imageUrls.map((imageUrl) => loadInputImageAsDataUrl(imageUrl)));
|
|
|
|
const { payload } = await requestWithEndpointFallback({
|
|
baseURL,
|
|
mode,
|
|
suffix: "/mj/submit/imagine",
|
|
method: "POST",
|
|
apiKey,
|
|
body: JSON.stringify({
|
|
base64Array,
|
|
prompt,
|
|
}),
|
|
contentType: "application/json",
|
|
});
|
|
|
|
const code = Number((payload as any)?.code ?? 0);
|
|
const taskId = String((payload as any)?.result || "").trim();
|
|
if (![1, 22].includes(code) || !taskId) {
|
|
throw new Error(extractErrorMessage(payload));
|
|
}
|
|
|
|
return taskId;
|
|
}
|
|
|
|
async function fetchMidjourneyTask(params: {
|
|
baseURL: string;
|
|
apiKey: string;
|
|
mode: "fast" | "relax";
|
|
taskId: string;
|
|
}) {
|
|
const { payload } = await requestWithEndpointFallback({
|
|
baseURL: params.baseURL,
|
|
mode: params.mode,
|
|
suffix: `/mj/task/${params.taskId}/fetch`,
|
|
method: "GET",
|
|
apiKey: params.apiKey,
|
|
});
|
|
|
|
return (payload || {}) as MidjourneyTaskResult;
|
|
}
|
|
|
|
async function waitForMidjourneyTask(params: {
|
|
baseURL: string;
|
|
apiKey: string;
|
|
mode: "fast" | "relax";
|
|
taskId: string;
|
|
}) {
|
|
const deadlineAt = Date.now() + DEFAULT_TIMEOUT_MS;
|
|
let lastTask: MidjourneyTaskResult | null = null;
|
|
|
|
while (Date.now() < deadlineAt) {
|
|
const task = await fetchMidjourneyTask(params);
|
|
lastTask = task;
|
|
|
|
const status = String(task.status || "").toUpperCase();
|
|
if (status === "SUCCESS" && task.imageUrl) {
|
|
return task;
|
|
}
|
|
if (status === "FAILURE") {
|
|
throw new Error(task.failReason || task.description || "Midjourney 任务失败");
|
|
}
|
|
|
|
await new Promise((resolve) => setTimeout(resolve, POLL_INTERVAL_MS));
|
|
}
|
|
|
|
throw new Error(
|
|
`Midjourney task timeout: ${params.taskId} (${lastTask?.status || "unknown"})`
|
|
);
|
|
}
|
|
|
|
function formatMarkdown(image: UploadedImage, imageSize: string) {
|
|
return ``;
|
|
}
|
|
|
|
export function createMidjourney({ apiKey, baseURL }: MidjourneyConfig) {
|
|
return {
|
|
provider: "midjourney",
|
|
chat: (modelId: string) => ({
|
|
provider: "midjourney",
|
|
doGenerate: async (options: {
|
|
messages: Array<{ role: string; content: any }>;
|
|
imageParams?: MidjourneyImageParams;
|
|
overrideModelId?: string;
|
|
}) => {
|
|
const { prompt, imageUrls } = extractPromptAndImages(options.messages);
|
|
const activeModelId = options.overrideModelId || modelId;
|
|
const mode = resolveMode(baseURL, activeModelId);
|
|
const imageSize = normalizeImageSizeLabel(options.imageParams?.imageSize);
|
|
const finalPrompt = appendAspectRatio(prompt, options.imageParams?.aspectRatio);
|
|
|
|
console.log("[Midjourney] Submitting request", {
|
|
model: activeModelId,
|
|
mode,
|
|
imageCount: imageUrls.length,
|
|
promptPreview: finalPrompt.slice(0, 120),
|
|
});
|
|
|
|
const taskId = await submitImagineTask({
|
|
baseURL,
|
|
apiKey,
|
|
mode,
|
|
prompt: finalPrompt,
|
|
imageUrls,
|
|
});
|
|
|
|
console.log("[Midjourney] Task submitted", { taskId, mode });
|
|
const taskResult = await waitForMidjourneyTask({
|
|
baseURL,
|
|
apiKey,
|
|
mode,
|
|
taskId,
|
|
});
|
|
|
|
if (!taskResult.imageUrl) {
|
|
throw new Error(taskResult.failReason || taskResult.description || "Midjourney 未返回图片地址");
|
|
}
|
|
|
|
const uploadedImage = await downloadAndUploadImage(taskResult.imageUrl);
|
|
const text = formatMarkdown(uploadedImage, imageSize);
|
|
|
|
return {
|
|
text,
|
|
finishReason: "stop" as const,
|
|
usage: {
|
|
promptTokens: 0,
|
|
completionTokens: text.length,
|
|
},
|
|
providerMetadata: {
|
|
midjourney: {
|
|
taskId,
|
|
mode,
|
|
progress: taskResult.progress || "",
|
|
buttons: taskResult.buttons || [],
|
|
imageUrl: taskResult.imageUrl,
|
|
},
|
|
},
|
|
};
|
|
},
|
|
doStream: async function* (options: {
|
|
messages: Array<{ role: string; content: any }>;
|
|
imageParams?: MidjourneyImageParams;
|
|
overrideModelId?: string;
|
|
}) {
|
|
const result = await this.doGenerate(options);
|
|
yield {
|
|
type: "text-delta",
|
|
textDelta: result.text,
|
|
};
|
|
},
|
|
}),
|
|
};
|
|
}
|