1
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@kigi316782

ブラウザだけで Stable Diffusion を動かす ── WebGPU × ONNX Runtime Web × 自作パイプラインで実現したローカル AI 画像生成ツール

1
Posted at

目次

  1. はじめに
  2. アーキテクチャ全体像
  3. WebGPU とは？ ONNX Runtime Web との関係
  4. Stable Diffusion パイプラインを TypeScript で自作した話
  5. LoRA マージをブラウザ内で実行する
  6. Worker 通信の型安全な設計
  7. OPFS でモデルをキャッシュする
  8. テスト戦略: ブラウザ API なしで 415 件のテスト
  9. AdSense + WebGPU の両立: COEP の罠
  10. デプロイ: CloudFormation + S3 + CloudFront
  11. まとめと今後の展望
  12. 参考リンク

はじめに

**「ブラウザだけで Stable Diffusion を動かす」**── この言葉を聞いてどう感じるでしょうか。

通常、Stable Diffusion を使った画像生成には Python 環境の構築、CUDA ドライバのインストール、ComfyUI や AUTOMATIC1111 のセットアップなど、少なくない準備が必要です。クラウド API を使えば手軽ですが、プロンプトや生成画像がサーバーに送信されることになります。

sd.keydrop.net は、これらの課題を解決するために開発したブラウザ完結型の AI 画像生成ツールです。

  • サーバーコストゼロ: 推論は全てユーザーの GPU 上で実行。静的ファイルのホスティングだけで運用可能
  • プライバシー完全保護: モデルファイル・プロンプト・生成画像はブラウザの外に出ない
  • インストール不要: Chrome を開いてモデルをドロップすれば即座に生成開始

ユーザーは自身が持つ ONNX 形式の SD モデル（SD-Turbo 等）をブラウザにアップロードし、テキストプロンプトから画像を生成できます。さらにブラウザ内で LoRA のマージまでできます。この記事では、このツールの技術的な実装詳細を深掘りします。

本ツールは「ユーザー持ち込み型」です。運営側はモデルを一切配布しません。モデルは Hugging Face 等から ONNX 形式のものをユーザー自身がダウンロードしてアップロードする方式です。

アーキテクチャ全体像

全体は以下のようなレイヤ構造になっています。

┌─────────────────────────────────────────────────────────────────┐
│                        Browser (Main Thread)                   │
│                                                                 │
│  ┌──────────────────────────────────────────────────────────┐  │
│  │                       UI Layer                           │  │
│  │  main.ts → header / generate-view / merge-view /         │  │
│  │             settings-view / about-view                   │  │
│  └──────────────────────────────────────────────────────────┘  │
│                          ▲            │                        │
│                          │            ▼                        │
│  ┌──────────────────────────────────────────────────────────┐  │
│  │                   Worker Client Layer                   │  │
│  │  workers/inference-client.ts  workers/merge-client.ts   │  │
│  └──────────────────────────────────────────────────────────┘  │
│                              │ postMessage                     │
└──────────────────────────────┼─────────────────────────────────┘
                               │
                    ┌──────────┴──────────┐
                    ▼                     ▼
           ┌─────────────────┐   ┌─────────────────┐
           │ Inference Worker│   │   Merge Worker  │
           │  (web worker)   │   │  (web worker)   │
           └─────────────────┘   └─────────────────┘
                    │                     │
                    ▼                     ▼
           ┌─────────────────┐   ┌─────────────────┐
           │ ONNX Runtime Web│   │ safetensors +   │
           │  + WebGPU EP    │   │ Pure TS matmul  │
           │  + SD pipeline  │   │                 │
           └─────────────────┘   └─────────────────┘
                    │
                    ▼
           ┌─────────────────┐
           │  GPU (WebGPU)   │
           └─────────────────┘

設計上の判断ポイント

なぜ React / Vue を使わないのか？

