"""
title: Manifold Pipe for FLUX.1.1 Pro of Black Forest Labs
author: nknickrehm
funding_url: https://github.com/open-webui
version: 0.0.1
license: MIT
requirements: pydantic, requests
environment_variables: BFL_API_KEY
"""
import base64
import os
import time
from typing import Any, Dict, Generator, Iterator, List, Union
import requests
from open_webui.utils.misc import get_last_user_message
from pydantic import BaseModel, Field
class Pipe:
class Valves(BaseModel):
BFL_API_KEY: str = Field(
default="", description="Your API Key for Black Forest Labs"
)
def __init__(self):
self.type = "manifold"
self.id = "BFL_FLUX_1_1_PRO"
self.name = "BFL_FLUX.1.1-pro: "
self.valves = self.Valves(
BFL_API_KEY=os.getenv("BFL_API_KEY", ""),
)
def send_image_generation_request(self, prompt: str) -> str:
request = requests.post(
"https://api.bfl.ml/v1/flux-pro-1.1",
headers={
"accept": "application/json",
"x-key": self.valves.BFL_API_KEY,
"Content-Type": "application/json",
},
json={"prompt": prompt, "width": 1024, "height": 768},
).json()
print(request)
return request["id"]
def poll_result(self, request_id: str) -> str:
while self.get_result(request_id)["status"] not in [
"Ready",
"Error",
"Content Moderated",
"Request Moderated",
"Task not found",
]:
time.sleep(1)
result = self.get_result(request_id)
status = result["status"]
if status == "Ready":
return result["result"]["sample"]
elif status in [
"Error",
"Content Moderated",
"Request Moderated",
"Task not found",
]:
raise RuntimeError(f"Image generation failed. Status: {status}")
def get_result(self, request_id: str):
result = requests.get(
"https://api.bfl.ml/v1/get_result",
headers={
"accept": "application/json",
"x-key": self.valves.BFL_API_KEY,
},
params={
"id": request_id,
},
).json()
return result
def pipes(self) -> List[Dict[str, str]]:
return [{"id": "flux-1-1-pro", "name": "flux-1.1-pro"}]
async def pipe(
self, body: Dict[str, Any]
) -> Union[str, Generator[str, None, None], Iterator[str]]:
prompt = get_last_user_message(body["messages"])
try:
bfl_task_id = self.send_image_generation_request(prompt)
image_url = self.poll_result(bfl_task_id)
return f""
except requests.exceptions.RequestException as e:
return f"Error: Request failed: {e}"
except Exception as e:
return f"Error: {e}"
