[inference provider] Add wavespeed.ai as an inference provider

packages/inference/README.md

@@ -67,6 +67,7 @@ Currently, we support the following providers:
 - [Cohere](https://cohere.com)
 - [Cerebras](https://cerebras.ai/)
 - [Groq](https://groq.com)
+- [Wavespeed.ai](https://wavespeed.ai/)
 - [Z.ai](https://z.ai/)
 
 To send requests to a third-party provider, you have to pass the `provider` parameter to the inference function. The default value of the `provider` parameter is "auto", which will select the first of the providers available for the model, sorted by your preferred order in https://hf.co/settings/inference-providers.
@@ -105,6 +106,7 @@ Only a subset of models are supported when requesting third-party providers. You
 - [Cerebras supported models](https://huggingface.co/api/partners/cerebras/models)
 - [Groq supported models](https://console.groq.com/docs/models)
 - [Novita AI supported models](https://huggingface.co/api/partners/novita/models)
+- [Wavespeed.ai supported models](https://huggingface.co/api/partners/wavespeed/models)
 - [Z.ai supported models](https://huggingface.co/api/partners/zai-org/models)
 
 ❗**Important note:** To be compatible, the third-party API must adhere to the "standard" shape API we expect on HF model pages for each pipeline task type.

packages/inference/src/lib/getProviderHelper.ts

@@ -52,6 +52,7 @@ import * as Replicate from "../providers/replicate.js";
 import * as Sambanova from "../providers/sambanova.js";
 import * as Scaleway from "../providers/scaleway.js";
 import * as Together from "../providers/together.js";
+import * as Wavespeed from "../providers/wavespeed.js";
 import * as Zai from "../providers/zai-org.js";
 import type { InferenceProvider, InferenceProviderOrPolicy, InferenceTask } from "../types.js";
 import { InferenceClientInputError } from "../errors.js";
@@ -173,6 +174,11 @@ export const PROVIDERS: Record<InferenceProvider, Partial<Record<InferenceTask,
 		conversational: new Together.TogetherConversationalTask(),
 		"text-generation": new Together.TogetherTextGenerationTask(),
 	},
+	wavespeed: {
+		"text-to-image": new Wavespeed.WavespeedAITextToImageTask(),
+		"text-to-video": new Wavespeed.WavespeedAITextToVideoTask(),
+		"image-to-image": new Wavespeed.WavespeedAIImageToImageTask(),
+	},
 	"zai-org": {
 		conversational: new Zai.ZaiConversationalTask(),
 	},

packages/inference/src/providers/consts.ts

@@ -39,5 +39,6 @@ export const HARDCODED_MODEL_INFERENCE_MAPPING: Record<
 	sambanova: {},
 	scaleway: {},
 	together: {},
+	wavespeed: {},
 	"zai-org": {},
 };

packages/inference/src/providers/wavespeed.ts

@@ -0,0 +1,185 @@
+import type { TextToImageArgs } from "../tasks/cv/textToImage.js";
+import type { ImageToImageArgs } from "../tasks/cv/imageToImage.js";
+import type { TextToVideoArgs } from "../tasks/cv/textToVideo.js";
+import type { BodyParams, RequestArgs, UrlParams } from "../types.js";
+import { delay } from "../utils/delay.js";
+import { omit } from "../utils/omit.js";
+import { base64FromBytes } from "../utils/base64FromBytes.js";
+import type { TextToImageTaskHelper, TextToVideoTaskHelper, ImageToImageTaskHelper } from "./providerHelper.js";
+import { TaskProviderHelper } from "./providerHelper.js";
+import {
+	InferenceClientInputError,
+	InferenceClientProviderApiError,
+	InferenceClientProviderOutputError,
+} from "../errors.js";
+
+const WAVESPEEDAI_API_BASE_URL = "https://api.wavespeed.ai";
+
+/**
+ * Response structure for task status and results
+ */
+interface WaveSpeedAITaskResponse {
+	id: string;
+	model: string;
+	outputs: string[];
+	urls: {
+		get: string;
+	};
+	has_nsfw_contents: boolean[];
+	status: "created" | "processing" | "completed" | "failed";
+	created_at: string;
+	error: string;
+	executionTime: number;
+	timings: {
+		inference: number;
+	};
+}
+
+/**
+ * Response structure for initial task submission
+ */
+interface WaveSpeedAISubmitResponse {
+	id: string;
+	urls: {
+		get: string;
+	};
+}
+
+/**
+ * Response structure for WaveSpeed AI API
+ */
+interface WaveSpeedAIResponse {
+	code: number;
+	message: string;
+	data: WaveSpeedAITaskResponse;
+}
+
+/**
+ * Response structure for WaveSpeed AI API with submit response data
+ */
+interface WaveSpeedAISubmitTaskResponse {
+	code: number;
+	message: string;
+	data: WaveSpeedAISubmitResponse;
+}
+
+abstract class WavespeedAITask extends TaskProviderHelper {
+	constructor(url?: string) {
+		super("wavespeed", url || WAVESPEEDAI_API_BASE_URL);
+	}
+
+	makeRoute(params: UrlParams): string {
+		return `/api/v3/${params.model}`;
+	}
+
+	preparePayload(params: BodyParams<ImageToImageArgs | TextToImageArgs | TextToVideoArgs>): Record<string, unknown> {
+		const payload: Record<string, unknown> = {
+			...omit(params.args, ["inputs", "parameters"]),
+			...params.args.parameters,
+			prompt: params.args.inputs,
+		};
+		// Add LoRA support if adapter is specified in the mapping
+		if (params.mapping?.adapter === "lora") {
+			payload.loras = [
+				{
+					path: params.mapping.hfModelId,
+					scale: 1, // Default scale value
+				},
+			];
+		}
+		return payload;
+	}
+
+	override async getResponse(
+		response: WaveSpeedAISubmitTaskResponse,
+		url?: string,
+		headers?: Record<string, string>
+	): Promise<Blob> {
+		if (!headers) {
+			throw new InferenceClientInputError("Headers are required for WaveSpeed AI API calls");
+		}
+
+		const resultUrl = response.data.urls.get;
+
+		// Poll for results until completion
+		while (true) {
+			const resultResponse = await fetch(resultUrl, { headers });
+
+			if (!resultResponse.ok) {
+				throw new InferenceClientProviderApiError(
+					"Failed to fetch response status from WaveSpeed AI API",
+					{ url: resultUrl, method: "GET" },
+					{
+						requestId: resultResponse.headers.get("x-request-id") ?? "",
+						status: resultResponse.status,
+						body: await resultResponse.text(),
+					}
+				);
+			}
+
+			const result: WaveSpeedAIResponse = await resultResponse.json();
+			const taskResult = result.data;
+
+			switch (taskResult.status) {
+				case "completed": {
+					// Get the media data from the first output URL
+					if (!taskResult.outputs?.[0]) {
+						throw new InferenceClientProviderOutputError(
+							"Received malformed response from WaveSpeed AI API: No output URL in completed response"
+						);
+					}
+					const mediaResponse = await fetch(taskResult.outputs[0]);
+					if (!mediaResponse.ok) {
+						throw new InferenceClientProviderApiError(
+							"Failed to fetch generation output from WaveSpeed AI API",
+							{ url: taskResult.outputs[0], method: "GET" },
+							{
+								requestId: mediaResponse.headers.get("x-request-id") ?? "",
+								status: mediaResponse.status,
+								body: await mediaResponse.text(),
+							}
+						);
+					}
+					return await mediaResponse.blob();
+				}
+				case "failed": {
+					throw new InferenceClientProviderOutputError(taskResult.error || "Task failed");
+				}
+
+				default: {
+					// Wait before polling again
+					await delay(500);
+					continue;
+				}
+			}
+		}
+	}
+}
+
+export class WavespeedAITextToImageTask extends WavespeedAITask implements TextToImageTaskHelper {
+	constructor() {
+		super(WAVESPEEDAI_API_BASE_URL);
+	}
+}
+
+export class WavespeedAITextToVideoTask extends WavespeedAITask implements TextToVideoTaskHelper {
+	constructor() {
+		super(WAVESPEEDAI_API_BASE_URL);
+	}
+}
+
+export class WavespeedAIImageToImageTask extends WavespeedAITask implements ImageToImageTaskHelper {
+	constructor() {
+		super(WAVESPEEDAI_API_BASE_URL);
+	}
+
+	async preparePayloadAsync(args: ImageToImageArgs): Promise<RequestArgs> {
+		return {
+			...args,
+			inputs: args.parameters?.prompt,
+			image: base64FromBytes(
+				new Uint8Array(args.inputs instanceof ArrayBuffer ? args.inputs : await (args.inputs as Blob).arrayBuffer())
+			),
+		};
+	}
+}

packages/inference/src/types.ts

@@ -66,6 +66,7 @@ export const INFERENCE_PROVIDERS = [
 	"sambanova",
 	"scaleway",
 	"together",
+	"wavespeed",
 	"zai-org",
 ] as const;
 

packages/inference/test/InferenceClient.spec.ts

@@ -2291,6 +2291,119 @@ describe.skip("InferenceClient", () => {
 		TIMEOUT
 	);
 
+	describe.concurrent(
+		"Wavespeed AI",
+		() => {
+			const client = new InferenceClient(env.HF_WAVESPEED_KEY ?? "dummy");
+
+			HARDCODED_MODEL_INFERENCE_MAPPING["wavespeed"] = {
+				"black-forest-labs/FLUX.1-schnell": {
+					provider: "wavespeed",
+					hfModelId: "black-forest-labs/FLUX.1-schnell",
+					providerId: "wavespeed-ai/flux-schnell",
+					status: "live",
+					task: "text-to-image",
+				},
+				"Wan-AI/Wan2.1-T2V-14B": {
+					provider: "wavespeed",
+					hfModelId: "wavespeed-ai/wan-2.1/t2v-480p",
+					providerId: "wavespeed-ai/wan-2.1/t2v-480p",
+					status: "live",
+					task: "text-to-video",
+				},
+				"HiDream-ai/HiDream-E1-Full": {
+					provider: "wavespeed",
+					hfModelId: "wavespeed-ai/hidream-e1-full",
+					providerId: "wavespeed-ai/hidream-e1-full",
+					status: "live",
+					task: "image-to-image",
+				},
+				"openfree/flux-chatgpt-ghibli-lora": {
+					provider: "wavespeed",
+					hfModelId: "openfree/flux-chatgpt-ghibli-lora",
+					providerId: "wavespeed-ai/flux-dev-lora",
+					status: "live",
+					task: "text-to-image",
+					adapter: "lora",
+					adapterWeightsPath: "flux-chatgpt-ghibli-lora.safetensors",
+				},
+				"linoyts/yarn_art_Flux_LoRA": {
+					provider: "wavespeed",
+					hfModelId: "linoyts/yarn_art_Flux_LoRA",
+					providerId: "wavespeed-ai/flux-dev-lora-ultra-fast",
+					status: "live",
+					task: "text-to-image",
+					adapter: "lora",
+					adapterWeightsPath: "pytorch_lora_weights.safetensors",
+				},
+			};
+			it(`textToImage - black-forest-labs/FLUX.1-schnell`, async () => {
+				const res = await client.textToImage({
+					model: "black-forest-labs/FLUX.1-schnell",
+					provider: "wavespeed",
+					inputs:
+						"Cute boy with a hat, exploring nature, holding a telescope, backpack, surrounded by flowers, cartoon style, vibrant colors.",
+				});
+				expect(res).toBeInstanceOf(Blob);
+			});
+
+			it(`textToImage - openfree/flux-chatgpt-ghibli-lora`, async () => {
+				const res = await client.textToImage({
+					model: "openfree/flux-chatgpt-ghibli-lora",
+					provider: "wavespeed",
+					inputs:
+						"Cute boy with a hat, exploring nature, holding a telescope, backpack, surrounded by flowers, cartoon style, vibrant colors.",
+				});
+				expect(res).toBeInstanceOf(Blob);
+			});
+
+			it(`textToImage - linoyts/yarn_art_Flux_LoRA`, async () => {
+				const res = await client.textToImage({
+					model: "linoyts/yarn_art_Flux_LoRA",
+					provider: "wavespeed",
+					inputs:
+						"Cute boy with a hat, exploring nature, holding a telescope, backpack, surrounded by flowers, cartoon style, vibrant colors.",
+				});
+				expect(res).toBeInstanceOf(Blob);
+			});
+
+			it(`textToVideo - Wan-AI/Wan2.1-T2V-14B`, async () => {
+				const res = await client.textToVideo({
+					model: "Wan-AI/Wan2.1-T2V-14B",
+					provider: "wavespeed",
+					inputs:
+						"A cool street dancer, wearing a baggy hoodie and hip-hop pants, dancing in front of a graffiti wall, night neon background, quick camera cuts, urban trends.",
+					parameters: {
+						guidance_scale: 5,
+						num_inference_steps: 30,
+						seed: -1,
+					},
+					duration: 5,
+					enable_safety_checker: true,
+					flow_shift: 2.9,
+					size: "480*832",
+				});
+				expect(res).toBeInstanceOf(Blob);
+			});
+
+			it(`imageToImage - HiDream-ai/HiDream-E1-Full`, async () => {
+				const res = await client.imageToImage({
+					model: "HiDream-ai/HiDream-E1-Full",
+					provider: "wavespeed",
+					inputs: new Blob([readTestFile("cheetah.png")], { type: "image/png" }),
+					parameters: {
+						prompt: "The leopard chases its prey",
+						guidance_scale: 5,
+						num_inference_steps: 30,
+						seed: -1,
+					},
+				});
+				expect(res).toBeInstanceOf(Blob);
+			});
+		},
+		TIMEOUT
+	);
+
 	describe.concurrent(
 		"PublicAI",
 		() => {