add image-to-image for replicate

Author: zeke

packages/inference/src/lib/getProviderHelper.ts

@@ -138,6 +138,7 @@ export const PROVIDERS: Record<InferenceProvider, Partial<Record<InferenceTask,
 		"text-to-image": new Replicate.ReplicateTextToImageTask(),
 		"text-to-speech": new Replicate.ReplicateTextToSpeechTask(),
 		"text-to-video": new Replicate.ReplicateTextToVideoTask(),
+		"image-to-image": new Replicate.ReplicateImageToImageTask(),
 	},
 	sambanova: {
 		conversational: new Sambanova.SambanovaConversationalTask(),

packages/inference/src/providers/replicate.ts

@@ -18,7 +18,9 @@ import { InferenceClientProviderOutputError } from "../errors.js";
 import { isUrl } from "../lib/isUrl.js";
 import type { BodyParams, HeaderParams, UrlParams } from "../types.js";
 import { omit } from "../utils/omit.js";
-import { TaskProviderHelper, type TextToImageTaskHelper, type TextToVideoTaskHelper } from "./providerHelper.js";
+import { TaskProviderHelper, type ImageToImageTaskHelper, type TextToImageTaskHelper, type TextToVideoTaskHelper } from "./providerHelper.js";
+import type { ImageToImageArgs } from "../tasks/cv/imageToImage.js";
+import { base64FromBytes } from "../utils/base64FromBytes.js";
 export interface ReplicateOutput {
 	output?: string | string[];
 }
@@ -152,3 +154,65 @@ export class ReplicateTextToVideoTask extends ReplicateTask implements TextToVid
 		throw new InferenceClientProviderOutputError("Received malformed response from Replicate text-to-video API");
 	}
 }
+
+export class ReplicateImageToImageTask extends ReplicateTask implements ImageToImageTaskHelper {
+	override preparePayload(params: BodyParams): Record<string, unknown> {
+		const inputs = params.args.inputs as Blob;
+		const parameters = params.args.parameters as Record<string, unknown> | undefined;
+		
+		return {
+			input: {
+				...omit(params.args, ["inputs", "parameters"]),
+				...(parameters || {}),
+				prompt: parameters?.prompt || "",
+				image: inputs, // This will be processed in preparePayloadAsync
+			},
+			version: params.model.includes(":") ? params.model.split(":")[1] : undefined,
+		};
+	}
+
+	async preparePayloadAsync(args: ImageToImageArgs): Promise<import("../types.js").RequestArgs> {
+		const { inputs, parameters, ...restArgs } = args;
+		
+		// Convert Blob to base64 data URL
+		const bytes = new Uint8Array(await inputs.arrayBuffer());
+		const base64 = base64FromBytes(bytes);
+		const imageInput = `data:${inputs.type || "image/jpeg"};base64,${base64}`;
+
+		return {
+			...restArgs,
+			inputs: imageInput,
+			parameters: {
+				...parameters,
+				image: imageInput,
+			},
+		};
+	}
+
+	override async getResponse(response: ReplicateOutput): Promise<Blob> {
+		if (
+			typeof response === "object" &&
+			!!response &&
+			"output" in response &&
+			Array.isArray(response.output) &&
+			response.output.length > 0 &&
+			typeof response.output[0] === "string"
+		) {
+			const urlResponse = await fetch(response.output[0]);
+			return await urlResponse.blob();
+		}
+
+		if (
+			typeof response === "object" &&
+			!!response &&
+			"output" in response &&
+			typeof response.output === "string" &&
+			isUrl(response.output)
+		) {
+			const urlResponse = await fetch(response.output);
+			return await urlResponse.blob();
+		}
+
+		throw new InferenceClientProviderOutputError("Received malformed response from Replicate image-to-image API");
+	}
+}

packages/inference/test/InferenceClient.spec.ts

@@ -7,6 +7,7 @@ import {
 	chatCompletion,
 	chatCompletionStream,
 	HfInference,
+	imageToImage,
 	InferenceClient,
 	textGeneration,
 	textToImage,
@@ -1277,6 +1278,18 @@ describe.skip("InferenceClient", () => {
 
 				expect(res).toBeInstanceOf(Blob);
 			});
+
+			it("imageToImage - FLUX Kontext Dev", async () => {
+				const res = await client.imageToImage({
+					model: "black-forest-labs/flux-kontext-dev",
+					provider: "replicate",
+					inputs: new Blob([readTestFile("stormtrooper_depth.png")], { type: "image/png" }),
+					parameters: {
+						prompt: "Change the stormtrooper armor to golden color while keeping the same pose and helmet design",
+					},
+				});
+				expect(res).toBeInstanceOf(Blob);
+			});
 		},
 		TIMEOUT
 	);