Tăng tốc ứng dụng client-side AI với WebGPU

Hiện nay, các ứng dụng AI đã trở nên phổ biến và việc sử dụng cũng như tích hợp các mô hình thông minh vào các trang web đang trở thành xu hướng tất yếu. Trước đây, hầu hết các tác vụ AI nặng nề đều phải thực hiện ở server-side (máy chủ) hoặc thiết lập phức tạp trên máy người dùng. Tuy nhiên, với sự phát triển của các công nghệ xử lý phía client như WebGPU hoặc WebAssembly (WASM), cách tiếp cận này đang dần thay đổi.

Bài viết sẽ giới thiệu về các công nghệ WASM, WebGPU để tăng tốc xử lý tác vụ AI trên trình duyệt của người dùng (client-side). Ngoài ra, chúng ta cũng sẽ xây dựng một ứng dụng nhận diện chữ viết (OCR) đơn giản bằng WebGPU kết hợp với model Vision Language Model (VLM) thế hệ mới.

WebGPU và WebAssembly: Nền tảng cho AI trên trình duyệt

Trước đây, trong các mô hình máy khách - máy chủ truyền thống, các máy chủ sẽ đóng vai trò xử lý chính, máy khách chỉ đóng vai trò hiển thị kết quả. Tuy nhiên, với sự phát triển của phần cứng máy tính, và lượng người dùng truy cập internet ngày càng tăng, nhu cầu chia tải xử lý về phía máy khách ngày càng cao. Điều này dẫn đến sự ra đời của các công nghệ xử lý phía client như WebAssembly và WebGPU.

WebAssembly - Nâng cấp hiệu năng tính toán

WebAssembly (WASM) là một định dạng mã nhị phân cho phép chạy code viết bằng các ngôn ngữ như C, C++, Rust ngay trên trình duyệt web với tốc độ gần tương đương ứng dụng gốc. WASM xuất hiện phiên bản đầu tiên từ năm 2017, trở thành tiêu chuẩn W3C (World Wide Web Consortium) từ năm 2024 và đã được hỗ trợ rộng rãi trên tất cả các trình duyệt hiện đại.

Trước khi WebGPU xuất hiện, WASM gần như là phương pháp duy nhất để chạy các thuật toán AI phức tạp trên trình duyệt web. Các thư viện như TensorFlow.js hay ONNX Runtime đời đầu đều dựa nhiều vào WASM (kết hợp với SIMD) để thực hiện tính toán trên CPU.

Hiện nay, WASM đã có thể chạy ổn định trên hầu hết mọi thiết bị máy tính từ cũ đến mới, thậm chí trên các điện thoại giá rẻ mà không yêu cầu card đồ họa rời.

WebGPU - Tối ưu sức mạnh đồ họa và tính toán song song

WebGPU là chuẩn API đồ họa thế hệ mới được thúc đẩy phát triển bởi Google và Mozilla nhằm thay thế WebGL đã già cỗi với nhiều hạn chế. Khác với WebGL chỉ tập trung vào vẽ 3D, WebGPU được xây dựng từ đầu để hỗ trợ mạnh mẽ cho Compute Shaders cho phép tính toán song song trên GPU - điều kiện tiên quyết để tăng tốc xử lý trên GPU.

WebGPU cho phép trình duyệt truy cập trực tiếp và hiệu quả vào GPU của thiết bị. Đặc biệt, nó hỗ trợ tính toán số thực 16-bit (fp16), giúp giảm một nửa băng thông bộ nhớ và tăng tốc độ suy luận (inference) của model AI lên đáng kể so với WebGL hay WASM.

Chính thức ra mắt vào năm 2021, WebGPU tuy còn khá mới mẻ nhưng đã chứng tỏ sức mạnh của nó với các ứng dụng AI và trò chơi 3D. Hiện nay các trình duyệt phổ biến đều đã hỗ trợ WebGPU, bao gồm Chrome, Firefox, Edge và Safari.

Ứng dụng WebGPU trong kỹ thuật xử lý hình ảnh

