@simulatte/doppler 0.1.3 → 0.1.5

package/src/inference/pipelines/diffusion/scheduler.d.ts

@@ -9,11 +9,27 @@ import type { DiffusionSchedulerConfig } from './types.js';
 export interface DiffusionScheduler {
   type: string;
   steps: number;
-  sigmas: Float32Array;
+  sigmas: Float32Array | null;
   timesteps: Float32Array;
+  predictionType?: string;
+  sigmaData?: number;
+}
+
+export interface DiffusionSchedulerStepResult {
+  prevSample: Float32Array;
+  predOriginalSample: Float32Array;
 }
 
 export declare function buildScheduler(
   config: DiffusionSchedulerConfig,
   stepsOverride?: number | null
 ): DiffusionScheduler;
+
+export declare function stepScmScheduler(
+  config: DiffusionScheduler,
+  modelOutput: Float32Array,
+  timestep: number,
+  sample: Float32Array,
+  stepIndex?: number,
+  noise?: Float32Array | null
+): DiffusionSchedulerStepResult;

package/src/inference/pipelines/diffusion/scheduler.js

@@ -34,6 +34,84 @@ function buildFlowMatchSchedule(config, steps) {
   return sigmas;
 }
 
+function buildScmTimesteps(steps, config) {
+  const maxTimesteps = Number.isFinite(config.maxTimesteps) ? config.maxTimesteps : 1.5708;
+  const intermediateTimesteps = Number.isFinite(config.intermediateTimesteps)
+    ? config.intermediateTimesteps
+    : 1.3;
+  const count = Math.max(1, steps);
+  if (count === 1) {
+    return new Float32Array([maxTimesteps, 0.0]);
+  }
+  if (count === 2) {
+    return new Float32Array([maxTimesteps, intermediateTimesteps, 0.0]);
+  }
+  return linspace(maxTimesteps, 0.0, count + 1);
+}
+
+export function stepScmScheduler(config, modelOutput, timestep, sample, stepIndex = 0, noise = null) {
+  if (!config || config.type !== 'scm') {
+    throw new Error('stepScmScheduler requires scheduler.type="scm".');
+  }
+  if (!(modelOutput instanceof Float32Array)) {
+    throw new Error('stepScmScheduler requires modelOutput as Float32Array.');
+  }
+  if (!(sample instanceof Float32Array)) {
+    throw new Error('stepScmScheduler requires sample as Float32Array.');
+  }
+  if (modelOutput.length !== sample.length) {
+    throw new Error(
+      `stepScmScheduler requires modelOutput and sample with matching sizes; got ${modelOutput.length} and ${sample.length}.`
+    );
+  }
+  if (!(config.timesteps instanceof Float32Array) || config.timesteps.length < 2) {
+    throw new Error('stepScmScheduler requires scheduler.timesteps with length >= 2.');
+  }
+  if (!Number.isInteger(stepIndex) || stepIndex < 0 || stepIndex + 1 >= config.timesteps.length) {
+    throw new Error(
+      `stepScmScheduler received invalid stepIndex=${stepIndex} for ${config.timesteps.length} timesteps.`
+    );
+  }
+
+  const parameterization = config.predictionType ?? 'trigflow';
+  if (parameterization !== 'trigflow') {
+    throw new Error(`Unsupported SCM predictionType "${parameterization}".`);
+  }
+
+  const s = config.timesteps[stepIndex];
+  const t = config.timesteps[stepIndex + 1];
+  const predOriginalSample = new Float32Array(sample.length);
+  const prevSample = new Float32Array(sample.length);
+
+  const cosS = Math.cos(s);
+  const sinS = Math.sin(s);
+  const cosT = Math.cos(t);
+  const sinT = Math.sin(t);
+
+  for (let i = 0; i < sample.length; i++) {
+    const predX0 = cosS * sample[i] - sinS * modelOutput[i];
+    predOriginalSample[i] = predX0;
+    prevSample[i] = predX0;
+  }
+
+  if (stepIndex + 1 < config.timesteps.length - 1) {
+    if (!(noise instanceof Float32Array) || noise.length !== sample.length) {
+      throw new Error(
+        'stepScmScheduler requires a Float32Array noise tensor for multi-step SCM updates.'
+      );
+    }
+    const sigmaData = Number.isFinite(config.sigmaData) ? config.sigmaData : 0.5;
+    for (let i = 0; i < prevSample.length; i++) {
+      prevSample[i] = cosT * predOriginalSample[i] + sinT * noise[i] * sigmaData;
+    }
+  }
+
+  return {
+    prevSample,
+    predOriginalSample,
+  };
+}
+
 export function buildScheduler(config, stepsOverride = null) {
   if (!config) {
     throw new Error('Scheduler config is required');
@@ -43,15 +121,25 @@ export function buildScheduler(config, stepsOverride = null) {
   if (typeof type !== 'string' || !type) {
     throw new Error('Diffusion scheduler requires a scheduler type.');
   }
-  const sigmas = type === 'flowmatch_euler'
-    ? buildFlowMatchSchedule(config, steps)
-    : buildLinearSigmaSchedule(steps);
   const trainSteps = Number.isFinite(config.numTrainTimesteps)
     ? config.numTrainTimesteps
     : null;
   if (!Number.isFinite(trainSteps) || trainSteps <= 0) {
     throw new Error('Diffusion scheduler requires valid numTrainTimesteps.');
   }
+  if (type === 'scm') {
+    return {
+      type,
+      steps,
+      sigmas: null,
+      timesteps: buildScmTimesteps(steps, config),
+      predictionType: config.predictionType ?? 'trigflow',
+      sigmaData: Number.isFinite(config.sigmaData) ? config.sigmaData : 0.5,
+    };
+  }
+  const sigmas = type === 'flowmatch_euler'
+    ? buildFlowMatchSchedule(config, steps)
+    : buildLinearSigmaSchedule(steps);
   const timesteps = new Float32Array(steps);
   for (let i = 0; i < steps; i++) {
     timesteps[i] = sigmas[i] * trainSteps;

package/src/inference/pipelines/diffusion/text-encoder-gpu.d.ts

@@ -16,25 +16,27 @@ export interface DiffusionTextEncoderWeightsEntry {
 
 export interface DiffusionTextEncoderWeights {
   text_encoder: DiffusionTextEncoderWeightsEntry;
-  text_encoder_2: DiffusionTextEncoderWeightsEntry;
-  text_encoder_3: DiffusionTextEncoderWeightsEntry;
+  text_encoder_2?: DiffusionTextEncoderWeightsEntry | null;
+  text_encoder_3?: DiffusionTextEncoderWeightsEntry | null;
   transformer?: DiffusionTextEncoderWeightsEntry;
 }
 
 export interface DiffusionTextTokens {
   text_encoder: number[];
-  text_encoder_2: number[];
-  text_encoder_3: number[];
+  text_encoder_2?: number[];
+  text_encoder_3?: number[];
 }
 
 export interface DiffusionTextConditioning {
   pooled: Float32Array;
   context: Tensor;
+  attentionMask?: Uint32Array | null;
   profile?: {
     totalMs?: number | null;
     clipMs?: number | null;
     clip2Ms?: number | null;
     t5Ms?: number | null;
+    gemmaMs?: number | null;
   } | null;
 }
 

package/src/inference/pipelines/diffusion/text-encoder-gpu.js

@@ -40,6 +40,8 @@ import {
   inferDiffusionMatmulDtypeFromBuffer,
   sumDiffusionProfileTimings,
 } from './helpers.js';
+import { initRoPEFrequencies } from '../text/init.js';
+import { processLayerGPU } from '../text/layer.js';
 
 const QUICK_GELU_ALPHA = 1.702;
 const SUPPORTED_CLIP_HIDDEN_ACTIVATIONS = new Set(['gelu', 'quick_gelu']);
@@ -56,6 +58,16 @@ function padTokens(tokens, maxLength, padTokenId) {
   return out;
 }
 
+function normalizeTokens(tokens, maxLength, fallbackTokenId) {
+  const source = Array.isArray(tokens) ? tokens : [];
+  const limit = Number.isFinite(maxLength) && maxLength > 0 ? Math.floor(maxLength) : source.length;
+  const trimmed = source.slice(0, limit);
+  if (trimmed.length > 0) {
+    return Uint32Array.from(trimmed);
+  }
+  return new Uint32Array([fallbackTokenId >>> 0]);
+}
+
 function findEosIndex(tokens, eosTokenId) {
   if (eosTokenId == null) return tokens.length - 1;
   for (let i = 0; i < tokens.length; i++) {
@@ -702,7 +714,264 @@ async function runT5Encoder(tokens, weightsEntry, config, runtime, options = {})
   };
 }
 
+function buildGemma2LayerTypes(layerCount, slidingWindow) {
+  if (!Number.isFinite(slidingWindow) || slidingWindow <= 0) {
+    return Array.from({ length: layerCount }, () => 'full_attention');
+  }
+  return Array.from({ length: layerCount }, (_, index) => (
+    index % 2 === 1 ? 'full_attention' : 'sliding_attention'
+  ));
+}
+
+function getGemma2LayerWeight(weights, prefix, layerIdx, suffix, required = true) {
+  const key = `${prefix}.model.layers.${layerIdx}.${suffix}`;
+  const weight = weights.get(key) || null;
+  if (!weight && required) {
+    throw new Error(`Missing Gemma2 diffusion weight "${key}".`);
+  }
+  return weight;
+}
+
+function resolveGemma2TextConfig(config) {
+  const hiddenSize = config.hidden_size;
+  const numHeads = config.num_attention_heads;
+  const numKVHeads = config.num_key_value_heads ?? numHeads;
+  const headDim = config.head_dim ?? (
+    Number.isFinite(hiddenSize) && Number.isFinite(numHeads) && numHeads > 0
+      ? Math.floor(hiddenSize / numHeads)
+      : null
+  );
+  const numLayers = config.num_hidden_layers;
+  const intermediateSize = config.intermediate_size;
+  const maxPositionEmbeddings = config.max_position_embeddings;
+  const rmsNormEps = config.rms_norm_eps ?? 1e-6;
+
+  if (!Number.isFinite(hiddenSize) || hiddenSize <= 0) {
+    throw new Error('Gemma2 diffusion text encoder requires hidden_size.');
+  }
+  if (!Number.isFinite(numHeads) || numHeads <= 0) {
+    throw new Error('Gemma2 diffusion text encoder requires num_attention_heads.');
+  }
+  if (!Number.isFinite(numKVHeads) || numKVHeads <= 0) {
+    throw new Error('Gemma2 diffusion text encoder requires num_key_value_heads.');
+  }
+  if (!Number.isFinite(headDim) || headDim <= 0) {
+    throw new Error('Gemma2 diffusion text encoder requires head_dim or hidden_size/num_attention_heads.');
+  }
+  if (!Number.isFinite(numLayers) || numLayers <= 0) {
+    throw new Error('Gemma2 diffusion text encoder requires num_hidden_layers.');
+  }
+  if (!Number.isFinite(intermediateSize) || intermediateSize <= 0) {
+    throw new Error('Gemma2 diffusion text encoder requires intermediate_size.');
+  }
+  if (!Number.isFinite(maxPositionEmbeddings) || maxPositionEmbeddings <= 0) {
+    throw new Error('Gemma2 diffusion text encoder requires max_position_embeddings.');
+  }
+
+  return {
+    hiddenSize,
+    numHeads,
+    numKVHeads,
+    headDim,
+    numLayers,
+    intermediateSize,
+    maxPositionEmbeddings,
+    rmsNormEps,
+    ropeTheta: config.rope_theta ?? 10000,
+    slidingWindow: config.sliding_window ?? 4096,
+    scaleEmbeddings: config.scale_embeddings !== false,
+  };
+}
+
+async function runGemma2TextEncoder(tokens, weightsEntry, config, runtime, options = {}) {
+  const device = getDevice();
+  if (!device) throw new Error('Gemma2 diffusion text encoder requires a WebGPU device.');
+  if (!weightsEntry?.weights || !weightsEntry?.shapes) {
+    throw new Error('Gemma2 diffusion text encoder requires loaded weights.');
+  }
+
+  const prefix = options.prefix ?? 'text_encoder';
+  const localRecorder = options.recorder
+    ? null
+    : (options.profile ? new CommandRecorder(device, `${prefix}_gemma2_encoder`, { profile: true }) : null);
+  const recorder = options.recorder ?? localRecorder;
+  const ops = createKernelOps(recorder);
+  const release = createDiffusionBufferReleaser(recorder);
+  const destroy = createDiffusionBufferDestroyer(recorder);
+  const weights = weightsEntry.weights;
+  const activationDtype = resolveDiffusionActivationDtype(runtime);
+  const resolved = resolveGemma2TextConfig(config);
+  const padTokenId = config.pad_token_id ?? config.bos_token_id ?? 0;
+  const tokenIds = normalizeTokens(tokens, options.maxLength ?? resolved.maxPositionEmbeddings, padTokenId);
+  const numTokens = tokenIds.length;
+  const tokenBuffer = createDiffusionIndexBuffer(device, tokenIds, `${prefix}_tokens`);
+
+  const embedKey = `${prefix}.model.embed_tokens.weight`;
+  const embedWeight = expectDiffusionWeight(
+    weights.get(embedKey),
+    embedKey
+  );
+  const embedDtype = resolveEmbeddingDtype(embedWeight, weightsEntry, embedKey, runtime);
+  let hidden = await ops.gather(
+    tokenBuffer,
+    getBuffer(embedWeight),
+    numTokens,
+    resolved.hiddenSize,
+    config.vocab_size,
+    {
+      embeddingDtype: embedDtype,
+      outputDtype: activationDtype,
+      transpose: false,
+    }
+  );
+  destroy(tokenBuffer);
+
+  if (resolved.scaleEmbeddings) {
+    const scaled = await ops.scale(hidden, Math.sqrt(resolved.hiddenSize), {
+      count: numTokens * resolved.hiddenSize,
+    });
+    release(hidden.buffer);
+    hidden = createTensor(scaled.buffer, scaled.dtype, [numTokens, resolved.hiddenSize], 'gemma2_embed');
+  }
+
+  const layerWeights = new Map();
+  for (let layerIdx = 0; layerIdx < resolved.numLayers; layerIdx++) {
+    layerWeights.set(`layer_${layerIdx}`, {
+      inputNorm: getGemma2LayerWeight(weights, prefix, layerIdx, 'input_layernorm.weight'),
+      qProj: getGemma2LayerWeight(weights, prefix, layerIdx, 'self_attn.q_proj.weight'),
+      kProj: getGemma2LayerWeight(weights, prefix, layerIdx, 'self_attn.k_proj.weight'),
+      vProj: getGemma2LayerWeight(weights, prefix, layerIdx, 'self_attn.v_proj.weight'),
+      oProj: getGemma2LayerWeight(weights, prefix, layerIdx, 'self_attn.o_proj.weight'),
+      postAttentionNorm: getGemma2LayerWeight(weights, prefix, layerIdx, 'post_attention_layernorm.weight'),
+      preFeedforwardNorm: getGemma2LayerWeight(weights, prefix, layerIdx, 'pre_feedforward_layernorm.weight'),
+      gate: getGemma2LayerWeight(weights, prefix, layerIdx, 'mlp.gate_proj.weight'),
+      up: getGemma2LayerWeight(weights, prefix, layerIdx, 'mlp.up_proj.weight'),
+      down: getGemma2LayerWeight(weights, prefix, layerIdx, 'mlp.down_proj.weight'),
+    });
+  }
+
+  const ropeFreqs = await initRoPEFrequencies({
+    headDim: resolved.headDim,
+    maxSeqLen: resolved.maxPositionEmbeddings,
+    ropeTheta: resolved.ropeTheta,
+    ropeLocalTheta: null,
+    ropeScale: 1,
+    ropeLocalScale: null,
+    ropeScalingType: null,
+    ropeLocalScalingType: null,
+    ropeScaling: null,
+    ropeLocalScaling: null,
+  }, true);
+
+  const context = {
+    useGPU: true,
+    activationDtype,
+    recorder,
+    currentSeqLen: 0,
+    kvCache: null,
+    weightConfig: {
+      rmsNormWeightOffset: true,
+    },
+    debugFlags: {},
+    weights: layerWeights,
+    ropeFreqsCos: ropeFreqs.cos,
+    ropeFreqsSin: ropeFreqs.sin,
+    ropeLocalCos: ropeFreqs.localCos,
+    ropeLocalSin: ropeFreqs.localSin,
+    config: {
+      hiddenSize: resolved.hiddenSize,
+      intermediateSize: resolved.intermediateSize,
+      numHeads: resolved.numHeads,
+      numKVHeads: resolved.numKVHeads,
+      headDim: resolved.headDim,
+      rmsNormEps: resolved.rmsNormEps,
+      slidingWindow: resolved.slidingWindow,
+      attnLogitSoftcapping: 50.0,
+      queryPreAttnScalar: resolved.headDim,
+      queryKeyNorm: false,
+      attentionOutputGate: false,
+      causalAttention: true,
+      hiddenActivation: 'gelu',
+      swigluLimit: null,
+      useMoE: false,
+      layerTypes: buildGemma2LayerTypes(resolved.numLayers, resolved.slidingWindow),
+      preFeedforwardNorm: true,
+      postFeedforwardNorm: false,
+      postAttentionNorm: true,
+    },
+  };
+
+  for (let layerIdx = 0; layerIdx < resolved.numLayers; layerIdx++) {
+    const output = await processLayerGPU(
+      layerIdx,
+      hidden.buffer,
+      numTokens,
+      true,
+      numTokens * resolved.hiddenSize,
+      context
+    );
+    hidden = createTensor(output.buffer, output.dtype, [numTokens, resolved.hiddenSize], `gemma2_layer_${layerIdx}`);
+  }
+
+  const finalNormKey = `${prefix}.model.norm.weight`;
+  const finalNorm = expectDiffusionWeight(weights.get(finalNormKey), finalNormKey);
+  const final = await ops.rmsNorm(hidden, getBuffer(finalNorm), resolved.rmsNormEps, {
+    batchSize: numTokens,
+    hiddenSize: resolved.hiddenSize,
+    rmsNormWeightOffset: true,
+  });
+  release(hidden.buffer);
+
+  let profile = null;
+  if (localRecorder) {
+    localRecorder.submit();
+    const timings = await localRecorder.resolveProfileTimings();
+    profile = timings ? { totalMs: sumDiffusionProfileTimings(timings) ?? 0, timings } : { totalMs: null };
+  }
+
+  return {
+    hidden: final,
+    attentionMask: Uint32Array.from({ length: numTokens }, () => 1),
+    maxLength: numTokens,
+    hiddenSize: resolved.hiddenSize,
+    profile,
+  };
+}
+
 export async function runTextEncodersForPrompt(tokensByEncoder, weightsByComponent, modelConfig, runtime, options = {}) {
+  const layout = modelConfig?.layout ?? 'sd3';
+  if (layout === 'sana') {
+    const gemmaConfig = modelConfig?.components?.text_encoder?.config || {};
+    const gemmaMaxLength = runtime?.textEncoder?.maxLength;
+    if (!Number.isFinite(gemmaMaxLength) || gemmaMaxLength <= 0) {
+      throw new Error('Sana Gemma2 encoder requires runtime.textEncoder.maxLength.');
+    }
+    const profileEnabled = options.profile === true;
+    const gemma = await runGemma2TextEncoder(
+      tokensByEncoder.text_encoder,
+      weightsByComponent.text_encoder,
+      gemmaConfig,
+      runtime,
+      {
+        prefix: 'text_encoder',
+        maxLength: gemmaMaxLength,
+        profile: profileEnabled,
+      }
+    );
+
+    return {
+      pooled: new Float32Array(0),
+      context: gemma.hidden,
+      attentionMask: gemma.attentionMask,
+      profile: profileEnabled
+        ? {
+            totalMs: gemma.profile?.totalMs ?? null,
+            gemmaMs: gemma.profile?.totalMs ?? null,
+          }
+        : null,
+    };
+  }
+
   const clipConfig = modelConfig?.components?.text_encoder?.config || {};
   const clip2Config = modelConfig?.components?.text_encoder_2?.config || {};
   const t5Config = modelConfig?.components?.text_encoder_3?.config || {};
@@ -749,6 +1018,7 @@ export async function runTextEncodersForPrompt(tokensByEncoder, weightsByCompone
   return {
     pooled,
     context: t5.hidden,
+    attentionMask: null,
     profile,
   };
 }

package/src/inference/pipelines/diffusion/text-encoder.js

@@ -1,5 +1,6 @@
 import { BPETokenizer } from '../../tokenizers/bpe.js';
 import { SentencePieceTokenizer } from '../../tokenizers/sentencepiece.js';
+import { BundledTokenizer } from '../../tokenizers/bundled.js';
 import { loadAuxText, loadAuxFile } from '../../../storage/shard-manager.js';
 
 function parseMerges(text) {
@@ -136,11 +137,27 @@ async function loadSentencePieceTokenizer(tokenizerConfig, options = {}) {
   return tokenizer;
 }
 
+async function loadBundledTokenizer(tokenizerConfig, options = {}) {
+  const { baseUrl } = options;
+  const tokenizerJsonText = await loadTextAsset(tokenizerConfig.tokenizerFile, baseUrl);
+  const tokenizerJson = JSON.parse(tokenizerJsonText);
+  const tokenizer = new BundledTokenizer({
+    vocabSize: 0,
+    deferSpecialTokens: true,
+  });
+  tokenizer.load(tokenizerJson);
+  return tokenizer;
+}
+
 export async function loadDiffusionTokenizers(diffusionConfig, options = {}) {
   const tokenizers = {};
   const config = diffusionConfig?.tokenizers || {};
   if (config.text_encoder) {
-    tokenizers.text_encoder = await loadBpeTokenizer(config.text_encoder, options);
+    if (config.text_encoder.type === 'bundled') {
+      tokenizers.text_encoder = await loadBundledTokenizer(config.text_encoder, options);
+    } else {
+      tokenizers.text_encoder = await loadBpeTokenizer(config.text_encoder, options);
+    }
   }
   if (config.text_encoder_2) {
     tokenizers.text_encoder_2 = await loadBpeTokenizer(config.text_encoder_2, options);

package/src/inference/pipelines/diffusion/types.d.ts

@@ -27,6 +27,10 @@ export interface DiffusionSchedulerConfig {
   eta: number;
   numTrainTimesteps: number;
   shift: number;
+  predictionType?: string;
+  sigmaData?: number;
+  maxTimesteps?: number;
+  intermediateTimesteps?: number;
 }
 
 export interface DiffusionLatentConfig {