@simulatte/doppler 0.1.8 → 0.2.0

package/src/inference/pipelines/diffusion/sd3-transformer.js

@@ -1,1194 +0,0 @@
-import { getDevice } from '../../../gpu/device.js';
-import { createTensor, dtypeBytes } from '../../../gpu/tensor.js';
-import { getBuffer } from '../../../gpu/weight-buffer.js';
-import { acquireBuffer } from '../../../memory/buffer-pool.js';
-import {
-  runConv2D,
-  runTranspose,
-  runGather,
-  runLayerNorm,
-  runRMSNorm,
-  runMatmul,
-  runAttention,
-  runGeLU,
-  runSiLURowSplit,
-  runResidualAdd,
-  runBiasAdd,
-  runModulate,
-  runPixelShuffle,
-  recordConv2D,
-  recordTranspose,
-  recordGather,
-  recordLayerNorm,
-  recordRMSNorm,
-  recordMatmul,
-  recordAttention,
-  recordGeLU,
-  recordSiLURowSplit,
-  recordResidualAdd,
-  recordBiasAdd,
-  recordModulate,
-  recordPixelShuffle,
-} from '../../../gpu/kernels/index.js';
-import { log } from '../../../debug/index.js';
-import { createSD3WeightResolver } from './sd3-weights.js';
-import {
-  resolveDiffusionActivationDtype,
-  createDiffusionBufferReleaser,
-  createDiffusionBufferDestroyer,
-  createDiffusionIndexBuffer,
-  expectDiffusionWeight,
-  normalizeDiffusionLocationDtype,
-  normalizeDiffusionMatmulLocationDtype,
-  inferDiffusionMatmulDtypeFromBuffer,
-} from './helpers.js';
-
-function reshapeTensor(tensor, shape, label) {
-  return createTensor(tensor.buffer, tensor.dtype, shape, label);
-}
-
-function createKernelOps(recorder) {
-  if (!recorder) {
-    return {
-      conv2d: runConv2D,
-      transpose: runTranspose,
-      gather: runGather,
-      layerNorm: runLayerNorm,
-      rmsNorm: runRMSNorm,
-      attention: runAttention,
-      gelu: runGeLU,
-      siluRowSplit: runSiLURowSplit,
-      residualAdd: runResidualAdd,
-      biasAdd: runBiasAdd,
-      modulate: runModulate,
-      pixelShuffle: runPixelShuffle,
-    };
-  }
-  return {
-    conv2d: (...args) => recordConv2D(recorder, ...args),
-    transpose: (...args) => recordTranspose(recorder, ...args),
-    gather: (...args) => recordGather(recorder, ...args),
-    layerNorm: (...args) => recordLayerNorm(recorder, ...args),
-    rmsNorm: (...args) => recordRMSNorm(recorder, ...args),
-    attention: (...args) => recordAttention(recorder, ...args),
-    gelu: (...args) => recordGeLU(recorder, ...args),
-    siluRowSplit: (...args) => recordSiLURowSplit(recorder, ...args),
-    residualAdd: (...args) => recordResidualAdd(recorder, ...args),
-    biasAdd: (...args) => recordBiasAdd(recorder, ...args),
-    modulate: (...args) => recordModulate(recorder, ...args),
-    pixelShuffle: (...args) => recordPixelShuffle(recorder, ...args),
-  };
-}
-
-function createVectorBuffer(device, data, label) {
-  const buffer = acquireBuffer(data.byteLength, undefined, label);
-  device.queue.writeBuffer(buffer, 0, data);
-  return buffer;
-}
-
-function resolveEmbeddingDtype(weight, weightsEntry, key, runtime) {
-  if (weight && weight.dtype) return weight.dtype;
-  const locationDtype = weightsEntry?.dtypes?.get(key);
-  const mapped = normalizeDiffusionLocationDtype(locationDtype);
-  if (!mapped) return null;
-  if (mapped !== 'f16') return mapped;
-  const allowUpcast = runtime?.loading?.allowF32UpcastNonMatmul !== false;
-  return allowUpcast ? 'f32' : 'f16';
-}
-
-function resolveMatmulDtype(weight, resolver, name) {
-  if (weight && weight.dtype) return weight.dtype;
-  if (!resolver || !name) return null;
-  const locationDtype = resolver.dtype(name);
-  return normalizeDiffusionMatmulLocationDtype(locationDtype);
-}
-
-function resolveBiasDtype(weight, resolver, name) {
-  if (weight && weight.dtype) return weight.dtype;
-  if (!resolver || !name) return 'f32';
-  const locationDtype = resolver.dtype(name);
-  const mapped = normalizeDiffusionLocationDtype(locationDtype);
-  return mapped || 'f32';
-}
-
-async function runMatmulResolved(input, weight, resolver, name, M, N, K, options = {}) {
-  const { recorder = null, ...rest } = options;
-  const resolved = resolveMatmulDtype(weight, resolver, name);
-  const bDtype = inferDiffusionMatmulDtypeFromBuffer(weight, N, K, resolved);
-  const nextOptions = bDtype ? { ...rest, bDtype } : rest;
-  if (recorder) {
-    return recordMatmul(recorder, input, weight, M, N, K, nextOptions);
-  }
-  return runMatmul(input, weight, M, N, K, nextOptions);
-}
-
-function createBiasTensorWithDtype(weight, size, label, resolver, name) {
-  if (!weight) return null;
-  const dtype = resolveBiasDtype(weight, resolver, name);
-  return createTensor(getBuffer(weight), dtype, [size], label);
-}
-
-function resolveTransformerLayerNormEps(config, runtime) {
-  const modelEps = Number(config?.norm_eps ?? config?.layer_norm_eps);
-  if (Number.isFinite(modelEps) && modelEps > 0) {
-    return modelEps;
-  }
-  const runtimeEps = Number(runtime?.backend?.layerNormEps);
-  if (Number.isFinite(runtimeEps) && runtimeEps > 0) {
-    return runtimeEps;
-  }
-  throw new Error(
-    'Diffusion transformer requires a positive layer norm epsilon from ' +
-    'transformer.config.norm_eps (or layer_norm_eps) or runtime.inference.diffusion.backend.layerNormEps.'
-  );
-}
-
-async function splitQKV(qkv, numTokens, hiddenSize, label, recorder) {
-  const device = getDevice();
-  const bytesPerElement = dtypeBytes(qkv.dtype);
-  const sliceBytes = numTokens * hiddenSize * bytesPerElement;
-  const qBuf = acquireBuffer(sliceBytes, undefined, `${label}_q`);
-  const kBuf = acquireBuffer(sliceBytes, undefined, `${label}_k`);
-  const vBuf = acquireBuffer(sliceBytes, undefined, `${label}_v`);
-
-  const encoder = recorder ? recorder.getEncoder() : device.createCommandEncoder();
-  encoder.copyBufferToBuffer(qkv.buffer, 0, qBuf, 0, sliceBytes);
-  encoder.copyBufferToBuffer(qkv.buffer, sliceBytes, kBuf, 0, sliceBytes);
-  encoder.copyBufferToBuffer(qkv.buffer, sliceBytes * 2, vBuf, 0, sliceBytes);
-  if (!recorder) {
-    device.queue.submit([encoder.finish()]);
-  }
-
-  return {
-    q: createTensor(qBuf, qkv.dtype, [numTokens, hiddenSize], `${label}_q`),
-    k: createTensor(kBuf, qkv.dtype, [numTokens, hiddenSize], `${label}_k`),
-    v: createTensor(vBuf, qkv.dtype, [numTokens, hiddenSize], `${label}_v`),
-  };
-}
-
-async function runFusedQKV(input, weight, biasTensor, numTokens, hiddenSize, outputDtype, label, matmul, weightName, ops, release, recorder) {
-  const qkv = await matmul(input, weight, weightName, numTokens, hiddenSize * 3, hiddenSize, {
-    outputDtype,
-    transposeB: 'auto',
-  });
-
-  let qkvTensor = qkv;
-  if (biasTensor) {
-    qkvTensor = await ops.biasAdd(qkv, biasTensor, numTokens, hiddenSize * 3);
-  }
-
-  const split = await splitQKV(qkvTensor, numTokens, hiddenSize, label, recorder);
-  release(qkvTensor.buffer);
-  return split;
-}
-
-async function runQKV(input, weights, bias, numTokens, hiddenSize, label, matmul, weightNames, ops, release, recorder) {
-  const outputDtype = input.dtype;
-  if (weights.qkv) {
-    return runFusedQKV(
-      input,
-      weights.qkv,
-      bias?.qkv ?? null,
-      numTokens,
-      hiddenSize,
-      outputDtype,
-      label,
-      matmul,
-      weightNames?.qkv ?? null,
-      ops,
-      release,
-      recorder
-    );
-  }
-
-  const qWeight = expectDiffusionWeight(weights.q, `${label}.q`);
-  const kWeight = expectDiffusionWeight(weights.k, `${label}.k`);
-  const vWeight = expectDiffusionWeight(weights.v, `${label}.v`);
-
-  let q = await matmul(input, qWeight, weightNames?.q ?? null, numTokens, hiddenSize, hiddenSize, {
-    outputDtype,
-    transposeB: 'auto',
-  });
-  let k = await matmul(input, kWeight, weightNames?.k ?? null, numTokens, hiddenSize, hiddenSize, {
-    outputDtype,
-    transposeB: 'auto',
-  });
-  let v = await matmul(input, vWeight, weightNames?.v ?? null, numTokens, hiddenSize, hiddenSize, {
-    outputDtype,
-    transposeB: 'auto',
-  });
-
-  if (bias?.q) q = await ops.biasAdd(q, bias.q, numTokens, hiddenSize);
-  if (bias?.k) k = await ops.biasAdd(k, bias.k, numTokens, hiddenSize);
-  if (bias?.v) v = await ops.biasAdd(v, bias.v, numTokens, hiddenSize);
-
-  return { q, k, v };
-}
-
-async function applyQKNorm(tensor, weight, numTokens, numHeads, headDim, eps, ops) {
-  const flattened = createTensor(tensor.buffer, tensor.dtype, [numTokens * numHeads, headDim], 'qk_norm_in');
-  const normed = await ops.rmsNorm(flattened, getBuffer(weight), eps, {
-    batchSize: numTokens * numHeads,
-    hiddenSize: headDim,
-  });
-  return reshapeTensor(normed, [numTokens, numHeads, headDim], 'qk_norm_out');
-}
-
-async function concatKV(a, b, numTokensA, numTokensB, hiddenSize, recorder) {
-  const device = getDevice();
-  const bytesPerElement = a.dtype === 'f16' ? 2 : 4;
-  const outputSize = (numTokensA + numTokensB) * hiddenSize * bytesPerElement;
-  const output = acquireBuffer(outputSize, undefined, 'kv_concat');
-  const encoder = recorder ? recorder.getEncoder() : device.createCommandEncoder();
-  encoder.copyBufferToBuffer(a.buffer, 0, output, 0, numTokensA * hiddenSize * bytesPerElement);
-  encoder.copyBufferToBuffer(b.buffer, 0, output, numTokensA * hiddenSize * bytesPerElement, numTokensB * hiddenSize * bytesPerElement);
-  if (!recorder) {
-    device.queue.submit([encoder.finish()]);
-  }
-  return createTensor(output, a.dtype, [numTokensA + numTokensB, hiddenSize], 'kv_concat');
-}
-
-async function runAttentionBlock(input, weights, bias, numTokens, hiddenSize, numHeads, headDim, normWeights, eps, matmul, weightNames, ops, release) {
-  let q = await matmul(input, weights.q, weightNames?.q ?? null, numTokens, hiddenSize, hiddenSize, {
-    outputDtype: input.dtype,
-    transposeB: 'auto',
-  });
-  let k = await matmul(input, weights.k, weightNames?.k ?? null, numTokens, hiddenSize, hiddenSize, {
-    outputDtype: input.dtype,
-    transposeB: 'auto',
-  });
-  let v = await matmul(input, weights.v, weightNames?.v ?? null, numTokens, hiddenSize, hiddenSize, {
-    outputDtype: input.dtype,
-    transposeB: 'auto',
-  });
-
-  if (bias?.q) q = await ops.biasAdd(q, bias.q, numTokens, hiddenSize);
-  if (bias?.k) k = await ops.biasAdd(k, bias.k, numTokens, hiddenSize);
-  if (bias?.v) v = await ops.biasAdd(v, bias.v, numTokens, hiddenSize);
-
-  if (normWeights?.q) {
-    const normed = await applyQKNorm(q, normWeights.q, numTokens, numHeads, headDim, eps, ops);
-    release(q.buffer);
-    q = normed;
-  }
-  if (normWeights?.k) {
-    const normed = await applyQKNorm(k, normWeights.k, numTokens, numHeads, headDim, eps, ops);
-    release(k.buffer);
-    k = normed;
-  }
-
-  const attn = await ops.attention(q, k, v, null, numHeads, headDim, {
-    seqLen: numTokens,
-    kvLen: numTokens,
-    numKVHeads: numHeads,
-    causal: false,
-  });
-
-  release(q.buffer);
-  release(k.buffer);
-  release(v.buffer);
-
-  return attn;
-}
-
-function resolveModulationSegments(weight, hiddenSize, fallbackSegments, resolver, name) {
-  const shape = weight?.shape || (resolver && name ? resolver.shape(name) : null);
-  const rows = Array.isArray(shape) ? shape[0] : null;
-  if (Number.isFinite(rows) && Number.isFinite(hiddenSize) && hiddenSize > 0) {
-    const segments = rows / hiddenSize;
-    if (Number.isInteger(segments) && segments > 0) {
-      return segments;
-    }
-    throw new Error(
-      `Modulation segments mismatch for ${name || 'unknown'}: rows=${rows}, hidden=${hiddenSize}, ` +
-      `expected an integer multiple instead of falling back to ${fallbackSegments}.`
-    );
-  }
-  throw new Error(
-    `Modulation tensor "${name || 'unknown'}" is missing shape metadata. ` +
-    `Runtime cannot fall back to ${fallbackSegments} segments.`
-  );
-}
-
-function resolveModulationOffsets(segments, hiddenSize) {
-  if (segments === 9) {
-    return {
-      attn: { scale: 0, shift: hiddenSize, gate: hiddenSize * 2 },
-      attn2: { scale: hiddenSize * 3, shift: hiddenSize * 4, gate: hiddenSize * 5 },
-      ff: { scale: hiddenSize * 6, shift: hiddenSize * 7, gate: hiddenSize * 8 },
-    };
-  }
-  if (segments === 6) {
-    const attn = { scale: 0, shift: hiddenSize, gate: hiddenSize * 2 };
-    return {
-      attn,
-      attn2: { ...attn },
-      ff: { scale: hiddenSize * 3, shift: hiddenSize * 4, gate: hiddenSize * 5 },
-    };
-  }
-  throw new Error(`Unsupported modulation segments=${segments} (expected 6 or 9).`);
-}
-
-async function buildModulation(timeText, weight, bias, hiddenSize, segments, runtime, matmul, weightName, ops) {
-  const device = getDevice();
-  const activationDtype = resolveDiffusionActivationDtype(runtime);
-  const outDim = hiddenSize * segments;
-  const bytesPerElement = activationDtype === 'f16' ? 2 : 4;
-  const bufferSize = (outDim + hiddenSize) * bytesPerElement;
-  const outputBuffer = acquireBuffer(bufferSize, undefined, 'sd3_modulate');
-
-  const mod = await matmul(timeText, weight, weightName, 1, outDim, hiddenSize, {
-    outputDtype: activationDtype,
-    transposeB: 'auto',
-    outputBuffer,
-  });
-
-  if (bias) {
-    await ops.biasAdd(mod, bias, 1, outDim);
-  }
-
-  const zeroOffset = outDim * bytesPerElement;
-  device.queue.writeBuffer(outputBuffer, zeroOffset, new Uint8Array(hiddenSize * bytesPerElement));
-
-  return {
-    tensor: createTensor(outputBuffer, activationDtype, [1, outDim], 'sd3_mod'),
-    zeroOffset: outDim,
-  };
-}
-
-async function applyAdaLayerNorm(input, weight, bias, eps, mod, offsets, runtime, ops, release, options = {}) {
-  const { numTokens, hiddenSize } = options;
-  const normed = await ops.layerNorm(input, weight, bias, eps, { batchSize: numTokens, hiddenSize });
-  const modulated = await ops.modulate(normed, mod.tensor, {
-    numTokens,
-    hiddenSize,
-    scaleOffset: offsets.scale,
-    shiftOffset: offsets.shift,
-    gateOffset: offsets.gate,
-    hasGate: false,
-    addOne: true,
-  });
-  release(normed.buffer);
-  return modulated;
-}
-
-async function applyGate(output, mod, offsets, ops, release, options = {}) {
-  const { numTokens, hiddenSize, zeroOffset } = options;
-  const gated = await ops.modulate(output, mod.tensor, {
-    numTokens,
-    hiddenSize,
-    scaleOffset: offsets.gate,
-    shiftOffset: zeroOffset,
-    gateOffset: offsets.gate,
-    hasGate: false,
-    addOne: false,
-  });
-  release(output.buffer);
-  return gated;
-}
-
-async function runFFN(input, weights, bias, numTokens, hiddenSize, runtime, matmul, weightNames, ops, release) {
-  const activationDtype = resolveDiffusionActivationDtype(runtime);
-  const upDim = weights.up.shape[0];
-  const downInput = weights.down.shape[1];
-  let up = await matmul(input, weights.up, weightNames?.up ?? null, numTokens, upDim, hiddenSize, {
-    outputDtype: activationDtype,
-    transposeB: 'auto',
-  });
-  if (bias?.up) up = await ops.biasAdd(up, bias.up, numTokens, upDim);
-
-  let act = null;
-  let intermediate = upDim;
-  if (Number.isFinite(downInput) && upDim === downInput * 2) {
-    act = await ops.siluRowSplit(up, {
-      numTokens,
-      dim: downInput,
-      activation: 'gelu',
-      swigluLimit: null,
-    });
-    intermediate = downInput;
-  } else {
-    act = await ops.gelu(up, { size: numTokens * upDim });
-  }
-  release(up.buffer);
-
-  let down = await matmul(act, weights.down, weightNames?.down ?? null, numTokens, hiddenSize, intermediate, {
-    outputDtype: activationDtype,
-    transposeB: 'auto',
-  });
-  if (bias?.down) down = await ops.biasAdd(down, bias.down, numTokens, hiddenSize);
-  release(act.buffer);
-  return down;
-}
-
-export async function runSD3Transformer(latents, context, timeText, weightsEntry, modelConfig, runtime, options = {}) {
-  const device = getDevice();
-  if (!device) {
-    throw new Error('SD3 transformer requires a WebGPU device.');
-  }
-
-  const resolver = createSD3WeightResolver(weightsEntry, modelConfig);
-  const recorder = options.recorder ?? null;
-  const ops = createKernelOps(recorder);
-  const release = createDiffusionBufferReleaser(recorder);
-  const destroy = createDiffusionBufferDestroyer(recorder);
-  const matmul = (input, weight, name, M, N, K, options = {}) =>
-    runMatmulResolved(input, weight, resolver, name, M, N, K, { ...options, recorder });
-  const config = modelConfig?.components?.transformer?.config || {};
-  const hiddenSize = config.num_attention_heads * config.attention_head_dim;
-  const numHeads = config.num_attention_heads;
-  const headDim = config.attention_head_dim;
-  const patchSize = config.patch_size;
-  const layerNormEps = resolveTransformerLayerNormEps(config, runtime);
-
-  const latentChannels = latents.shape[0];
-  const latentHeight = latents.shape[1];
-  const latentWidth = latents.shape[2];
-  const gridHeight = Math.floor(latentHeight / patchSize);
-  const gridWidth = Math.floor(latentWidth / patchSize);
-  const tokenCount = gridHeight * gridWidth;
-
-  const projWeight = expectDiffusionWeight(resolver.get('pos_embed.proj.weight'), 'pos_embed.proj.weight');
-  const projBias = resolver.get('pos_embed.proj.bias');
-
-  const conv = await ops.conv2d(latents, projWeight, projBias, {
-    inChannels: latentChannels,
-    outChannels: hiddenSize,
-    height: latentHeight,
-    width: latentWidth,
-    kernelH: patchSize,
-    kernelW: patchSize,
-    stride: patchSize,
-    pad: 0,
-  });
-
-  const tokens = await ops.transpose(conv, hiddenSize, tokenCount);
-  release(conv.buffer);
-
-  const posEmbed = expectDiffusionWeight(resolver.get('pos_embed.pos_embed'), 'pos_embed.pos_embed');
-  const posShape = resolver.shape('pos_embed.pos_embed') || [1, tokenCount, hiddenSize];
-  const maxTokens = posShape[1];
-  const maxGrid = Math.floor(Math.sqrt(maxTokens));
-
-  if (maxGrid * maxGrid !== maxTokens) {
-    log.warn('Diffusion', 'pos_embed size is not square; using sequential indices.');
-  }
-
-  const posIndices = new Uint32Array(tokenCount);
-  for (let y = 0; y < gridHeight; y++) {
-    const srcY = maxGrid * (y / Math.max(1, gridHeight));
-    const srcYIdx = Math.min(maxGrid - 1, Math.floor(srcY));
-    for (let x = 0; x < gridWidth; x++) {
-      const srcX = maxGrid * (x / Math.max(1, gridWidth));
-      const srcXIdx = Math.min(maxGrid - 1, Math.floor(srcX));
-      posIndices[y * gridWidth + x] = srcYIdx * maxGrid + srcXIdx;
-    }
-  }
-
-  const posBuffer = createDiffusionIndexBuffer(device, posIndices, 'sd3_pos_idx');
-  const posEmbedKey = resolver.key('pos_embed.pos_embed');
-  const posEmbedDtype = resolveEmbeddingDtype(posEmbed, weightsEntry, posEmbedKey, runtime);
-  const pos = await ops.gather(
-    posBuffer,
-    getBuffer(posEmbed),
-    tokenCount,
-    hiddenSize,
-    maxTokens,
-    {
-      embeddingDtype: posEmbedDtype,
-      outputDtype: tokens.dtype,
-      transpose: false,
-    }
-  );
-  destroy(posBuffer);
-
-  const xCombined = await ops.residualAdd(tokens, pos, tokenCount * hiddenSize, { useVec4: true });
-  release(tokens.buffer);
-  release(pos.buffer);
-
-  let x = createTensor(xCombined.buffer, xCombined.dtype, [tokenCount, hiddenSize], 'sd3_tokens');
-  let ctx = context;
-  let ctxOwned = false;
-
-  const ones = new Float32Array(hiddenSize).fill(1.0);
-  const zeros = new Float32Array(hiddenSize);
-  const onesBuf = createVectorBuffer(device, ones, 'sd3_ln_weight');
-  const zerosBuf = createVectorBuffer(device, zeros, 'sd3_ln_bias');
-
-  const dualLayers = new Set(config.dual_attention_layers || []);
-  const attn2Layers = Array.isArray(config.attn2_layers)
-    ? new Set(config.attn2_layers)
-    : null;
-  const numLayers = config.num_layers;
-
-  for (let layerIdx = 0; layerIdx < numLayers; layerIdx++) {
-    const modWeightName = `transformer_blocks.${layerIdx}.norm1.linear.weight`;
-    const modBiasName = `transformer_blocks.${layerIdx}.norm1.linear.bias`;
-    const modWeight = expectDiffusionWeight(
-      resolver.get(modWeightName),
-      modWeightName
-    );
-    const modBias = resolver.get(modBiasName);
-    const modSegments = resolveModulationSegments(modWeight, hiddenSize, 9, resolver, modWeightName);
-    if (modSegments < 6) {
-      throw new Error(`Unsupported modulation segments=${modSegments} for ${modWeightName}`);
-    }
-    const modBiasTensor = createBiasTensorWithDtype(
-      modBias,
-      hiddenSize * modSegments,
-      'sd3_mod_bias',
-      resolver,
-      modBiasName
-    );
-    const mod = await buildModulation(timeText, modWeight, modBiasTensor, hiddenSize, modSegments, runtime, matmul, modWeightName, ops);
-
-    const offsets = resolveModulationOffsets(modSegments, hiddenSize);
-    const attnOffsets = offsets.attn;
-    const attn2Offsets = offsets.attn2;
-    const ffOffsets = offsets.ff;
-
-    let ctxMod = null;
-    let ctxOffsets = null;
-    let ctxAttnOffsets = null;
-    let ctxFfOffsets = null;
-    if (dualLayers.has(layerIdx)) {
-      const ctxWeightName = `transformer_blocks.${layerIdx}.norm1_context.linear.weight`;
-      const ctxBiasName = `transformer_blocks.${layerIdx}.norm1_context.linear.bias`;
-      const ctxWeight = expectDiffusionWeight(
-        resolver.get(ctxWeightName),
-        ctxWeightName
-      );
-      const ctxBias = resolver.get(ctxBiasName);
-      const ctxSegments = resolveModulationSegments(ctxWeight, hiddenSize, 6, resolver, ctxWeightName);
-      if (ctxSegments < 6) {
-        throw new Error(`Unsupported modulation segments=${ctxSegments} for ${ctxWeightName}`);
-      }
-      const ctxBiasTensor = createBiasTensorWithDtype(
-        ctxBias,
-        hiddenSize * ctxSegments,
-        'sd3_ctx_mod_bias',
-        resolver,
-        ctxBiasName
-      );
-      ctxMod = await buildModulation(timeText, ctxWeight, ctxBiasTensor, hiddenSize, ctxSegments, runtime, matmul, ctxWeightName, ops);
-      ctxOffsets = resolveModulationOffsets(ctxSegments, hiddenSize);
-      ctxAttnOffsets = ctxOffsets.attn;
-      ctxFfOffsets = ctxOffsets.ff;
-    }
-
-    const xAttnIn = await applyAdaLayerNorm(
-      x,
-      onesBuf,
-      zerosBuf,
-      layerNormEps,
-      mod,
-      attnOffsets,
-      runtime,
-      ops,
-      release,
-      { numTokens: tokenCount, hiddenSize }
-    );
-
-    if (dualLayers.has(layerIdx)) {
-      const ctxAttnIn = await applyAdaLayerNorm(
-        ctx,
-        onesBuf,
-        zerosBuf,
-        layerNormEps,
-        ctxMod,
-        ctxAttnOffsets,
-        runtime,
-        ops,
-        release,
-        { numTokens: ctx.shape[0], hiddenSize }
-      );
-
-      const attnWeightNames = {
-        q: `transformer_blocks.${layerIdx}.attn.to_q.weight`,
-        k: `transformer_blocks.${layerIdx}.attn.to_k.weight`,
-        v: `transformer_blocks.${layerIdx}.attn.to_v.weight`,
-        qkv: `transformer_blocks.${layerIdx}.attn.qkv.weight`,
-      };
-      const attnWeights = {
-        q: resolver.get(attnWeightNames.q),
-        k: resolver.get(attnWeightNames.k),
-        v: resolver.get(attnWeightNames.v),
-        qkv: resolver.get(attnWeightNames.qkv),
-      };
-      const attnBiasNames = {
-        q: `transformer_blocks.${layerIdx}.attn.to_q.bias`,
-        k: `transformer_blocks.${layerIdx}.attn.to_k.bias`,
-        v: `transformer_blocks.${layerIdx}.attn.to_v.bias`,
-        qkv: `transformer_blocks.${layerIdx}.attn.qkv.bias`,
-      };
-      const attnBias = {
-        q: createBiasTensorWithDtype(
-          resolver.get(attnBiasNames.q),
-          hiddenSize,
-          'sd3_attn_q_bias',
-          resolver,
-          attnBiasNames.q
-        ),
-        k: createBiasTensorWithDtype(
-          resolver.get(attnBiasNames.k),
-          hiddenSize,
-          'sd3_attn_k_bias',
-          resolver,
-          attnBiasNames.k
-        ),
-        v: createBiasTensorWithDtype(
-          resolver.get(attnBiasNames.v),
-          hiddenSize,
-          'sd3_attn_v_bias',
-          resolver,
-          attnBiasNames.v
-        ),
-        qkv: createBiasTensorWithDtype(
-          resolver.get(attnBiasNames.qkv),
-          hiddenSize * 3,
-          'sd3_attn_qkv_bias',
-          resolver,
-          attnBiasNames.qkv
-        ),
-      };
-      const addWeightNames = {
-        q: `transformer_blocks.${layerIdx}.attn.add_q_proj.weight`,
-        k: `transformer_blocks.${layerIdx}.attn.add_k_proj.weight`,
-        v: `transformer_blocks.${layerIdx}.attn.add_v_proj.weight`,
-        qkv: `transformer_blocks.${layerIdx}.attn.add_qkv.weight`,
-      };
-      const addWeights = {
-        q: resolver.get(addWeightNames.q),
-        k: resolver.get(addWeightNames.k),
-        v: resolver.get(addWeightNames.v),
-        qkv: resolver.get(addWeightNames.qkv),
-      };
-      const addBiasNames = {
-        q: `transformer_blocks.${layerIdx}.attn.add_q_proj.bias`,
-        k: `transformer_blocks.${layerIdx}.attn.add_k_proj.bias`,
-        v: `transformer_blocks.${layerIdx}.attn.add_v_proj.bias`,
-        qkv: `transformer_blocks.${layerIdx}.attn.add_qkv.bias`,
-      };
-      const addBias = {
-        q: createBiasTensorWithDtype(
-          resolver.get(addBiasNames.q),
-          hiddenSize,
-          'sd3_attn_add_q_bias',
-          resolver,
-          addBiasNames.q
-        ),
-        k: createBiasTensorWithDtype(
-          resolver.get(addBiasNames.k),
-          hiddenSize,
-          'sd3_attn_add_k_bias',
-          resolver,
-          addBiasNames.k
-        ),
-        v: createBiasTensorWithDtype(
-          resolver.get(addBiasNames.v),
-          hiddenSize,
-          'sd3_attn_add_v_bias',
-          resolver,
-          addBiasNames.v
-        ),
-        qkv: createBiasTensorWithDtype(
-          resolver.get(addBiasNames.qkv),
-          hiddenSize * 3,
-          'sd3_attn_add_qkv_bias',
-          resolver,
-          addBiasNames.qkv
-        ),
-      };
-
-      const normWeights = {
-        q: resolver.get(`transformer_blocks.${layerIdx}.attn.norm_q.weight`),
-        k: resolver.get(`transformer_blocks.${layerIdx}.attn.norm_k.weight`),
-        qAdd: resolver.get(`transformer_blocks.${layerIdx}.attn.norm_added_q.weight`),
-        kAdd: resolver.get(`transformer_blocks.${layerIdx}.attn.norm_added_k.weight`),
-      };
-
-      let { q: qx, k: kx, v: vx } = await runQKV(
-        xAttnIn,
-        attnWeights,
-        attnBias,
-        tokenCount,
-        hiddenSize,
-        `sd3_attn_${layerIdx}`,
-        matmul,
-        attnWeightNames,
-        ops,
-        release,
-        recorder
-      );
-
-      let { q: qc, k: kc, v: vc } = await runQKV(
-        ctxAttnIn,
-        addWeights,
-        addBias,
-        ctx.shape[0],
-        hiddenSize,
-        `sd3_attn_add_${layerIdx}`,
-        matmul,
-        addWeightNames,
-        ops,
-        release,
-        recorder
-      );
-
-      if (normWeights.q) {
-        const normed = await applyQKNorm(qx, normWeights.q, tokenCount, numHeads, headDim, layerNormEps, ops);
-        release(qx.buffer);
-        qx = normed;
-      }
-      if (normWeights.k) {
-        const normed = await applyQKNorm(kx, normWeights.k, tokenCount, numHeads, headDim, layerNormEps, ops);
-        release(kx.buffer);
-        kx = normed;
-      }
-      if (normWeights.qAdd) {
-        const normed = await applyQKNorm(qc, normWeights.qAdd, ctx.shape[0], numHeads, headDim, layerNormEps, ops);
-        release(qc.buffer);
-        qc = normed;
-      }
-      if (normWeights.kAdd) {
-        const normed = await applyQKNorm(kc, normWeights.kAdd, ctx.shape[0], numHeads, headDim, layerNormEps, ops);
-        release(kc.buffer);
-        kc = normed;
-      }
-
-      const kAll = await concatKV(kx, kc, tokenCount, ctx.shape[0], hiddenSize, recorder);
-      const vAll = await concatKV(vx, vc, tokenCount, ctx.shape[0], hiddenSize, recorder);
-
-      const attnX = await ops.attention(qx, kAll, vAll, null, numHeads, headDim, {
-        seqLen: tokenCount,
-        kvLen: tokenCount + ctx.shape[0],
-        numKVHeads: numHeads,
-        causal: false,
-      });
-
-      const attnC = await ops.attention(qc, kAll, vAll, null, numHeads, headDim, {
-        seqLen: ctx.shape[0],
-        kvLen: tokenCount + ctx.shape[0],
-        numKVHeads: numHeads,
-        causal: false,
-      });
-
-      const outWeightName = `transformer_blocks.${layerIdx}.attn.to_out.0.weight`;
-      const outWeight = expectDiffusionWeight(
-        resolver.get(outWeightName),
-        outWeightName
-      );
-      const outBiasName = `transformer_blocks.${layerIdx}.attn.to_out.0.bias`;
-      const outBias = resolver.get(outBiasName);
-      const outAddWeightName = `transformer_blocks.${layerIdx}.attn.to_add_out.weight`;
-      const outAddWeight = expectDiffusionWeight(
-        resolver.get(outAddWeightName),
-        outAddWeightName
-      );
-      const outAddBiasName = `transformer_blocks.${layerIdx}.attn.to_add_out.bias`;
-      const outAddBias = resolver.get(outAddBiasName);
-
-      let attnOutX = await matmul(attnX, outWeight, outWeightName, tokenCount, hiddenSize, hiddenSize, {
-        outputDtype: attnX.dtype,
-        transposeB: 'auto',
-      });
-      if (outBias) {
-        attnOutX = await ops.biasAdd(
-          attnOutX,
-          createBiasTensorWithDtype(outBias, hiddenSize, 'sd3_attn_out_bias', resolver, outBiasName),
-          tokenCount,
-          hiddenSize
-        );
-      }
-
-      let attnOutC = await matmul(attnC, outAddWeight, outAddWeightName, ctx.shape[0], hiddenSize, hiddenSize, {
-        outputDtype: attnC.dtype,
-        transposeB: 'auto',
-      });
-      if (outAddBias) {
-        attnOutC = await ops.biasAdd(
-          attnOutC,
-          createBiasTensorWithDtype(outAddBias, hiddenSize, 'sd3_attn_out_add_bias', resolver, outAddBiasName),
-          ctx.shape[0],
-          hiddenSize
-        );
-      }
-
-      const gatedX = await applyGate(attnOutX, mod, attnOffsets, ops, release, { numTokens: tokenCount, hiddenSize, zeroOffset: mod.zeroOffset });
-      const gatedC = await applyGate(attnOutC, ctxMod, ctxAttnOffsets, ops, release, { numTokens: ctx.shape[0], hiddenSize, zeroOffset: ctxMod.zeroOffset });
-
-      const xRes = await ops.residualAdd(x, gatedX, tokenCount * hiddenSize, { useVec4: true });
-      const cRes = await ops.residualAdd(ctx, gatedC, ctx.shape[0] * hiddenSize, { useVec4: true });
-
-      release(xAttnIn.buffer);
-      release(ctxAttnIn.buffer);
-      release(qx.buffer);
-      release(kx.buffer);
-      release(vx.buffer);
-      release(qc.buffer);
-      release(kc.buffer);
-      release(vc.buffer);
-      release(kAll.buffer);
-      release(vAll.buffer);
-      release(attnX.buffer);
-      release(attnC.buffer);
-      release(gatedX.buffer);
-      release(gatedC.buffer);
-      release(x.buffer);
-      if (ctxOwned) {
-        release(ctx.buffer);
-      }
-
-      x = createTensor(xRes.buffer, xRes.dtype, [tokenCount, hiddenSize], 'sd3_x');
-      ctx = createTensor(cRes.buffer, cRes.dtype, [ctx.shape[0], hiddenSize], 'sd3_ctx');
-      ctxOwned = true;
-
-      const ctxFfIn = await applyAdaLayerNorm(
-        ctx,
-        onesBuf,
-        zerosBuf,
-        layerNormEps,
-        ctxMod,
-        ctxFfOffsets,
-        runtime,
-        ops,
-        release,
-        { numTokens: ctx.shape[0], hiddenSize }
-      );
-
-      const ffCtxWeightNames = {
-        up: `transformer_blocks.${layerIdx}.ff_context.net.0.proj.weight`,
-        down: `transformer_blocks.${layerIdx}.ff_context.net.2.weight`,
-      };
-      const ffCtxWeights = {
-        up: expectDiffusionWeight(
-          resolver.get(ffCtxWeightNames.up),
-          ffCtxWeightNames.up
-        ),
-        down: expectDiffusionWeight(
-          resolver.get(ffCtxWeightNames.down),
-          ffCtxWeightNames.down
-        ),
-      };
-      const ffCtxBiasNames = {
-        up: `transformer_blocks.${layerIdx}.ff_context.net.0.proj.bias`,
-        down: `transformer_blocks.${layerIdx}.ff_context.net.2.bias`,
-      };
-      const ffCtxBias = {
-        up: createBiasTensorWithDtype(
-          resolver.get(ffCtxBiasNames.up),
-          ffCtxWeights.up.shape[0],
-          'sd3_ff_ctx_up_bias',
-          resolver,
-          ffCtxBiasNames.up
-        ),
-        down: createBiasTensorWithDtype(
-          resolver.get(ffCtxBiasNames.down),
-          hiddenSize,
-          'sd3_ff_ctx_down_bias',
-          resolver,
-          ffCtxBiasNames.down
-        ),
-      };
-      const ffCtxOut = await runFFN(
-        ctxFfIn,
-        ffCtxWeights,
-        ffCtxBias,
-        ctx.shape[0],
-        hiddenSize,
-        runtime,
-        matmul,
-        ffCtxWeightNames,
-        ops,
-        release
-      );
-      const ffCtxGated = await applyGate(ffCtxOut, ctxMod, ctxFfOffsets, ops, release, { numTokens: ctx.shape[0], hiddenSize, zeroOffset: ctxMod.zeroOffset });
-      const ctxRes2 = await ops.residualAdd(ctx, ffCtxGated, ctx.shape[0] * hiddenSize, { useVec4: true });
-
-      release(ctxFfIn.buffer);
-      release(ffCtxGated.buffer);
-      if (ctxOwned) {
-        release(ctx.buffer);
-      }
-      ctx = createTensor(ctxRes2.buffer, ctxRes2.dtype, [ctx.shape[0], hiddenSize], 'sd3_ctx');
-      ctxOwned = true;
-
-    } else {
-      release(xAttnIn.buffer);
-    }
-
-    const hasAttn2 = attn2Layers ? attn2Layers.has(layerIdx) : config.dual_attention_layers ? dualLayers.has(layerIdx) : true;
-    if (hasAttn2) {
-      const xAttn2In = await applyAdaLayerNorm(
-        x,
-        onesBuf,
-        zerosBuf,
-        layerNormEps,
-        mod,
-        attn2Offsets,
-        runtime,
-        ops,
-        release,
-        { numTokens: tokenCount, hiddenSize }
-      );
-
-      const attn2WeightNames = {
-        q: `transformer_blocks.${layerIdx}.attn2.to_q.weight`,
-        k: `transformer_blocks.${layerIdx}.attn2.to_k.weight`,
-        v: `transformer_blocks.${layerIdx}.attn2.to_v.weight`,
-        qkv: `transformer_blocks.${layerIdx}.attn2.qkv.weight`,
-      };
-      const attn2Weights = {
-        q: resolver.get(attn2WeightNames.q),
-        k: resolver.get(attn2WeightNames.k),
-        v: resolver.get(attn2WeightNames.v),
-        qkv: resolver.get(attn2WeightNames.qkv),
-      };
-      const attn2BiasNames = {
-        q: `transformer_blocks.${layerIdx}.attn2.to_q.bias`,
-        k: `transformer_blocks.${layerIdx}.attn2.to_k.bias`,
-        v: `transformer_blocks.${layerIdx}.attn2.to_v.bias`,
-        qkv: `transformer_blocks.${layerIdx}.attn2.qkv.bias`,
-      };
-      const attn2Bias = {
-        q: createBiasTensorWithDtype(
-          resolver.get(attn2BiasNames.q),
-          hiddenSize,
-          'sd3_attn2_q_bias',
-          resolver,
-          attn2BiasNames.q
-        ),
-        k: createBiasTensorWithDtype(
-          resolver.get(attn2BiasNames.k),
-          hiddenSize,
-          'sd3_attn2_k_bias',
-          resolver,
-          attn2BiasNames.k
-        ),
-        v: createBiasTensorWithDtype(
-          resolver.get(attn2BiasNames.v),
-          hiddenSize,
-          'sd3_attn2_v_bias',
-          resolver,
-          attn2BiasNames.v
-        ),
-        qkv: createBiasTensorWithDtype(
-          resolver.get(attn2BiasNames.qkv),
-          hiddenSize * 3,
-          'sd3_attn2_qkv_bias',
-          resolver,
-          attn2BiasNames.qkv
-        ),
-      };
-
-      let { q: q2, k: k2, v: v2 } = await runQKV(
-        xAttn2In,
-        attn2Weights,
-        attn2Bias,
-        tokenCount,
-        hiddenSize,
-        `sd3_attn2_${layerIdx}`,
-        matmul,
-        attn2WeightNames,
-        ops,
-        release,
-        recorder
-      );
-
-      const normQ2 = resolver.get(`transformer_blocks.${layerIdx}.attn2.norm_q.weight`);
-      const normK2 = resolver.get(`transformer_blocks.${layerIdx}.attn2.norm_k.weight`);
-      if (normQ2) {
-        const normed = await applyQKNorm(q2, normQ2, tokenCount, numHeads, headDim, layerNormEps, ops);
-        release(q2.buffer);
-        q2 = normed;
-      }
-      if (normK2) {
-        const normed = await applyQKNorm(k2, normK2, tokenCount, numHeads, headDim, layerNormEps, ops);
-        release(k2.buffer);
-        k2 = normed;
-      }
-
-      const attn2 = await ops.attention(q2, k2, v2, null, numHeads, headDim, {
-        seqLen: tokenCount,
-        kvLen: tokenCount,
-        numKVHeads: numHeads,
-        causal: false,
-      });
-
-      const attn2OutWeightName = `transformer_blocks.${layerIdx}.attn2.to_out.0.weight`;
-      const attn2OutWeight = expectDiffusionWeight(
-        resolver.get(attn2OutWeightName),
-        attn2OutWeightName
-      );
-      const attn2OutBiasName = `transformer_blocks.${layerIdx}.attn2.to_out.0.bias`;
-      const attn2OutBias = resolver.get(attn2OutBiasName);
-      let attn2Out = await matmul(attn2, attn2OutWeight, attn2OutWeightName, tokenCount, hiddenSize, hiddenSize, {
-        outputDtype: attn2.dtype,
-        transposeB: 'auto',
-      });
-      if (attn2OutBias) {
-        attn2Out = await ops.biasAdd(
-          attn2Out,
-          createBiasTensorWithDtype(attn2OutBias, hiddenSize, 'sd3_attn2_out_bias', resolver, attn2OutBiasName),
-          tokenCount,
-          hiddenSize
-        );
-      }
-
-      const gated2 = await applyGate(attn2Out, mod, attn2Offsets, ops, release, { numTokens: tokenCount, hiddenSize, zeroOffset: mod.zeroOffset });
-      const xRes2 = await ops.residualAdd(x, gated2, tokenCount * hiddenSize, { useVec4: true });
-
-      release(xAttn2In.buffer);
-      release(q2.buffer);
-      release(k2.buffer);
-      release(v2.buffer);
-      release(attn2.buffer);
-      release(attn2Out.buffer);
-      release(gated2.buffer);
-      release(x.buffer);
-
-      x = createTensor(xRes2.buffer, xRes2.dtype, [tokenCount, hiddenSize], 'sd3_x');
-    }
-
-    const xFfIn = await applyAdaLayerNorm(
-      x,
-      onesBuf,
-      zerosBuf,
-      layerNormEps,
-      mod,
-      ffOffsets,
-      runtime,
-      ops,
-      release,
-      { numTokens: tokenCount, hiddenSize }
-    );
-
-    const ffWeightNames = {
-      up: `transformer_blocks.${layerIdx}.ff.net.0.proj.weight`,
-      down: `transformer_blocks.${layerIdx}.ff.net.2.weight`,
-    };
-    const ffWeights = {
-      up: expectDiffusionWeight(
-        resolver.get(ffWeightNames.up),
-        ffWeightNames.up
-      ),
-      down: expectDiffusionWeight(
-        resolver.get(ffWeightNames.down),
-        ffWeightNames.down
-      ),
-    };
-    const ffBiasNames = {
-      up: `transformer_blocks.${layerIdx}.ff.net.0.proj.bias`,
-      down: `transformer_blocks.${layerIdx}.ff.net.2.bias`,
-    };
-    const ffBias = {
-      up: createBiasTensorWithDtype(
-        resolver.get(ffBiasNames.up),
-        ffWeights.up.shape[0],
-        'sd3_ff_up_bias',
-        resolver,
-        ffBiasNames.up
-      ),
-      down: createBiasTensorWithDtype(
-        resolver.get(ffBiasNames.down),
-        hiddenSize,
-        'sd3_ff_down_bias',
-        resolver,
-        ffBiasNames.down
-      ),
-    };
-
-    const ffOut = await runFFN(
-      xFfIn,
-      ffWeights,
-      ffBias,
-      tokenCount,
-      hiddenSize,
-      runtime,
-      matmul,
-      ffWeightNames,
-      ops,
-      release
-    );
-    const ffGated = await applyGate(ffOut, mod, ffOffsets, ops, release, { numTokens: tokenCount, hiddenSize, zeroOffset: mod.zeroOffset });
-    const xRes3 = await ops.residualAdd(x, ffGated, tokenCount * hiddenSize, { useVec4: true });
-
-    release(xFfIn.buffer);
-    release(ffGated.buffer);
-    release(x.buffer);
-
-    x = createTensor(xRes3.buffer, xRes3.dtype, [tokenCount, hiddenSize], 'sd3_x');
-
-    release(mod.tensor.buffer);
-    if (ctxMod?.tensor?.buffer) {
-      release(ctxMod.tensor.buffer);
-    }
-  }
-
-  const normOutWeightName = 'norm_out.linear.weight';
-  const normOutWeight = expectDiffusionWeight(resolver.get(normOutWeightName), normOutWeightName);
-  const normOutBias = resolver.get('norm_out.linear.bias');
-  const normOutSegments = resolveModulationSegments(normOutWeight, hiddenSize, 2, resolver, normOutWeightName);
-  const normOutBiasTensor = createBiasTensorWithDtype(
-    normOutBias,
-    hiddenSize * normOutSegments,
-    'sd3_norm_out_bias',
-    resolver,
-    'norm_out.linear.bias'
-  );
-  const normOut = await buildModulation(timeText, normOutWeight, normOutBiasTensor, hiddenSize, normOutSegments, runtime, matmul, normOutWeightName, ops);
-
-  const xNorm = await ops.layerNorm(x, onesBuf, zerosBuf, layerNormEps, { batchSize: tokenCount, hiddenSize });
-  const xMod = await ops.modulate(xNorm, normOut.tensor, {
-    numTokens: tokenCount,
-    hiddenSize,
-    scaleOffset: 0,
-    shiftOffset: hiddenSize,
-    gateOffset: 0,
-    hasGate: false,
-    addOne: true,
-  });
-
-  release(xNorm.buffer);
-  release(x.buffer);
-  release(normOut.tensor.buffer);
-  if (ctxOwned) {
-    release(ctx.buffer);
-  }
-  release(onesBuf);
-  release(zerosBuf);
-
-  const projOutWeightName = 'proj_out.weight';
-  const projOutWeight = expectDiffusionWeight(resolver.get(projOutWeightName), projOutWeightName);
-  const projOutBiasName = 'proj_out.bias';
-  const projOutBias = resolver.get(projOutBiasName);
-  let patch = await matmul(xMod, projOutWeight, projOutWeightName, tokenCount, projOutWeight.shape[0], hiddenSize, {
-    outputDtype: xMod.dtype,
-    transposeB: 'auto',
-  });
-  if (projOutBias) {
-    patch = await ops.biasAdd(
-      patch,
-      createBiasTensorWithDtype(projOutBias, projOutWeight.shape[0], 'sd3_proj_out_bias', resolver, projOutBiasName),
-      tokenCount,
-      projOutWeight.shape[0]
-    );
-  }
-
-  release(xMod.buffer);
-
-  const patchChannels = projOutWeight.shape[0];
-  const output = await ops.pixelShuffle(patch, {
-    outChannels: latentChannels,
-    outHeight: latentHeight,
-    outWidth: latentWidth,
-    gridWidth,
-    gridHeight,
-    patchSize,
-    patchChannels,
-  });
-
-  release(patch.buffer);
-
-  return output;
-}