Với sự ra đời của các công nghệ WebGPU và WASM, cách xử lý các bài toán phức tạp về thị giác máy tính như phân loại hình ảnh hoặc nhận diện ký tự quang học (OCR) đã thay đổi đáng kể.

Trong thời gian dài, thư viện Tesseract được xem là tiêu chuẩn vàng cho bài toán OCR. Tesseract hoạt động dựa trên phương pháp nhận diện mẫu (Pattern Matching) và Machine Learning truyền thống, nhận diện tốt các ký tự đơn giản, rõ nét. Tuy nhiên, Tesseract hoạt động kém hiệu quả với một số loại hình ảnh như ảnh chụp thực tế, chữ viết tay ngoằn ngoèo, nền nhiễu hoặc bố cục phức tạp, ... dẫn đến độ chính xác bị giảm đáng kể cũng như thiếu hoàn toàn ngữ cảnh của hình ảnh.

Sự xuất hiện của các mô hình đa phương thức (Vision Language Models - VLM) như SmolVLM, Florence-2 hoặc Qwen3-VL đã thay đổi hoàn toàn cách tiếp cận này. Nhờ khả năng xử lý đầu cuối (end-to-end) và hiểu ngữ cảnh sâu sắc, các mô hình này hoạt động hiệu quả trên cả các bố cục phức tạp, đồng thời mở rộng khả năng sang việc mô tả và phân tích nội dung hình ảnh thay vì chỉ trích xuất văn bản thuần túy.

Xem thêm bài viết đánh giá về các mô hình AI nhỏ tại đây.

Trong phần tiếp theo, chúng ta sẽ kết hợp sức mạnh của WebGPU và mô hình SmolVLM để xây dựng một ứng dụng OCR đơn giản nhưng không kém phần mạnh mẽ.

Demo: Xây dựng ứng dụng OCR đơn giản với SmolVLM

Đầu tiên, ứng dụng OCR của chúng ta sẽ cần chuẩn bị thư mục webgpu-ocr với 2 file mã nguồn đơn giản:

webgpu-ocr/
├── webgpu-ocr-demo.html
└── webgpu-ocr-demo.js

Giao diện (HTML)

Đầu tiên, các bạn tạo file webgpu-ocr-demo.html để upload và xử lý ảnh, với thư viện Bootstrap thân thiện.

File: webgpu-ocr-demo.html

<!DOCTYPE html>
<html lang="en">

<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Simple AI OCR (SmolVLM)</title>
    <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.0/dist/css/bootstrap.min.css" rel="stylesheet">
    <link href="https://cdn.jsdelivr.net/npm/bootstrap-icons@1.7.2/font/bootstrap-icons.css" rel="stylesheet">
</head>

<body>
    <div class="container py-5">
        <div class="text-center mb-5">
            <h1 class="mb-2"><i class="bi bi-robot me-2"></i>Simple AI OCR</h1>
            <p class="text-secondary">Client-side OCR with SmolVLM & WebGPU</p>
        </div>

        <div id="status" class="alert alert-info text-center fw-semibold mx-auto" style="max-width: 800px;">
            <i class="bi bi-gear-fill me-2 animate-spin"></i>Initializing...
        </div>

        <div class="row justify-content-center">
            <div class="col-lg-8">
                <!-- Controls -->
                <div class="card shadow-sm mb-4">
                    <div class="card-body p-4 text-center">
                        <label class="btn btn-primary btn-lg px-5">
                            <i class="bi bi-upload me-2"></i> Select Image
                            <input type="file" id="file-input" accept="image/*" class="d-none" disabled>
                        </label>

                    </div>
                </div>

                <!-- Result Area -->
                <div class="row">
                    <div class="col-md-6 mb-4">
                        <div class="card shadow-sm h-100">
                            <div class="card-header bg-white">
                                <h5 class="mb-0">Image Preview</h5>
                            </div>
                            <div class="card-body p-0">
                                <div id="preview-container" class="h-100 p-3">
                                    <span class="placeholder-text">No image selected</span>
                                </div>
                            </div>
                        </div>
                    </div>
                    <div class="col-md-6 mb-4">
                        <div class="card shadow-sm h-100">
                            <div class="card-header bg-white">
                                <h5 class="mb-0">Extracted Text</h5>
                            </div>
                            <div class="card-body p-0">
                                <textarea id="result" class="form-control border-0 h-100 p-3"
                                    style="resize: none; min-height: 300px;"
                                    placeholder="Text will appear here..."></textarea>
                            </div>
                        </div>
                    </div>
                </div>
            </div>
        </div>
    </div>
    <script type="module" src="webgpu-ocr-demo.js"></script>

