diffusers.js library for running diffusion models on GPU/WebGPU
See demo here https://islamov.ai/diffusers.js/
Installation
npm i @aislamov/diffusers.jsUsage
Browser (see examples/react)
import { DiffusionPipeline } from '@aislamov/diffusers.js'
const pipe = DiffusionPipeline.fromPretrained('aislamov/stable-diffusion-2-1-base-onnx')
const images = pipe.run({
prompt: "an astronaut running a horse",
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 (see 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: "an astronaut running a horse",
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('Image saved as output.png');
})'aislamov/stable-diffusion-2-1-base-onnx' model is optimized for GPU and will fail to load without CUDA/DML/WebGPU support. Please use 'cpu' revision on a machine without GPU.
const pipe = DiffusionPipeline.fromPretrained('aislamov/stable-diffusion-2-1-base-onnx', { revision: 'cpu' })Running examples
Browser/React
npm install && npm run build
cd examples/react && npm install
npm run startNode
npm install && npm run build
cd examples/node && npm install
node src/txt2img.mjs --prompt "a dog" --rev cpuNode.js GPU support
DML and CUDA support for node.js is already merged to onnxruntime microsoft/onnxruntime#16050 but it is not released in current onnxruntime-node package. I will build my own version if they will not update it in the next release.
How does it work
It uses my own modified build of onnx runtime for web with 64bit and other changes. You can see list of contributions here https://github.com/search?q=is%3Apr%20author%3Adakenf%20&type=pullrequests
I'm trying to get all changes to be merged to official version.
Also, i've had to do a lot of fixes to emscripten like this emscripten-core/emscripten#19737, WebAssembly spec and V8 engine like this and this in order to make this all work