[Inference] text-to-image: option to return a URL instead of Blob (#1207)

Context (internal):
https://huggingface.slack.com/archives/C07KX53FZTK/p1739562014823989?thread_ts=1739561786.364929&cid=C07KX53FZTK

Author: SBrandeis

packages/inference/src/tasks/cv/textToImage.ts

@@ -24,6 +24,10 @@ interface BlackForestLabsResponse {
 	polling_url: string;
 }
 
+interface TextToImageOptions extends Options {
+	outputType?: "url" | "blob";
+}
+
 function getResponseFormatArg(provider: InferenceProvider) {
 	switch (provider) {
 		case "fal-ai":
@@ -43,7 +47,15 @@ function getResponseFormatArg(provider: InferenceProvider) {
  * This task reads some text input and outputs an image.
  * Recommended model: stabilityai/stable-diffusion-2
  */
-export async function textToImage(args: TextToImageArgs, options?: Options): Promise<Blob> {
+export async function textToImage(
+	args: TextToImageArgs,
+	options?: TextToImageOptions & { outputType: "url" }
+): Promise<string>;
+export async function textToImage(
+	args: TextToImageArgs,
+	options?: TextToImageOptions & { outputType?: undefined | "blob" }
+): Promise<Blob>;
+export async function textToImage(args: TextToImageArgs, options?: TextToImageOptions): Promise<Blob | string> {
 	const payload =
 		!args.provider || args.provider === "hf-inference" || args.provider === "sambanova"
 			? args
@@ -66,11 +78,18 @@ export async function textToImage(args: TextToImageArgs, options?: Options): Pro
 
 	if (res && typeof res === "object") {
 		if (args.provider === "black-forest-labs" && "polling_url" in res && typeof res.polling_url === "string") {
+			if (options?.outputType === "url") {
+				return res.polling_url;
+			}
 			return await pollBflResponse(res.polling_url);
 		}
 		if (args.provider === "fal-ai" && "images" in res && Array.isArray(res.images) && res.images[0].url) {
-			const image = await fetch(res.images[0].url);
-			return await image.blob();
+			if (options?.outputType === "url") {
+				return res.images[0].url;
+			} else {
+				const image = await fetch(res.images[0].url);
+				return await image.blob();
+			}
 		}
 		if (
 			args.provider === "hyperbolic" &&
@@ -79,17 +98,24 @@ export async function textToImage(args: TextToImageArgs, options?: Options): Pro
 			res.images[0] &&
 			typeof res.images[0].image === "string"
 		) {
+			if (options?.outputType === "url") {
+				return `data:image/jpeg;base64,${res.images[0].image}`;
+			}
 			const base64Response = await fetch(`data:image/jpeg;base64,${res.images[0].image}`);
-			const blob = await base64Response.blob();
-			return blob;
+			return await base64Response.blob();
 		}
 		if ("data" in res && Array.isArray(res.data) && res.data[0].b64_json) {
 			const base64Data = res.data[0].b64_json;
+			if (options?.outputType === "url") {
+				return `data:image/jpeg;base64,${base64Data}`;
+			}
 			const base64Response = await fetch(`data:image/jpeg;base64,${base64Data}`);
-			const blob = await base64Response.blob();
-			return blob;
+			return await base64Response.blob();
 		}
 		if ("output" in res && Array.isArray(res.output)) {
+			if (options?.outputType === "url") {
+				return res.output[0];
+			}
 			const urlResponse = await fetch(res.output[0]);
 			const blob = await urlResponse.blob();
 			return blob;
@@ -99,6 +125,10 @@ export async function textToImage(args: TextToImageArgs, options?: Options): Pro
 	if (!isValidOutput) {
 		throw new InferenceOutputError("Expected Blob");
 	}
+	if (options?.outputType === "url") {
+		const b64 = await res.arrayBuffer().then((buf) => Buffer.from(buf).toString("base64"));
+		return `data:image/jpeg;base64,${b64}`;
+	}
 	return res;
 }
 

packages/inference/test/HfInference.spec.ts

@@ -8,6 +8,7 @@ import { textToVideo } from "../src/tasks/cv/textToVideo";
 import { readTestFile } from "./test-files";
 import "./vcr";
 import { HARDCODED_MODEL_ID_MAPPING } from "../src/providers/consts";
+import { isUrl } from "../src/lib/isUrl";
 
 const TIMEOUT = 60000 * 3;
 const env = import.meta.env;
@@ -1326,6 +1327,26 @@ describe.concurrent("HfInference", () => {
 				});
 				expect(res).toBeInstanceOf(Blob);
 			});
+
+			it("textToImage URL", async () => {
+				const res = await textToImage(
+					{
+						model: "black-forest-labs/FLUX.1-dev",
+						provider: "black-forest-labs",
+						accessToken: env.HF_BLACK_FOREST_LABS_KEY,
+						inputs: "A raccoon driving a truck",
+						parameters: {
+							height: 256,
+							width: 256,
+							num_inference_steps: 4,
+							seed: 8817,
+						},
+					},
+					{ outputType: "url" }
+				);
+				expect(res).toBeTypeOf("string");
+				expect(isUrl(res)).toBeTruthy();
+			});
 		},
 		TIMEOUT
 	);