Bản demo Stable Diffusion WebGPU

• Bản demo tạo ảnh Stable Diffusion chạy trực tiếp trong trình duyệt, hoạt động dựa trên WebGPU và WebAssembly mà không cần máy chủ riêng
• Cần bật các flag "Experimental WebAssembly" và "Experimental WebAssembly JavaScript Promise Integration (JSPI)" trên Chrome mới nhất để sử dụng
• Tạo ảnh bằng cách nhập Prompt, Negative Prompt và số inference step; các tệp mô hình được cache nên không cần tải lại mỗi lần
• Là kết quả của việc patch onnxruntime, emscripten, binaryen để cấp phát bộ nhớ vượt quá 4GB và port StableDiffusionPipeline của Python sang JS
• Hiện tốc độ còn chậm do chưa hỗ trợ đa luồng và các phép toán WebGPU chưa được triển khai, nhưng cho thấy khả năng cải thiện trong tương lai nhờ triển khai kernel JS và hỗ trợ memory64

---

Yêu cầu để chạy

• Cần bật các flag Experimental WebAssembly và Experimental WebAssembly JavaScript Promise Integration (JSPI) trên trình duyệt Chrome mới nhất
• Model files được cache nên khi truy cập lại không cần tải xuống lại

Đầu vào và cách hoạt động

• Cung cấp các mục đầu vào gồm Prompt, Negative Prompt và số inference step
  • Do sử dụng PNDM Scheduler, step thực tế được xử lý từ giá trị nhập i thành i+1
• Mỗi step mất khoảng 1 phút; việc chạy VAE decoder để tạo ảnh mất thêm khoảng 10 giây
• Khi mở DevTools, tốc độ tổng thể chậm hơn khoảng 2 lần
• UNET chỉ chạy trên CPU, nhanh hơn GPU 10%; trên GPU nó không tạo ra kết quả đúng và khiến tab trình duyệt bị treo
• Số step tối thiểu để có kết quả chấp nhận được là 20; cho mục đích demo thì 3 step cũng không vấn đề

FAQ

• Lỗi protobuf parsing failed

  • Trong DevTools, vào Application → Storage rồi chạy "Clear site data"

• Lỗi sbox_fatal_memory_exceeded

  • Đây là trạng thái thiếu RAM cần thiết để chạy SD; hãy thử khắc phục bằng cách tải lại tab hoặc trình duyệt

• Cách triển khai

  • Port StableDiffusionPipeline của Python sang JS
  • Patch onnxruntime và emscripten+binaryen (toolchain biên dịch WebAssembly) để hỗ trợ cấp phát và sử dụng bộ nhớ vượt quá 4GB
  • Khi các pull request liên quan được đưa vào bản phát hành, bất kỳ ai cũng có thể biên dịch và chạy mã sử dụng bộ nhớ trên 4GB trong trình duyệt

• Nguyên nhân tốc độ chậm

  • Hiện chưa hỗ trợ đa luồng, nên chỉ dùng một lõi CPU
  • Không thể tạo bộ nhớ 64-bit bằng SharedArrayBuffer thông qua constructor WebAssembly.Memory
  • Đã đề xuất thay đổi đặc tả của flag "memory64"; sau khi được chấp nhận, dự kiến sẽ patch engine V8 để hỗ trợ

• Có chạy trên GPU không

  • Có chạy trên GPU, nhưng WebGPU của onnxruntime còn ở giai đoạn đầu nên nhiều phép toán chưa được triển khai
  • Dữ liệu liên tục được truyền giữa CPU và JS
  • Khi kernel JS được triển khai cho phần lớn phép toán, tốc độ có thể được cải thiện đáng kể

• Có thể chạy local không

  • Có thể; mã của trang này được cung cấp trong repository stable-diffusion-webgpu-minimal

• Có thể chạy LLM lớn bằng transformers.js không

  • Có thể dùng gói onnxruntime đã được patch (@aislamov/onnxruntime-web64), nhưng không đảm bảo hoạt động trong mọi trường hợp
  • Bản build này bị giới hạn ở 8GB bộ nhớ, nên có thể tải trọng số khoảng 4GB

