add image-segmentation support for fal (#1602)

adds `image-segmentation` task for fal, following internal discussion
[here](https://huggingface.slack.com/archives/C0664PDFGSJ/p1751531970106289?thread_ts=1748529796.374799&cid=C0664PDFGSJ)

cc @Vaibhavs10

---------

Co-authored-by: Celina Hanouti <[email protected]>

Author: linoytsaban

packages/inference/src/lib/getProviderHelper.ts

@@ -67,6 +67,7 @@ export const PROVIDERS: Record<InferenceProvider, Partial<Record<InferenceTask,
 		"text-to-video": new FalAI.FalAITextToVideoTask(),
 		"image-to-image": new FalAI.FalAIImageToImageTask(),
 		"automatic-speech-recognition": new FalAI.FalAIAutomaticSpeechRecognitionTask(),
+		"image-segmentation": new FalAI.FalAIImageSegmentationTask(),
 	},
 	"featherless-ai": {
 		conversational: new FeatherlessAI.FeatherlessAIConversationalTask(),

packages/inference/src/providers/fal-ai.ts

@@ -16,12 +16,12 @@
  */
 import { base64FromBytes } from "../utils/base64FromBytes.js";
 
-import type { AutomaticSpeechRecognitionOutput } from "@huggingface/tasks";
+import type { AutomaticSpeechRecognitionOutput, ImageSegmentationOutput } from "@huggingface/tasks";
 import { isUrl } from "../lib/isUrl.js";
 import type { BodyParams, HeaderParams, InferenceTask, ModelId, RequestArgs, UrlParams } from "../types.js";
 import { delay } from "../utils/delay.js";
 import { omit } from "../utils/omit.js";
-import type { ImageToImageTaskHelper } from "./providerHelper.js";
+import type { ImageSegmentationTaskHelper, ImageToImageTaskHelper } from "./providerHelper.js";
 import {
 	type AutomaticSpeechRecognitionTaskHelper,
 	TaskProviderHelper,
@@ -36,6 +36,7 @@ import {
 	InferenceClientProviderOutputError,
 } from "../errors.js";
 import type { ImageToImageArgs } from "../tasks/index.js";
+import type { ImageSegmentationArgs } from "../tasks/cv/imageSegmentation.js";
 
 export interface FalAiQueueOutput {
 	request_id: string;
@@ -406,3 +407,87 @@ export class FalAITextToSpeechTask extends FalAITask {
 		}
 	}
 }
+export class FalAIImageSegmentationTask extends FalAiQueueTask implements ImageSegmentationTaskHelper {
+	task: InferenceTask;
+	constructor() {
+		super("https://queue.fal.run");
+		this.task = "image-segmentation";
+	}
+
+	override makeRoute(params: UrlParams): string {
+		if (params.authMethod !== "provider-key") {
+			return `/${params.model}?_subdomain=queue`;
+		}
+		return `/${params.model}`;
+	}
+
+	override preparePayload(params: BodyParams): Record<string, unknown> {
+		return {
+			...omit(params.args, ["inputs", "parameters"]),
+			...(params.args.parameters as Record<string, unknown>),
+			sync_mode: true,
+		};
+	}
+
+	async preparePayloadAsync(args: ImageSegmentationArgs): Promise<RequestArgs> {
+		const blob = "data" in args && args.data instanceof Blob ? args.data : "inputs" in args ? args.inputs : undefined;
+		const mimeType = blob instanceof Blob ? blob.type : "image/png";
+		const base64Image = base64FromBytes(
+			new Uint8Array(blob instanceof ArrayBuffer ? blob : await (blob as Blob).arrayBuffer())
+		);
+		return {
+			...omit(args, ["inputs", "parameters", "data"]),
+			...args.parameters,
+			...args,
+			image_url: `data:${mimeType};base64,${base64Image}`,
+			sync_mode: true,
+		};
+	}
+
+	override async getResponse(
+		response: FalAiQueueOutput,
+		url?: string,
+		headers?: Record<string, string>
+	): Promise<ImageSegmentationOutput> {
+		const result = await this.getResponseFromQueueApi(response, url, headers);
+		if (
+			typeof result === "object" &&
+			result !== null &&
+			"image" in result &&
+			typeof result.image === "object" &&
+			result.image !== null &&
+			"url" in result.image &&
+			typeof result.image.url === "string"
+		) {
+			const maskResponse = await fetch(result.image.url);
+			if (!maskResponse.ok) {
+				throw new InferenceClientProviderApiError(
+					`Failed to fetch segmentation mask from ${result.image.url}`,
+					{ url: result.image.url, method: "GET" },
+					{
+						requestId: maskResponse.headers.get("x-request-id") ?? "",
+						status: maskResponse.status,
+						body: await maskResponse.text(),
+					}
+				);
+			}
+			const maskBlob = await maskResponse.blob();
+			const maskArrayBuffer = await maskBlob.arrayBuffer();
+			const maskBase64 = base64FromBytes(new Uint8Array(maskArrayBuffer));
+
+			return [
+				{
+					label: "mask", // placeholder label, as Fal does not provide labels in the response(?)
+					score: 1.0, // placeholder score, as Fal does not provide scores in the response(?)
+					mask: maskBase64,
+				},
+			];
+		}
+
+		throw new InferenceClientProviderOutputError(
+			`Received malformed response from Fal.ai image-segmentation API: expected { image: { url: string } } format, got instead: ${JSON.stringify(
+				response
+			)}`
+		);
+	}
+}

packages/inference/src/providers/hf-inference.ts

@@ -76,6 +76,7 @@ import { base64FromBytes } from "../utils/base64FromBytes.js";
 import type { ImageToImageArgs } from "../tasks/cv/imageToImage.js";
 import type { AutomaticSpeechRecognitionArgs } from "../tasks/audio/automaticSpeechRecognition.js";
 import { omit } from "../utils/omit.js";
+import { ImageSegmentationArgs } from "../tasks/cv/imageSegmentation.js";
 interface Base64ImageGeneration {
 	data: Array<{
 		b64_json: string;
@@ -345,6 +346,15 @@ export class HFInferenceImageSegmentationTask extends HFInferenceTask implements
 			"Received malformed response from HF-Inference image-segmentation API: expected Array<{label: string, mask: string, score: number}>"
 		);
 	}
+
+	async preparePayloadAsync(args: ImageSegmentationArgs): Promise<RequestArgs> {
+		return {
+			...args,
+			inputs: base64FromBytes(
+				new Uint8Array(args.inputs instanceof ArrayBuffer ? args.inputs : await (args.inputs as Blob).arrayBuffer())
+			),
+		};
+	}
 }
 
 export class HFInferenceImageToTextTask extends HFInferenceTask implements ImageToTextTaskHelper {

packages/inference/src/providers/providerHelper.ts

@@ -54,6 +54,7 @@ import { toArray } from "../utils/toArray.js";
 import type { ImageToImageArgs } from "../tasks/cv/imageToImage.js";
 import type { AutomaticSpeechRecognitionArgs } from "../tasks/audio/automaticSpeechRecognition.js";
 import type { ImageToVideoArgs } from "../tasks/cv/imageToVideo.js";
+import { ImageSegmentationArgs } from "../tasks/cv/imageSegmentation.js";
 
 /**
  * Base class for task-specific provider helpers
@@ -161,6 +162,7 @@ export interface ImageToVideoTaskHelper {
 export interface ImageSegmentationTaskHelper {
 	getResponse(response: unknown, url?: string, headers?: HeadersInit): Promise<ImageSegmentationOutput>;
 	preparePayload(params: BodyParams<ImageSegmentationInput & BaseArgs>): Record<string, unknown> | BodyInit;
+	preparePayloadAsync(args: ImageSegmentationArgs): Promise<RequestArgs>;
 }
 
 export interface ImageClassificationTaskHelper {

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

@@ -3,9 +3,9 @@ import { resolveProvider } from "../../lib/getInferenceProviderMapping.js";
 import { getProviderHelper } from "../../lib/getProviderHelper.js";
 import type { BaseArgs, Options } from "../../types.js";
 import { innerRequest } from "../../utils/request.js";
-import { preparePayload, type LegacyImageInput } from "./utils.js";
+import { makeRequestOptions } from "../../lib/makeRequestOptions.js";
 
-export type ImageSegmentationArgs = BaseArgs & (ImageSegmentationInput | LegacyImageInput);
+export type ImageSegmentationArgs = BaseArgs & ImageSegmentationInput;
 
 /**
  * This task reads some image input and outputs the likelihood of classes & bounding boxes of detected objects.
@@ -17,10 +17,11 @@ export async function imageSegmentation(
 ): Promise<ImageSegmentationOutput> {
 	const provider = await resolveProvider(args.provider, args.model, args.endpointUrl);
 	const providerHelper = getProviderHelper(provider, "image-segmentation");
-	const payload = preparePayload(args);
+	const payload = await providerHelper.preparePayloadAsync(args);
 	const { data: res } = await innerRequest<ImageSegmentationOutput>(payload, providerHelper, {
 		...options,
 		task: "image-segmentation",
 	});
-	return providerHelper.getResponse(res);
+	const { url, info } = await makeRequestOptions(args, providerHelper, { ...options, task: "image-segmentation" });
+	return providerHelper.getResponse(res, url, info.headers as Record<string, string>);
 }