推論パイプライン自体が重い処理であり、フレームワークの仮想 DOM オーバーヘッドは不要と判断しました。Vanilla TypeScript + Tailwind CSS v4 の組み合わせで、バンドルサイズを最小に保ちつつ、strict モードの型安全性を確保しています。

なぜ Web Worker？

ONNX Runtime Web の InferenceSession.run() は非同期ですが、WebGPU のコマンドバッファ構築はメインスレッドをブロックする可能性があります。Worker に分離することで UI のフリーズを防ぎます。

なぜエンジン抽象化？

InferenceEngine というインターフェースを定義し、ONNX Runtime Web / Transformers.js / Ratchet の各バックエンドを差し替え可能にしています。現時点では ONNX Runtime Web のみ実装済みですが、将来の拡張に備えた設計です。

WebGPU とは？ ONNX Runtime Web との関係

WebGL との違い

WebGL はグラフィックス描画用の API であり、汎用計算（GPGPU）には不向きでした。WebGPU は Vulkan / Metal / DirectX 12 を抽象化した新しい API で、コンピュートシェーダによる汎用 GPU 計算を第一級でサポートしています。

ONNX Runtime Web は、この WebGPU のコンピュートシェーダを使ってニューラルネットワークの演算をGPU上で実行します。テンソルの行列積や畳み込みなどの重い演算がブラウザ内で GPU アクセラレーションされるわけです。

SharedArrayBuffer と COEP の要件

ONNX Runtime Web の WebGPU EP は内部で SharedArrayBuffer を使用します。これを有効にするには、以下のレスポンスヘッダが必須です。

Cross-Origin-Opener-Policy: same-origin
Cross-Origin-Embedder-Policy: credentialless

credentialless と require-corp の選択は重要です。require-corp を選ぶとサードパーティスクリプト（AdSense 等）がブロックされます。credentialless なら SharedArrayBuffer を有効にしつつ、広告も読み込めます。この罠については後述の「COEP の罠」セクションで詳しく解説します。

ONNX Runtime Web の構成

ONNX Runtime Web の WASM パスとスレッド数は Worker 内で明示的に設定しています。

export function configureOrtWasm(ort: OrtModule): void {
  try {
    const base =
      typeof self !== 'undefined' && self.location
        ? new URL('/ort-wasm/', self.location.origin).toString()
        : '/ort-wasm/';
    (ort.env.wasm as unknown as { wasmPaths: string }).wasmPaths = base;
    ort.env.wasm.numThreads = 1;
    (ort.env as unknown as { logLevel: string }).logLevel = 'error';
  } catch (e) {
    console.warn('[shared] failed to configure wasm env', e);
  }
}

ONNX セッションの作成時には WebGPU EP を指定します。

const sessionOptions: InferenceSession.SessionOptions = {
  executionProviders: ['webgpu'],
  graphOptimizationLevel: 'all',
  enableMemPattern: false,
};

Stable Diffusion パイプラインを TypeScript で自作した話

Stable Diffusion の推論パイプラインは通常 Python（Hugging Face diffusers）で実装されていますが、ブラウザで動かすためにすべて TypeScript で再実装しました。パイプラインは以下の流れで動作します。

プロンプト → CLIP Tokenizer → Text Encoder → UNet ループ → VAE Decoder → PNG

パイプラインの run() メソッド

全体のフローを見てみましょう。