Sau khi tạo xong, giao diện của ứng dụng sẽ như sau:

Logic xử lý (JavaScript)

Tiếp theo, chúng ta viết toàn bộ logic xử lý ở webgpu-ocr-demo.js. Ở đây, chúng ta sẽ sử dụng thư viện Transformers.js để tải và chạy model SmolVLM.

File: webgpu-ocr-demo.js

import { AutoModelForVision2Seq, AutoProcessor, RawImage, env } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers';

// Cấu hình môi trường
env.allowLocalModels = false;
env.useBrowserCache = true;

const MODEL_ID = 'HuggingFaceTB/SmolVLM-256M-Instruct';

let model = null;
let processor = null;

const statusEl = document.getElementById('status');
const fileInput = document.getElementById('file-input');
const previewContainer = document.getElementById('preview-container');
const resultTextarea = document.getElementById('result');

// Hàm hiển thị trạng thái
function setStatus(message, type = 'info') {
    statusEl.innerHTML = message;
    statusEl.className = `alert alert-${type} text-center fw-semibold mx-auto`;
    statusEl.style.maxWidth = '800px';
}

// 1. Khởi tạo Model
async function init() {
    try {
        setStatus('<i class="bi bi-download me-2"></i>Loading SmolVLM model...', 'info');

        const device = 'webgpu';
        processor = await AutoProcessor.from_pretrained(MODEL_ID);

        model = await AutoModelForVision2Seq.from_pretrained(MODEL_ID, {
            device: device,
            dtype: {
                embed_tokens: 'fp32',
                vision_encoder: 'fp32',
                encoder_model: 'q4', // Quantization 4-bit để tối ưu bộ nhớ
                decoder_model_merged: 'q4',
            },
        });

        setStatus('<i class="bi bi-check-circle-fill me-2"></i>Ready! Select an image.', 'success');
        document.getElementById('file-input').disabled = false;

    } catch (e) {
        console.error(e);
        setStatus(`<i class="bi bi-exclamation-triangle-fill me-2"></i>Error: ${e.message}`, 'danger');
    }
}

// 2. Xử lý OCR
async function runOCR(imageUrl) {
    if (!model || !processor) return;

    try {
        setStatus('<i class="bi bi-cpu me-2 animate-spin"></i>Processing...', 'warning');
        resultTextarea.value = '';

        const image = await RawImage.fromURL(imageUrl);

        // Chuẩn bị hội thoại cho SmolVLM
        const messages = [
            {
                role: "user",
                content: [
                    { type: "image" },
                    { type: "text", text: "Extract all text from this image." }
                ]
            }
        ];

        // Định dạng đầu vào
        const text_inputs = processor.apply_chat_template(messages, { render_bos_token: false });
        const inputs = await processor(text_inputs, [image]);

        // Sinh kết quả
        const outputs = await model.generate({
            ...inputs,
            max_new_tokens: 512,
            do_sample: false,
            repetition_penalty: 1.1,
        });

        // Giải mã kết quả
        const generatedFullText = processor.decode(outputs[0], { skip_special_tokens: true });
        const promptText = processor.decode(inputs.input_ids[0], { skip_special_tokens: true });

        const cleanText = generatedFullText.replace(promptText, '').replace(/^A:\s*/, '').trim();
        resultTextarea.value = cleanText;
        setStatus('<i class="bi bi-check-circle-fill me-2"></i>Done!', 'success');
    } catch (e) {
        console.error(e);
        setStatus(`Error: ${e.message}`, 'danger');
    }
}

