Use httpx

lib/api.py

@@ -1,7 +1,7 @@
 import base64
 import io
 
-import requests
+import httpx
 import streamlit as st
 from openai import APIError, OpenAI
 from PIL import Image
@@ -47,13 +47,19 @@ def txt2img_generate(api_key, service, model, inputs, parameters, **kwargs):
     base_url = f"{Config.SERVICES[service]}/{model}"
 
     try:
-        response = requests.post(base_url, headers=headers, json=json)
+        response = httpx.post(base_url, headers=headers, json=json, timeout=Config.TXT2IMG_TIMEOUT)
         if response.status_code // 100 == 2:  # 2xx
             if service == "Hugging Face":
                 return Image.open(io.BytesIO(response.content))
             if service == "Fal":
-                bytes = base64.b64decode(response.json()["images"][0]["url"].split(",")[-1])
-                return Image.open(io.BytesIO(bytes))
+                # sync_mode means wait for image base64 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))
+                else:
+                    url = response.json()["images"][0]["url"]
+                    image = httpx.get(url, headers=headers, timeout=Config.TXT2IMG_TIMEOUT)
+                    return Image.open(io.BytesIO(image.content))
         else:
             return f"Error: {response.status_code} {response.text}"
     except Exception as e:

lib/config.py

@@ -9,6 +9,7 @@ Config = SimpleNamespace(
         "Perplexity": "https://api.perplexity.ai",
         "Fal": "https://fal.run",
     },
+    TXT2IMG_TIMEOUT=120,
     TXT2IMG_HIDDEN_PARAMETERS=[
         # sent to API but not shown in generation parameters accordion
         "enable_safety_checker",
@@ -27,7 +28,7 @@ Config = SimpleNamespace(
     TXT2IMG_DEFAULT_MODEL={
         # index of model in below lists
         "Hugging Face": 2,
-        "Fal": 1,
+        "Fal": 2,
     },
     TXT2IMG_MODELS={
         "Hugging Face": [
@@ -36,9 +37,8 @@ Config = SimpleNamespace(
             "stabilityai/stable-diffusion-xl-base-1.0",
         ],
         "Fal": [
-            # TODO: fix these models
-            # "fal-ai/aura-flow",
-            # "fal-ai/flux-pro",
+            "fal-ai/aura-flow",
+            "fal-ai/flux-pro",
             "fal-ai/fooocus",
             "fal-ai/kolors",
             "fal-ai/pixart-sigma",

lib/presets.py

@@ -24,11 +24,11 @@ ModelPresets = SimpleNamespace(
         "guidance_scale": 3.5,
         "guidance_scale_min": 1.0,
         "guidance_scale_max": 10.0,
-        "num_inference_steps": 50,
+        "num_inference_steps": 28,
         "num_inference_steps_min": 10,
         "num_inference_steps_max": 50,
-        "parameters": ["seed", "guidance_scale", "num_inference_steps", "expand_prompt"],
-        "kwargs": {"num_images": 1},
+        "parameters": ["seed", "num_inference_steps", "guidance_scale", "expand_prompt"],
+        "kwargs": {"num_images": 1, "sync_mode": False},
     },
     FLUX_DEV={
         "name": "FLUX.1 Dev",
@@ -50,7 +50,7 @@ ModelPresets = SimpleNamespace(
         "guidance_scale_min": 1.0,
         "guidance_scale_max": 10.0,
         "parameters": ["seed", "image_size", "num_inference_steps", "guidance_scale"],
-        "kwargs": {"num_images": 1, "sync_mode": True, "safety_tolerance": 6},
+        "kwargs": {"num_images": 1, "sync_mode": False, "safety_tolerance": 6},
     },
     FLUX_SCHNELL={
         "name": "FLUX.1 Schnell",

pages/2_🎨_Text_to_Image.py

@@ -19,6 +19,8 @@ PRESET_MODEL = {
     "black-forest-labs/flux.1-dev": ModelPresets.FLUX_DEV,
     "black-forest-labs/flux.1-schnell": ModelPresets.FLUX_SCHNELL,
     "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/fooocus": ModelPresets.FOOOCUS,
     "fal-ai/kolors": ModelPresets.KOLORS,
     "fal-ai/pixart-sigma": ModelPresets.PIXART_SIGMA,
@@ -179,14 +181,12 @@ for message in st.session_state.txt2img_messages:
                         div[data-testid="stMarkdownContainer"] p:not(:last-of-type) { margin-bottom: 0 }
                     </style>
                     """)
-                    filtered_parameters = {
-                        k: v
+                    filtered_parameters = [
+                        f"`{k}`: {v}"
                         for k, v in message["parameters"].items()
                         if k not in Config.TXT2IMG_HIDDEN_PARAMETERS
-                    }
-                    md = f"`model`: {message['model']}\n\n"
-                    md += "\n\n".join([f"`{k}`: {v}" for k, v in filtered_parameters.items()])
-                    st.markdown(md)
+                    ]
+                    st.markdown(f"`model`: {message['model']}\n\n" + "\n\n".join(filtered_parameters))
 
             if role == "assistant":
                 # image is full width when _not_ in full-screen mode

requirements.txt

@@ -1,4 +1,5 @@
+h2
+httpx
 openai==1.41.0
 pillow
-requests
 streamlit==1.37.1