diffusers.js 库用于在 GPU/WebGPU 上运行扩散模型

查看演示：https://islamov.ai/diffusers.js/

安装

npm i @aislamov/diffusers.js

使用方法

浏览器（参见 examples/react）

import { DiffusionPipeline } from '@aislamov/diffusers.js'

const pipe = DiffusionPipeline.fromPretrained('aislamov/stable-diffusion-2-1-base-onnx')
const images = pipe.run({
  prompt: "一个宇航员骑马",
  numInferenceSteps: 30,
})

const canvas = document.getElementById('canvas')
const data = await images[0].toImageData({ tensorLayout: 'NCWH', format: 'RGB' });
canvas.getContext('2d').putImageData(data, 0, 0);

Node.js（参见 examples/node）

import { DiffusionPipeline } from '@aislamov/diffusers.js'
import { PNG } from 'pngjs'

const pipe = DiffusionPipeline.fromPretrained('aislamov/stable-diffusion-2-1-base-onnx')
const images = pipe.run({
  prompt: "一个宇航员骑马",
  numInferenceSteps: 30,
})

const data = await images[0].mul(255).round().clipByValue(0, 255).transpose(0, 2, 3, 1)

const p = new PNG({ width: 512, height: 512, inputColorType: 2 })
p.data = Buffer.from(data.data)
p.pack().pipe(fs.createWriteStream('output.png')).on('finish', () => {
  console.log('图像已保存为 output.png');
})

'aislamov/stable-diffusion-2-1-base-onnx' 模型针对 GPU 进行了优化，如果没有 CUDA/DML/WebGPU 支持将无法加载。请在没有 GPU 的机器上使用 'cpu' 版本。

const pipe = DiffusionPipeline.fromPretrained('aislamov/stable-diffusion-2-1-base-onnx', { revision: 'cpu' })

运行示例

浏览器/React

npm install && npm run build
cd examples/react && npm install
npm run start

Node

npm install && npm run build
cd examples/node && npm install
node src/txt2img.mjs --prompt "一只狗" --rev cpu

Node.js GPU 支持

DML 和 CUDA 对 node.js 的支持已合并到 onnxruntime https://github.com/microsoft/onnxruntime/pull/16050，但在当前的 onnxruntime-node 包中尚未发布。如果他们在下一个版本中不更新，我将构建自己的版本。

工作原理

它使用了我自己修改的 onnx runtime 网页版本，包含 64 位和其他更改。你可以在这里查看贡献列表：https://github.com/search?q=is%3Apr%20author%3Adakenf%20&type=pullrequests

我正在尝试将所有更改合并到官方版本中。

此外，为了使这一切正常运行，我不得不对 emscripten 进行大量修复，如 https://github.com/emscripten-core/emscripten/pull/19737，以及对 WebAssembly 规范和 V8 引擎进行修改，如这个和这个。