async run(
  params: RunParams,
  onProgress?: (p: PipelineProgress) => void,
  signal?: AbortSignal,
): Promise<Blob> {
  const { width, height, steps, seed, prompt, negativePrompt = '' } = params;
  const cfgScale = this.config.useCfg ? params.cfgScale : 0;
  signal?.throwIfAborted();

  // 1) トークナイズ
  const condIds = this.tokenizer.encode(prompt);
  const uncondIds = cfgScale > 1 ? this.tokenizer.encode(negativePrompt) : null;

  // 2) テキストエンコーダ
  const condEmbed = await this.runTextEncoder(condIds);
  const uncondEmbed = uncondIds ? await this.runTextEncoder(uncondIds) : null;
  signal?.throwIfAborted();

  // 3) latent 初期化
  const latentH = Math.floor(height / this.config.vaeScaleFactor);
  const latentW = Math.floor(width / this.config.vaeScaleFactor);
  const latentSize = this.config.latentChannels * latentH * latentW;
  const rng = new SeededRandom(seed);
  let latent = rng.gaussianArray(latentSize);

  // 4) scheduler
  const scheduler = new EulerAncestralDiscreteScheduler({
    numInferenceSteps: steps,
    predictionType: this.config.predictionType,
  });
  latent = scheduler.scaleInitialLatent(latent);

  // 5) 拡散ループ
  const totalSteps = scheduler.timesteps.length;
  for (let i = 0; i < totalSteps; i++) {
    signal?.throwIfAborted();
    onProgress?.({ step: i + 1, totalSteps });

    const timestep = scheduler.timesteps[i];
    const scaled = scheduler.scaleModelInput(latent, i);

    const noisePredCond = await this.runUnet(scaled, timestep, condEmbed, latentH, latentW);
    let noisePred = noisePredCond;
    if (cfgScale > 1 && uncondEmbed) {
      const noisePredUncond = await this.runUnet(
        scaled, timestep, uncondEmbed, latentH, latentW,
      );
      noisePred = new Float32Array(noisePredCond.length);
      for (let k = 0; k < noisePred.length; k++) {
        noisePred[k] = noisePredUncond[k] + cfgScale * (noisePredCond[k] - noisePredUncond[k]);
      }
    }

    latent = scheduler.step(noisePred, i, latent, rng);
  }

  // 6) VAE scaling を外す
  const vaeInput = new Float32Array(latent.length);
  const invScale = 1 / this.config.vaeScalingFactor;
  for (let k = 0; k < latent.length; k++) vaeInput[k] = latent[k] * invScale;

  // 7) VAE decode
  signal?.throwIfAborted();
  const image = await this.runVaeDecoder(vaeInput, latentH, latentW);

  // 8) テンソル → PNG
  return await tensorToPng(image, { width, height });
}

SD-Turbo（v_prediction、CFG なし、1〜4 ステップ）と SD 1.5（epsilon、CFG あり）の両方に対応しています。PipelineConfig でモデルの特性を切り替えます。

export interface PipelineConfig {
  latentChannels: number;       // SD は 4
  vaeScaleFactor: number;       // SD は 8
  vaeScalingFactor: number;     // SD-Turbo: 0.13025, SD1.5: 0.18215
  dtype: 'float16' | 'float32';
  predictionType: 'epsilon' | 'v_prediction';
  useCfg: boolean;              // SD-Turbo は false
}

CLIP BPE Tokenizer ── なぜ自作したか

当初は Transformers.js の利用を検討しましたが、v4 が Stable Diffusion の CLIP Tokenizer に対応していなかったため、OpenAI CLIP の Python 実装を参考に TypeScript で自作しました。

まず、BPE トークナイザの要となるのが bytesToUnicode マップです。これは UTF-8 バイト列を可視化可能な Unicode 文字に変換するための OpenAI 独自のハックです。

function bytesToUnicode(): Map<number, string> {
  const bs: number[] = [];
  for (let b = 33; b <= 126; b++) bs.push(b);
  for (let b = 161; b <= 172; b++) bs.push(b);
  for (let b = 174; b <= 255; b++) bs.push(b);
  const cs = bs.slice();
  let n = 0;
  for (let b = 0; b < 256; b++) {
    if (!bs.includes(b)) {
      bs.push(b);
      cs.push(256 + n);
      n++;
    }
  }
  const map = new Map<number, string>();
  for (let i = 0; i < bs.length; i++) {
    map.set(bs[i], String.fromCodePoint(cs[i]));
  }
  return map;
}

BPE アルゴリズム本体は、1つの単語を受け取り、マージルールに基づいてサブワードに分割します。

