Spaces:
Paused
Paused
File size: 2,831 Bytes
08808ae | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | /**
* 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;
}
}
|