Black Forest Labs Image Models (#1193)

Co-authored-by: Julien Chaumond <[email protected]>
Co-authored-by: SBrandeis <[email protected]>

Author: 3 people

.github/workflows/test.yml

@@ -48,6 +48,7 @@ jobs:
           HF_TOGETHER_KEY: dummy
           HF_NOVITA_KEY: dummy
           HF_FIREWORKS_KEY: dummy
+          HF_BLACK_FOREST_LABS_KEY: dummy
 
   browser:
     runs-on: ubuntu-latest
@@ -91,6 +92,7 @@ jobs:
           HF_TOGETHER_KEY: dummy
           HF_NOVITA_KEY: dummy
           HF_FIREWORKS_KEY: dummy
+          HF_BLACK_FOREST_LABS_KEY: dummy
 
   e2e:
     runs-on: ubuntu-latest
@@ -161,3 +163,4 @@ jobs:
           HF_TOGETHER_KEY: dummy
           HF_NOVITA_KEY: dummy
           HF_FIREWORKS_KEY: dummy
+          HF_BLACK_FOREST_LABS_KEY: dummy

packages/inference/README.md

@@ -54,6 +54,7 @@ Currently, we support the following providers:
 - [Replicate](https://replicate.com)
 - [Sambanova](https://sambanova.ai)
 - [Together](https://together.xyz)
+- [Blackforestlabs](https://blackforestlabs.ai)
 
 To send requests to a third-party provider, you have to pass the `provider` parameter to the inference function. Make sure your request is authenticated with an access token.
 ```ts

packages/inference/src/lib/makeRequestOptions.ts

@@ -6,6 +6,7 @@ import { SAMBANOVA_API_BASE_URL } from "../providers/sambanova";
 import { TOGETHER_API_BASE_URL } from "../providers/together";
 import { NOVITA_API_BASE_URL } from "../providers/novita";
 import { FIREWORKS_AI_API_BASE_URL } from "../providers/fireworks-ai";
+import { BLACKFORESTLABS_AI_API_BASE_URL } from "../providers/black-forest-labs";
 import type { InferenceProvider } from "../types";
 import type { InferenceTask, Options, RequestArgs } from "../types";
 import { isUrl } from "./isUrl";
@@ -80,8 +81,13 @@ export async function makeRequestOptions(
 
 	const headers: Record<string, string> = {};
 	if (accessToken) {
-		headers["Authorization"] =
-			provider === "fal-ai" && authMethod === "provider-key" ? `Key ${accessToken}` : `Bearer ${accessToken}`;
+		if (provider === "fal-ai" && authMethod === "provider-key") {
+			headers["Authorization"] = `Key ${accessToken}`;
+		} else if (provider === "black-forest-labs" && authMethod === "provider-key") {
+			headers["X-Key"] = accessToken;
+		} else {
+			headers["Authorization"] = `Bearer ${accessToken}`;
+		}
 	}
 
 	// e.g. @huggingface/inference/3.1.3
@@ -148,6 +154,12 @@ function makeUrl(params: {
 
 	const shouldProxy = params.provider !== "hf-inference" && params.authMethod !== "provider-key";
 	switch (params.provider) {
+		case "black-forest-labs": {
+			const baseUrl = shouldProxy
+				? HF_HUB_INFERENCE_PROXY_TEMPLATE.replace("{{PROVIDER}}", params.provider)
+				: BLACKFORESTLABS_AI_API_BASE_URL;
+			return `${baseUrl}/${params.model}`;
+		}
 		case "fal-ai": {
 			const baseUrl = shouldProxy
 				? HF_HUB_INFERENCE_PROXY_TEMPLATE.replace("{{PROVIDER}}", params.provider)

packages/inference/src/providers/black-forest-labs.ts

@@ -0,0 +1,18 @@
+export const BLACKFORESTLABS_AI_API_BASE_URL = "https://api.us1.bfl.ai/v1";
+
+/**
+ * See the registered mapping of HF model ID => Black Forest Labs model ID here:
+ *
+ * https://huggingface.co/api/partners/blackforestlabs/models
+ *
+ * This is a publicly available mapping.
+ *
+ * If you want to try to run inference for a new model locally before it's registered on huggingface.co,
+ * you can add it to the dictionary "HARDCODED_MODEL_ID_MAPPING" in consts.ts, for dev purposes.
+ *
+ * - If you work at Black Forest Labs and want to update this mapping, please use the model mapping API we provide on huggingface.co
+ * - If you're a community member and want to add a new supported HF model to Black Forest Labs, please open an issue on the present repo
+ * and we will tag Black Forest Labs team members.
+ *
+ * Thanks!
+ */

packages/inference/src/providers/consts.ts

@@ -16,6 +16,7 @@ export const HARDCODED_MODEL_ID_MAPPING: Record<InferenceProvider, Record<ModelI
 	 * Example:
 	 * "Qwen/Qwen2.5-Coder-32B-Instruct": "Qwen2.5-Coder-32B-Instruct",
 	 */
+	"black-forest-labs": {},
 	"fal-ai": {},
 	"fireworks-ai": {},
 	"hf-inference": {},

packages/inference/src/providers/novita.ts

@@ -15,4 +15,4 @@ export const NOVITA_API_BASE_URL = "https://api.novita.ai/v3/openai";
  * and we will tag Novita team members.
  *
  * Thanks!
- */ 
+ */

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

@@ -3,6 +3,7 @@ import { InferenceOutputError } from "../../lib/InferenceOutputError";
 import type { BaseArgs, InferenceProvider, Options } from "../../types";
 import { omit } from "../../utils/omit";
 import { request } from "../custom/request";
+import { delay } from "../../utils/delay";
 
 export type TextToImageArgs = BaseArgs & TextToImageInput;
 
@@ -14,6 +15,10 @@ interface Base64ImageGeneration {
 interface OutputUrlImageGeneration {
 	output: string[];
 }
+interface BlackForestLabsResponse {
+	id: string;
+	polling_url: string;
+}
 
 function getResponseFormatArg(provider: InferenceProvider) {
 	switch (provider) {
@@ -44,12 +49,17 @@ export async function textToImage(args: TextToImageArgs, options?: Options): Pro
 					...getResponseFormatArg(args.provider),
 					prompt: args.inputs,
 			  };
-	const res = await request<TextToImageOutput | Base64ImageGeneration | OutputUrlImageGeneration>(payload, {
+	const res = await request<
+		TextToImageOutput | Base64ImageGeneration | OutputUrlImageGeneration | BlackForestLabsResponse
+	>(payload, {
 		...options,
 		taskHint: "text-to-image",
 	});
 
 	if (res && typeof res === "object") {
+		if (args.provider === "black-forest-labs" && "polling_url" in res && typeof res.polling_url === "string") {
+			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();
@@ -72,3 +82,33 @@ export async function textToImage(args: TextToImageArgs, options?: Options): Pro
 	}
 	return res;
 }
+
+async function pollBflResponse(url: string): Promise<Blob> {
+	const urlObj = new URL(url);
+	for (let step = 0; step < 5; step++) {
+		await delay(1000);
+		console.debug(`Polling Black Forest Labs API for the result... ${step + 1}/5`);
+		urlObj.searchParams.set("attempt", step.toString(10));
+		const resp = await fetch(urlObj, { headers: { "Content-Type": "application/json" } });
+		if (!resp.ok) {
+			throw new InferenceOutputError("Failed to fetch result from black forest labs API");
+		}
+		const payload = await resp.json();
+		if (
+			typeof payload === "object" &&
+			payload &&
+			"status" in payload &&
+			typeof payload.status === "string" &&
+			payload.status === "Ready" &&
+			"result" in payload &&
+			typeof payload.result === "object" &&
+			payload.result &&
+			"sample" in payload.result &&
+			typeof payload.result.sample === "string"
+		) {
+			const image = await fetch(payload.result.sample);
+			return await image.blob();
+		}
+	}
+	throw new InferenceOutputError("Failed to fetch result from black forest labs API");
+}

packages/inference/src/types.ts

@@ -29,14 +29,15 @@ export interface Options {
 export type InferenceTask = Exclude<PipelineType, "other">;
 
 export const INFERENCE_PROVIDERS = [
+	"black-forest-labs",
 	"fal-ai",
 	"fireworks-ai",
-	"nebius",
 	"hf-inference",
+	"nebius",
+	"novita",
 	"replicate",
 	"sambanova",
 	"together",
-	"novita",
 ] as const;
 
 export type InferenceProvider = (typeof INFERENCE_PROVIDERS)[number];

packages/inference/src/utils/delay.ts

@@ -0,0 +1,5 @@
+export function delay(ms: number): Promise<void> {
+	return new Promise((resolve) => {
+		setTimeout(() => resolve(), ms);
+	});
+}

packages/inference/test/HfInference.spec.ts

@@ -2,7 +2,7 @@ import { assert, describe, expect, it } from "vitest";
 
 import type { ChatCompletionStreamOutput } from "@huggingface/tasks";
 
-import { chatCompletion, HfInference } from "../src";
+import { chatCompletion, HfInference, textToImage } from "../src";
 import { textToVideo } from "../src/tasks/cv/textToVideo";
 import { readTestFile } from "./test-files";
 import "./vcr";
@@ -1214,4 +1214,30 @@ describe.concurrent("HfInference", () => {
 		},
 		TIMEOUT
 	);
+	describe.concurrent(
+		"Black Forest Labs",
+		() => {
+			HARDCODED_MODEL_ID_MAPPING["black-forest-labs"] = {
+				"black-forest-labs/FLUX.1-dev": "flux-dev",
+				// "black-forest-labs/FLUX.1-schnell": "flux-pro",
+			};
+
+			it("textToImage", 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,
+					},
+				});
+				expect(res).toBeInstanceOf(Blob);
+			});
+		},
+		TIMEOUT
+	);
 });