トップ › 技術ブログ › WebGPU で AI 超解像をブラウザの中で動かす
WebGPU で AI 超解像をブラウザの中で動かす ― 手箱イメージの waifu2x 実装
この記事は、当社(Gratias合同会社)が開発・運営している無料ツール「手箱(Tebako)」の画像ツールについて、開発チームが仕組みを解説するものです。自社のサービスの紹介を含みます。
手箱イメージには「AI で超解像度処理」という機能があります。小さい画像を AI で 2倍・4倍に引き伸ばしたり、大きさを変えずに輪郭を締めたりするものです。手箱では画像をサーバに送らずに、利用者の PC の GPU をブラウザから使って計算しています。
この記事では、使っている AI モデルとライセンス、ブラウザの中での実行のしかた、大きな画像を扱うためのタイル分割とメモリの見積もり、そして WebGPU が使えない PC での扱いを、開発の記録をもとに紹介します。手箱全体の「送らない」作りについては別の記事にまとめています。
使っているモデルとライセンス
拡大に使っているのは、nagadomi さんが公開している waifu2x(nunif)の学習済みモデルです。手箱では次のものを、配布物のまま改変せずに自社のサイトから配信しています。
- 写真向け(swin_unet/photo)とイラスト向け(swin_unet/art)の、それぞれ 2倍・4倍
- ノイズ除去の強さ違い(noise0〜3)と、2倍・4倍の組み合わせ
- タイルの継ぎ目をなじませる重みを作る小さなモデル
nunif はコードだけでなく重みまで MIT ライセンスで公開されています。設計の段階では写真向けに別のモデルも候補にしていましたが、重みの配布条件を個別に確かめる必要がありました。waifu2x なら写真・イラストの両方を同じ仕組みで扱え、ライセンスも明確なので、こちらを選びました。推論を動かすのは Microsoft の ONNX Runtime Web(MIT)で、これも改変せずに使っています。どちらもツールの「使用しているオープンソースソフトウェア」の表示に載せています。
ブラウザの中での実行のしかた
AI の処理は、画面とは別のスレッド(Web Worker)の中で丸ごと行います。画面のスレッドからは画像の画素と設定だけを渡し、Worker は進み具合を返しながら、モデルの取得・タイルの切り出し・推論・貼り合わせまでを済ませます。推論には ONNX Runtime Web の WebGPU 版を使い、計算は GPU で行います。
// 実行プロバイダは WebGPU だけ。
// ['webgpu','wasm'] と書くと、GPU が無いときに黙って CPU に落ちる
let EP = ['webgpu']
ort.env.wasm.numThreads = 1
ここには二つの判断があります。一つめは、GPU が使えないときに CPU での計算へ黙って切り替えないことです。CPU で動かすと極端に遅くなり、利用者には止まったように見えます。どれほど違うかは、後述の「WebGPU が使えない PC では」で数字を出します。
二つめは、ONNX Runtime の WebAssembly 部分を 1 スレッドで動かしていることです。複数のスレッドを使うには、ページを「クロスオリジン分離」の状態にする必要があります。そうすると、手箱の運営を支えている広告のスクリプトが読み込めなくなります。計算の本体は GPU なので、ここは 1 スレッドで割り切りました。
タイルに分けて推論し、継ぎ目を消す
大きな画像を一度にモデルへ渡すと、GPU の作業領域が足りなくなるおそれがあります。そこで画像を正方形のタイルに分けて順に推論し、結果を貼り合わせます。手順は nunif の Web 版の実装にならっています。
- 画像の外周を少し広げ、タイルの格子に合うまで右と下を埋める
- 重なりを持たせながらタイルを切り出し、1 枚ずつ推論する
- 出てきたタイルを、継ぎ目用の重みで加重平均しながら貼り合わせる
- 左上から目的の大きさを切り出す
外周の埋め方は、写真向けは端で折り返し、イラスト向けは端の画素を伸ばします。nunif はこれもモデルで行いますが、手箱では同じことを JavaScript で書き、GPU とのやり取りを 1 回減らし、取得するファイルも減らしました。イラスト向けでは、1 色だけのタイルは推論しても 1 色にしかならないので、推論を省いています。
継ぎ目は数字でも確かめています。出力の列ごとに隣の列との差を取り、境目の列だけが跳ねていないかを、分割しない方法で拡大した対照と比べました。差の比は境目の列で 1.30〜1.63、対照で 1.69 で、境目が目立つことはありませんでした。
この検査で見落とした不具合もあります。貼り合わせの重みを色のチャンネルごとに進めていたため、赤を書いた時点で重みが進み、緑と青が半分の明るさになっていました。継ぎ目の検査は列どうしの差しか見ないので、全体が一様に暗くなる誤りは素通りします。いまは RGB のすべてを分割しない結果と突き合わせ、ずれが 62 から 0.5 に下がったことを確かめています。
タイルの大きさは測って決めた
タイルは大きいほど速そうに思えますが、実際はそうなりませんでした。1000×1000 の画像を 4倍にする時間を、タイルの大きさを変えて当社の PC で測った結果の一部です(RTX 2060 と RTX 5060 で動作を確認しています)。
| タイル | 枚数 | 時間 | 1 画素あたり |
|---|---|---|---|
| 64 | 529 | 14.17 秒 | 6.6 µs |
| 160 | 64 | 4.71 秒 | 3.0 µs |
| 304 | 16 | 4.51 秒 | 3.1 µs |
| 400 | 9 | 6.68 秒 | 4.6 µs |
| 592 | 4 | 32.51 秒 | 23.2 µs |
1 画素あたりの時間は 112〜304 の範囲でほぼ平らで、400 で悪化し、544 以上で大きく崩れました。小さすぎると 1 枚ごとの呼び出しの手間が勝ち、大きすぎても遅くなります。最初に作った「作業領域が最小のものの 1.5 倍以内で、いちばん大きいタイル」という決め方は 592 を選び、いちばん速いタイルより 621% 遅くなっていました。GPU の無い環境(ソフトウェアで WebGPU をまねる仕組み)で測った数字から、GPU での動きを推し量ったのが誤りのもとでした。
いまは、画像ごとに「枚数 ×(1 枚あたりの手間 + タイルの画素数)」が最小になるタイルを、304 以下・メモリの上限の内側から選んでいます。1 枚あたりの手間は、上の実測から逆算して画素 5,000 個ぶんとしました。
メモリの上限を先に計算する
タブがメモリを使いすぎて落ちると、利用者には何が起きたか分かりません。そこで手箱では、始める前に貼り合わせに要るメモリを見積もり、上限を超えるときは実行せずに理由を伝えます。
見積もりは出力 1 画素あたり 28 バイトです。色の累積に 12 バイト、重みに 4 バイト、書き出す RGBA に 4 バイト、それに透過のある画像で不透明度を別に処理するときの山の分を足しています。透過があるかは読み込むまで分からないので、あるものとして見積もります。上限は 600MB にしました。4倍だと 1000px 角でおよそ 320MB、1400px 角では 630MB になり、後者は断って「2倍にするか、先に縮小する」よう案内します。
大きさを変えずに高画質化する「大きさはそのまま」の仕上がりでは、AI に 2倍で描かせてから元の大きさに戻し、最後に細部を立てます。この場合はスマートフォンの写真(1200 万画素)でも断らずに済むよう、画像を横の帯に分け、帯の上下に 24 画素ずつ余分に通してから切り落としています。分けない場合と比べた試験では、差は 255 階調で平均 0.063・最大 4 で、帯の継ぎ目の行だけが目立つことはありませんでした。
透過のある PNG は、色と不透明度を別々に同じモデルへ通します。透明な画素の色は多くの形式で黒なので、不透明な側の色を外へ広げてから通し、輪郭に黒がにじまないようにしています。
モデルの取得とキャッシュ
AI モデルは、利用者が「AI で超解像度処理」を選ぶまで 1 バイトも取りに行きません。ONNX Runtime を含む部分は別の塊にしたので、画像ツール本体の JavaScript は 38.7KB から 44.6KB に増えただけです。
初回に取得するのは、処理用のプログラム(約 6.4MB)と、選んだ種類・倍率のモデル 1 つ(16〜18MB)です。写真向けを選べばイラスト向けは取りに行きません。待機画面には取得の進み具合を出します。取得したモデルはブラウザの Cache Storage に入れ、2 回目以降は通信せずに使います。ブラウザの通常のキャッシュ任せにすると、「2 回目以降は待たない」と約束できないためです。
この取得については、画面に「通信が発生するのはこの超解像度処理AIモデルの取得だけで、あなたの画像は送信されません。」と明記しています。モデルは外部の配信サービスを使わず、手箱と同じサイトから配っています。
WebGPU が使えない PC では
WebGPU は Linux や一部のスマートフォンでは使えないことがあり、手箱ではこれを一定の割合で起きる前提として扱っています。
使えるかどうかは、アプリを開いたときではなく、AI を選んだその時に調べます。起動時に調べると、AI を使わない人の環境でも GPU に触ることになるからです。調べるのは、安全な接続(https)か、ブラウザに WebGPU の機能があるか、GPU のアダプタを取得できるか、の 3 点で、だめなときはどれに当たったかを画面に出し、「通常の拡大に切り替える」ボタンを置きます。通常の拡大は自前の Lanczos 法で、Worker の中で計算します。
CPU で推論する道を本番の形で測ったわけではありませんが、目安になる数字があります。GPU の無い試験環境で、CPU が GPU をまねる仕組み(SwiftShader)を通すと、タイル 1 枚の推論に 538 秒かかりました。実機の GPU では 1000×1000 の 4倍がまるごと 4.51 秒です。桁が違うので、黙って切り替えて「遅いけれど動く」とは言えないと判断しました。
まとめ
手箱イメージの AI 超解像度処理は、waifu2x のモデルを ONNX Runtime Web と WebGPU で動かし、画像を送らずに利用者の GPU で拡大しています。そのために、タイルの大きさを実測で選び、メモリを先に見積もって危ないときは断り、モデルは必要なときだけ取得して手元に残す、という形にしました。WebGPU が無い環境では理由を示して通常の拡大へ案内し、黙って遅い道に切り替えることはしません。
ブラウザの中の AI は、端末の性能やブラウザの対応に左右されます。それでも、画像を預けずに済むことには、利用者とこちらの双方に価値があると考えています。
手箱イメージ
写真の位置情報の削除、リサイズと形式の変換、書類の影の除去、AI による超解像度処理(2倍・4倍・大きさはそのままの高画質化)を、画像をアップロードせずにブラウザの中で行える無料ツールです。AI の処理には WebGPU に対応したブラウザ(新しい Chrome・Edge)が必要です。
手箱イメージを開く