private bpe(token: string): string[] {
  if (this.cache.has(token)) {
    return this.cache.get(token)!.split(' ');
  }

  // 末尾に </w> を付けて単語境界を表現
  const chars = Array.from(token);
  if (chars.length === 0) return [];
  chars[chars.length - 1] = chars[chars.length - 1] + '</w>';

  let word = chars;
  let pairs = getPairs(word);
  if (pairs.size === 0) {
    this.cache.set(token, word.join(' '));
    return word;
  }

  while (true) {
    // 最もランクが若いペアを探す
    let bestPair: string | null = null;
    let bestRank = Infinity;
    for (const pair of pairs) {
      const rank = this.bpeRanks.get(pair);
      if (rank !== undefined && rank < bestRank) {
        bestRank = rank;
        bestPair = pair;
      }
    }
    if (bestPair === null) break;

    const [a, b] = bestPair.split(' ');
    const newWord: string[] = [];
    let i = 0;
    while (i < word.length) {
      const j = word.indexOf(a, i);
      if (j === -1) {
        newWord.push(...word.slice(i));
        break;
      }
      newWord.push(...word.slice(i, j));
      if (j < word.length - 1 && word[j + 1] === b) {
        newWord.push(a + b);
        i = j + 2;
      } else {
        newWord.push(word[j]);
        i = j + 1;
      }
    }
    word = newWord;
    if (word.length === 1) break;
    pairs = getPairs(word);
  }

  const result = word.join(' ');
  this.cache.set(token, result);
  return word;
}

encode() メソッドでは、CLIP の正規表現パターンで pre-tokenize してから各トークンに BPE を適用し、SOS/EOS トークンとパディングを付けて長さ 77 の固定長配列を返します。

Euler Ancestral Discrete Scheduler

diffusers の Python 実装を TypeScript に移植しました。スケジューラの核心は step() メソッドです。

step(
  modelOutput: Float32Array,
  stepIndex: number,
  sample: Float32Array,
  rng: SeededRandom,
): Float32Array {
  const sigma = this.sigmas[stepIndex];
  const sigmaNext = this.sigmas[stepIndex + 1];

  // 1) epsilon / v_prediction → pred_original_sample (x0)
  const predOriginal = new Float32Array(sample.length);
  if (this.predictionType === 'epsilon') {
    for (let i = 0; i < sample.length; i++) {
      predOriginal[i] = sample[i] - sigma * modelOutput[i];
    }
  } else {
    // v_prediction: pred_x0 = sample / (σ²+1) - σ * v / √(σ²+1)
    const sig2p1 = sigma * sigma + 1;
    const invSq = 1 / sig2p1;
    const invS = 1 / Math.sqrt(sig2p1);
    for (let i = 0; i < sample.length; i++) {
      predOriginal[i] = sample[i] * invSq - sigma * modelOutput[i] * invS;
    }
  }

  // 2) ancestral sigma 計算
  const sigmaUpSq = Math.max(
    (sigmaNext * sigmaNext * (sigma * sigma - sigmaNext * sigmaNext)) /
      Math.max(sigma * sigma, 1e-12),
    0,
  );
  const sigmaUp = Math.sqrt(sigmaUpSq);
  const sigmaDown = Math.sqrt(Math.max(sigmaNext * sigmaNext - sigmaUp * sigmaUp, 0));

  // 3) 勾配
  const dt = sigmaDown - sigma;
  const prevSample = new Float32Array(sample.length);
  for (let i = 0; i < sample.length; i++) {
    const derivative = (sample[i] - predOriginal[i]) / sigma;
    prevSample[i] = sample[i] + derivative * dt;
  }

  // 4) ancestral ノイズを加算
  if (sigmaNext > 0 && sigmaUp > 0) {
    for (let i = 0; i < prevSample.length; i++) {
      prevSample[i] += rng.gaussian() * sigmaUp;
    }
  }
  return prevSample;
}

epsilon 予測と v_prediction の両方に対応しています。SD-Turbo は v_prediction、SD 1.5 は epsilon を使います。この分岐はスケジューラの中で吸収されるため、パイプライン側は意識する必要がありません。

