feat(providers): add BFL Images API package

packages/capabilities/image-generation/pyproject.toml

@@ -27,6 +27,7 @@ Issues = "https://github.com/withceleste/celeste-python/issues"
 
 [tool.uv.sources]
 celeste-ai = { workspace = true }
+celeste-bfl = { workspace = true }
 celeste-google = { workspace = true }
 
 [project.entry-points."celeste.packages"]

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

@@ -32,6 +32,8 @@ class ImageGenerationUsage(Usage):
     reasoning_tokens: int | None = None
     num_images: int | None = None
     billed_units: float | None = None
+    input_mp: float | None = None
+    output_mp: float | None = None
 
 
 class ImageGenerationOutput(Output[ImageArtifact | list[ImageArtifact]]):

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

@@ -1,14 +1,11 @@
-"""BFL (Black Forest Labs) client implementation for FLUX.2 image generation."""
+"""BFL client implementation for image generation."""
 
-import asyncio
-import json
-import time
 from typing import Any, Unpack
 
-import httpx
+from celeste_bfl.images.client import BFLImagesClient
 
 from celeste.artifacts import ImageArtifact
-from celeste.mime_types import ApplicationMimeType
+from celeste.exceptions import ValidationError
 from celeste.parameters import ParameterMapper
 from celeste_image_generation.client import ImageGenerationClient
 from celeste_image_generation.io import (
@@ -18,12 +15,11 @@
 )
 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."""
+class BFLImageGenerationClient(BFLImagesClient, ImageGenerationClient):
+    """BFL client for image generation."""
 
     @classmethod
     def parameter_mappers(cls) -> list[ParameterMapper]:
@@ -37,12 +33,8 @@ def _init_request(self, inputs: ImageGenerationInput) -> dict[str, Any]:
 
     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,
-        )
+        usage = super()._parse_usage(response_data)
+        return ImageGenerationUsage(**usage)
 
     def _parse_content(
         self,
@@ -55,95 +47,21 @@ def _parse_content(
 
         if not sample_url:
             msg = f"No image URL in {self.provider} response"
-            raise ValueError(msg)
+            raise ValidationError(msg)
 
         return ImageArtifact(url=sample_url)
 
     def _parse_finish_reason(
         self, response_data: dict[str, Any]
-    ) -> ImageGenerationFinishReason | None:
+    ) -> ImageGenerationFinishReason:
         """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 = {
-            **self.auth.get_headers(),
-            "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 = {
-            **self.auth.get_headers(),
-            "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)
+        return ImageGenerationFinishReason(reason=None)
 
 
 __all__ = ["BFLImageGenerationClient"]

packages/capabilities/image-generation/src/celeste_image_generation/providers/bfl/parameters.py

@@ -1,14 +1,47 @@
-"""BFL parameter mappers for image generation."""
+"""BFL Images parameter mappers for image generation.
+
+AspectRatioMapper is defined locally (handles "WxH" → width/height transformation).
+Other mappers subclass from provider and add capability-specific `name` attribute.
+"""
 
 from typing import Any
 
-from celeste import Model
+from celeste_bfl.images.parameters import (
+    GuidanceMapper as _GuidanceMapper,
+)
+from celeste_bfl.images.parameters import (
+    HeightMapper as _HeightMapper,
+)
+from celeste_bfl.images.parameters import (
+    OutputFormatMapper as _OutputFormatMapper,
+)
+from celeste_bfl.images.parameters import (
+    PromptUpsamplingMapper as _PromptUpsamplingMapper,
+)
+from celeste_bfl.images.parameters import (
+    SafetyToleranceMapper as _SafetyToleranceMapper,
+)
+from celeste_bfl.images.parameters import (
+    SeedMapper as _SeedMapper,
+)
+from celeste_bfl.images.parameters import (
+    StepsMapper as _StepsMapper,
+)
+from celeste_bfl.images.parameters import (
+    WidthMapper as _WidthMapper,
+)
+
+from celeste.models import Model
 from celeste.parameters import ParameterMapper
 from celeste_image_generation.parameters import ImageGenerationParameter
 
 
 class AspectRatioMapper(ParameterMapper):
-    """Map aspect_ratio to BFL width/height parameters."""
+    """Map aspect_ratio to BFL width/height parameters.
+
+    Converts 'WxH' string to width/height, rounded to nearest multiple of 16.
+    Delegates to provider's WidthMapper and HeightMapper for the actual mapping.
+    """
 
     name = ImageGenerationParameter.ASPECT_RATIO
 
@@ -31,130 +64,35 @@ def map(
         width = ((width + 8) // 16) * 16
         height = ((height + 8) // 16) * 16
 
-        request["width"] = width
-        request["height"] = height
+        # Delegate to provider mappers
+        request = _WidthMapper().map(request, width, model)
+        request = _HeightMapper().map(request, height, model)
         return request
 
 
-class PromptUpsamplingMapper(ParameterMapper):
-    """Map prompt_upsampling parameter to BFL request format."""
-
+class PromptUpsamplingMapper(_PromptUpsamplingMapper):
     name = ImageGenerationParameter.PROMPT_UPSAMPLING
 
-    def map(
-        self,
-        request: dict[str, Any],
-        value: object,
-        model: Model,
-    ) -> dict[str, Any]:
-        """Transform prompt_upsampling into provider request."""
-        validated_value = self._validate_value(value, model)
-        if validated_value is None:
-            return request
-
-        request["prompt_upsampling"] = validated_value
-        return request
-
-
-class SeedMapper(ParameterMapper):
-    """Map seed parameter to BFL request format."""
 
+class SeedMapper(_SeedMapper):
     name = ImageGenerationParameter.SEED
 
-    def map(
-        self,
-        request: dict[str, Any],
-        value: object,
-        model: Model,
-    ) -> dict[str, Any]:
-        """Transform seed into provider request."""
-        validated_value = self._validate_value(value, model)
-        if validated_value is None:
-            return request
-
-        request["seed"] = validated_value
-        return request
-
-
-class SafetyToleranceMapper(ParameterMapper):
-    """Map safety_tolerance parameter to BFL request format."""
 
+class SafetyToleranceMapper(_SafetyToleranceMapper):
     name = ImageGenerationParameter.SAFETY_TOLERANCE
 
-    def map(
-        self,
-        request: dict[str, Any],
-        value: object,
-        model: Model,
-    ) -> dict[str, Any]:
-        """Transform safety_tolerance into provider request."""
-        validated_value = self._validate_value(value, model)
-        if validated_value is None:
-            return request
-
-        request["safety_tolerance"] = validated_value
-        return request
-
-
-class OutputFormatMapper(ParameterMapper):
-    """Map output_format parameter to BFL request format."""
 
+class OutputFormatMapper(_OutputFormatMapper):
     name = ImageGenerationParameter.OUTPUT_FORMAT
 
-    def map(
-        self,
-        request: dict[str, Any],
-        value: object,
-        model: Model,
-    ) -> dict[str, Any]:
-        """Transform output_format into provider request."""
-        validated_value = self._validate_value(value, model)
-        if validated_value is None:
-            return request
-
-        request["output_format"] = validated_value
-        return request
-
-
-class StepsMapper(ParameterMapper):
-    """Map steps parameter to BFL request format (flex only)."""
 
+class StepsMapper(_StepsMapper):
     name = ImageGenerationParameter.STEPS
 
-    def map(
-        self,
-        request: dict[str, Any],
-        value: object,
-        model: Model,
-    ) -> dict[str, Any]:
-        """Transform steps into provider request."""
-        validated_value = self._validate_value(value, model)
-        if validated_value is None:
-            return request
-
-        request["steps"] = validated_value
-        return request
-
-
-class GuidanceMapper(ParameterMapper):
-    """Map guidance parameter to BFL request format (flex only)."""
 
+class GuidanceMapper(_GuidanceMapper):
     name = ImageGenerationParameter.GUIDANCE
 
-    def map(
-        self,
-        request: dict[str, Any],
-        value: object,
-        model: Model,
-    ) -> dict[str, Any]:
-        """Transform guidance into provider request."""
-        validated_value = self._validate_value(value, model)
-        if validated_value is None:
-            return request
-
-        request["guidance"] = validated_value
-        return request
-
 
 BFL_PARAMETER_MAPPERS: list[ParameterMapper] = [
     AspectRatioMapper(),

packages/capabilities/image-generation/tests/integration_tests/test_image_generation/test_generate.py

@@ -35,6 +35,7 @@ async def test_generate(provider: Provider, model: str, parameters: dict) -> Non
     client = create_client(
         capability=Capability.IMAGE_GENERATION,
         provider=provider,
+        model=model,
     )
     prompt = "A red apple on a white background"
 

packages/providers/bfl/pyproject.toml

@@ -0,0 +1,17 @@
+[project]
+name = "celeste-bfl"
+version = "0.3.0"
+description = "BFL (Black Forest Labs) provider package for Celeste AI"
+authors = [{name = "Kamilbenkirane", email = "kamil@withceleste.ai"}]
+license = {text = "Apache-2.0"}
+requires-python = ">=3.12"
+
+[tool.uv.sources]
+celeste-ai = { workspace = true }
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/celeste_bfl"]

packages/providers/bfl/src/celeste_bfl/__init__.py

@@ -0,0 +1 @@
+"""BFL (Black Forest Labs) provider package for Celeste AI."""

packages/providers/bfl/src/celeste_bfl/images/__init__.py

@@ -0,0 +1 @@
+"""BFL Images API provider package."""