如何使用 QVAC 和 Socket.io 在 Node.js 中构建一个离线 AI 图像生成器

const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);

const app = express();
const server = http.createServer(app);
const io = new Server(server);

const PORT = process.env.PORT || 3000;

app.use(express.json());
app.use(express.static(path.join(__dirname, 'public');

const CONFIG_PATH = path.join(__dirname, '.device-preference.json');

function getPreferredDevice() {
try {
if (fs.existsSync(CONFIG_PATH)) {
const data = JSON.parse(fs.readFileSyncCONFIG_PATH, 'utf8'));
return data.device || null;
}
} catch (err) {
console.error('读取设备偏好设置失败:', err.message);
}
return null;
}

function setPreferredDevice(device) {
try {
fs.writeFileSync(CONFIG_PATH, JSON.stringify({ device }), 'utf8');
} catch (err) {
console.error('写入设备偏好设置失败:', err.message);
}
}

// 全局模型状态
let loadedModelId = process.modelId || null;
let modelLoadPercent = 0;
let modelLoadStatus = '等待触发...';
let isModelLoading = false;

const modelSize = (SD_V2_1_1B_Q8_0.expectedSize / (1024 * 1024 * 1024)).toFixed(2) + ' GB';

function broadcastModelProgresspercent, status) {
io.emit('model-download-progress', { percent, status, size: modelSize });
}