fp16 変換 ── IEEE 754 半精度の実装

SD-Turbo の ONNX モデルは fp16 量子化されています。ONNX Runtime Web は float16 テンソルを Uint16Array で受け渡しする仕様のため、JavaScript 側で Float32Array との相互変換が必要です。

// f32 → u16 変換用の共有バッファ
const f32Buf = new ArrayBuffer(4);
const f32View = new Float32Array(f32Buf);
const u32View = new Uint32Array(f32Buf);

export function f32ToF16(value: number): number {
  f32View[0] = value;
  const x = u32View[0];

  const sign = (x >>> 16) & 0x8000;
  let exp = ((x >>> 23) & 0xff) - (127 - 15);
  const mant = x & 0x7fffff;

  if (exp >= 0x1f) {
    // Infinity / NaN
    return sign | 0x7c00 | (mant ? 0x0200 : 0);
  }
  if (exp <= 0) {
    if (exp < -10) {
      return sign; // 極小値はゼロに
    }
    // サブノーマル
    const mantWithLead = mant | 0x800000;
    const shift = 14 - exp;
    const m = mantWithLead >>> shift;
    const round = ((mantWithLead >>> (shift - 1)) & 1) &&
      (exp < 0 || (mantWithLead & ((1 << (shift - 1)) - 1)) !== 0);
    return sign | m | (round ? 1 : 0);
  }
  // 通常範囲
  let f16 = sign | (exp << 10) | (mant >>> 13);
  if (mant & 0x1000) {
    f16 += 1; // 四捨五入
  }
  return f16 & 0xffff;
}

SharedArrayBuffer を使ったビット操作による変換は、IEEE 754 の仕様を直接扱うため理解しやすく、外部ライブラリへの依存も不要です。

セッション入力型の自動判別

モデルによって入力テンソルの dtype が異なります（fp16 / fp32 / int32 / int64）。ONNX セッションの inputMetadata からモデルが期待する型を自動判定して、適切な型のテンソルを生成します。

type InputTypeMap = Record<string, Tensor.Type>;

function buildInputTypeMap(session: InferenceSession): InputTypeMap {
  const map: InputTypeMap = {};
  const meta = session.inputMetadata as readonly {
    name: string;
    isTensor: boolean;
    type?: Tensor.Type;
  }[];
  for (const m of meta) {
    if (m.isTensor && m.type) map[m.name] = m.type;
  }
  return map;
}

これにより、fp16 モデルでも fp32 モデルでも同じパイプラインコードで動作します。

LoRA マージをブラウザ内で実行する

LoRA（Low-Rank Adaptation）のマージ機能もブラウザ内で完結します。Pure TypeScript で safetensors フォーマットの読み書きと行列積を実装しています。

safetensors フォーマットの解説

safetensors はシンプルなバイナリフォーマットです。

[8 bytes: header length (LE u64)][header JSON][binary data...]

ヘッダ読み込みの実装は以下の通りです。ファイル全体をメモリに展開せず、必要な部分だけを Blob.slice() で読み取ります。

export async function readHeader(file: Blob): Promise<SafetensorsMetadata> {
  const headerLenBuf = await file.slice(0, 8).arrayBuffer();
  const view = new DataView(headerLenBuf);
  const headerLen = Number(view.getBigUint64(0, true));
  if (headerLen <= 0 || headerLen > 100 * 1024 * 1024) {
    throw new Error(`Invalid safetensors header length: ${headerLen}`);
  }
  const headerBytes = await file.slice(8, 8 + headerLen).arrayBuffer();
  const text = new TextDecoder().decode(headerBytes);
  let header: Record<string, unknown>;
  try {
    header = JSON.parse(text) as Record<string, unknown>;
  } catch (e) {
    throw new Error(`Invalid safetensors header JSON: ${(e as Error).message}`);
  }

  const metadata = (header.__metadata__ ?? {}) as Record<string, string>;
  const tensors: TensorDescriptor[] = [];
  for (const [key, rawDef] of Object.entries(header)) {
    if (key === '__metadata__') continue;
    const def = rawDef as { dtype?: string; shape?: number[]; data_offsets?: [number, number] };
    if (!def.dtype || !def.shape || !def.data_offsets) continue;
    tensors.push({
      name: key,
      dtype: def.dtype as SafetensorsDtype,
      shape: def.shape,
      dataOffsets: [def.data_offsets[0], def.data_offsets[1]],
    });
  }

  return { metadata, tensors, headerTotalBytes: 8 + headerLen };
}

