Bridge reference-to-video requests to Python providers

Author: georgi

packages/protocol/src/bridge-protocol.ts

@@ -46,7 +46,7 @@
  *   3. Update `MIN_NODETOOL_CORE_VERSION` to that new release.
  */
 
-export const BRIDGE_PROTOCOL_VERSION = 5;
+export const BRIDGE_PROTOCOL_VERSION = 6;
 
 /**
  * Hard floor: the JS runtime rejects (at `discover`) any worker reporting a

packages/protocol/tests/error-and-protocol-coverage.test.ts

@@ -79,7 +79,7 @@ describe("error-helpers.isTRPCErrorWithCode", () => {
 
 describe("bridge-protocol constants", () => {
   it("BRIDGE_PROTOCOL_VERSION is the current speaking version", () => {
-    expect(BRIDGE_PROTOCOL_VERSION).toBe(5);
+    expect(BRIDGE_PROTOCOL_VERSION).toBe(6);
   });
 
   it("MIN_BRIDGE_PROTOCOL_VERSION is the hard floor at 1", () => {

packages/runtime/src/providers/python-provider.ts

@@ -25,6 +25,8 @@ import type {
   TextToSpeechParams,
   TextToVideoParams,
   ImageToVideoParams,
+  ReferenceToVideoInputs,
+  ReferenceToVideoParams,
   TextToMusicParams,
   EncodedAudioResult
 } from "./types.js";
@@ -79,6 +81,7 @@ export class PythonProvider extends BaseProvider {
   private _secrets: Record<string, string>;
   private _supportsStreamingTTS = true;
   private _supportsEncodedTTS = true;
+  private _workerCapabilities = new Set<string>();
 
   constructor(
     providerId: string,
@@ -99,19 +102,26 @@ export class PythonProvider extends BaseProvider {
       this._bridge = bridge;
       this._pythonProviderId = providerIdOrOptions;
       this._secrets = secrets;
+      this.referenceToVideo = BaseProvider.prototype.referenceToVideo;
       return;
     }
 
     const { _id, _bridge, _bridgeProviderId, _capabilities, ...rawSecrets } =
       providerIdOrOptions;
     super(_id);
+    const wrappedReferenceToVideo = this.referenceToVideo;
     this._bridge = _bridge;
     this._pythonProviderId = _bridgeProviderId ?? _id;
+    this.referenceToVideo = BaseProvider.prototype.referenceToVideo;
     if (Array.isArray(_capabilities)) {
+      this._workerCapabilities = new Set(_capabilities.map(String));
       this._supportsStreamingTTS = _capabilities.includes("text_to_speech");
       this._supportsEncodedTTS = _capabilities.includes(
         "text_to_speech_encoded"
       );
+      if (this._workerCapabilities.has("reference_to_video")) {
+        this.referenceToVideo = wrappedReferenceToVideo;
+      }
     }
     this._secrets = Object.fromEntries(
       Object.entries(rawSecrets).filter(
@@ -187,10 +197,25 @@ export class PythonProvider extends BaseProvider {
     // The public provider id may be an alias (notably `huggingface-local`) so
     // selections route back through this bridge adapter instead of colliding
     // with a built-in remote provider that uses the worker's original id.
-    return models.map((model) => ({
-      ...model,
-      provider: this.provider
-    }));
+    return models.map((model) => {
+      const supportedTasks = Array.isArray(model.supportedTasks)
+        ? model.supportedTasks.map(String)
+        : Array.isArray(model.supported_tasks)
+          ? model.supported_tasks.map(String)
+          : undefined;
+      const normalizedModel: Record<string, unknown> = {
+        ...model,
+        provider: this.provider
+      };
+      if (modelType === "video" && supportedTasks) {
+        normalizedModel.supportedTasks = this._workerCapabilities.has(
+          "reference_to_video"
+        )
+          ? supportedTasks
+          : supportedTasks.filter((task) => task !== "reference_to_video");
+      }
+      return normalizedModel;
+    });
   }
 
   // ── Chat completion ───────────────────────────────────────────────
@@ -327,6 +352,47 @@ export class PythonProvider extends BaseProvider {
     );
   }
 
+  async referenceToVideo(
+    inputs: ReferenceToVideoInputs,
+    params: ReferenceToVideoParams
+  ): Promise<Uint8Array> {
+    if (!this._workerCapabilities.has("reference_to_video")) {
+      throw new Error("Python worker does not support reference_to_video");
+    }
+    if (
+      Array.isArray(params.model.supportedTasks) &&
+      !params.model.supportedTasks.includes("reference_to_video")
+    ) {
+      throw new Error(
+        `Video model ${params.model.id} does not support reference_to_video`
+      );
+    }
+    if (!Array.isArray(params.model.supportedTasks)) {
+      const models = await this._getModels("video");
+      const model = models.find(
+        (candidate) =>
+          isRecord(candidate) && String(candidate.id ?? "") === params.model.id
+      );
+      if (
+        !isRecord(model) ||
+        !Array.isArray(model.supportedTasks) ||
+        !model.supportedTasks.includes("reference_to_video")
+      ) {
+        throw new Error(
+          `Video model ${params.model.id} does not support reference_to_video`
+        );
+      }
+    }
+    const { signal, ...wireParams } = params;
+    return this._bridge.providerReferenceToVideo(
+      this._pythonProviderId,
+      inputs,
+      { ...wireParams, model: params.model.id },
+      this._secrets,
+      signal
+    );
+  }
+
   async *textToSpeech(
     args: TextToSpeechParams
   ): AsyncGenerator<StreamingAudioChunk> {

packages/runtime/src/python-bridge-base.ts

@@ -113,7 +113,8 @@ import type {
   BlenderEvent,
   BlenderExecuteOptions,
   BlenderExecuteResult,
-  PythonBridge
+  PythonBridge,
+  ReferenceToVideoInputs
 } from "./python-bridge-types.js";
 import {
   comfyStatusInfoSchema,
@@ -1160,6 +1161,56 @@ export abstract class PythonBridgeBase
     return result.blobs.video;
   }
 
+  async providerReferenceToVideo(
+    providerId: string,
+    inputs: ReferenceToVideoInputs,
+    params: Record<string, unknown>,
+    secrets?: Record<string, string>,
+    signal?: AbortSignal
+  ): Promise<Uint8Array> {
+    const { images, videos } = inputs;
+    if (
+      !Array.isArray(images) ||
+      !Array.isArray(videos) ||
+      (images.length === 0 && videos.length === 0)
+    ) {
+      throw new Error(
+        "reference_to_video requires at least one reference image or video"
+      );
+    }
+    if (
+      [...images, ...videos].some((bytes) => !(bytes instanceof Uint8Array))
+    ) {
+      throw new Error(
+        "reference_to_video reference media must be binary buffers"
+      );
+    }
+    if ([...images, ...videos].some((bytes) => bytes.byteLength === 0)) {
+      throw new Error("reference_to_video reference media must be non-empty");
+    }
+    const totalBytes = [...images, ...videos].reduce(
+      (total, bytes) => total + bytes.byteLength,
+      0
+    );
+    if (totalBytes > 192 * 1024 * 1024) {
+      throw new Error(
+        `reference_to_video reference media is ${totalBytes} bytes, exceeding the 201326592 byte limit`
+      );
+    }
+    const result = await this._providerBlobCall(
+      "provider.reference_to_video",
+      {
+        provider: providerId,
+        reference_images: images,
+        reference_videos: videos,
+        params,
+        secrets: secrets ?? {}
+      },
+      signal
+    );
+    return result.blobs.video;
+  }
+
   async providerTextToAudio(
     providerId: string,
     params: Record<string, unknown>,

packages/runtime/src/python-bridge-types.ts

@@ -2,6 +2,8 @@ import { EventEmitter } from "node:events";
 import { z } from "zod";
 
 import type { ASRResult } from "./providers/types.js";
+import type { ReferenceToVideoInputs } from "./providers/types.js";
+export type { ReferenceToVideoInputs } from "./providers/types.js";
 
 interface NodeMetadataProperty {
   name: string;
@@ -590,6 +592,13 @@ export interface PythonBridge extends EventEmitter {
     secrets?: Record<string, string>,
     signal?: AbortSignal
   ): Promise<Uint8Array>;
+  providerReferenceToVideo(
+    providerId: string,
+    inputs: ReferenceToVideoInputs,
+    params: Record<string, unknown>,
+    secrets?: Record<string, string>,
+    signal?: AbortSignal
+  ): Promise<Uint8Array>;
   providerTextToAudio(
     providerId: string,
     params: Record<string, unknown>,

packages/runtime/src/swappable-python-bridge.ts

@@ -278,6 +278,22 @@ export class SwappableBridge extends EventEmitter implements PythonBridge {
     );
   }
 
+  providerReferenceToVideo(
+    providerId: string,
+    inputs: import("./providers/types.js").ReferenceToVideoInputs,
+    params: Record<string, unknown>,
+    secrets?: Record<string, string>,
+    signal?: AbortSignal
+  ): Promise<Uint8Array> {
+    return this._target.providerReferenceToVideo(
+      providerId,
+      inputs,
+      params,
+      secrets,
+      signal
+    );
+  }
+
   providerTextToAudio(
     providerId: string,
     params: Record<string, unknown>,

packages/runtime/tests/providers/provider-registry-extended.test.ts

@@ -214,6 +214,110 @@ describe("provider-registry — extended coverage", () => {
     );
   });
 
+  it("gates reference video on worker capability and model support", async () => {
+    const oldWorkerModels = vi.fn(async () => [
+      {
+        id: "misreported",
+        name: "Misreported",
+        supported_tasks: ["image_to_video", "reference_to_video"]
+      }
+    ]);
+    const oldWorker = new PythonProvider({
+      _id: "wangp-old",
+      _bridge: { getProviderModels: oldWorkerModels }
+    } as any);
+    expect(oldWorker.getCapabilities()).not.toContain("reference_to_video");
+    await expect(oldWorker.getAvailableVideoModels()).resolves.toEqual([
+      expect.objectContaining({
+        supportedTasks: ["image_to_video"],
+        provider: "wangp-old"
+      })
+    ]);
+    await expect(
+      oldWorker.referenceToVideo(
+        { images: [new Uint8Array([1])], videos: [] },
+        {
+          model: {
+            id: "misreported",
+            name: "Misreported",
+            provider: "wangp-old"
+          }
+        }
+      )
+    ).rejects.toThrow("does not support referenceToVideo");
+
+    const referenceToVideo = vi.fn(async () => new Uint8Array([9]));
+    const models = vi.fn(async () => [
+      {
+        id: "start-only",
+        name: "Start only",
+        supportedTasks: ["image_to_video"]
+      },
+      {
+        id: "reference-model",
+        name: "Reference",
+        supportedTasks: ["reference_to_video"]
+      }
+    ]);
+    const worker = new PythonProvider({
+      _id: "wangp",
+      _capabilities: ["reference_to_video"],
+      _bridge: {
+        getProviderModels: models,
+        providerReferenceToVideo: referenceToVideo
+      },
+      API_KEY: "secret"
+    } as any);
+    expect(worker.getCapabilities()).toContain("reference_to_video");
+    await expect(
+      worker.referenceToVideo(
+        {
+          images: [new Uint8Array([1]), new Uint8Array([2])],
+          videos: [new Uint8Array([3])]
+        },
+        {
+          model: { id: "start-only", name: "Start only", provider: "wangp" },
+          prompt: "reject"
+        }
+      )
+    ).rejects.toThrow("start-only does not support reference_to_video");
+    expect(referenceToVideo).not.toHaveBeenCalled();
+
+    const signal = new AbortController().signal;
+    await expect(
+      worker.referenceToVideo(
+        {
+          images: [new Uint8Array([1]), new Uint8Array([2])],
+          videos: [new Uint8Array([3])]
+        },
+        {
+          model: {
+            id: "reference-model",
+            name: "Reference",
+            provider: "wangp"
+          },
+          prompt: "forward",
+          useReferenceVideoAudio: true,
+          signal
+        }
+      )
+    ).resolves.toEqual(new Uint8Array([9]));
+    expect(referenceToVideo).toHaveBeenCalledWith(
+      "wangp",
+      {
+        images: [new Uint8Array([1]), new Uint8Array([2])],
+        videos: [new Uint8Array([3])]
+      },
+      {
+        model: "reference-model",
+        prompt: "forward",
+        useReferenceVideoAudio: true
+      },
+      { API_KEY: "secret" },
+      signal
+    );
+  });
+
   it("discovers music and routes encoded audio through the Python bridge", async () => {
     const wav = new Uint8Array([
       0x52, 0x49, 0x46, 0x46, 0, 0, 0, 0, 0x57, 0x41, 0x56, 0x45