Add BFL API

lib/api.py

@@ -1,5 +1,6 @@
 import base64
 import io
+import time
 
 import httpx
 import streamlit as st
@@ -20,6 +21,7 @@ def txt2txt_generate(api_key, service, model, parameters, **kwargs):
         return st.write_stream(stream)
     except APIError as e:
         # OpenAI uses this message for streaming errors and attaches response.error to error.body
+        # https://github.com/openai/openai-python/blob/v1.0.0/src/openai/_streaming.py#L59
         return e.body if e.message == "An error occurred during streaming" else e.message
     except Exception as e:
         return str(e)
@@ -31,19 +33,28 @@ def txt2img_generate(api_key, service, model, inputs, parameters, **kwargs):
         headers["Authorization"] = f"Bearer {api_key}"
         headers["X-Wait-For-Model"] = "true"
         headers["X-Use-Cache"] = "false"
+
     if service == "Fal":
         headers["Authorization"] = f"Key {api_key}"
 
+    if service == "BFL":
+        headers["x-key"] = api_key
+
     json = {}
     if service == "Hugging Face":
         json = {
             "inputs": inputs,
             "parameters": {**parameters, **kwargs},
         }
+
     if service == "Fal":
         json = {**parameters, **kwargs}
         json["prompt"] = inputs
 
+    if service == "BFL":
+        json = {**parameters, **kwargs}
+        json["prompt"] = inputs
+
     base_url = f"{Config.SERVICES[service]}/{model}"
 
     try:
@@ -51,8 +62,9 @@ def txt2img_generate(api_key, service, model, inputs, parameters, **kwargs):
         if response.status_code // 100 == 2:  # 2xx
             if service == "Hugging Face":
                 return Image.open(io.BytesIO(response.content))
+
             if service == "Fal":
-                # sync_mode means wait for image base64 instead of CDN link
+                # Sync mode means wait for image base64 string instead of CDN link
                 if parameters.get("sync_mode", True):
                     bytes = base64.b64decode(response.json()["images"][0]["url"].split(",")[-1])
                     return Image.open(io.BytesIO(bytes))
@@ -60,6 +72,32 @@ def txt2img_generate(api_key, service, model, inputs, parameters, **kwargs):
                     url = response.json()["images"][0]["url"]
                     image = httpx.get(url, headers=headers, timeout=Config.TXT2IMG_TIMEOUT)
                     return Image.open(io.BytesIO(image.content))
+
+            # BFL is async so we need to poll for result
+            # https://api.bfl.ml/docs
+            if service == "BFL":
+                id = response.json()["id"]
+                url = f"{Config.SERVICES[service]}/get_result?id={id}"
+
+                retries = 0
+                while retries < Config.TXT2IMG_TIMEOUT:
+                    response = httpx.get(url, timeout=Config.TXT2IMG_TIMEOUT)
+                    if response.status_code // 100 != 2:
+                        return f"Error: {response.status_code} {response.text}"
+
+                    if response.json()["status"] == "Ready":
+                        image = httpx.get(
+                            response.json()["result"]["sample"],
+                            headers=headers,
+                            timeout=Config.TXT2IMG_TIMEOUT,
+                        )
+                        return Image.open(io.BytesIO(image.content))
+
+                    retries += 1
+                    time.sleep(1)
+
+                return "Error: API timeout"
+
         else:
             return f"Error: {response.status_code} {response.text}"
     except Exception as e:

lib/config.py