Pure TypeScript での行列積

LoRA マージの核心は applyLoraDelta 関数です。LoRA の up / down 行列の積を計算し、スケーリングしてベースモデルの重みに加算します。

export function applyLoraDelta(
  baseF32: FloatLike,
  baseShape: number[],
  down: FloatLike,
  downShape: number[],
  up: FloatLike,
  upShape: number[],
  scale: number,
): void {
  // up @ down → [out, in_total]
  const outDim = upShape[0];
  const rank = upShape[1];
  if (downShape[0] !== rank) {
    throw new Error(`rank mismatch: down[${downShape}] vs up[${upShape}]`);
  }
  const inTotal = downShape.slice(1).reduce((a, b) => a * b, 1);
  const baseTotal = baseShape.reduce((a, b) => a * b, 1);
  if (baseTotal !== outDim * inTotal) {
    throw new Error(
      `base tensor (${baseShape.join('x')}) shape mismatch with up/down (${outDim}x${inTotal})`,
    );
  }

  // matmul: delta[o, i] = sum_r up[o, r] * down[r, i]
  for (let o = 0; o < outDim; o++) {
    const upRow = o * rank;
    const baseRow = o * inTotal;
    for (let i = 0; i < inTotal; i++) {
      let acc = 0;
      const downCol = i;
      for (let r = 0; r < rank; r++) {
        acc += up[upRow + r] * down[r * inTotal + downCol];
      }
      baseF32[baseRow + i] += scale * acc;
    }
  }
}

マージ式は merged = base + (alpha / rank) * weight * (up @ down) です。Kohya 形式の LoRA に対応しており、テンソル単位で逐次処理することで OOM を回避しています。

Pure TypeScript での行列積は GPU に比べると遅いですが、ブラウザ内で完結するメリットは大きいです。将来的には WebGPU コンピュートシェーダで高速化する計画があります。

Worker 通信の型安全な設計

推論 Worker とメインスレッド間の通信には、TypeScript の discriminated union を使った型付きメッセージを採用しています。

export type MainToWorkerMessage =
  | { type: 'init'; requestId: string; engineId: EngineId }
  | { type: 'load'; requestId: string; modelFiles: ModelFiles; hubModelId?: string }
  | { type: 'generate'; requestId: string; params: GenerationParams }
  | { type: 'cancel'; requestId: string }
  | { type: 'dispose'; requestId: string };

export type WorkerToMainMessage =
  | { type: 'ready'; requestId: string }
  | { type: 'loaded'; requestId: string; modelName: string }
  | {
      type: 'generated';
      requestId: string;
      pngBuffer: ArrayBuffer;  // transferable
      width: number;
      height: number;
    }
  | { type: 'disposed'; requestId: string }
  | { type: 'canceled'; requestId: string }
  | { type: 'error'; requestId: string; message: string; name?: string }
  | {
      type: 'progress';
      requestId: string;
      phase: 'loading' | 'compiling' | 'sampling' | 'decoding';
      step?: number;
      totalSteps?: number;
      ratio?: number;   // 0..1
      message?: string;
    };

すべてのメッセージは requestId を持ち、メインスレッド側のクライアントラッパが Promise の解決キーとして使います。AbortSignal にも対応しており、生成途中のキャンセルが可能です。