io.on('connection', (socket) => {
console.log('客户端已连接:', socket.id);

socket.on('disconnect', () => {
console.log('客户端已断开连接:', socket.id);
});

// 触发模型下载
socket.on('trigger-model-download', async () => {
// 如果模型已经加载完成，检查它是否仍在工作进程中
if (loadedModelId) {
try {
await getLoadedModelInfo({ modelId: loadedModelId });
socket.emit('model-download-progress', {
percent: 100,
status: '模型已成功在本地加载。',
size: modelSize
});
return;
} catch (err) {
console.log('模型ID无效或未找到，正在重置状态并重新加载..., ' + err.message);
loadedModelId = null;
process.modelId = null;
}
}

// 如果模型当前正在加载中，报告当前的加载进度
if (isModelLoading) {
socket.emit('model-download-progress', {
percent: Math.round(modelLoadPercent),
status: modelLoadStatus,
size: modelSize
});
return;
}

isModelLoading = true;
modelLoadPercent = 0;
modelLoadStatus = '开始加载模型...';
broadcastModelProgress(modelLoadPercent, modelLoadStatus);

try {
console.log('开始加载模型...'
const preferredDevice = getPreferredDevice();
const loadConfig = { prediction: "v" };
if (preferredDevice) {
loadConfig.device = preferredDevice;
if (preferredDevice === 'cpu') {
loadConfig.threads = 4;
}
console.log(`使用缓存的设备偏好设置：${preferredDevice}`);
}

loadedModelId = await loadModel({
modelSrc: SD_V2_1_1B_Q8_0,
modelType: "sdcpp-generation",
modelConfig: loadConfig,
onProgress: (p) => {
modelLoadPercent = p.percentage;
modelLoadStatus = p.percentage >= 100 ? '模型已成功在本地加载.' : `正在下载模型权重... (${p.percentage.toFixed(1)}%)`;
broadcastModelProgress(Math.round(modelLoadPercent), modelLoadStatus);
}
});
process.modelId = loadedModelId;

isModelLoading = false;
console.log('模型已成功加载。ID：', loadedModelId);
} catch (err) {
isModelLoading = false;
modelLoadPercent = 0;
modelLoadStatus = '加载模型失败: ' + err.message;
console.error('加载模型失败:', err);
broadcastModelProgress(0, modelLoadStatus);
socket.emit('error_event', { message: '加载模型失败: ' + err.message });
}
});

socket.on('generate', async (data) => {
const { prompt, ratio } = data;
if (!prompt || prompt.trim() === '') {
socket.emit('error_event', { message: '需要提供提示信息' });
return;
}

if (!loadedModelId) {
socket.emit('error_event', { message: '模型尚未加载完成' });
return;
}

const runDiffusion = async (modelIdToUse) => {
socket.emit('progress', {
percent: 0,
status: '开始执行扩散过程...'
sub: 'DIFFUSION INITIALIZING'
});

console.log(`正在使用模型ID ${modelIdToUse}，为提示信息 „\({prompt}"“ 和比率 \){ratio} 生成图像`);

const { progressStream, outputs, stats } = diffusion({
modelId: modelIdToUse,
prompt,
});

// 输出每一步的进展情况
for await (const { step, totalSteps } of progressStream) {
const percent = Math.round((step / totalSteps) * 100);
socket.emit('progress', {
percent,
status: `当前步骤为 \({step}/\){totalSteps}…`,
sub: 'RUNNING DIFFUSION'
});
}

// 处理输出结果
const buffers = await outputs;
if (!buffers || buffers.length === 0) {
throw new Error('扩散模型未返回任何图像缓冲区');
}

// 将图像缓冲区转换为base64 Data URL格式，而不是将其保存到磁盘上
const base64Data = Buffer.from(buffers[0]).toString('base64');
const dataUrl = `data:image/png;base64,${base64Data}`;

// 发送成功信息
socket.emit('success', {
url: dataUrl,
prompt,
seed: (await stats).seed || -1
});

console.log(`图像已生成，并以base64 Data URL格式发送成功。`);
};

try {
await runDiffusion(loadedModelId);
} catch (err) {
console.error('图像生成失败:', err);

const isCrash = err.code === 50205 || (err.message && err.message.includes('WORKER_CRASHED'));
if (isCrash) {
console.log('在GPU上执行时工作进程崩溃了。正在尝试切换到CPU模式…');

// 保存设备偏好设置，以便下次直接使用CPU进行加载，避免重复加载
setPreferredDevice('cpu');

// 重置模型状态
loadedModelId = null;
process.modelId = null;

socket.emit('progress', {
percent: 0,
status: 'GPU驱动程序崩溃了。正在自动切换到CPU模式…',
sub: 'CPU FALLBACK LOADING'
});

try {
console.log('正在使用CPU加载模型…');
isModelLoading = true;
modelLoadPercent = 0;
modelLoadStatus = '正在加载CPU模型的权重…';
broadcastModelProgress(modelLoadPercent, modelLoadStatus);

loadedModelId = await loadModel({
modelSrc: SD_V2_1_1B_Q8_0,
modelType: "sdcpp-generation",
modelConfig: { prediction: "v", device: 'cpu', threads: 4 },
onProgress: (p) => {
modelLoadPercent = p.percentage;
modelLoadStatus = `正在加载CPU模型的权重… (${p.percentage.toFixed(1)}%)`;
broadcastModelProgress(Math.round(modelLoadPercent), modelLoadStatus);
}
});
process.modelId = loadedModelId;
isModelLoading = false;
console.log('模型已成功在CPU上加载。ID：', loadedModelId);

// 尝试再次使用CPU执行扩散过程
await runDiffusion(loadedModelId);
} catch (cpuErr) {
console.error('在CPU上执行切换失败:', cpuErr);
isModelLoading = false;
socket.emit('error_event', { message: '在CPU上生成图像失败：' + cpuErr.message });
}
} else {
if (err.message && (err.message.includes('MODEL_NOT_FOUND') || err.message.includes('未找到')) {
loadedModelId = null;
process.modelId = null;
broadcastModelProgress(0, '模型状态丢失。请重新触发加载过程.');
}
socket.emit('error_event', { message: '图像生成失败: ' + err.message });
}
}
});
});

app.get '*', (req, res) => {
res.sendFile(path.join(__dirname, 'public', 'index.html'));
});

server.listen(PORT, () => {
console.log(`服务器正在运行，地址为 http://localhost:${PORT}`);
});

// 清理退出处理程序
async function handleCleanup() {
const modelId = process.modelId || loadedModelId;
if (modelId && modelId !== 'mock-model-id') {
try {
await unloadModel({ modelId, clearStorage: false });
} catch (err) {}
}
process.exit(0);
}