decoder.js

/**
 * MediaTok neural decoder — runs ONNX models via WebGPU
 *
 * Two-stage decode:
 *   1. TokenProjector: token IDs (1, 444) -> latent grid (1, 4, 8, 8)
 *   2. TAESDDecoder: latent grid -> RGB frame (1, 3, 64, 64)
 */

export class GigaTokenDecoder {
  constructor() {
    this.projSession = null;
    this.decSession = null;
    this.ort = null;
  }

  async load(projPath, decPath) {
    // Default paths for HF Space
    if (!projPath) projPath = 'https://huggingface.co/Daankular/mediatok-models/resolve/main/projector.onnx';
    if (!decPath) decPath = 'https://huggingface.co/Daankular/mediatok-models/resolve/main/taesd_decoder.onnx';
    this.ort = window.ort;
    if (!this.ort) {
      throw new Error('ONNX Runtime not loaded — include onnxruntime-web');
    }

    // Try WebGPU first, fallback to wasm
    let backend = 'webgpu';
    try {
      await this.ort.env.init();
    } catch (_) {}

    try {
      this.projSession = await this.ort.InferenceSession.create(projPath, {
        executionProviders: ['webgpu', 'wasm'],
      });
    } catch (e) {
      try {
        this.projSession = await this.ort.InferenceSession.create(projPath, {
          executionProviders: ['wasm'],
        });
      } catch (e2) {
        throw new Error(`Failed to load projector: ${e2.message}`);
      }
    }

    try {
      this.decSession = await this.ort.InferenceSession.create(decPath, {
        executionProviders: ['webgpu', 'wasm'],
      });
    } catch (e) {
      try {
        this.decSession = await this.ort.InferenceSession.create(decPath, {
          executionProviders: ['wasm'],
        });
      } catch (e2) {
        throw new Error(`Failed to load decoder: ${e2.message}`);
      }
    }
  }

  async decode(tokenIds) {
    if (!this.projSession || !this.decSession) {
      throw new Error('Decoder not loaded — call load() first');
    }

    // Prepare input tensor (1, 444) int64
    const inputTensor = new this.ort.Tensor(
      'int64',
      BigInt64Array.from(tokenIds.map(BigInt)),
      [1, 444]
    );

    // Stage 1: projector
    const projResult = await this.projSession.run({ token_ids: inputTensor });
    const latent = projResult.latent;

    // Stage 2: decoder
    const decResult = await this.decSession.run({ latent });
    const frame = decResult.frame;

    // frame is (1, 3, 64, 64) float32
    // Convert to RGB array for canvas rendering
    const data = frame.data;
    const H = 64, W = 64;
    const rgb = Array.from({ length: H }, () => Array.from({ length: W }, () => [0, 0, 0]));

    for (let y = 0; y < H; y++) {
      for (let x = 0; x < W; x++) {
        const i = y * W + x;
        rgb[y][x][0] = data[i];           // R
        rgb[y][x][1] = data[H*W + i];     // G
        rgb[y][x][2] = data[2*H*W + i];   // B
      }
    }

    return rgb;
  }
}