// Lắng nghe sự kiện chọn file
fileInput.addEventListener('change', (e) => {
    const file = e.target.files[0];
    if (!file) return;

    const reader = new FileReader();
    reader.onload = (event) => {
        const url = event.target.result;
        previewContainer.innerHTML = `<img id="preview" src="${url}" style="max-width:100%; max-height:500px">`;
        runOCR(url);
    };
    reader.readAsDataURL(file);
});

// Bắt đầu khởi tạo
init();

Hướng dẫn chạy

Do cơ chế bảo mật CORS và ES Modules của trình duyệt, bạn sẽ không thể mở trực tiếp file HTML trên trình duyệt. Bạn sẽ cần một Local Server để chạy.

Để tạo một máy chủ web đơn giản, bạn có thể chạy lệnh Python sau từ Terminal của thư mục dự án:

python -m http.server 8080

Sau đó truy cập: http://localhost:8080/webgpu-ocr-demo.html

Ngoài ra, nếu có sẵn Node.js, các bạn có thể chạy lệnh sau: npx serve -l 8080

Kết quả và tiềm năng mở rộng

Khi chạy lần đầu, trình duyệt sẽ tải dữ liệu model và lưu cache lại, những lần sau sẽ khởi động tức thì. Kết quả nhận diện của SmolVLM thường rất ấn tượng, vượt trội hoàn toàn so với các giải pháp cũ.

Kết quả demo OCR

Các use case thực tế

Với nền tảng WebGPU và VLM đã xây dựng, bạn có thể dễ dàng mở rộng cho các bài toán thú vị khác chỉ bằng cách đổi câu lệnh prompt (trong code):

• Đọc chữ viết tay: Thử nghiệm với các ghi chú ngoằn ngoèo, bản thảo cũ khó đọc

• Trích xuất dữ liệu (Data Extraction): Đưa vào hình ảnh hóa đơn và yêu cầu trích xuất dữ liệu

• Mô tả hình ảnh (Image Captioning): Thay vì yêu cầu "Extract all text", hãy thử "Describe this image in detail". Ứng dụng sẽ trở thành công cụ mô tả hình ảnh hỗ trợ người khiếm thị.

Tip: Nếu máy bạn không hỗ trợ WebGPU, bạn có thể đổi dòng device: 'webgpu' thành device: 'wasm' trong file JS. Khi đó model sẽ chạy bằng CPU (chậm hơn đáng kể nhưng tương thích mọi thiết bị).

Thử nghiệm với các model khác

Ngoài SmolVLM, bạn cũng có thể thử nghiệm với các model khác bằng cách thay đổi biến MODEL_ID trong file webgpu-ocr-demo.js. Dưới đây là một số model tương thích tốt:

• Florence-2: onnx-community/Florence-2-base-ft (Nhẹ hơn, chuyên dụng cho OCR/Captioning)
• Qwen2-VL: onnx-community/Qwen2-VL-2B-Instruct (Mô hình 2B tham số, mạnh mẽ hơn nhưng yêu cầu phần cứng cao hơn)

// Ví dụ chuyển sang Florence-2
const MODEL_ID = 'onnx-community/Florence-2-base-ft';

Ngoài ra, bạn cũng có thể trực tiếp xem mã nguồn của dự án tại Github và trang demo trực tiếp.

Lời kết

Việc đưa AI xuống client-side với WebGPU không chỉ giúp giảm chi phí server mà còn nâng cao quyền riêng tư dữ liệu cho người dùng. Sự chuyển dịch từ các thuật toán cổ điển (như Tesseract) sang các mô hình Generative AI (như SmolVLM) đánh dấu một bước tiến lớn về chất lượng và khả năng hiểu của máy tính.

Hy vọng bài viết này giúp bạn có cái nhìn tổng quan và bước đầu làm chủ được công nghệ thú vị này!

Tài liệu tham khảo

• Transformers.js Documentation
• WebGPU API trên MDN
• Lịch sử WebAssembly
• SmolVLM Model Card