type フィールドによる discriminated union のおかげで、Worker 側の onmessage ハンドラでは switch (msg.type) で型安全にディスパッチできます。progress メッセージには phase と ratio が含まれるため、UI 側でフェーズ別のプログレスバーを表示できます。

OPFS でモデルをキャッシュする

SD モデルは数 GB あり、毎回アップロードするのは非現実的です。Origin Private File System（OPFS）を使ってブラウザ内に永続化しています。

const ROOT_DIR = 'models';

async function getRoot(): Promise<FileSystemDirectoryHandle> {
  if (!('storage' in navigator) || typeof navigator.storage.getDirectory !== 'function') {
    throw new Error('OPFS is not supported in this browser');
  }
  const root = await navigator.storage.getDirectory();
  return root.getDirectoryHandle(ROOT_DIR, { create: true });
}

export async function saveFile(name: string, file: File | Blob): Promise<void> {
  const dir = await getRoot();
  const handle = await dir.getFileHandle(name, { create: true });
  const writable = await handle.createWritable();
  await writable.write(file);
  await writable.close();
}

export async function readFile(name: string): Promise<File> {
  const dir = await getRoot();
  const handle = await dir.getFileHandle(name);
  const file = await handle.getFile();
  return new File([file], name, { type: file.type, lastModified: file.lastModified });
}

OPFS は通常のファイルシステム API とは異なり、ユーザーにパーミッションダイアログを出さずにアクセスできます。navigator.storage.getDirectory() で取得したルートの下に models ディレクトリを作り、モデルファイルを保存します。2回目以降はアップロードなしで即座にモデルを読み込めます。

OPFS のストレージ容量はブラウザとデバイスに依存しますが、Chrome では通常、ディスク容量の最大 60% まで利用可能です。SD-Turbo の 5 ファイル合計で約 2GB なので、ほとんどの環境で問題なく動作します。

テスト戦略: ブラウザ API なしで 415 件のテスト

ブラウザ依存の API（WebGPU、OPFS 等）を使うプロジェクトでありながら、テストは Node.js 上で実行しています。

[PASS] fp16:         23 passed, 0 failed   (IEEE 754 半精度変換)
[PASS] tokenizer:   231 passed, 0 failed   (OpenAI CLIP 参照出力と一致)
[PASS] scheduler:   132 passed, 0 failed   (Euler Ancestral, PRNG 統計)
[PASS] safetensors:  29 passed, 0 failed   (読込/書込 roundtrip、型変換)

テスト戦略のポイント

  1. パイプラインの各コンポーネントを分離テスト: Tokenizer、Scheduler、fp16 変換、safetensors パーサはすべて純粋な計算ロジックなので、ブラウザ API に依存しない
  2. 実データでの検証: OpenAI CLIP の実際の vocab.json と merges.txt をフィクスチャとして使い、Python の参照実装と同一の出力を確認
  3. Node スタブ: DOM API や localStorage は最小限のスタブで代替

Tokenizer のテストでは 231 件のケースを網羅しており、英語・日本語・特殊文字・境界条件を含む OpenAI CLIP の参照出力と完全一致することを検証しています。

AdSense + WebGPU の両立: COEP の罠

WebGPU 推論に必要な SharedArrayBuffer と Google AdSense を同じページで動かすには、COEP（Cross-Origin-Embedder-Policy）の設定で工夫が必要です。

問題

SharedArrayBuffer を有効にするには COEP ヘッダが必要ですが、require-corp を指定すると、Cross-Origin-Resource-Policy ヘッダを返さないサードパーティリソースがすべてブロックされます。Google AdSense のスクリプトやフレームはこのヘッダを返さないため、広告が表示されなくなります。

解決策

credentialless を使います。

Cross-Origin-Embedder-Policy: credentialless

credentialless は require-corp と同様に SharedArrayBuffer を有効にしつつ、クロスオリジンのサブリソースを Cross-Origin-Resource-Policy ヘッダなしでも読み込めるようにします。Cookie はクロスオリジンリクエストに付与されませんが、AdSense はコンテキスト広告の配信に支障がないため問題ありません。

