Merge pull request #4 from black-forest-labs/saksham/add-kontext

- Add Kontext Pro API support
- Support to use load balanced endpoint `api.bfl.ai`, with use of `polling_url`
- Add lazy cleanup of `task_id: polling_url` to reduce memory overhead on local client
- Add Kontext unit test case
- Updated readme to use dev mode by default

Author: TheNoise2Signal

README.md

@@ -18,6 +18,11 @@ from blackforest import BFLClient
 # or
 from blackforestlabs import BFLClient
 
+import os
+
+# For synchronous API call (great for testing, but please use "production" for async calls, for faster throughput)
+os.environ["BFL_ENV"] = "dev"  
+
 # Initialize the client
 client = BFLClient(api_key="your-api-key")
 
@@ -29,13 +34,39 @@ inputs = {
         "output_format": "jpeg"
     }
 response = client.generate("flux-pro-1.1", inputs)
+
+# For Flux Kontext Pro with reference images
+kontext_inputs = {
+        "prompt": "A beautiful landscape in the style of the reference image",
+        "input_image": "path/to/reference/image.jpg",  # File path (auto-encoded) or base64
+        "input_image_2": "path/to/another/image.png",  # Optional multiref (experimental)
+        "aspect_ratio": "16:9",  # Between 1:4 and 4:1
+        "output_format": "png",
+        "seed": 42,  # Optional for reproducibility
+        "safety_tolerance": 2,  # 0-6, higher = less strict
+        "prompt_upsampling": True  # Enhanced prompt processing
+    }
+response = client.generate("flux-kontext-pro", kontext_inputs)
 ```
 
 ## Features
 
 - Official Python interface for Black Forest Labs API
 - Automatic request handling and response parsing
 - Type hints for better IDE support
+- Support for all Flux models including Flux Kontext Pro with multi-reference capabilities
+
+## Supported Models
+
+- `flux-dev` - Development model
+- `flux-pro` - Professional model 
+- `flux-pro-1.1` - Enhanced professional model
+- `flux-pro-1.1-ultra` - Ultra high-quality model
+- `flux-kontext-pro` - Context-aware model with reference image support and experimental multi-reference capabilities
+- `flux-pro-1.0-fill` - Image inpainting model
+- `flux-pro-1.0-expand` - Image expansion model
+- `flux-pro-1.0-canny` - Canny edge-guided model
+- `flux-pro-1.0-depth` - Depth-guided model
 
 ## Requirements
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "blackforest"
-version = "0.1.0"
+version = "0.1.1"
 description = "The official Python library for the Black Forest Labs API"
 readme = "README.md"
 authors = [

src/blackforest/client.py

@@ -37,7 +37,7 @@ class BFLClient:
     def __init__(
         self,
         api_key: str,
-        base_url: str = "https://api.us1.bfl.ai",
+        base_url: str = "https://api.bfl.ai",
         timeout: int = 30,
     ):
         """
@@ -57,6 +57,50 @@ def __init__(
             'Content-Type': 'application/json',
             'Accept': 'application/json',
         })
+        # Map to store task_id -> (polling_url, timestamp)
+        self._task_polling_urls = {}
+
+    def _cleanup_expired_polling_urls(self) -> None:
+        """
+        Remove polling URL entries that are older than 30 minutes.
+        """
+        current_time = time.time()
+        expired_keys = []
+        
+        for task_id, (polling_url, timestamp) in self._task_polling_urls.items():
+            # Remove entries older than 30 minutes (1800 seconds)
+            if current_time - timestamp > 1800:
+                expired_keys.append(task_id)
+        
+        for key in expired_keys:
+            del self._task_polling_urls[key]
+
+    def clear_polling_urls(self) -> None:
+        """
+        Manually clear all stored polling URLs.
+        """
+        self._task_polling_urls.clear()
+
+    def _get_polling_endpoint(self, task_id: str) -> str:
+        """
+        Get the appropriate polling endpoint for a task.
+        
+        Args:
+            task_id: The task ID to get the endpoint for
+            
+        Returns:
+            The endpoint (relative path or full URL) to use for polling
+        """
+        # Clean up expired entries
+        self._cleanup_expired_polling_urls()
+        
+        if task_id in self._task_polling_urls:
+            polling_url, _ = self._task_polling_urls[task_id]  # Extract URL from tuple
+            # Return the full polling URL as-is since _request() now handles full URLs
+            return polling_url
+        else:
+            # Fallback to constructed URL
+            return f"/v1/get_result?id={task_id}"
 
     def _request(
         self,
@@ -71,7 +115,7 @@ def _request(
 
         Args:
             method: HTTP method
-            endpoint: API endpoint
+            endpoint: API endpoint (can be a relative path or full URL)
             params: URL parameters
             data: Form data
             json: JSON data
@@ -82,7 +126,11 @@ def _request(
         Raises:
             BFLError: If the API request fails
         """
-        url = urljoin(self.base_url, endpoint)
+        # If endpoint is already a full URL, use it directly
+        if endpoint.startswith(('http://', 'https://')):
+            url = endpoint
+        else:
+            url = urljoin(self.base_url, endpoint)
 
         try:
             response = self.session.request(
@@ -112,6 +160,38 @@ def _encode_image(self, image_path: str) -> str:
         with open(image_path, 'rb') as image_file:
             return base64.b64encode(image_file.read()).decode('utf-8')
 
+    def _is_file_path(self, value: str) -> bool:
+        """Check if a string is a file path (not base64 or URL)."""
+        if value.startswith(('http://', 'https://')):
+            return False
+        # Check if it looks like base64 (no path separators, mostly alphanumeric with +/=)
+        if '/' not in value and '\\' not in value and len(value) > 100:
+            try:
+                base64.b64decode(value)
+                return False  # It's valid base64
+            except Exception:
+                pass
+        # Check if file exists
+        return os.path.exists(value)
+
+    def _process_kontext_inputs(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
+        """Process flux-kontext-pro inputs to automatically encode image paths."""
+        processed_inputs = inputs.copy()
+        
+        # Image fields that might need encoding
+        image_fields = ['input_image', 'input_image_2', 'input_image_3', 'input_image_4']
+        
+        for field in image_fields:
+            if field in processed_inputs and processed_inputs[field]:
+                value = processed_inputs[field]
+                if isinstance(value, str) and self._is_file_path(value):
+                    try:
+                        processed_inputs[field] = self._encode_image(value)
+                    except Exception as e:
+                        raise BFLError(f"Error encoding image file '{value}' for field '{field}': {str(e)}")
+        
+        return processed_inputs
+
     def _process_folder(self, folder_path: str) -> List[str]:
         """Process all images in a folder and return list of base64 encoded images."""
         folder = Path(folder_path)
@@ -198,8 +278,15 @@ def process_image(
         # Make the API request
         response = self._request("POST", endpoint, json=payload)
 
+        # Store the polling URL for this task if available
+        task_id = response.get("id")
+        if task_id and "polling_url" in response:
+            # Clean up expired entries before adding new one
+            self._cleanup_expired_polling_urls()
+            self._task_polling_urls[task_id] = (response["polling_url"], time.time())
+
         return ImageProcessingResponse(
-            task_id=response.get("id"),
+            task_id=task_id,
             status="submitted",
             result=response
         )
@@ -214,7 +301,8 @@ def get_task_status(self, task_id: str) -> ImageProcessingResponse:
         Returns:
             ImageProcessingResponse containing current status and result if available
         """
-        response = self._request("GET", f"/v1/get_result?id={task_id}")
+        endpoint = self._get_polling_endpoint(task_id)
+        response = self._request("GET", endpoint)
 
         return ImageProcessingResponse(
             task_id=task_id,
@@ -250,13 +338,15 @@ def get_polling_result(
         while attempts < config.max_retries:
             print(f"Polling task {task_id} for result. \
                   Attempt {attempts + 1} of {config.max_retries}")
-            response = self._request("GET", f"/v1/get_result?id={task_id}")
+            
+            endpoint = self._get_polling_endpoint(task_id)
+            response = self._request("GET", endpoint)
 
             # Check if the task is complete
             if response.get("status") in ["Ready", "completed", "failed"]:
                 if response.get("status") == "failed":
-                    raise BFLError(f"Task failed: {response.get('error', \
-                                                                'Unknown error')}")
+                    error_msg = response.get('error', 'Unknown error')
+                    raise BFLError(f"Task failed: {error_msg}")
                 return response.get("result", {})
 
             # Check for timeout
@@ -319,30 +409,42 @@ def generate(
             raise BFLError(f"Model {model} not supported. \
                            Supported models: {list(self.model_input_registry.keys())}")
 
+        # Process inputs for automatic image encoding if using flux-kontext-pro
+        processed_inputs = inputs
+        if model == "flux-kontext-pro":
+            processed_inputs = self._process_kontext_inputs(inputs)
+
         # Convert the inputs to the appropriate type
-        typed_inputs = input_cls(**inputs)
+        typed_inputs = input_cls(**processed_inputs)
 
         response = self._request("POST", f"{self.api_version}/{model}",
                                   json=typed_inputs.model_dump(exclude_none=True))
 
+        # Store the polling URL for this task
+        task_id = response["id"]
+        if "polling_url" in response:
+            # Clean up expired entries before adding new one
+            self._cleanup_expired_polling_urls()
+            self._task_polling_urls[task_id] = (response["polling_url"], time.time())
+
         # Track usage if requested
         if track_usage:
             self.track_usage_via_api(model, 1)
 
         # If sync is True, poll for results
         if config.sync:
             try:
-                result = self.get_polling_result(response["id"], config)
+                result = self.get_polling_result(task_id, config)
                 return SyncResponse(
-                    id=response["id"],
+                    id=task_id,
                     result=result
                 )
             except Exception as e:
                 raise BFLError(f"Error getting synchronous result: {str(e)}")
 
         else:
             return AsyncResponse(
-                id=response["id"],
+                id=task_id,
                 polling_url=response["polling_url"]
             )
 
@@ -383,5 +485,4 @@ def track_usage_via_api(self, name: str, n: int = 1) -> None:
             self._request("POST", endpoint, json=payload)
             print(f"Successfully tracked usage for {name} with {n} generations")
         except BFLError as e:
-            raise BFLError(f"Failed to track usage for {name}: {str(e)}")
-q
+            raise BFLError(f"Failed to track usage for {name}: {str(e)}")

src/blackforest/resources/mapping/model_input_registry.py

@@ -1,4 +1,5 @@
 from blackforest.types.inputs.flux_dev import FluxDevInputs
+from blackforest.types.inputs.flux_kontext_pro import FluxKontextProInputs
 from blackforest.types.inputs.flux_pro import FluxProInputs
 from blackforest.types.inputs.flux_pro_1_1 import FluxPro11Inputs
 from blackforest.types.inputs.flux_pro_canny import FluxProCannyInputs
@@ -16,4 +17,5 @@
     "flux-pro-1.0-expand": FluxProExpandInputs,
     "flux-pro-1.0-canny": FluxProCannyInputs,
     "flux-pro-1.0-depth": FluxProDepthInputs,
+    "flux-kontext-pro": FluxKontextProInputs,
 }

src/blackforest/types/inputs/flux_kontext_pro.py

@@ -0,0 +1,67 @@
+import re
+from typing import Optional
+
+from pydantic import Field, HttpUrl, model_validator
+
+from blackforest.types.base.output_format import OutputFormat
+from blackforest.types.inputs.generic import GenericImageInput
+
+
+class FluxKontextProInputs(GenericImageInput):
+    """Inputs for the Flux Kontext Pro model."""
+    
+    # Override prompt to make it required
+    prompt: str = Field(
+        example="ein fantastisches bild",
+        description="Text prompt for image generation.",
+    )
+
+    # Override output_format default to png
+    output_format: Optional[OutputFormat] = Field(
+        default=OutputFormat.png,
+        description="Output format for the generated image. Can be 'jpeg' or 'png'.",
+    )
+
+    # Kontext-specific input image fields
+    input_image: Optional[str] = Field(
+        default=None,
+        description="Base64 encoded image or URL to use with Kontext.",
+    )
+    input_image_2: Optional[str] = Field(
+        default=None,
+        description="Base64 encoded image or URL to use with Kontext. *Experimental Multiref*",
+    )
+    input_image_3: Optional[str] = Field(
+        default=None,
+        description="Base64 encoded image or URL to use with Kontext. *Experimental Multiref*",
+    )
+    input_image_4: Optional[str] = Field(
+        default=None,
+        description="Base64 encoded image or URL to use with Kontext. *Experimental Multiref*",
+    )
+
+    # Custom aspect ratio field with specific validation
+    aspect_ratio: Optional[str] = Field(
+        default=None, description="Aspect ratio of the image between 21:9 and 9:21"
+    )
+
+    @model_validator(mode="after")
+    def validate_aspect_ratio(self):
+        try:
+            if self.aspect_ratio is not None:
+                # ensure proper format (1:1) and ratio is between 21:9 and 9:21
+                if not re.match(r"^\d+:\d+$", self.aspect_ratio):
+                    raise ValueError(
+                        "Aspect ratio must be in the format of 'width:height'"
+                    )
+                width, height = map(int, self.aspect_ratio.split(":"))
+                ratio = width / height
+                min_ratio = 1 / 4
+                max_ratio = 4 / 1
+                if not (min_ratio <= ratio <= max_ratio):
+                    raise ValueError(
+                        f"Aspect ratio {self.aspect_ratio} ({ratio:.3f}) must be between 1:4 and 4:1"
+                    )
+        except Exception as e:
+            raise ValueError(f"Invalid aspect ratio: {e}")
+        return self