[Inference] Black Forest: poll for result even if output is url (#1209)

huggingface · Feb 18, 2025 · 0b602a5 · 0b602a5
1 parent 081a6ab
commit 0b602a5
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/packages/inference/src/tasks/cv/textToImage.ts b/packages/inference/src/tasks/cv/textToImage.ts
@@ -78,10 +78,7 @@ export async function textToImage(args: TextToImageArgs, options?: TextToImageOp
 
 	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);
+			return await pollBflResponse(res.polling_url, options?.outputType);
 		}
 		if (args.provider === "fal-ai" && "images" in res && Array.isArray(res.images) && res.images[0].url) {
 			if (options?.outputType === "url") {
@@ -132,7 +129,7 @@ export async function textToImage(args: TextToImageArgs, options?: TextToImageOp
 	return res;
 }
 
-async function pollBflResponse(url: string): Promise<Blob> {
+async function pollBflResponse(url: string, outputType?: "url" | "blob"): Promise<Blob> {
 	const urlObj = new URL(url);
 	for (let step = 0; step < 5; step++) {
 		await delay(1000);
@@ -155,6 +152,9 @@ async function pollBflResponse(url: string): Promise<Blob> {
 			"sample" in payload.result &&
 			typeof payload.result.sample === "string"
 		) {
+			if (outputType === "url") {
+				return payload.result.sample;
+			}
 			const image = await fetch(payload.result.sample);
 			return await image.blob();
 		}