463 lines
18 KiB
TypeScript
463 lines
18 KiB
TypeScript
import { createHash, randomUUID } from 'node:crypto';
|
||
import { mkdir, readFile, rename, rm, writeFile } from 'node:fs/promises';
|
||
import { join } from 'node:path';
|
||
import type {
|
||
ImageWorkspaceCapabilities,
|
||
ImageWorkspaceGenerationSettings,
|
||
ImageWorkspaceImage,
|
||
ImageWorkspaceProject,
|
||
ImageWorkspaceReferenceUploadInput,
|
||
ImageWorkspaceReferenceUploadResult,
|
||
ImageWorkspaceSendMessageInput,
|
||
ImageWorkspaceSnapshot,
|
||
} from '../../shared/image-workspace';
|
||
|
||
const LOCAL_WORKSPACE_SCHEMA_VERSION = 1;
|
||
const LOCAL_WORKSPACE_DIRECTORY = 'image-workspace-development';
|
||
const LOCAL_WORKSPACE_FILE = 'workspace.json';
|
||
const MAX_PROJECT_NAME_LENGTH = 80;
|
||
const MAX_PROMPT_LENGTH = 4_000;
|
||
const MAX_AGENT_COUNT = 50;
|
||
const MAX_REFERENCE_BYTES = 5 * 1024 * 1024;
|
||
|
||
const LOCAL_CAPABILITIES: ImageWorkspaceCapabilities = {
|
||
modes: [
|
||
{ id: 'generate', label: '图片生成', description: '从文字或参考图生成新的图片' },
|
||
{ id: 'edit', label: '参考图编辑', description: '使用参考图继续调整画面' },
|
||
],
|
||
models: [
|
||
{ id: 'local-placeholder-v1', label: '标准创作模型', description: '适用于图片生成与参考图编辑' },
|
||
],
|
||
aspectRatios: [
|
||
{ id: '1:1', label: '1:1' },
|
||
{ id: '3:4', label: '3:4' },
|
||
{ id: '4:3', label: '4:3' },
|
||
{ id: '16:9', label: '16:9' },
|
||
],
|
||
resolutions: [
|
||
{ id: '1024', label: '1K' },
|
||
{ id: '2048', label: '2K' },
|
||
],
|
||
outputCounts: [
|
||
{ id: '1', label: '1 张' },
|
||
{ id: '2', label: '2 张' },
|
||
{ id: '4', label: '4 张' },
|
||
],
|
||
defaultModeId: 'generate',
|
||
defaultModelId: 'local-placeholder-v1',
|
||
defaultAspectRatioId: '1:1',
|
||
defaultResolutionId: '1024',
|
||
defaultOutputCountId: '1',
|
||
maxReferenceImages: 4,
|
||
referenceUpload: {
|
||
enabled: true,
|
||
acceptedMimeTypes: ['image/png', 'image/jpeg', 'image/webp'],
|
||
maxBytes: MAX_REFERENCE_BYTES,
|
||
},
|
||
};
|
||
|
||
type PersistedImageWorkspace = {
|
||
schemaVersion: typeof LOCAL_WORKSPACE_SCHEMA_VERSION;
|
||
activeProjectId: string | null;
|
||
projects: ImageWorkspaceProject[];
|
||
referencesByProjectId: Record<string, ImageWorkspaceImage[]>;
|
||
};
|
||
|
||
type LocalImageWorkspaceOptions = {
|
||
userDataDir: string;
|
||
now?: () => Date;
|
||
createId?: () => string;
|
||
};
|
||
|
||
export class LocalImageWorkspaceError extends Error {
|
||
readonly status: number;
|
||
readonly code: string;
|
||
|
||
constructor(status: number, code: string, message: string) {
|
||
super(message);
|
||
this.name = 'LocalImageWorkspaceError';
|
||
this.status = status;
|
||
this.code = code;
|
||
}
|
||
}
|
||
|
||
function createEmptyState(): PersistedImageWorkspace {
|
||
return {
|
||
schemaVersion: LOCAL_WORKSPACE_SCHEMA_VERSION,
|
||
activeProjectId: null,
|
||
projects: [],
|
||
referencesByProjectId: {},
|
||
};
|
||
}
|
||
|
||
function clone<T>(value: T): T {
|
||
return JSON.parse(JSON.stringify(value)) as T;
|
||
}
|
||
|
||
function isRecord(value: unknown): value is Record<string, unknown> {
|
||
return typeof value === 'object' && value !== null && !Array.isArray(value);
|
||
}
|
||
|
||
function isPersistedWorkspace(value: unknown): value is PersistedImageWorkspace {
|
||
return isRecord(value)
|
||
&& value.schemaVersion === LOCAL_WORKSPACE_SCHEMA_VERSION
|
||
&& (typeof value.activeProjectId === 'string' || value.activeProjectId === null)
|
||
&& Array.isArray(value.projects)
|
||
&& isRecord(value.referencesByProjectId);
|
||
}
|
||
|
||
function escapeXml(value: string): string {
|
||
return value
|
||
.replaceAll('&', '&')
|
||
.replaceAll('<', '<')
|
||
.replaceAll('>', '>')
|
||
.replaceAll('"', '"')
|
||
.replaceAll("'", ''');
|
||
}
|
||
|
||
function wrapPrompt(prompt: string, maxCharacters = 24, maxLines = 3): string[] {
|
||
const characters = Array.from(prompt.replace(/\s+/g, ' ').trim());
|
||
const lines: string[] = [];
|
||
for (let index = 0; index < characters.length && lines.length < maxLines; index += maxCharacters) {
|
||
const line = characters.slice(index, index + maxCharacters).join('');
|
||
const truncated = index + maxCharacters < characters.length && lines.length === maxLines - 1;
|
||
lines.push(truncated ? `${line.slice(0, Math.max(0, line.length - 1))}…` : line);
|
||
}
|
||
return lines.length > 0 ? lines : ['未填写提示词'];
|
||
}
|
||
|
||
function resolveDimensions(settings: Required<ImageWorkspaceGenerationSettings>): { width: number; height: number } {
|
||
const longEdge = settings.resolutionId === '2048' ? 2048 : 1024;
|
||
const [widthRatio, heightRatio] = (settings.aspectRatioId ?? '1:1')
|
||
.split(':')
|
||
.map((value) => Number(value));
|
||
if (!widthRatio || !heightRatio) return { width: longEdge, height: longEdge };
|
||
if (widthRatio >= heightRatio) {
|
||
return { width: longEdge, height: Math.round(longEdge * heightRatio / widthRatio) };
|
||
}
|
||
return { width: Math.round(longEdge * widthRatio / heightRatio), height: longEdge };
|
||
}
|
||
|
||
function optionId(
|
||
value: string | undefined,
|
||
defaultValue: string,
|
||
options: Array<{ id: string }>,
|
||
label: string,
|
||
): string {
|
||
const resolved = value || defaultValue;
|
||
if (!options.some((option) => option.id === resolved)) {
|
||
throw new LocalImageWorkspaceError(400, 'IMAGE_WORKSPACE_INVALID_SETTING', `当前创作空间不支持这个${label}`);
|
||
}
|
||
return resolved;
|
||
}
|
||
|
||
function normalizeSettings(settings: ImageWorkspaceGenerationSettings): Required<ImageWorkspaceGenerationSettings> {
|
||
return {
|
||
modeId: optionId(settings.modeId, LOCAL_CAPABILITIES.defaultModeId!, LOCAL_CAPABILITIES.modes, '创作模式'),
|
||
modelId: optionId(settings.modelId, LOCAL_CAPABILITIES.defaultModelId!, LOCAL_CAPABILITIES.models, '模型'),
|
||
aspectRatioId: optionId(
|
||
settings.aspectRatioId,
|
||
LOCAL_CAPABILITIES.defaultAspectRatioId!,
|
||
LOCAL_CAPABILITIES.aspectRatios,
|
||
'画幅',
|
||
),
|
||
resolutionId: optionId(
|
||
settings.resolutionId,
|
||
LOCAL_CAPABILITIES.defaultResolutionId!,
|
||
LOCAL_CAPABILITIES.resolutions,
|
||
'分辨率',
|
||
),
|
||
outputCountId: optionId(
|
||
settings.outputCountId,
|
||
LOCAL_CAPABILITIES.defaultOutputCountId!,
|
||
LOCAL_CAPABILITIES.outputCounts,
|
||
'生成数量',
|
||
),
|
||
};
|
||
}
|
||
|
||
function createPlaceholderImage(
|
||
id: string,
|
||
prompt: string,
|
||
settings: Required<ImageWorkspaceGenerationSettings>,
|
||
referenceImageIds: string[],
|
||
outputIndex: number,
|
||
): ImageWorkspaceImage {
|
||
const seed = createHash('sha256')
|
||
.update(JSON.stringify({ prompt, settings, referenceImageIds, outputIndex }))
|
||
.digest('hex');
|
||
const hue = Number.parseInt(seed.slice(0, 4), 16) % 360;
|
||
const accentHue = (hue + 58 + Number.parseInt(seed.slice(4, 6), 16) % 90) % 360;
|
||
const { width, height } = resolveDimensions(settings);
|
||
const promptLines = wrapPrompt(prompt);
|
||
const promptText = promptLines.map((line, index) => (
|
||
`<tspan x="${Math.round(width * 0.08)}" dy="${index === 0 ? 0 : Math.round(height * 0.075)}">${escapeXml(line)}</tspan>`
|
||
)).join('');
|
||
const metadata = [
|
||
LOCAL_CAPABILITIES.models.find((model) => model.id === settings.modelId)?.label ?? '标准创作模型',
|
||
settings.aspectRatioId,
|
||
`${width}×${height}`,
|
||
referenceImageIds.length > 0 ? `${referenceImageIds.length} 张参考图` : '纯文字创作',
|
||
].join(' · ');
|
||
const svg = `
|
||
<svg xmlns="http://www.w3.org/2000/svg" width="${width}" height="${height}" viewBox="0 0 ${width} ${height}">
|
||
<rect width="${width}" height="${height}" fill="hsl(${hue} 72% 88%)"/>
|
||
<circle cx="${Math.round(width * 0.82)}" cy="${Math.round(height * 0.2)}" r="${Math.round(Math.min(width, height) * 0.28)}" fill="hsl(${accentHue} 78% 68%)" opacity="0.72"/>
|
||
<rect x="${Math.round(width * 0.06)}" y="${Math.round(height * 0.08)}" width="${Math.round(width * 0.88)}" height="${Math.round(height * 0.84)}" rx="${Math.round(Math.min(width, height) * 0.04)}" fill="white" fill-opacity="0.78" stroke="#202638" stroke-width="${Math.max(4, Math.round(Math.min(width, height) * 0.008))}"/>
|
||
<text x="${Math.round(width * 0.08)}" y="${Math.round(height * 0.19)}" fill="#202638" font-family="Arial, sans-serif" font-size="${Math.max(24, Math.round(Math.min(width, height) * 0.055))}" font-weight="800">CREATIVE · ${outputIndex + 1}</text>
|
||
<text x="${Math.round(width * 0.08)}" y="${Math.round(height * 0.39)}" fill="#202638" font-family="Arial, sans-serif" font-size="${Math.max(22, Math.round(Math.min(width, height) * 0.047))}" font-weight="700">${promptText}</text>
|
||
<text x="${Math.round(width * 0.08)}" y="${Math.round(height * 0.82)}" fill="#5f4638" font-family="Arial, sans-serif" font-size="${Math.max(16, Math.round(Math.min(width, height) * 0.027))}" font-weight="700">${escapeXml(metadata)}</text>
|
||
</svg>
|
||
`.trim();
|
||
const url = `data:image/svg+xml;base64,${Buffer.from(svg, 'utf8').toString('base64')}`;
|
||
return {
|
||
id,
|
||
url,
|
||
thumbnailUrl: url,
|
||
alt: `生成图片:${prompt}`,
|
||
};
|
||
}
|
||
|
||
function validateBase64(contentBase64: string): Buffer {
|
||
const normalized = contentBase64.trim();
|
||
if (!normalized || normalized.length % 4 !== 0 || !/^[A-Za-z0-9+/]*={0,2}$/.test(normalized)) {
|
||
throw new LocalImageWorkspaceError(400, 'IMAGE_WORKSPACE_INVALID_REFERENCE', '参考图内容不是有效的 Base64');
|
||
}
|
||
const bytes = Buffer.from(normalized, 'base64');
|
||
if (bytes.length === 0 || bytes.length > MAX_REFERENCE_BYTES) {
|
||
throw new LocalImageWorkspaceError(400, 'IMAGE_WORKSPACE_INVALID_REFERENCE', '参考图大小必须在 5 MB 以内');
|
||
}
|
||
return bytes;
|
||
}
|
||
|
||
export function isLocalImageWorkspaceDevelopmentEnabled(options: {
|
||
isPackaged: boolean;
|
||
configuredMode?: string;
|
||
isDevelopmentServer?: boolean;
|
||
}): boolean {
|
||
if (options.isPackaged) return false;
|
||
|
||
const configuredMode = options.configuredMode?.trim().toLowerCase();
|
||
return configuredMode === 'local'
|
||
|| (!configuredMode && options.isDevelopmentServer === true);
|
||
}
|
||
|
||
export function getLocalImageWorkspaceDirectory(userDataDir: string): string {
|
||
return join(userDataDir, LOCAL_WORKSPACE_DIRECTORY);
|
||
}
|
||
|
||
export class LocalImageWorkspace {
|
||
private readonly rootDirectory: string;
|
||
private readonly stateFile: string;
|
||
private readonly now: () => Date;
|
||
private readonly createId: () => string;
|
||
private state: PersistedImageWorkspace | null = null;
|
||
private operation: Promise<void> = Promise.resolve();
|
||
|
||
constructor(options: LocalImageWorkspaceOptions) {
|
||
this.rootDirectory = getLocalImageWorkspaceDirectory(options.userDataDir);
|
||
this.stateFile = join(this.rootDirectory, LOCAL_WORKSPACE_FILE);
|
||
this.now = options.now ?? (() => new Date());
|
||
this.createId = options.createId ?? randomUUID;
|
||
}
|
||
|
||
getSnapshot(): Promise<ImageWorkspaceSnapshot> {
|
||
return this.enqueue(async () => this.snapshot(await this.load()));
|
||
}
|
||
|
||
createProject(name: string): Promise<ImageWorkspaceSnapshot> {
|
||
return this.mutate(async (state) => {
|
||
const projectName = name.trim();
|
||
if (!projectName || projectName.length > MAX_PROJECT_NAME_LENGTH) {
|
||
throw new LocalImageWorkspaceError(400, 'IMAGE_WORKSPACE_INVALID_PROJECT_NAME', '项目名称需为 1–80 个字符');
|
||
}
|
||
const timestamp = this.now().toISOString();
|
||
const projectId = `local-project-${this.createId()}`;
|
||
state.projects.push({
|
||
id: projectId,
|
||
name: projectName,
|
||
agents: [],
|
||
activeAgentId: null,
|
||
messages: [],
|
||
updatedAt: timestamp,
|
||
});
|
||
state.referencesByProjectId[projectId] = [];
|
||
state.activeProjectId = projectId;
|
||
});
|
||
}
|
||
|
||
addAgent(projectId: string): Promise<ImageWorkspaceSnapshot> {
|
||
return this.mutate(async (state) => {
|
||
const project = this.requireProject(state, projectId);
|
||
if (project.agents.length >= MAX_AGENT_COUNT) {
|
||
throw new LocalImageWorkspaceError(409, 'IMAGE_WORKSPACE_AGENT_LIMIT', '每个项目最多添加 50 个 Agent');
|
||
}
|
||
const agentId = `local-agent-${this.createId()}`;
|
||
project.agents.push({ id: agentId, name: `创作 Agent ${project.agents.length + 1}` });
|
||
project.updatedAt = this.now().toISOString();
|
||
});
|
||
}
|
||
|
||
sendMessage(input: ImageWorkspaceSendMessageInput): Promise<ImageWorkspaceSnapshot> {
|
||
return this.mutate(async (state) => {
|
||
const project = this.requireProject(state, input.projectId);
|
||
if (!project.agents.some((agent) => agent.id === input.agentId)) {
|
||
throw new LocalImageWorkspaceError(404, 'IMAGE_WORKSPACE_AGENT_NOT_FOUND', '当前项目中不存在这个 Agent');
|
||
}
|
||
const prompt = input.prompt.trim();
|
||
if (!prompt || prompt.length > MAX_PROMPT_LENGTH) {
|
||
throw new LocalImageWorkspaceError(400, 'IMAGE_WORKSPACE_INVALID_PROMPT', '创作描述需为 1–4000 个字符');
|
||
}
|
||
const referenceIds = [...new Set(input.referenceImageIds)];
|
||
if (referenceIds.length > LOCAL_CAPABILITIES.maxReferenceImages) {
|
||
throw new LocalImageWorkspaceError(400, 'IMAGE_WORKSPACE_REFERENCE_LIMIT', '参考图数量超过创作空间限制');
|
||
}
|
||
const availableImages = new Map<string, ImageWorkspaceImage>();
|
||
for (const image of state.referencesByProjectId[project.id] ?? []) availableImages.set(image.id, image);
|
||
for (const message of project.messages) {
|
||
for (const image of message.images) availableImages.set(image.id, image);
|
||
}
|
||
const references = referenceIds.map((id) => {
|
||
const image = availableImages.get(id);
|
||
if (!image) {
|
||
throw new LocalImageWorkspaceError(400, 'IMAGE_WORKSPACE_REFERENCE_NOT_FOUND', '所选参考图不存在');
|
||
}
|
||
return image;
|
||
});
|
||
const settings = normalizeSettings(input.settings);
|
||
const outputCount = Number(settings.outputCountId);
|
||
const timestamp = this.now().toISOString();
|
||
const images = Array.from({ length: outputCount }, (_, index) => createPlaceholderImage(
|
||
`local-image-${this.createId()}`,
|
||
prompt,
|
||
settings,
|
||
referenceIds,
|
||
index,
|
||
));
|
||
project.messages.push(
|
||
{
|
||
id: `local-message-${this.createId()}`,
|
||
role: 'user',
|
||
agentId: input.agentId,
|
||
text: prompt,
|
||
images: clone(references),
|
||
status: 'succeeded',
|
||
createdAt: timestamp,
|
||
},
|
||
{
|
||
id: `local-message-${this.createId()}`,
|
||
role: 'assistant',
|
||
agentId: input.agentId,
|
||
text: `已生成 ${outputCount} 张图片。`,
|
||
images,
|
||
status: 'succeeded',
|
||
createdAt: timestamp,
|
||
},
|
||
);
|
||
project.activeAgentId = input.agentId;
|
||
project.updatedAt = timestamp;
|
||
state.activeProjectId = project.id;
|
||
});
|
||
}
|
||
|
||
uploadReference(input: ImageWorkspaceReferenceUploadInput): Promise<ImageWorkspaceReferenceUploadResult> {
|
||
return this.enqueue(async () => {
|
||
const state = await this.load();
|
||
const project = this.requireProject(state, input.projectId);
|
||
if (!LOCAL_CAPABILITIES.referenceUpload.acceptedMimeTypes.includes(input.mimeType)) {
|
||
throw new LocalImageWorkspaceError(400, 'IMAGE_WORKSPACE_INVALID_REFERENCE_TYPE', '创作空间仅支持 PNG、JPEG 和 WebP');
|
||
}
|
||
const bytes = validateBase64(input.contentBase64);
|
||
const fileName = input.fileName.trim().slice(0, 180) || '本地参考图';
|
||
const reference: ImageWorkspaceImage = {
|
||
id: `local-reference-${this.createId()}`,
|
||
url: `data:${input.mimeType};base64,${bytes.toString('base64')}`,
|
||
alt: fileName,
|
||
};
|
||
state.referencesByProjectId[project.id] ??= [];
|
||
state.referencesByProjectId[project.id].push(reference);
|
||
project.updatedAt = this.now().toISOString();
|
||
await this.persist(state);
|
||
return { workspace: this.snapshot(state), reference: clone(reference) };
|
||
});
|
||
}
|
||
|
||
reset(): Promise<ImageWorkspaceSnapshot> {
|
||
return this.enqueue(async () => {
|
||
await rm(this.rootDirectory, { recursive: true, force: true });
|
||
this.state = createEmptyState();
|
||
return this.snapshot(this.state);
|
||
});
|
||
}
|
||
|
||
private async mutate(
|
||
operation: (state: PersistedImageWorkspace) => Promise<void>,
|
||
): Promise<ImageWorkspaceSnapshot> {
|
||
return this.enqueue(async () => {
|
||
const state = await this.load();
|
||
await operation(state);
|
||
await this.persist(state);
|
||
return this.snapshot(state);
|
||
});
|
||
}
|
||
|
||
private enqueue<T>(operation: () => Promise<T>): Promise<T> {
|
||
const result = this.operation.then(operation, operation);
|
||
this.operation = result.then(() => undefined, () => undefined);
|
||
return result;
|
||
}
|
||
|
||
private async load(): Promise<PersistedImageWorkspace> {
|
||
if (this.state) return this.state;
|
||
try {
|
||
const parsed = JSON.parse(await readFile(this.stateFile, 'utf8')) as unknown;
|
||
if (!isPersistedWorkspace(parsed)) {
|
||
throw new LocalImageWorkspaceError(
|
||
500,
|
||
'IMAGE_WORKSPACE_LOCAL_DATA_INVALID',
|
||
'创作空间数据暂时无法读取',
|
||
);
|
||
}
|
||
this.state = parsed;
|
||
} catch (error) {
|
||
if ((error as NodeJS.ErrnoException).code === 'ENOENT') {
|
||
this.state = createEmptyState();
|
||
} else {
|
||
throw error;
|
||
}
|
||
}
|
||
return this.state;
|
||
}
|
||
|
||
private async persist(state: PersistedImageWorkspace): Promise<void> {
|
||
await mkdir(this.rootDirectory, { recursive: true });
|
||
const temporaryFile = `${this.stateFile}.${process.pid}.tmp`;
|
||
await writeFile(temporaryFile, `${JSON.stringify(state, null, 2)}\n`, 'utf8');
|
||
try {
|
||
await rename(temporaryFile, this.stateFile);
|
||
} catch (error) {
|
||
const code = (error as NodeJS.ErrnoException).code;
|
||
if (code !== 'EEXIST' && code !== 'EPERM') throw error;
|
||
await rm(this.stateFile, { force: true });
|
||
await rename(temporaryFile, this.stateFile);
|
||
}
|
||
this.state = state;
|
||
}
|
||
|
||
private requireProject(state: PersistedImageWorkspace, projectId: string): ImageWorkspaceProject {
|
||
const project = state.projects.find((item) => item.id === projectId);
|
||
if (!project) {
|
||
throw new LocalImageWorkspaceError(404, 'IMAGE_WORKSPACE_PROJECT_NOT_FOUND', '项目不存在');
|
||
}
|
||
return project;
|
||
}
|
||
|
||
private snapshot(state: PersistedImageWorkspace): ImageWorkspaceSnapshot {
|
||
return {
|
||
capabilities: clone(LOCAL_CAPABILITIES),
|
||
projects: clone(state.projects),
|
||
activeProjectId: state.activeProjectId,
|
||
};
|
||
}
|
||
}
|