• Kế hoạch pull request lên repository onnxruntime

  • Dự kiến sẽ tiến hành; đây là công việc thứ hai sau lần trước bổ sung tăng tốc GPU cho binding node.js

[GN⁺]

Ý kiến trên Hacker News

• Nhóm MLC đã chạy được cái này từ tháng 3 rồi: https://github.com/mlc-ai/web-stable-diffusion
  Điều ấn tượng hơn là sau đó họ còn bổ sung hỗ trợ cho nhiều mô hình ngôn ngữ lớn: https://webllm.mlc.ai/

  • Thật sự rất ấn tượng và hiệu năng cũng có vẻ tốt hơn nhiều. Tôi đang theo một hướng tiếp cận khác là chạy mô hình ONNX tùy ý mà không cần chỉnh sửa trước

• “Tải 3.5GB và dùng 8GB RAM”
  Thật thú vị khi trình duyệt đã thành ra như vậy. Web đã nuốt chửng hệ điều hành

  • Tôi không thật sự hiểu lời chỉ trích này. Bạn đang chạy Stable Diffusion trên máy tính của mình thông qua trình duyệt, vậy thì làm sao có thể làm mà không tải xuống và đưa vào RAM được?
    Kể cả tải về và chạy trực tiếp không qua trình duyệt thì kích thước tải xuống và lượng RAM dùng cũng sẽ gần như vậy thôi
  • Web đã nuốt luôn tiền đề ban đầu của Java
  • Nên mới có công ty bán laptop mà hệ điều hành về cơ bản là trình duyệt, và cũng có công ty làm điều tương tự trên smart TV
  • Khi băng thông tăng lên và sandbox của web trưởng thành hơn, thật thú vị khi thấy ứng dụng tiến hóa thành thứ cứ mở lên là dùng, thay vì tải về, cài đặt và bảo trì. Có người sẽ ghét điều này, nhưng với số đông nó mở ra rất nhiều cánh cửa
  • Giờ tôi nghĩ WebGPU cũng hữu ích cho cả những mục đích như thế này, chứ không chỉ cho các trang WebXR tương tác thời gian thực, trạng thái multiplayer streaming hay nhiều draw call.
    Cung cấp trải nghiệm kiểu này qua trình duyệt có những mặt đơn giản hơn, và có lẽ cũng không có cách nào dễ hơn để chạy Stable Diffusion, nên hy vọng các dự án như vậy tiếp tục được hỗ trợ

• Cần Chrome mới nhất với cờ Experimental WebAssembly và Experimental WebAssembly JavaScript Promise Integration (JSPI) được bật, vậy chắc tôi phải chờ đến khi nó vào Firefox thôi

  • Ít ra thì Firefox giờ đã cho phép tùy biến kiểu thanh cuộn CSS, nên đó cũng là tin tốt
  • Tôi đã bật chrome:flags như yêu cầu trên Brave nhưng vẫn không chạy được. Tôi chưa từng tải Chrome trên mấy máy Mac M1 của mình, và giờ cũng không định bắt đầu
  • Tôi đã thử bật trên Canary nhưng vẫn không làm nó chạy được
  • Ý là bạn sẽ không thèm thử trên Chrome luôn à?
  • Cờ thứ hai dùng để làm gì vậy?

• Cái này cũng đã được triển khai ở đây rồi: https://websd.mlc.ai/#text-to-image-generation-demo
  MLC dùng Apache TVM để sinh mã WebGPU và tự động tinh chỉnh, hiệu năng cũng khá ổn

• Tôi đã tích hợp Stable Diffusion của MLC-AI vào website của mình như một trình tạo nền tùy chỉnh. Dữ liệu mô hình lấy từ Hugging Face
  https://dustinbrett.com/

  • Nhìn mức độ mượt và nhanh thế này lại càng thấy buồn khi mọi app trên điện thoại chỉ để hiện vài màn hình và gọi vài API mà cũng đòi tới 200MB
  • Tải tốt trên iPhone 12 mini và điện thoại cũng không bị quá nhiệt. Thật sự rất ấn tượng
  • Website rất ấn tượng và mọi thứ đều chạy cực kỳ mượt. Rốt cuộc tôi đã chơi Doom lâu hơn mức cần thiết