@@ -5,11 +5,12 @@ Config = SimpleNamespace(
     ICON="⚡",
     LAYOUT="wide",
     SERVICES={
+        "BFL": "https://api.bfl.ml/v1",
+        "Fal": "https://fal.run",
         "Hugging Face": "https://api-inference.huggingface.co/models",
         "Perplexity": "https://api.perplexity.ai",
-        "Fal": "https://fal.run",
     },
-    TXT2IMG_TIMEOUT=120,
+    TXT2IMG_TIMEOUT=60,
     TXT2IMG_HIDDEN_PARAMETERS=[
         # sent to API but not shown in generation parameters accordion
         "enable_safety_checker",
@@ -26,23 +27,33 @@ Config = SimpleNamespace(
     ],
     TXT2IMG_NEGATIVE_PROMPT="ugly, unattractive, disfigured, deformed, mutated, malformed, blurry, grainy, noisy, oversaturated, undersaturated, overexposed, underexposed, worst quality, low details, lowres, watermark, signature, autograph, trademark, sloppy, cluttered",
     TXT2IMG_DEFAULT_MODEL={
-        # index of model in below lists
+        # The index of model in below lists
+        "BFL": 2,
+        "Fal": 0,
         "Hugging Face": 2,
-        "Fal": 2,
     },
     TXT2IMG_MODELS={
-        "Hugging Face": [
-            "black-forest-labs/flux.1-dev",
-            "black-forest-labs/flux.1-schnell",
-            "stabilityai/stable-diffusion-xl-base-1.0",
+        # Model IDs referenced in Text_to_Image.py
+        "BFL": [
+            "flux-dev",
+            "flux-pro",
+            "flux-pro-1.1",
         ],
         "Fal": [
             "fal-ai/aura-flow",
+            "fal-ai/flux/dev",
+            "fal-ai/flux/schnell",
             "fal-ai/flux-pro",
+            "fal-ai/flux-pro/v1.1",
             "fal-ai/fooocus",
             "fal-ai/kolors",
             "fal-ai/stable-diffusion-v3-medium",
         ],
+        "Hugging Face": [
+            "black-forest-labs/flux.1-dev",
+            "black-forest-labs/flux.1-schnell",
+            "stabilityai/stable-diffusion-xl-base-1.0",
+        ],
     },
     TXT2IMG_DEFAULT_IMAGE_SIZE="square_hd",  # fal image sizes
     TXT2IMG_IMAGE_SIZES=[
@@ -76,7 +87,6 @@ Config = SimpleNamespace(
         "1344x704",  # 21:11
         "1408x704",  # 2:1
     ],
-    # TODO: txt2img fooocus styles like "Fooocus V2" and "Fooocus Enhance" (use multiselect in UI)
     TXT2TXT_DEFAULT_SYSTEM="You are a helpful assistant. Be precise and concise.",
     TXT2TXT_DEFAULT_MODEL={
         "Hugging Face": 4,

lib/presets.py

@@ -1,9 +1,9 @@
 from types import SimpleNamespace
 
-# txt2txt services
+# txt2txt
 ServicePresets = SimpleNamespace(
+    # Every service has model and system messages
     HUGGING_FACE={
-        # every service has model and system messages
         "frequency_penalty": 0.0,
         "frequency_penalty_min": -2.0,
         "frequency_penalty_max": 2.0,
@@ -17,7 +17,7 @@ ServicePresets = SimpleNamespace(
     },
 )
 
-# txt2img models
+# txt2img
 ModelPresets = SimpleNamespace(
     AURA_FLOW={
         "name": "AuraFlow",
@@ -30,33 +30,89 @@ ModelPresets = SimpleNamespace(
         "parameters": ["seed", "num_inference_steps", "guidance_scale", "expand_prompt"],
         "kwargs": {"num_images": 1, "sync_mode": False},
     },
-    FLUX_DEV={
+    FLUX_1_1_PRO_BFL={
+        "name": "FLUX1.1 Pro",
+        "parameters": ["seed", "width", "height", "prompt_upsampling"],
+        "kwargs": {"safety_tolerance": 6},
+    },
+    FLUX_PRO_BFL={
+        "name": "FLUX.1 Pro",
+        "guidance_scale": 2.5,
+        "guidance_scale_min": 1.5,
+        "guidance_scale_max": 5.0,
+        "num_inference_steps": 40,
+        "num_inference_steps_min": 10,
+        "num_inference_steps_max": 50,
+        "parameters": ["seed", "width", "height", "steps", "guidance", "prompt_upsampling"],
+        "kwargs": {"safety_tolerance": 6, "interval": 1},
+    },
+    FLUX_DEV_BFL={
         "name": "FLUX.1 Dev",
         "num_inference_steps": 28,
         "num_inference_steps_min": 10,
         "num_inference_steps_max": 50,
-        "guidance_scale": 3.5,
-        "guidance_scale_min": 1.0,
-        "guidance_scale_max": 10.0,
-        "parameters": ["width", "height", "guidance_scale", "num_inference_steps"],
-        "kwargs": {"max_sequence_length": 512},
+        "guidance_scale": 3.0,
+        "guidance_scale_min": 1.5,
+        "guidance_scale_max": 5.0,
+        "parameters": ["seed", "width", "height", "steps", "guidance", "prompt_upsampling"],
+        "kwargs": {"safety_tolerance": 6},
     },
-    FLUX_PRO={
+    FLUX_1_1_PRO_FAL={
+        "name": "FLUX1.1 Pro",
+        "parameters": ["seed", "image_size"],
+        "kwargs": {
+            "num_images": 1,
+            "sync_mode": False,
+            "safety_tolerance": 6,
+            "enable_safety_checker": False,
+        },
+    },
+    FLUX_PRO_FAL={
         "name": "FLUX.1 Pro",
+        "guidance_scale": 2.5,
+        "guidance_scale_min": 1.5,
+        "guidance_scale_max": 5.0,
+        "num_inference_steps": 40,
+        "num_inference_steps_min": 10,
+        "num_inference_steps_max": 50,
+        "parameters": ["seed", "image_size", "num_inference_steps", "guidance_scale"],
+        "kwargs": {"num_images": 1, "sync_mode": False, "safety_tolerance": 6},
+    },
+    FLUX_DEV_FAL={
+        "name": "FLUX.1 Dev",
         "num_inference_steps": 28,
         "num_inference_steps_min": 10,
         "num_inference_steps_max": 50,
-        "guidance_scale": 3.5,
-        "guidance_scale_min": 1.0,
-        "guidance_scale_max": 10.0,
+        "guidance_scale": 3.0,
+        "guidance_scale_min": 1.5,
+        "guidance_scale_max": 5.0,
         "parameters": ["seed", "image_size", "num_inference_steps", "guidance_scale"],
         "kwargs": {"num_images": 1, "sync_mode": False, "safety_tolerance": 6},
     },
-    FLUX_SCHNELL={
+    FLUX_SCHNELL_FAL={
+        "name": "FLUX.1 Schnell",
+        "num_inference_steps": 4,
+        "num_inference_steps_min": 1,
+        "num_inference_steps_max": 12,
+        "parameters": ["seed", "image_size", "num_inference_steps"],
+        "kwargs": {"num_images": 1, "sync_mode": False, "enable_safety_checker": False},
+    },
+    FLUX_DEV_HF={
+        "name": "FLUX.1 Dev",
+        "num_inference_steps": 28,
+        "num_inference_steps_min": 10,
+        "num_inference_steps_max": 50,
+        "guidance_scale": 3.0,
+        "guidance_scale_min": 1.5,
+        "guidance_scale_max": 5.0,
+        "parameters": ["width", "height", "guidance_scale", "num_inference_steps"],
+        "kwargs": {"max_sequence_length": 512},
+    },
+    FLUX_SCHNELL_HF={
         "name": "FLUX.1 Schnell",
         "num_inference_steps": 4,
         "num_inference_steps_min": 1,
-        "num_inference_steps_max": 8,
+        "num_inference_steps_max": 12,
         "parameters": ["width", "height", "num_inference_steps"],
         "kwargs": {"guidance_scale": 0.0, "max_sequence_length": 256},
     },

pages/2_🎨_Text_to_Image.py

@@ -6,34 +6,45 @@ import streamlit as st
 from lib import Config, ModelPresets, txt2img_generate
 
 SERVICE_SESSION = {
+    "BFL": "api_key_bfl",
     "Fal": "api_key_fal",
     "Hugging Face": "api_key_hugging_face",
 }
 
 SESSION_TOKEN = {
+    "api_key_bfl": os.environ.get("BFL_API_KEY") or None,
     "api_key_fal": os.environ.get("FAL_KEY") or None,
     "api_key_hugging_face": os.environ.get("HF_TOKEN") or None,
 }
 
+# Model IDs in lib/config.py
 PRESET_MODEL = {
-    "black-forest-labs/flux.1-dev": ModelPresets.FLUX_DEV,
-    "black-forest-labs/flux.1-schnell": ModelPresets.FLUX_SCHNELL,
+    "black-forest-labs/flux.1-dev": ModelPresets.FLUX_DEV_HF,
+    "black-forest-labs/flux.1-schnell": ModelPresets.FLUX_SCHNELL_HF,
     "stabilityai/stable-diffusion-xl-base-1.0": ModelPresets.STABLE_DIFFUSION_XL,
     "fal-ai/aura-flow": ModelPresets.AURA_FLOW,
-    "fal-ai/flux-pro": ModelPresets.FLUX_PRO,
+    "fal-ai/flux/dev": ModelPresets.FLUX_DEV_FAL,
+    "fal-ai/flux/schnell": ModelPresets.FLUX_SCHNELL_FAL,
+    "fal-ai/flux-pro": ModelPresets.FLUX_PRO_FAL,
+    "fal-ai/flux-pro/v1.1": ModelPresets.FLUX_1_1_PRO_FAL,
     "fal-ai/fooocus": ModelPresets.FOOOCUS,
     "fal-ai/kolors": ModelPresets.KOLORS,
     "fal-ai/stable-diffusion-v3-medium": ModelPresets.STABLE_DIFFUSION_3,
+    "flux-pro-1.1": ModelPresets.FLUX_1_1_PRO_BFL,
+    "flux-pro": ModelPresets.FLUX_PRO_BFL,
+    "flux-dev": ModelPresets.FLUX_DEV_BFL,
 }
 
-# config
 st.set_page_config(
     page_title=f"{Config.TITLE} | Text to Image",
     page_icon=Config.ICON,
     layout=Config.LAYOUT,
 )
 
-# initialize state
+# Initialize Streamlit session state
+if "api_key_bfl" not in st.session_state:
+    st.session_state.api_key_bfl = ""
+
 if "api_key_fal" not in st.session_state:
     st.session_state.api_key_fal = ""
 
@@ -49,7 +60,6 @@ if "txt2img_messages" not in st.session_state:
 if "txt2img_seed" not in st.session_state:
     st.session_state.txt2img_seed = 0
 
-# sidebar
 st.logo("logo.svg")
 st.sidebar.header("Settings")
 service = st.sidebar.selectbox(
@@ -59,7 +69,7 @@ service = st.sidebar.selectbox(
     disabled=st.session_state.running,
 )
 
-# disable API key input and hide value if set by environment variable (handle empty string value later)
+# Disable API key input and hide value if set by environment variable; handle empty string value later.
 for display_name, session_key in SERVICE_SESSION.items():
     if service == display_name:
         st.session_state[session_key] = st.sidebar.text_input(
@@ -75,7 +85,6 @@ model = st.sidebar.selectbox(
     options=Config.TXT2IMG_MODELS[service],
     index=Config.TXT2IMG_DEFAULT_MODEL[service],
     disabled=st.session_state.running,
-    format_func=lambda x: x.split("/")[1],
 )
 
 # heading
@@ -84,7 +93,7 @@ st.html("""
     <p>Generate an image from a text prompt.</p>
 """)
 
-# build parameters from preset
+# Build parameters from preset by rendering the appropriate input widgets
 parameters = {}
 preset = PRESET_MODEL[model]
 for param in preset["parameters"]:
@@ -134,7 +143,7 @@ for param in preset["parameters"]:
             value=Config.TXT2IMG_DEFAULT_ASPECT_RATIO,
             disabled=st.session_state.running,
         )
-    if param == "guidance_scale":
+    if param in ["guidance_scale", "guidance"]:
         parameters[param] = st.sidebar.slider(
             "Guidance Scale",
             preset["guidance_scale_min"],
@@ -143,7 +152,7 @@ for param in preset["parameters"]:
             0.1,
             disabled=st.session_state.running,
         )
-    if param == "num_inference_steps":
+    if param in ["num_inference_steps", "steps"]:
         parameters[param] = st.sidebar.slider(
             "Inference Steps",
             preset["num_inference_steps_min"],
@@ -152,20 +161,20 @@ for param in preset["parameters"]:
             1,
             disabled=st.session_state.running,
         )
-    if param == "expand_prompt":
+    if param in ["expand_prompt", "prompt_expansion"]:
         parameters[param] = st.sidebar.checkbox(
-            "Expand Prompt",
+            "Prompt Expansion",
             value=False,
             disabled=st.session_state.running,
         )
-    if param == "prompt_expansion":
+    if param == "prompt_upsampling":
         parameters[param] = st.sidebar.checkbox(
-            "Prompt Expansion",
+            "Prompt Upsampling",
             value=False,
             disabled=st.session_state.running,
         )
 
-# wrap the prompt in an expander to display additional parameters
+# Wrap the prompt in an accordion to display additional parameters
 for message in st.session_state.txt2img_messages:
     role = message["role"]
     with st.chat_message(role):
@@ -202,7 +211,7 @@ for message in st.session_state.txt2img_messages:
                 """)
                 st.write(message["content"])  # success will be image, error will be text
 
-# button row
+# Buttons for deleting last generation or clearing all generations
 if st.session_state.txt2img_messages:
     button_container = st.empty()
     with button_container.container():
@@ -235,7 +244,8 @@ if st.session_state.txt2img_messages:
 else:
     button_container = None
 
-# show the prompt and spinner while loading then update state and re-render
+# Set running state to True and show spinner while loading.
+# Update state and refresh on response; errors will be displayed as chat messages.
 if prompt := st.chat_input(
     "What do you want to see?",
     on_submit=lambda: setattr(st.session_state, "running", True),