feat(image-generation): add BFL (Black Forest Labs) provider (#54)

* feat(image-generation): add BFL (Black Forest Labs) provider

Add support for BFL's FLUX.2 image generation models:
- flux-2-pro: Professional model with automatic prompt upsampling
- flux-2-flex: Flexible model with configurable steps and guidance

Changes:
- Add BFL provider with polling-based async generation
- Add BFL to Provider enum and credential system
- Add new ImageGenerationParameter values (guidance, steps, seed, etc.)
- Add billed_units field to ImageGenerationUsage
- Add BFL to integration tests and CI workflow
- Support BFL_API_KEY environment variable

* chore: bump versions for BFL provider release

- celeste-ai: 0.2.11 → 0.2.12
- celeste-image-generation: 0.2.9 → 0.2.10

* fix: use content param for httpx.Response construction

httpx.Response doesn't accept json parameter - use content with JSON-encoded bytes instead.

Author: Kamilbenkirane

.env.example

@@ -7,6 +7,9 @@ OPENAI_API_KEY=your-openai-api-key-here
 # Anthropic
 ANTHROPIC_API_KEY=your-anthropic-api-key-here
 
+# BFL (Black Forest Labs)
+BFL_API_KEY=your-bfl-api-key-here
+
 # Google
 GOOGLE_API_KEY=your-google-api-key-here
 

.github/workflows/publish.yml

@@ -55,6 +55,7 @@ jobs:
         env:
           OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
           ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
+          BFL_API_KEY: ${{ secrets.BFL_API_KEY }}
           GOOGLE_API_KEY: ${{ secrets.GOOGLE_API_KEY }}
           MISTRAL_API_KEY: ${{ secrets.MISTRAL_API_KEY }}
           COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}

packages/image-generation/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "celeste-image-generation"
-version = "0.2.9"
+version = "0.2.10"
 description = "Image generation package for Celeste AI. Unified interface for all providers"
 authors = [{name = "Kamilbenkirane", email = "kamil@withceleste.ai"}]
 readme = "README.md"

packages/image-generation/src/celeste_image_generation/io.py

@@ -31,6 +31,7 @@ class ImageGenerationUsage(Usage):
     output_tokens: int | None = None
     reasoning_tokens: int | None = None
     generated_images: int | None = None
+    billed_units: float | None = None
 
 
 class ImageGenerationOutput(Output[ImageArtifact]):

packages/image-generation/src/celeste_image_generation/models.py

@@ -1,13 +1,15 @@
 """Model definitions for image generation."""
 
 from celeste import Model
