"""
title: FLUX.1.1 Pro Manifold Function for Black Forest Lab Image Generation Models
author: mobilestack, credit to bgeneto
author_url: https://github.com/mobilestack/open-webui-flux-1.1-pro
forked from https://github.com/bgeneto/open-webui-flux-image-gen
funding_url: https://github.com/open-webui
version: 0.2
license: MIT
requirements: pydantic, requests
environment_variables: FLUX_PRO_API_BASE_URL, FLUX_PRO_API_KEY
supported providers: together.ai
https://api.together.xyz/v1/images/generations
"""
import base64
import os
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 representing the FLUX.1.1-pro Manifold Function.
"""
class Valves(BaseModel):
"""
Pydantic model for storing API keys and base URLs.
"""
FLUX_PRO_API_KEY: str = Field(
default="", description="Your API Key for Flux.1.1 Pro"
)
FLUX_PRO_API_BASE_URL: str = Field(
default="https://api.together.xyz/v1/images/generations",
description="Base URL for the API",
)
def __init__(self):
"""
Initialize the Pipe class with default values and environment variables.
"""
self.type = "manifold"
self.id = "FLUX_1_1_PRO"
self.name = "FLUX.1.1-pro: "
self.valves = self.Valves(
FLUX_PRO_API_KEY=os.getenv("FLUX_PRO_API_KEY", ""),
FLUX_PRO_API_BASE_URL=os.getenv(
"FLUX_PRO_API_BASE_URL",
"https://api.together.xyz/v1/images/generations",
),
)
def url_to_img_data(self, url: str) -> str:
"""
Convert a URL to base64-encoded image data.
Args:
url (str): The URL of the image.
Returns:
str: Base64-encoded image data.
"""
headers = {"Authorization": f"Bearer {self.valves.FLUX_PRO_API_KEY}"}
response = requests.get(url, headers=headers)
response.raise_for_status()
content_type = response.headers.get("Content-Type", "application/octet-stream")
encoded_content = base64.b64encode(response.content).decode("utf-8")
return f"data:{content_type};base64,{encoded_content}"
def stream_response(
self, headers: Dict[str, str], payload: Dict[str, Any]
) -> Generator[str, None, None]:
yield self.non_stream_response(headers, payload)
def get_img_extension(self, img_data: str) -> Union[str, None]:
"""
Get the image extension based on the base64-encoded data.
Args:
img_data (str): Base64-encoded image data.
Returns:
Union[str, None]: The image extension or None if unsupported.
"""
if img_data.startswith("/9j/"):
return "jpeg"
elif img_data.startswith("iVBOR"):
return "png"
elif img_data.startswith("R0lG"):
return "gif"
elif img_data.startswith("UklGR"):
return "webp"
return None
def handle_json_response(self, response: requests.Response) -> str:
"""
Handle JSON response from the API.
Args:
response (requests.Response): The response object.
Returns:
str: The formatted image data or an error message.
"""
resp = response.json()
if "output" in resp:
img_data = resp["output"][0]
elif "data" in resp and "b64_json" in resp["data"][0]:
img_data = resp["data"][0]["b64_json"]
else:
return "Error: Unexpected response format for the image provider!"
# split ;base64, from img_data
try:
img_data = img_data.split(";base64,")[1]
except IndexError:
pass
img_ext = self.get_img_extension(img_data[:9])
if not img_ext:
return "Error: Unsupported image format!"
# rebuild img_data with proper format
img_data = f"data:image/{img_ext};base64,{img_data}"
return f"\n`GeneratedImage.{img_ext}`"
def handle_image_response(self, response: requests.Response) -> str:
"""
Handle image response from the API.
Args:
response (requests.Response): The response object.
Returns:
str: The formatted image data.
"""
content_type = response.headers.get("Content-Type", "")
# check image type in the content type
img_ext = "png"
if "image/" in content_type:
img_ext = content_type.split("/")[-1]
image_base64 = base64.b64encode(response.content).decode("utf-8")
return f"\n`GeneratedImage.{img_ext}`"
def non_stream_response(
self, headers: Dict[str, str], payload: Dict[str, Any]
) -> str:
"""
Get a non-streaming response from the API.
Args:
headers (Dict[str, str]): The headers for the request.
payload (Dict[str, Any]): The payload for the request.
Returns:
str: The response from the API.
"""
try:
response = requests.post(
url=self.valves.FLUX_PRO_API_BASE_URL,
headers=headers,
json=payload,
stream=False,
timeout=(3.05, 60),
)
response.raise_for_status()
content_type = response.headers.get("Content-Type", "")
if "application/json" in content_type:
return self.handle_json_response(response)
elif "image/" in content_type:
return self.handle_image_response(response)
else:
return f"Error: Unsupported content type {content_type}"
except requests.exceptions.RequestException as e:
return f"Error: Request failed: {e}"
except Exception as e:
return f"Error: {e}"
def pipes(self) -> List[Dict[str, str]]:
"""
Get the list of available pipes.
Returns:
List[Dict[str, str]]: The list of pipes.
"""
return [{"id": "flux_1_1_pro", "name": "Flux 1.1 PRO"}]
def pipe(
self, body: Dict[str, Any]
) -> Union[str, Generator[str, None, None], Iterator[str]]:
"""
Process the pipe request.
Args:
body (Dict[str, Any]): The request body.
Returns:
Union[str, Generator[str, None, None], Iterator[str]]: The response from the API.
"""
headers = {
"Authorization": f"Bearer {self.valves.FLUX_PRO_API_KEY}",
"Content-Type": "application/json",
}
body["stream"] = False
prompt = get_last_user_message(body["messages"])
headers_map = {
"huggingface.co": {"x-wait-for-model": "true"},
"replicate.com": {"Prefer": "wait"},
"together.xyz": {},
}
payload_map = {
"huggingface.co": {"inputs": prompt},
"replicate.com": {
"input": {
"prompt": prompt,
"go_fast": True,
"num_outputs": 1,
"aspect_ratio": "1:1",
"output_format": "webp",
"output_quality": 90,
}
},
"together.xyz": {
"model": "black-forest-labs/FLUX.1.1-pro",
"prompt": prompt,
"width": 1024,
"height": 1024,
"steps": 4,
"n": 1,
"response_format": "b64_json",
},
}
payload = None
for key in payload_map:
if key in self.valves.FLUX_PRO_API_BASE_URL:
payload = payload_map[key]
headers.update(headers_map.get(key, {}))
break
if payload is None:
return "Error: Unsupported API base URL! Remember, that's the beauty of open-source: you can add your own..."
try:
if body.get("stream", False):
return self.stream_response(headers, payload)
else:
return self.non_stream_response(headers, payload)
except requests.exceptions.RequestException as e:
return f"Error: Request failed: {e}"
except Exception as e:
return f"Error: {e}"