최신 버전의 Chrome에서 "실험적인 WebAssembly" 및 "실험적인 WebAssembly JavaScript Promise 통합 (JSPI)" 플래그가 활성화된 상태로 Stable Diffusion WebGPU 데모를 실행해야 합니다.
데모의 각 추론 단계는 약 1분이 걸리며, 이미지를 생성하기 위해 VAE 디코더를 실행하는 데 추가로 약 10초가 소요됩니다.
DevTools를 열어두면 데모 속도가 약 2배로 느려집니다.
데모에서 사용되는 UNET 모델은 CPU에서만 실행되며, 브라우저 탭이 멈출 수 있지만 GPU에서 실행하는 것보다 10% 더 빠릅니다.
데모에서 수용 가능한 결과를 얻기 위한 최소한의 단계 수는 20입니다. 그러나 데모 목적으로는 3단계가 충분합니다.
모델 파일은 캐시되므로 데모를 실행할 때마다 다운로드할 필요가 없습니다.
"protobuf parsing failed" 오류가 발생하면 DevTools에서 Application -> Storage로 이동하여 사이트 데이터를 지울 수 있습니다.
"sbox_fatal_memory_exceeded" 오류가 발생하면 충분한 RAM이 없어 데모를 실행할 수 없는 것입니다. 탭이나 브라우저를 다시 로드해 볼 수 있습니다.
데모는 StableDiffusionPipeline을 Python에서 JS로 이식하고, onnxruntime과 emscripten+binaryen을 패치하여 4GB 이상의 메모리를 할당하고 사용할 수 있도록 지원함으로써 가능해졌습니다.
현재 데모는 멀티스레딩을 지원하지 않고 하나의 CPU 코어만 사용하기 때문에 느립니다. WebAssembly.Memory 생성자를 통한 64비트 메모리의 지원 부족도 제한 사항입니다.
데모는 GPU에서 실행되지만 webgpu와 onnxruntime은 초기 단계이므로 많은 작업이 아직 구현되지 않았습니다. 데이터는 JS를 통해 GPU와 CPU 사이에서 지속적으로 전송되어 처리 속도가 느려집니다. 더 많은 작업이 JS 커널을 갖게 되면 데모는 더 빨라질 것입니다.
GitHub에서 사용 가능한 코드를 사용하여 데모를 로컬에서 실행할 수 있습니다.
패치된 onnxruntime을 사용하여 transformers.js로 큰 LLMs를 실행할 수 있지만, 8GB의 메모리로 제한됩니다. 해당 패키지를 사용하여 최대 약 4GB의 가중치를 로드할 수 있습니다.
저자는 이전에 node.js 바인딩에 GPU 가속을 추가한 경험이 있으므로 onnxruntime 리포지토리에 풀 리퀘스트를 할 계획입니다.