• Với các tác vụ bị nghẽn ở CPU, có thể cân nhắc dùng service worker để luồng chính không bị khựng

• Tiếp theo sẽ là chatbot WebGPU à? Kiểu chat như ChatGPT trên GPU của tôi ngay trong trình duyệt?

  • Thực ra đã làm được rồi. Có một vài hạn chế
    Mô hình ngôn ngữ lớn càng tốt thì càng ngốn nhiều RAM/VRAM, nên đặc biệt khó trên WebGPU
  • https://webllm.mlc.ai/
    Đã có triển khai rồi

• Tôi tò mò là cái này dùng vào đâu?
  Điều tôi hiểu là nó tạo ảnh trong trình duyệt thay vì trên server. Điều duy nhất tôi nghĩ ra là không cần phải refresh trang để chỉnh sửa hoặc tạo ảnh mới.
  Vậy thì có thể có những website mà thiết kế hình ảnh thay đổi theo thời gian thực, và nếu cả chức năng cũng thay đổi đáng kể theo thì có thể sẽ khá hay. Chỉ là tôi không chắc Stable Diffusion hữu ích đến đâu trong việc tạo ra các component UI hay yếu tố hình ảnh của website

  • Nó hữu ích vì chi phí tạo ảnh hàng loạt là rất đắt. Về cơ bản là chuyển chi phí tính toán sang phía client
  • Mục đích là để chạy dễ dàng ngay trên máy của chính mình. Không cần phải là lập trình viên và cũng không cần cài ứng dụng
  • Chạy mô hình hoàn toàn cục bộ thì quyền riêng tư là một lợi thế lớn
    Các prompt nhạy cảm sẽ không bị lộ cho ai đó ở xa

• Có chỉ số nào so sánh hiệu năng trên WebGPU với hiệu năng chạy native không?

  • UNET mất khoảng 1 phút 10 giây trên WebGPU, còn trên CPU một luồng thì khoảng 1 phút. VAE mất 2 phút trên CPU, còn trên GPU thì khoảng 10 giây.
    Có vẻ là do phần lớn phép tính GPU cho VAE đã được triển khai, còn cho UNET thì chưa. Với trường hợp sau, trình duyệt vẫn liên tục ném dữ liệu qua lại giữa GPU và CPU ở mỗi bước

• Nếu đủ nhanh thì có thể dùng để render ảnh cục bộ cho mục đích cá nhân. Website chỉ cần gửi prompt, rồi mỗi người dùng có thể render ra ảnh khác nhau.
  Đến lúc đó thì bản quyền sẽ được tính thế nào nhỉ? Bản thân mô hình có thuộc đối tượng bản quyền không, hay cả hệ thống sẽ sụp đổ?

  • Khả năng đó khá thú vị, nhưng hiện tại vẫn còn rất xa thế giới đó. Ở chỗ khác trong chủ đề này cũng có nói là nó đang dùng tích cực 8GB RAM.
    Và tôi nghĩ không nhiều nhà thiết kế web sẵn sàng chấp nhận rủi ro mô hình hiểu sai prompt, tạo ra kết quả méo mó với số ngón tay sai, hay vô tình sinh ra nội dung tình dục/bạo lực trong ngữ cảnh không mong muốn.
    Ngày nay, với nhiều mô hình sinh ảnh, thường phải tạo hơn chục ảnh rồi chọn cái tốt nhất, còn những cái bị bỏ đi thật ra có thể khá tệ.
    Để website có thể minh họa động một cách thường xuyên trong đời sống hằng ngày thì chất lượng và tính dự đoán được của mô hình phải cao hơn bây giờ rất nhiều.
    Dù vậy, tôi cũng không muốn nói là nó tuyệt đối không thể đạt tới mức đó. Các mô hình gần đây đã làm được những điều mà chỉ vài năm trước người ta còn cho là không tưởng. So với https://xkcd.com/1425/, có lẽ giờ còn khó cả việc giải thích bối cảnh đằng sau trò đùa đó cho độc giả trẻ!