Automatically transform binary inputs into data uris

Author: aron

README.md

@@ -67,24 +67,15 @@ console.log(prediction.output);
 // ['https://replicate.delivery/pbxt/RoaxeXqhL0xaYyLm6w3bpGwF5RaNBjADukfFnMbhOyeoWBdhA/out-0.png']
 ```
 
-To run a model that takes a file input, pass a URL to a publicly accessible file. Or, for smaller files (<10MB), you can convert file data into a base64-encoded data URI and pass that directly:
+To run a model that takes a file input, pass a URL to a publicly accessible file. Or, for smaller files (<10MB), you can pass the data directly.
 
 
 ```js
 import { promises as fs } from "fs";
 
-// Read the file into a buffer
-const data = await fs.readFile("path/to/image.png");
-// Convert the buffer into a base64-encoded string
-const base64 = data.toString("base64");
-// Set MIME type for PNG image
-const mimeType = "image/png";
-// Create the data URI
-const dataURI = `data:${mimeType};base64,${base64}`;
-
 const model = "nightmareai/real-esrgan:42fed1c4974146d4d2414e2be2c5277c7fcf05fcc3a73abf41610695738c1d7b";
 const input = {
-  image: dataURI,
+  image: await fs.readFile("path/to/image.png"),
 };
 const output = await replicate.run(model, { input });
 // ['https://replicate.delivery/mgxm/e7b0e122-9daa-410e-8cde-006c7308ff4d/output.png']

index.test.ts

@@ -200,6 +200,46 @@ describe("Replicate client", () => {
       expect(prediction.id).toBe("ufawqhfynnddngldkgtslldrkq");
     });
 
+    test.each([
+      {
+        type: "file",
+        value: new File(["hello world"], "hello.txt", { type: "text/plain" }),
+        expected: "data:text/plain;base64,aGVsbG8gd29ybGQ=",
+      },
+      {
+        type: "blob",
+        value: new Blob(["hello world"], { type: "text/plain" }),
+        expected: "data:text/plain;base64,aGVsbG8gd29ybGQ=",
+      },
+      {
+        type: "buffer",
+        value: Buffer.from("hello world"),
+        expected: "data:application/octet-stream;base64,aGVsbG8gd29ybGQ=",
+      },
+    ])(
+      "converts a $type input into a base64 encoded string",
+      async ({ value: data, expected }) => {
+        let body: Record<string, any>;
+        nock(BASE_URL)
+          .post("/predictions")
+          .reply(201, (_uri, _body: Record<string, any>) => {
+            return (body = _body);
+          });
+
+        await client.predictions.create({
+          version:
+            "5c7d5dc6dd8bf75c1acaa8565735e7986bc5b66206b55cca93cb72c9bf15ccaa",
+          input: {
+            prompt: "Tell me a story",
+            data,
+          },
+          stream: true,
+        });
+
+        expect(body!.input.data).toEqual(expected);
+      }
+    );
+
     test("Passes stream parameter to API endpoint", async () => {
       nock(BASE_URL)
         .post("/predictions")

lib/deployments.js

@@ -1,3 +1,5 @@
+const { transformFileInputs } = require("./util");
+
 /**
  * Create a new prediction with a deployment
  *
@@ -11,7 +13,7 @@
  * @returns {Promise<object>} Resolves with the created prediction data
  */
 async function createPrediction(deployment_owner, deployment_name, options) {
-  const { stream, ...data } = options;
+  const { stream, input: _input, ...data } = options;
 
   if (data.webhook) {
     try {
@@ -22,11 +24,14 @@ async function createPrediction(deployment_owner, deployment_name, options) {
     }
   }
 
+  // Transform any file looking fields into strings (either data uri or urls).
+  const input = await transformFileInputs(_input);
+
   const response = await this.request(
     `/deployments/${deployment_owner}/${deployment_name}/predictions`,
     {
       method: "POST",
-      data: { ...data, stream },
+      data: { ...data, input, stream },
     }
   );
 

lib/predictions.js

@@ -1,3 +1,5 @@
+const { transformFileInputs } = require("./util");
+
 /**
  * Create a new prediction
  *
@@ -11,7 +13,7 @@
  * @returns {Promise<object>} Resolves with the created prediction
  */
 async function createPrediction(options) {
-  const { model, version, stream, ...data } = options;
+  const { model, version, stream, input: _input, ...data } = options;
 
   if (data.webhook) {
     try {
@@ -22,16 +24,19 @@ async function createPrediction(options) {
     }
   }
 
+  // Transform any file looking fields into strings (either data uri or urls).
+  const input = await transformFileInputs(_input);
+
   let response;
   if (version) {
     response = await this.request("/predictions", {
       method: "POST",
-      data: { ...data, stream, version },
+      data: { ...data, stream, input, version },
     });
   } else if (model) {
     response = await this.request(`/models/${model}/predictions`, {
       method: "POST",
-      data: { ...data, stream },
+      data: { ...data, stream, input },
     });
   } else {
     throw new Error("Either model or version must be specified");