+from celeste_image_generation.providers.bfl.models import MODELS as BFL_MODELS
 from celeste_image_generation.providers.bytedance.models import (
     MODELS as BYTEDANCE_MODELS,
 )
 from celeste_image_generation.providers.google.models import MODELS as GOOGLE_MODELS
 from celeste_image_generation.providers.openai.models import MODELS as OPENAI_MODELS
 
 MODELS: list[Model] = [
+    *BFL_MODELS,
     *BYTEDANCE_MODELS,
     *GOOGLE_MODELS,
     *OPENAI_MODELS,

packages/image-generation/src/celeste_image_generation/parameters.py

@@ -9,15 +9,27 @@ class ImageGenerationParameter(StrEnum):
     """Unified parameter names for image generation capability."""
 
     ASPECT_RATIO = "aspect_ratio"
+    GUIDANCE = "guidance"
+    OUTPUT_FORMAT = "output_format"
     PARTIAL_IMAGES = "partial_images"
+    PROMPT_UPSAMPLING = "prompt_upsampling"
     QUALITY = "quality"
+    SAFETY_TOLERANCE = "safety_tolerance"
+    SEED = "seed"
+    STEPS = "steps"
     WATERMARK = "watermark"
 
 
 class ImageGenerationParameters(Parameters):
     """Parameters for image generation."""
 
     aspect_ratio: str | None
+    guidance: float | None
+    output_format: str | None
     partial_images: int | None
+    prompt_upsampling: bool | None
     quality: str | None
+    safety_tolerance: int | None
+    seed: int | None
+    steps: int | None
     watermark: bool | None

packages/image-generation/src/celeste_image_generation/providers/__init__.py

@@ -8,6 +8,9 @@
 def _get_providers() -> list[tuple[Provider, type[Client]]]:
     """Lazy-load providers."""
     # Import clients directly from .client modules to avoid __init__.py imports
+    from celeste_image_generation.providers.bfl.client import (
+        BFLImageGenerationClient,
+    )
     from celeste_image_generation.providers.bytedance.client import (
         ByteDanceImageGenerationClient,
     )
@@ -19,6 +22,7 @@ def _get_providers() -> list[tuple[Provider, type[Client]]]:
     )
 
     return [
+        (Provider.BFL, BFLImageGenerationClient),
         (Provider.BYTEDANCE, ByteDanceImageGenerationClient),
         (Provider.GOOGLE, GoogleImageGenerationClient),
         (Provider.OPENAI, OpenAIImageGenerationClient),

packages/image-generation/src/celeste_image_generation/providers/bfl/__init__.py

@@ -0,0 +1,6 @@
+"""BFL (Black Forest Labs) provider for image generation."""
+
+from .client import BFLImageGenerationClient
+from .models import MODELS
+
+__all__ = ["MODELS", "BFLImageGenerationClient"]

packages/image-generation/src/celeste_image_generation/providers/bfl/client.py

@@ -0,0 +1,149 @@
+"""BFL (Black Forest Labs) client implementation for FLUX.2 image generation."""
+
+import asyncio
+import json
+import time
+from typing import Any, Unpack
+
+import httpx
+
+from celeste.artifacts import ImageArtifact
+from celeste.mime_types import ApplicationMimeType
+from celeste.parameters import ParameterMapper
+from celeste_image_generation.client import ImageGenerationClient
+from celeste_image_generation.io import (
+    ImageGenerationFinishReason,
+    ImageGenerationInput,
+    ImageGenerationUsage,
+)
+from celeste_image_generation.parameters import ImageGenerationParameters
+
+from . import config
+from .parameters import BFL_PARAMETER_MAPPERS
+
+
+class BFLImageGenerationClient(ImageGenerationClient):
+    """Black Forest Labs client for image generation."""
+
+    @classmethod
+    def parameter_mappers(cls) -> list[ParameterMapper]:
+        return BFL_PARAMETER_MAPPERS
+
+    def _init_request(self, inputs: ImageGenerationInput) -> dict[str, Any]:
+        """Initialize request for BFL API format."""
+        return {
+            "prompt": inputs.prompt,
+        }
+
+    def _parse_usage(self, response_data: dict[str, Any]) -> ImageGenerationUsage:
+        """Parse usage from response."""
+        submit_metadata = response_data.get("_submit_metadata", {})
+        cost = submit_metadata.get("cost")
+
+        return ImageGenerationUsage(
+            billed_units=float(cost) if cost is not None else None,
+        )
+
+    def _parse_content(
+        self,
+        response_data: dict[str, Any],
+        **parameters: Unpack[ImageGenerationParameters],
+    ) -> ImageArtifact:
+        """Parse content from response."""
+        result = response_data.get("result", {})
+        sample_url = result.get("sample")
+
+        if not sample_url:
+            msg = f"No image URL in {self.provider} response"
+            raise ValueError(msg)
+
+        return ImageArtifact(url=sample_url)
+
+    def _parse_finish_reason(
+        self, response_data: dict[str, Any]
+    ) -> ImageGenerationFinishReason | None:
+        """Parse finish reason from response."""
+        status = response_data.get("status")
+        if status == "Ready":
+            return ImageGenerationFinishReason(reason="COMPLETE")
+        elif status in ("Error", "Failed"):
+            error_msg = response_data.get("error", "Generation failed")
+            return ImageGenerationFinishReason(reason="ERROR", message=error_msg)
+        return None
+
+    async def _make_request(
+        self,
+        request_body: dict[str, Any],
+        **parameters: Unpack[ImageGenerationParameters],
+    ) -> httpx.Response:
+        """Make HTTP request(s) and return response object."""
+        headers = {
+            config.AUTH_HEADER_NAME: self.api_key.get_secret_value(),
+            "Content-Type": ApplicationMimeType.JSON,
+            "Accept": ApplicationMimeType.JSON,
+        }
+
+        endpoint = config.ENDPOINT.format(model_id=self.model.id)
+
+        submit_response = await self.http_client.post(
+            f"{config.BASE_URL}{endpoint}",
+            headers=headers,
+            json_body=request_body,
+        )
+
+        if submit_response.status_code != 200:
+            return submit_response
+
+        submit_data = submit_response.json()
+        polling_url = submit_data.get("polling_url")
+
+        if not polling_url:
+            msg = f"No polling_url in {self.provider} response"
+            raise ValueError(msg)
+
+        start_time = time.monotonic()
+        poll_headers = {
+            config.AUTH_HEADER_NAME: self.api_key.get_secret_value(),
+            "Accept": ApplicationMimeType.JSON,
+        }
+
+        while True:
+            elapsed = time.monotonic() - start_time
+            if elapsed >= config.POLLING_TIMEOUT:
+                msg = f"{self.provider} polling timed out after {config.POLLING_TIMEOUT} seconds"
+                raise TimeoutError(msg)
+
+            poll_response = await self.http_client.get(
+                polling_url,
+                headers=poll_headers,
+            )
+
+            if poll_response.status_code != 200:
+                return poll_response
+
+            poll_data = poll_response.json()
+            status = poll_data.get("status")
+
+            if status == "Ready":
+                final_data = {
+                    **poll_data,
+                    "_submit_metadata": submit_data,
+                }
+                return httpx.Response(
+                    status_code=200,
+                    content=json.dumps(final_data).encode("utf-8"),
+                    headers={"content-type": "application/json"},
+                    request=httpx.Request("GET", polling_url),
+                )
+            elif status in ("Error", "Failed"):
+                return httpx.Response(
+                    status_code=400,
+                    content=json.dumps(poll_data).encode("utf-8"),
+                    headers={"content-type": "application/json"},
+                    request=httpx.Request("GET", polling_url),
+                )
+
+            await asyncio.sleep(config.POLLING_INTERVAL)
+
+
+__all__ = ["BFLImageGenerationClient"]

packages/image-generation/src/celeste_image_generation/providers/bfl/config.py

@@ -0,0 +1,13 @@
+"""BFL (Black Forest Labs) provider configuration for image generation."""
+
+# HTTP Configuration
+BASE_URL = "https://api.bfl.ai"
+ENDPOINT = "/v1/{model_id}"
+
+# Authentication
+AUTH_HEADER_NAME = "x-key"
+AUTH_HEADER_PREFIX = ""  # Direct API key, no prefix
+
+# Polling Configuration
+POLLING_INTERVAL = 0.5  # seconds between polling attempts
+POLLING_TIMEOUT = 120.0  # 2 minutes for image generation