credentialless は比較的新しい仕様で、2024 年時点で Chrome / Edge でサポートされていますが、Safari は未対応です。WebGPU 自体も Safari では実験的サポートのため、実質的に Chromium 系ブラウザをターゲットとするこのプロジェクトでは問題になりません。

さらに、Cookie 同意バナーを実装して Google Consent Mode v2 に対応しています。

const STORAGE_KEY = 'sd-keydrop.consent';

export interface ConsentState {
  necessary: true;
  analytics: boolean;
  advertising: boolean;
  timestamp: string;
}

ユーザーは「すべて許可 / 必要最小限のみ / カスタマイズ」の 3 段階から選択でき、gtag('consent', 'update', ...) で Google 側に同意状態を伝達します。

デプロイ: CloudFormation + S3 + CloudFront

完全な静的サイトなので、サーバーレス構成でデプロイしています。

インフラ構成

  • S3: ビルド成果物のホスティング
  • CloudFront: CDN + HTTPS + レスポンスヘッダーポリシー（COEP/COOP）
  • CloudFormation: インフラの IaC 管理

キャッシュ戦略

デプロイスクリプトではファイル種別ごとに異なるキャッシュ戦略を適用しています。

# HTML ファイル: キャッシュなし（即座に反映）
aws s3 cp "${html_file}" "s3://${BUCKET_NAME}/${s3_key}" \
  --cache-control "no-cache, no-store, must-revalidate" \
  --content-type "text/html; charset=utf-8"

# ハッシュ付き静的アセット (JS/CSS): 長期キャッシュ
aws s3 sync dist/assets/ "s3://${BUCKET_NAME}/assets/" \
  --cache-control "public, max-age=31536000, immutable"

# ORT WASM ファイル: 1 日キャッシュ（ハッシュなし）
aws s3 sync dist/ort-wasm/ "s3://${BUCKET_NAME}/ort-wasm/" \
  --cache-control "public, max-age=86400"
ファイル種別 Cache-Control 理由
HTML no-cache, no-store デプロイ直後にユーザーへ反映
JS/CSS（ハッシュ付き） immutable, 1 year コンテンツハッシュで一意に識別
ORT WASM 1 day ファイル名にハッシュがないため中期キャッシュ
設定ファイル (robots.txt 等) 1 hour 頻繁には変わらないが更新可能に

CloudFront のレスポンスヘッダーポリシーで COOP / COEP を付与しているため、S3 側で個別にヘッダを設定する必要はありません。

まとめと今後の展望

実現したこと

  • ブラウザ内で Stable Diffusion の推論パイプライン全体を TypeScript で実装
  • WebGPU × ONNX Runtime Web でGPUアクセラレーション
  • CLIP BPE Tokenizer、Euler Ancestral Scheduler、fp16 変換を自作
  • safetensors パーサ + LoRA マージを Pure TypeScript で実装
  • OPFS による数 GB のモデルキャッシュ
  • 型安全な Worker 通信プロトコル
  • サーバーコストゼロで運用可能な静的サイト

現在の制約

  • SD-Turbo のみ検証済み: SD 1.5 系もコード上は対応しているが、十分な検証が行われていない
  • Safari 非対応: WebGPU の実装が実験的であり、credentialless COEP もサポートされていない
  • VRAM 要件: 4GB 以上の GPU が推奨。統合 GPU では厳しい場合がある

今後の展望

  • SDXL 対応: より高品質な画像生成
  • WebGPU コンピュートシェーダでの LoRA マージ高速化: 現在の Pure TypeScript 行列積を GPU で実行
  • OSS として公開予定: コミュニティからのフィードバックを受けて改善

ブラウザの機能は日々進化しており、WebGPU の成熟とともに「ブラウザで AI」はますます現実的になっています。この記事が、WebGPU を使った AI アプリケーション開発の参考になれば幸いです。

参考リンク

1
1
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
1
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?