@qvac/sdk 0.8.3 → 0.9.0

package/dist/server/rpc/handlers/index.d.ts

@@ -3,6 +3,7 @@ import { handleDownloadAsset } from "./download-asset";
 import { handleLoadModel } from "./load-model/index";
 import { handleUnloadModel } from "./unload-model";
 import { handleEmbed } from "./embed";
+import { handleTranscribe } from "./transcribe";
 import { handleTranscribeStream } from "./transcribe-stream";
 import { provideHandler } from "./provideHandler";
 import { stopProvideHandler } from "./stopProvideHandler";
@@ -13,10 +14,14 @@ import { cancelHandler } from "./cancelHandler";
 import { handleDeleteCache } from "./delete-cache";
 import { handleTextToSpeech } from "./text-to-speech";
 import { handleGetModelInfo } from "./get-model-info";
+import { handleFinetune } from "./finetune";
 import { handleOCRStream } from "./ocr-stream";
 import { handleHeartbeat } from "./heartbeat";
+import { handleDiffusionStream } from "./diffusion-stream";
 import { handlePluginInvoke, handlePluginInvokeStream } from "./plugin-invoke";
 import { handleModelRegistryList, handleModelRegistrySearch, handleModelRegistryGetModel } from "./registry";
+import { handleSuspend } from "./suspend";
+import { handleResume } from "./resume";
 export declare const handlers: {
     heartbeat: typeof handleHeartbeat;
     completionStream: typeof handleCompletionStream;
@@ -25,6 +30,7 @@ export declare const handlers: {
     loadModel: typeof handleLoadModel;
     unloadModel: typeof handleUnloadModel;
     embed: typeof handleEmbed;
+    transcribe: typeof handleTranscribe;
     transcribeStream: typeof handleTranscribeStream;
     provide: typeof provideHandler;
     stopProvide: typeof stopProvideHandler;
@@ -34,11 +40,15 @@ export declare const handlers: {
     cancel: typeof cancelHandler;
     textToSpeech: typeof handleTextToSpeech;
     getModelInfo: typeof handleGetModelInfo;
+    finetune: typeof handleFinetune;
     ocrStream: typeof handleOCRStream;
+    diffusionStream: typeof handleDiffusionStream;
     pluginInvoke: typeof handlePluginInvoke;
     pluginInvokeStream: typeof handlePluginInvokeStream;
     modelRegistryList: typeof handleModelRegistryList;
     modelRegistrySearch: typeof handleModelRegistrySearch;
     modelRegistryGetModel: typeof handleModelRegistryGetModel;
+    suspend: typeof handleSuspend;
+    resume: typeof handleResume;
 };
 //# sourceMappingURL=index.d.ts.map

package/dist/server/rpc/handlers/index.js

@@ -3,6 +3,7 @@ import { handleDownloadAsset } from "./download-asset.js";
 import { handleLoadModel } from "./load-model/index.js";
 import { handleUnloadModel } from "./unload-model.js";
 import { handleEmbed } from "./embed.js";
+import { handleTranscribe } from "./transcribe.js";
 import { handleTranscribeStream } from "./transcribe-stream.js";
 import { provideHandler } from "./provideHandler.js";
 import { stopProvideHandler } from "./stopProvideHandler.js";
@@ -13,10 +14,14 @@ import { cancelHandler } from "./cancelHandler.js";
 import { handleDeleteCache } from "./delete-cache.js";
 import { handleTextToSpeech } from "./text-to-speech.js";
 import { handleGetModelInfo } from "./get-model-info.js";
+import { handleFinetune } from "./finetune.js";
 import { handleOCRStream } from "./ocr-stream.js";
 import { handleHeartbeat } from "./heartbeat.js";
+import { handleDiffusionStream } from "./diffusion-stream.js";
 import { handlePluginInvoke, handlePluginInvokeStream } from "./plugin-invoke.js";
 import { handleModelRegistryList, handleModelRegistrySearch, handleModelRegistryGetModel, } from "./registry.js";
+import { handleSuspend } from "./suspend.js";
+import { handleResume } from "./resume.js";
 export const handlers = {
     heartbeat: handleHeartbeat,
     completionStream: handleCompletionStream,
@@ -25,6 +30,7 @@ export const handlers = {
     loadModel: handleLoadModel,
     unloadModel: handleUnloadModel,
     embed: handleEmbed,
+    transcribe: handleTranscribe,
     transcribeStream: handleTranscribeStream,
     provide: provideHandler,
     stopProvide: stopProvideHandler,
@@ -34,10 +40,14 @@ export const handlers = {
     cancel: cancelHandler,
     textToSpeech: handleTextToSpeech,
     getModelInfo: handleGetModelInfo,
+    finetune: handleFinetune,
     ocrStream: handleOCRStream,
+    diffusionStream: handleDiffusionStream,
     pluginInvoke: handlePluginInvoke,
     pluginInvokeStream: handlePluginInvokeStream,
     modelRegistryList: handleModelRegistryList,
     modelRegistrySearch: handleModelRegistrySearch,
     modelRegistryGetModel: handleModelRegistryGetModel,
+    suspend: handleSuspend,
+    resume: handleResume,
 };

package/dist/server/rpc/handlers/load-model/hyperdrive.js

@@ -12,6 +12,7 @@ import { getActiveDownload, registerDownload, unregisterDownload, createHyperdri
 import { getSDKConfig } from "../../../bare/registry/config-registry.js";
 import { FileNotFoundError, ChecksumValidationFailedError, DownloadCancelledError, HyperdriveDownloadFailedError, ModelLoadFailedError, NoBlobFoundError, } from "../../../../utils/errors-server.js";
 import { getServerLogger } from "../../../../logging/index.js";
+import { registerSwarm, unregisterSwarm, registerCorestore, unregisterCorestore, } from "../../../bare/runtime-lifecycle.js";
 const logger = getServerLogger();
 function getCorestoreDir(hyperdriveKey) {
     return path.join(getEnv().HOME_DIR, ".qvac", "corestore", hyperdriveKey);
@@ -40,21 +41,35 @@ async function setupHyperdrive(hyperdriveKey, corestoreDir) {
     });
     swarm.join(drive.discoveryKey, { server: false, client: true });
     swarm.flush().then(doneFindingPeers, doneFindingPeers);
+    registerCorestore(corestore, {
+        label: `hyperdrive:${hyperdriveKey}`,
+        createdAt: Date.now(),
+    });
+    registerSwarm(swarm, {
+        label: `hyperdrive:${hyperdriveKey}`,
+        createdAt: Date.now(),
+    });
     return { corestore, drive, swarm, corestoreDir };
 }
 async function cleanupHyperdrive(setup) {
     logger.debug("🧹 Cleaning up hyperdrive setup");
-    if (!setup.swarm.suspended) {
-        try {
-            await setup.swarm.suspend();
-        }
-        catch {
-            // If suspend fails, continue with cleanup
+    try {
+        if (!setup.swarm.suspended) {
+            try {
+                await setup.swarm.suspend();
+            }
+            catch (error) {
+                logger.debug("⚠️ Swarm suspend failed during cleanup:", error instanceof Error ? error.message : String(error));
+            }
         }
+        await setup.swarm.destroy();
+        await setup.drive.close();
+        await setup.corestore.close();
+    }
+    finally {
+        unregisterSwarm(setup.swarm);
+        unregisterCorestore(setup.corestore);
     }
-    await setup.swarm.destroy();
-    await setup.drive.close();
-    await setup.corestore.close();
 }
 export async function deleteCorestoreDirectory(corestoreDir) {
     logger.info(`🗑️ DELETING corestore directory... ${corestoreDir}`);

package/dist/server/rpc/handlers/load-model/registry.d.ts

@@ -1,5 +1,16 @@
 import type { ModelProgressUpdate } from "../../../../schemas/index";
+import { type RegistryItem } from "../../../../models/registry/index";
 import type { DownloadMetricsHooks } from "./types";
+/**
+ * Find companion ONNX data file in registry.
+ * ONNX models with external data may use either .onnx_data or .onnx.data.
+ */
+export declare function findOnnxCompanionDataFile(registryPath: string): RegistryItem | undefined;
+/**
+ * Resolve the main ONNX registry path from a companion data registry path.
+ * Supports both .onnx_data and .onnx.data naming conventions.
+ */
+export declare function getPairedOnnxRegistryPath(registryPath: string): string | undefined;
 /**
  * Download a model from the QVAC Registry.
  *

package/dist/server/rpc/handlers/load-model/registry.js

@@ -277,13 +277,26 @@ async function downloadShardedFilesFromRegistry(registryPath, registrySource, ca
 }
 /**
  * Find companion ONNX data file in registry.
- * ONNX models with external data have a .onnx file and a .onnx_data file.
+ * ONNX models with external data may use either .onnx_data or .onnx.data.
  */
-function findOnnxCompanionDataFile(registryPath) {
+export function findOnnxCompanionDataFile(registryPath) {
     if (!registryPath.endsWith(".onnx"))
         return undefined;
-    const dataPath = registryPath + "_data";
-    return getModelByPath(dataPath);
+    return (getModelByPath(`${registryPath}_data`) ??
+        getModelByPath(`${registryPath}.data`));
+}
+/**
+ * Resolve the main ONNX registry path from a companion data registry path.
+ * Supports both .onnx_data and .onnx.data naming conventions.
+ */
+export function getPairedOnnxRegistryPath(registryPath) {
+    if (registryPath.endsWith(".onnx_data")) {
+        return registryPath.slice(0, -"_data".length);
+    }
+    if (registryPath.endsWith(".onnx.data")) {
+        return registryPath.slice(0, -".data".length);
+    }
+    return undefined;
 }
 /**
  * Download ONNX model with companion data file from registry.
@@ -430,6 +443,20 @@ export async function downloadModelFromRegistry(registryPath, registrySource, pr
     const shardInfo = detectShardedModel(filename);
     // Look up model metadata from our generated models.ts
     const modelMetadata = getModelByPath(registryPath);
+    // ONNX external data: check if already present in paired ONNX cache directory.
+    // Avoids redundant single-file downloads when the companion .onnx download already placed it.
+    const pairedOnnxRegistryPath = getPairedOnnxRegistryPath(registryPath);
+    if (pairedOnnxRegistryPath) {
+        const onnxRegistryPath = pairedOnnxRegistryPath;
+        const onnxCacheKey = generateShortHash(onnxRegistryPath);
+        const pairedPath = getOnnxModelPath(onnxCacheKey, filename);
+        const validated = await validateCachedFile(pairedPath, filename, modelMetadata?.expectedSize || 0, modelMetadata?.sha256Checksum, hooks);
+        if (validated) {
+            logger.info(`✅ ONNX data file found in paired cache: ${validated}`);
+            hooks?.markCacheHit();
+            return validated;
+        }
+    }
     if (shardInfo.isSharded) {
         const cacheKey = generateShortHash(registryPath);
         const localShardMeta = modelMetadata?.shardMetadata;

package/dist/server/rpc/handlers/plugin-dispatch.d.ts

@@ -1,15 +1,3 @@
-/**
- * Dispatches a request to a plugin handler and returns a Promise.
- * Use for non-streaming (reply) handlers.
- *
- * @throws PluginHandlerTypeMismatchError if the handler is streaming
- */
 export declare function dispatchPluginReply<TRequest, TResponse>(modelId: string, handlerName: string, request: TRequest): Promise<TResponse>;
-/**
- * Dispatches a request to a plugin handler and returns an AsyncGenerator.
- * Use for streaming handlers.
- *
- * @throws PluginHandlerTypeMismatchError if the handler is not streaming
- */
-export declare function dispatchPluginStream<TRequest, TResponse>(modelId: string, handlerName: string, request: TRequest): AsyncGenerator<TResponse>;
+export declare function dispatchPluginStream<TRequest, TResponse>(modelId: string, handlerName: string, request: TRequest, inputStream?: AsyncIterable<Buffer>): AsyncGenerator<TResponse>;
 //# sourceMappingURL=plugin-dispatch.d.ts.map

package/dist/server/rpc/handlers/plugin-dispatch.js

@@ -2,10 +2,7 @@ import { getModelEntry } from "../../bare/registry/model-registry.js";
 import { getPlugin } from "../../plugins/index.js";
 import { profileReplyHandler, profileStreamHandler, } from "../profiling/index.js";
 import { ModelNotFoundError, ModelIsDelegatedError, PluginNotFoundError, PluginHandlerNotFoundError, PluginHandlerTypeMismatchError, } from "../../../utils/errors-server.js";
-/**
- * Resolves the plugin handler and returns both the result and streaming flag.
- */
-function resolvePluginHandler(modelId, handlerName, request) {
+function resolvePluginHandlerDef(modelId, handlerName) {
     const entry = getModelEntry(modelId);
     if (!entry) {
         throw new ModelNotFoundError(modelId);
@@ -22,17 +19,15 @@ function resolvePluginHandler(modelId, handlerName, request) {
         const availableHandlers = Object.keys(plugin.handlers);
         throw new PluginHandlerNotFoundError(entry.local.modelType, handlerName, availableHandlers);
     }
+    return handlerDef;
+}
+function resolvePluginHandler(modelId, handlerName, request) {
+    const handlerDef = resolvePluginHandlerDef(modelId, handlerName);
     return {
         result: handlerDef.handler(request),
         streaming: handlerDef.streaming,
     };
 }
-/**
- * Dispatches a request to a plugin handler and returns a Promise.
- * Use for non-streaming (reply) handlers.
- *
- * @throws PluginHandlerTypeMismatchError if the handler is streaming
- */
 export async function dispatchPluginReply(modelId, handlerName, request) {
     return profileReplyHandler({ op: handlerName, request }, async () => {
         const { result, streaming } = resolvePluginHandler(modelId, handlerName, request);
@@ -42,18 +37,23 @@ export async function dispatchPluginReply(modelId, handlerName, request) {
         return result;
     });
 }
-/**
- * Dispatches a request to a plugin handler and returns an AsyncGenerator.
- * Use for streaming handlers.
- *
- * @throws PluginHandlerTypeMismatchError if the handler is not streaming
- */
-export async function* dispatchPluginStream(modelId, handlerName, request) {
+export async function* dispatchPluginStream(modelId, handlerName, request, inputStream) {
     yield* profileStreamHandler({ op: handlerName, request }, async function* () {
-        const { result, streaming } = resolvePluginHandler(modelId, handlerName, request);
-        if (!streaming) {
-            throw new PluginHandlerTypeMismatchError(handlerName, "streaming", "reply");
+        const handlerDef = resolvePluginHandlerDef(modelId, handlerName);
+        if (inputStream) {
+            if (!handlerDef.duplex) {
+                throw new PluginHandlerTypeMismatchError(handlerName, "duplex", handlerDef.streaming ? "streaming" : "reply");
+            }
+            yield* handlerDef.handler(request, inputStream);
+        }
+        else {
+            if (handlerDef.duplex) {
+                throw new PluginHandlerTypeMismatchError(handlerName, "streaming", "duplex");
+            }
+            if (!handlerDef.streaming) {
+                throw new PluginHandlerTypeMismatchError(handlerName, "streaming", "reply");
+            }
+            yield* handlerDef.handler(request);
         }
-        yield* result;
     });
 }

package/dist/server/rpc/handlers/resume.d.ts

@@ -0,0 +1,3 @@
+import type { ResumeResponse } from "../../../schemas/index";
+export declare function handleResume(): Promise<ResumeResponse>;
+//# sourceMappingURL=resume.d.ts.map

package/dist/server/rpc/handlers/resume.js

@@ -0,0 +1,11 @@
+import { resumeRuntime } from "../../bare/runtime-lifecycle.js";
+import { LifecycleResumeFailedError } from "../../../utils/errors-server.js";
+export async function handleResume() {
+    try {
+        await resumeRuntime();
+        return { type: "resume" };
+    }
+    catch (error) {
+        throw new LifecycleResumeFailedError(error instanceof Error ? error.message : String(error), error);
+    }
+}

package/dist/server/rpc/handlers/suspend.d.ts

@@ -0,0 +1,3 @@
+import type { SuspendResponse } from "../../../schemas/index";
+export declare function handleSuspend(): Promise<SuspendResponse>;
+//# sourceMappingURL=suspend.d.ts.map

package/dist/server/rpc/handlers/suspend.js

@@ -0,0 +1,11 @@
+import { suspendRuntime } from "../../bare/runtime-lifecycle.js";
+import { LifecycleSuspendFailedError } from "../../../utils/errors-server.js";
+export async function handleSuspend() {
+    try {
+        await suspendRuntime();
+        return { type: "suspend" };
+    }
+    catch (error) {
+        throw new LifecycleSuspendFailedError(error instanceof Error ? error.message : String(error), error);
+    }
+}

package/dist/server/rpc/handlers/transcribe-stream.d.ts

@@ -1,3 +1,3 @@
 import type { TranscribeStreamRequest, TranscribeStreamResponse } from "../../../schemas/index";
-export declare function handleTranscribeStream(request: TranscribeStreamRequest): AsyncGenerator<TranscribeStreamResponse>;
+export declare function handleTranscribeStream(request: TranscribeStreamRequest, inputStream: AsyncIterable<Buffer>): AsyncGenerator<TranscribeStreamResponse>;
 //# sourceMappingURL=transcribe-stream.d.ts.map

package/dist/server/rpc/handlers/transcribe-stream.js

@@ -1,4 +1,4 @@
 import { dispatchPluginStream } from "./plugin-dispatch.js";
-export async function* handleTranscribeStream(request) {
-    yield* dispatchPluginStream(request.modelId, "transcribeStream", request);
+export async function* handleTranscribeStream(request, inputStream) {
+    yield* dispatchPluginStream(request.modelId, "transcribeStream", request, inputStream);
 }

package/dist/server/rpc/handlers/transcribe.d.ts

@@ -0,0 +1,3 @@
+import type { TranscribeRequest, TranscribeResponse } from "../../../schemas/index";
+export declare function handleTranscribe(request: TranscribeRequest): AsyncGenerator<TranscribeResponse>;
+//# sourceMappingURL=transcribe.d.ts.map

package/dist/server/rpc/handlers/transcribe.js

@@ -0,0 +1,4 @@
+import { dispatchPluginStream } from "./plugin-dispatch.js";
+export async function* handleTranscribe(request) {
+    yield* dispatchPluginStream(request.modelId, "transcribe", request);
+}

package/dist/server/rpc/profiling/operation-metrics.js

@@ -2,6 +2,8 @@
  * Declarative extraction of operation-level metrics from request/response data.
  * Used by operation wrappers to capture handler-specific profiling events.
  */
+import {} from "../../../schemas/index.js";
+import { readModelExecutionMs } from "../../../profiling/model-execution.js";
 const metricsRegistry = new Map();
 function extractDownloadStatsGauges(stats, gauges) {
     if (stats?.downloadTimeMs !== undefined) {
@@ -72,15 +74,16 @@ registerOperationMetrics({
     kind: "handler",
     getTags: (req) => (req.modelId ? { modelId: req.modelId } : {}),
     fromFinalChunk: (res) => {
-        if (!res.stats)
-            return undefined;
         const gauges = {};
-        if (res.stats.timeToFirstToken !== undefined)
+        if (res.stats?.timeToFirstToken !== undefined)
             gauges["timeToFirstToken"] = res.stats.timeToFirstToken;
-        if (res.stats.tokensPerSecond !== undefined)
+        if (res.stats?.tokensPerSecond !== undefined)
             gauges["tokensPerSecond"] = res.stats.tokensPerSecond;
-        if (res.stats.cacheTokens !== undefined)
+        if (res.stats?.cacheTokens !== undefined)
             gauges["cacheTokens"] = res.stats.cacheTokens;
+        const modelExecMs = readModelExecutionMs(res);
+        if (modelExecMs !== undefined)
+            gauges["modelExecutionTime"] = modelExecMs;
         return Object.keys(gauges).length > 0 ? gauges : undefined;
     },
 });
@@ -89,13 +92,27 @@ registerOperationMetrics({
     kind: "handler",
     getTags: (req) => (req.modelId ? { modelId: req.modelId } : {}),
     fromFinalChunk: (res) => {
-        if (!res.stats)
-            return undefined;
         const gauges = {};
-        if (res.stats.processedTokens !== undefined)
-            gauges["processedTokens"] = res.stats.processedTokens;
-        if (res.stats.processingTime !== undefined)
-            gauges["processingTime"] = res.stats.processingTime;
+        const modelExecMs = readModelExecutionMs(res);
+        if (modelExecMs !== undefined)
+            gauges["modelExecutionTime"] = modelExecMs;
+        // Common stats
+        if (res.stats?.totalTime !== undefined)
+            gauges["totalTime"] = res.stats.totalTime;
+        if (res.stats?.totalTokens !== undefined)
+            gauges["totalTokens"] = res.stats.totalTokens;
+        if (res.stats?.tokensPerSecond !== undefined)
+            gauges["tokensPerSecond"] = res.stats.tokensPerSecond;
+        if (res.stats?.timeToFirstToken !== undefined)
+            gauges["timeToFirstToken"] = res.stats.timeToFirstToken;
+        // NMT-specific
+        if (res.stats?.decodeTime !== undefined)
+            gauges["decodeTime"] = res.stats.decodeTime;
+        if (res.stats?.encodeTime !== undefined)
+            gauges["encodeTime"] = res.stats.encodeTime;
+        // LLM-specific
+        if (res.stats?.cacheTokens !== undefined)
+            gauges["cacheTokens"] = res.stats.cacheTokens;
         return Object.keys(gauges).length > 0 ? gauges : undefined;
     },
 });
@@ -103,31 +120,105 @@ registerOperationMetrics({
     op: "transcribeStream",
     kind: "handler",
     getTags: (req) => (req.modelId ? { modelId: req.modelId } : {}),
+    fromFinalChunk: (res) => {
+        const gauges = {};
+        const modelExecMs = readModelExecutionMs(res);
+        if (modelExecMs !== undefined)
+            gauges["modelExecutionTime"] = modelExecMs;
+        // Common stats
+        if (res.stats?.audioDuration !== undefined)
+            gauges["audioDuration"] = res.stats.audioDuration;
+        if (res.stats?.realTimeFactor !== undefined)
+            gauges["realTimeFactor"] = res.stats.realTimeFactor;
+        if (res.stats?.tokensPerSecond !== undefined)
+            gauges["tokensPerSecond"] = res.stats.tokensPerSecond;
+        if (res.stats?.totalTokens !== undefined)
+            gauges["totalTokens"] = res.stats.totalTokens;
+        if (res.stats?.totalSegments !== undefined)
+            gauges["totalSegments"] = res.stats.totalSegments;
+        // whisper-specific timings
+        if (res.stats?.whisperEncodeTime !== undefined)
+            gauges["whisperEncodeTime"] = res.stats.whisperEncodeTime;
+        if (res.stats?.whisperDecodeTime !== undefined)
+            gauges["whisperDecodeTime"] = res.stats.whisperDecodeTime;
+        // parakeet-specific timings
+        if (res.stats?.encoderTime !== undefined)
+            gauges["encoderTime"] = res.stats.encoderTime;
+        if (res.stats?.decoderTime !== undefined)
+            gauges["decoderTime"] = res.stats.decoderTime;
+        if (res.stats?.melSpecTime !== undefined)
+            gauges["melSpecTime"] = res.stats.melSpecTime;
+        return Object.keys(gauges).length > 0 ? gauges : undefined;
+    },
 });
 registerOperationMetrics({
     op: "textToSpeech",
     kind: "handler",
     getTags: (req) => (req.modelId ? { modelId: req.modelId } : {}),
+    fromFinalChunk: (res) => {
+        const gauges = {};
+        const modelExecMs = readModelExecutionMs(res);
+        if (modelExecMs !== undefined)
+            gauges["modelExecutionTime"] = modelExecMs;
+        if (res.stats?.audioDuration !== undefined)
+            gauges["audioDuration"] = res.stats.audioDuration;
+        if (res.stats?.totalSamples !== undefined)
+            gauges["totalSamples"] = res.stats.totalSamples;
+        return Object.keys(gauges).length > 0 ? gauges : undefined;
+    },
 });
 registerOperationMetrics({
     op: "embed",
     kind: "handler",
     getTags: (req) => (req.modelId ? { modelId: req.modelId } : {}),
+    fromResult: (res) => {
+        const gauges = {};
+        const modelExecMs = readModelExecutionMs(res);
+        if (modelExecMs !== undefined)
+            gauges["modelExecutionTime"] = modelExecMs;
+        if (res.stats?.totalTime !== undefined)
+            gauges["totalTime"] = res.stats.totalTime;
+        if (res.stats?.tokensPerSecond !== undefined)
+            gauges["tokensPerSecond"] = res.stats.tokensPerSecond;
+        if (res.stats?.totalTokens !== undefined)
+            gauges["totalTokens"] = res.stats.totalTokens;
+        return Object.keys(gauges).length > 0 ? gauges : undefined;
+    },
 });
 registerOperationMetrics({
     op: "ocrStream",
     kind: "handler",
     getTags: (req) => (req.modelId ? { modelId: req.modelId } : {}),
     fromFinalChunk: (res) => {
-        if (!res.stats)
-            return undefined;
         const gauges = {};
-        if (res.stats.detectionTime !== undefined)
+        if (res.stats?.detectionTime !== undefined)
             gauges["detectionTime"] = res.stats.detectionTime;
-        if (res.stats.recognitionTime !== undefined)
+        if (res.stats?.recognitionTime !== undefined)
             gauges["recognitionTime"] = res.stats.recognitionTime;
-        if (res.stats.totalTime !== undefined)
+        if (res.stats?.totalTime !== undefined)
             gauges["totalTime"] = res.stats.totalTime;
+        const modelExecMs = readModelExecutionMs(res);
+        if (modelExecMs !== undefined)
+            gauges["modelExecutionTime"] = modelExecMs;
+        return Object.keys(gauges).length > 0 ? gauges : undefined;
+    },
+});
+registerOperationMetrics({
+    op: "diffusionStream",
+    kind: "handler",
+    getTags: (req) => (req.modelId ? { modelId: req.modelId } : {}),
+    fromFinalChunk: (res) => {
+        if (!res.stats)
+            return undefined;
+        const gauges = {};
+        if (res.stats.generationMs !== undefined)
+            gauges["generationMs"] = res.stats.generationMs;
+        if (res.stats.totalSteps !== undefined)
+            gauges["totalSteps"] = res.stats.totalSteps;
+        if (res.stats.totalImages !== undefined)
+            gauges["totalImages"] = res.stats.totalImages;
+        if (res.stats.totalPixels !== undefined)
+            gauges["totalPixels"] = res.stats.totalPixels;
         return Object.keys(gauges).length > 0 ? gauges : undefined;
     },
 });

package/dist/server/rpc/profiling/operation-wrappers.d.ts

@@ -14,5 +14,5 @@ export interface ProfiledStreamOptions<TRequest> {
     perCall?: PerCallProfiling;
 }
 export declare function profileReplyHandler<TRequest, TResponse>(options: ProfiledReplyOptions<TRequest>, handler: () => Promise<TResponse>): Promise<TResponse>;
-export declare function profileStreamHandler<TRequest, TResponse>(options: ProfiledStreamOptions<TRequest>, handler: () => AsyncGenerator<TResponse>): AsyncGenerator<TResponse>;
+export declare function profileStreamHandler<TRequest, TResponse, TReturn = unknown>(options: ProfiledStreamOptions<TRequest>, handler: () => AsyncGenerator<TResponse, TReturn>): AsyncGenerator<TResponse, TReturn>;
 //# sourceMappingURL=operation-wrappers.d.ts.map