Add initial Fal txt2img models

lib/api.py

@@ -1,3 +1,4 @@
+import base64
 import io
 
 import requests
@@ -24,25 +25,35 @@ def txt2txt_generate(api_key, service, model, parameters, **kwargs):
 
 
 def txt2img_generate(api_key, service, model, inputs, parameters, **kwargs):
-    headers = {
-        "Authorization": f"Bearer {api_key}",
-        "X-Wait-For-Model": "true",
-        "X-Use-Cache": "false",
-    }
+    headers = {}
+    if service == "Huggingface":
+        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}"
+
+    json = {}
+    if service == "Huggingface":
+        json = {
+            "inputs": inputs,
+            "parameters": {**parameters, **kwargs},
+        }
+    if service == "Fal":
+        json = {**parameters, **kwargs}
+        json["prompt"] = inputs
+
     base_url = f"{Config.SERVICES[service]}/{model}"
 
     try:
-        response = requests.post(
-            base_url,
-            headers=headers,
-            json={
-                "inputs": inputs,
-                "parameters": {**parameters, **kwargs},
-            },
-        )
-        if response.status_code == 200:
-            return Image.open(io.BytesIO(response.content))
+        response = requests.post(base_url, headers=headers, json=json)
+        if response.status_code // 100 == 2:  # 2xx
+            if service == "Huggingface":
+                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))
         else:
-            return f"Error: {response.status_code} - {response.text}"
+            return f"Error: {response.status_code} {response.text}"
     except Exception as e:
         return str(e)

lib/config.py

@@ -7,22 +7,78 @@ Config = SimpleNamespace(
     SERVICES={
         "Huggingface": "https://api-inference.huggingface.co/models",
         "Perplexity": "https://api.perplexity.ai",
+        "Fal": "https://fal.run",
     },
-    TXT2IMG_NEGATIVE_PROMPT="ugly, unattractive, malformed, mutated, disgusting, blurry, grainy, noisy, oversaturated, undersaturated, overexposed, underexposed, worst quality, low details, lowres, watermark, signature, autograph, trademark, sloppy, cluttered",
-    TXT2IMG_DEFAULT_MODEL=2,
-    TXT2IMG_MODELS=[
-        "black-forest-labs/flux.1-dev",
-        "black-forest-labs/flux.1-schnell",
-        "stabilityai/stable-diffusion-xl-base-1.0",
+    TXT2IMG_HIDDEN_PARAMETERS=[
+        "enable_safety_checker",
+        "max_sequence_length",
+        "num_images",
+        "output_format",
+        "performance",
+        "safety_tolerance",
+        "scheduler",
+        "sharpness",
+        "style",
+        "styles",
+        "sync_mode",
     ],
-    TXT2IMG_DEFAULT_AR="1:1",
-    TXT2IMG_AR={
-        "7:4": (1344, 768),
-        "9:7": (1152, 896),
-        "1:1": (1024, 1024),
-        "7:9": (896, 1152),
-        "4:7": (768, 1344),
+    TXT2IMG_NEGATIVE_PROMPT="ugly, unattractive, malformed, mutated, disgusting, blurry, grainy, noisy, oversaturated, undersaturated, overexposed, underexposed, worst quality, low details, lowres, watermark, signature, autograph, trademark, sloppy, cluttered",
+    TXT2IMG_DEFAULT_MODEL={
+        "Huggingface": 2,
+        "Fal": 1,
+    },
+    TXT2IMG_MODELS={
+        "Huggingface": [
+            "black-forest-labs/flux.1-dev",
+            "black-forest-labs/flux.1-schnell",
+            "stabilityai/stable-diffusion-xl-base-1.0",
+        ],
+        "Fal": [
+            # "fal-ai/aura-flow",
+            # "fal-ai/flux-pro",
+            "fal-ai/fooocus",
+            "fal-ai/kolors",
+            "fal-ai/pixart-sigma",
+            "fal-ai/stable-diffusion-v3-medium",
+        ],
     },
+    TXT2IMG_DEFAULT_IMAGE_SIZE="square_hd",  # fal image sizes
+    TXT2IMG_IMAGE_SIZES=[
+        "landscape_16_9",
+        "landscape_4_3",
+        "square_hd",
+        "portrait_4_3",
+        "portrait_16_9",
+    ],
+    TXT2IMG_DEFAULT_ASPECT_RATIO="1024x1024",  # fooocus aspect ratios
+    TXT2IMG_ASPECT_RATIOS=[
+        "704x1408",
+        "704x1344",
+        "768x1344",
+        "768x1280",
+        "832x1216",
+        "832x1152",
+        "896x1152",
+        "896x1088",
+        "960x1088",
+        "960x1024",
+        "1024x1024",
+        "1024x960",
+        "1088x960",
+        "1088x896",
+        "1152x896",
+        "1152x832",
+        "1216x832",
+        "1280x768",
+        "1344x768",
+        "1344x704",
+        "1408x704",
+        "1472x704",
+        "1536x640",
+        "1600x640",
+        "1664x576",
+        "1728x576",
+    ],
     TXT2TXT_DEFAULT_SYSTEM="You are a helpful assistant. Be precise and concise.",
     TXT2TXT_DEFAULT_MODEL={
         "Huggingface": 4,

lib/presets.py

@@ -19,18 +19,40 @@ ServicePresets = SimpleNamespace(
 
 # txt2img models
 ModelPresets = SimpleNamespace(
-    FLUX_1_DEV={
+    AURA_FLOW={
+        "name": "AuraFlow",
+        "guidance_scale": 3.5,
+        "guidance_scale_min": 1.0,
+        "guidance_scale_max": 10.0,
+        "num_inference_steps": 50,
+        "num_inference_steps_min": 10,
+        "num_inference_steps_max": 50,
+        "parameters": ["seed", "guidance_scale", "num_inference_steps", "expand_prompt"],
+        "kwargs": {"num_images": 1},
+    },
+    FLUX_DEV={
         "name": "FLUX.1 Dev",
-        "num_inference_steps": 30,
+        "num_inference_steps": 28,
         "num_inference_steps_min": 10,
-        "num_inference_steps_max": 40,
+        "num_inference_steps_max": 50,
         "guidance_scale": 3.5,
         "guidance_scale_min": 1.0,
-        "guidance_scale_max": 7.0,
+        "guidance_scale_max": 10.0,
         "parameters": ["width", "height", "guidance_scale", "num_inference_steps"],
         "kwargs": {"max_sequence_length": 512},
     },
-    FLUX_1_SCHNELL={
+    FLUX_PRO={
+        "name": "FLUX.1 Pro",
+        "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": ["seed", "image_size", "num_inference_steps", "guidance_scale"],
+        "kwargs": {"num_images": 1, "sync_mode": True, "safety_tolerance": 6},
+    },
+    FLUX_SCHNELL={
         "name": "FLUX.1 Schnell",
         "num_inference_steps": 4,
         "num_inference_steps_min": 1,
@@ -38,14 +60,87 @@ ModelPresets = SimpleNamespace(
         "parameters": ["width", "height", "num_inference_steps"],
         "kwargs": {"guidance_scale": 0.0, "max_sequence_length": 256},
     },
+    FOOOCUS={
+        "name": "Fooocus",
+        "guidance_scale": 4.0,
+        "guidance_scale_min": 1.0,
+        "guidance_scale_max": 10.0,
+        "parameters": ["seed", "negative_prompt", "aspect_ratio", "guidance_scale"],
+        "kwargs": {
+            "num_images": 1,
+            "sync_mode": True,
+            "enable_safety_checker": False,
+            "output_format": "png",
+            "sharpness": 2,
+            "styles": ["Fooocus Enhance", "Fooocus V2", "Fooocus Sharp"],
+            "performance": "Quality",
+        },
+    },
+    KOLORS={
+        "name": "Kolors",
+        "guidance_scale": 5.0,
+        "guidance_scale_min": 1.0,
+        "guidance_scale_max": 10.0,
+        "num_inference_steps": 50,
+        "num_inference_steps_min": 10,
+        "num_inference_steps_max": 50,
+        "parameters": [
+            "seed",
+            "negative_prompt",
+            "image_size",
+            "guidance_scale",
+            "num_inference_steps",
+        ],
+        "kwargs": {
+            "num_images": 1,
+            "sync_mode": True,
+            "enable_safety_checker": False,
+            "scheduler": "EulerDiscreteScheduler",
+        },
+    },
+    PIXART_SIGMA={
+        "name": "PixArt-Σ",
+        "guidance_scale": 4.5,
+        "guidance_scale_min": 1.0,
+        "guidance_scale_max": 10.0,
+        "num_inference_steps": 35,
+        "num_inference_steps_min": 10,
+        "num_inference_steps_max": 50,
+        "parameters": ["seed", "negative_prompt", "image_size", "guidance_scale", "num_inference_steps"],
+        "kwargs": {
+            "num_images": 1,
+            "sync_mode": True,
+            "enable_safety_checker": False,
+            "style": "(No style)",
+            "scheduler": "SA-SOLVER",
+        },
+    },
+    STABLE_DIFFUSION_3={
+        "name": "SD3",
+        "guidance_scale": 5.0,
+        "guidance_scale_min": 1.0,
+        "guidance_scale_max": 10.0,
+        "num_inference_steps": 28,
+        "num_inference_steps_min": 10,
+        "num_inference_steps_max": 50,
+        "parameters": [
+            "seed",
+            "negative_prompt",
+            "image_size",
+            "guidance_scale",
+            "num_inference_steps",
+            "prompt_expansion",
+        ],
+        "kwargs": {"num_images": 1, "sync_mode": True, "enable_safety_checker": False},
+    },
     STABLE_DIFFUSION_XL={
         "name": "SDXL",
         "guidance_scale": 7.0,
         "guidance_scale_min": 1.0,
-        "guidance_scale_max": 15.0,
+        "guidance_scale_max": 10.0,
         "num_inference_steps": 40,
         "num_inference_steps_min": 10,
         "num_inference_steps_max": 50,
-        "parameters": ["width", "height", "guidance_scale", "num_inference_steps", "seed", "negative_prompt"],
+        "parameters": ["seed", "negative_prompt", "width", "height", "guidance_scale", "num_inference_steps"],
     },
 )

pages/2_🎨_Text_to_Image.py

@@ -5,16 +5,16 @@ import streamlit as st
 
 from lib import Config, ModelPresets, txt2img_generate
 
-HF_TOKEN = None
-FAL_KEY = None
-# HF_TOKEN = os.environ.get("HF_TOKEN") or None
-# FAL_KEY = os.environ.get("FAL_KEY") or None
-API_URL = "https://api-inference.huggingface.co/models"
-HEADERS = {"Authorization": f"Bearer {HF_TOKEN}", "X-Wait-For-Model": "true", "X-Use-Cache": "false"}
+HF_TOKEN = os.environ.get("HF_TOKEN") or None
+FAL_KEY = os.environ.get("FAL_KEY") or None
 PRESET_MODEL = {
-    "black-forest-labs/flux.1-dev": ModelPresets.FLUX_1_DEV,
-    "black-forest-labs/flux.1-schnell": ModelPresets.FLUX_1_SCHNELL,
+    "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/fooocus": ModelPresets.FOOOCUS,
+    "fal-ai/kolors": ModelPresets.KOLORS,
+    "fal-ai/pixart-sigma": ModelPresets.PIXART_SIGMA,
+    "fal-ai/stable-diffusion-v3-medium": ModelPresets.STABLE_DIFFUSION_3,
 }
 
 # config
@@ -34,8 +34,8 @@ if "api_key_huggingface" not in st.session_state:
 if "txt2img_messages" not in st.session_state:
     st.session_state.txt2img_messages = []
 
-if "txt2img_running" not in st.session_state:
-    st.session_state.txt2img_running = False
+if "running" not in st.session_state:
+    st.session_state.running = False
 
 if "txt2img_seed" not in st.session_state:
     st.session_state.txt2img_seed = 0
@@ -45,9 +45,9 @@ st.logo("logo.svg")
 st.sidebar.header("Settings")
 service = st.sidebar.selectbox(
     "Service",
-    options=["Huggingface"],
-    index=0,
-    disabled=st.session_state.txt2img_running,
+    options=["Fal", "Huggingface"],
+    index=1,
+    disabled=st.session_state.running,
 )
 
 if service == "Huggingface" and HF_TOKEN is None:
@@ -56,7 +56,7 @@ if service == "Huggingface" and HF_TOKEN is None:
         type="password",
         help="Cleared on page refresh",
         value=st.session_state.api_key_huggingface,
-        disabled=st.session_state.txt2img_running,
+        disabled=st.session_state.running,
     )
 else:
     st.session_state.api_key_huggingface = st.session_state.api_key_huggingface
@@ -67,7 +67,7 @@ if service == "Fal" and FAL_KEY is None:
         type="password",
         help="Cleared on page refresh",
         value=st.session_state.api_key_fal,
-        disabled=st.session_state.txt2img_running,
+        disabled=st.session_state.running,
     )
 else:
     st.session_state.api_key_fal = st.session_state.api_key_fal
@@ -80,17 +80,11 @@ if service == "Fal" and FAL_KEY is not None:
 
 model = st.sidebar.selectbox(
     "Model",
-    options=Config.TXT2IMG_MODELS,
-    index=Config.TXT2IMG_DEFAULT_MODEL,
-    disabled=st.session_state.txt2img_running,
+    options=Config.TXT2IMG_MODELS[service],
+    index=Config.TXT2IMG_DEFAULT_MODEL[service],
+    disabled=st.session_state.running,
     format_func=lambda x: x.split("/")[1],
 )
-aspect_ratio = st.sidebar.select_slider(
-    "Aspect Ratio",
-    options=list(Config.TXT2IMG_AR.keys()),
-    value=Config.TXT2IMG_DEFAULT_AR,
-    disabled=st.session_state.txt2img_running,
-)
 
 # heading
 st.html("""
@@ -102,10 +96,52 @@ st.html("""
 parameters = {}
 preset = PRESET_MODEL[model]
 for param in preset["parameters"]:
+    if param == "seed":
+        parameters[param] = st.sidebar.number_input(
+            "Seed",
+            min_value=-1,
+            max_value=(1 << 53) - 1,
+            value=-1,
+            disabled=st.session_state.running,
+        )
+    if param == "negative_prompt":
+        parameters[param] = st.sidebar.text_area(
+            "Negative Prompt",
+            value=Config.TXT2IMG_NEGATIVE_PROMPT,
+            disabled=st.session_state.running,
+        )
     if param == "width":
-        parameters[param] = Config.TXT2IMG_AR[aspect_ratio][0]
+        parameters[param] = st.sidebar.slider(
+            "Width",
+            step=64,
+            value=1024,
+            min_value=512,
+            max_value=2048,
+            disabled=st.session_state.running,
+        )
     if param == "height":
-        parameters[param] = Config.TXT2IMG_AR[aspect_ratio][1]
+        parameters[param] = st.sidebar.slider(
+            "Height",
+            step=64,
+            value=1024,
+            min_value=512,
+            max_value=2048,
+            disabled=st.session_state.running,
+        )
+    if param == "image_size":
+        parameters[param] = st.sidebar.select_slider(
+            "Image Size",
+            options=Config.TXT2IMG_IMAGE_SIZES,
+            value=Config.TXT2IMG_DEFAULT_IMAGE_SIZE,
+            disabled=st.session_state.running,
+        )
+    if param == "aspect_ratio":
+        parameters[param] = st.sidebar.select_slider(
+            "Aspect Ratio",
+            options=Config.TXT2IMG_ASPECT_RATIOS,
+            value=Config.TXT2IMG_DEFAULT_ASPECT_RATIO,
+            disabled=st.session_state.running,
+        )
     if param == "guidance_scale":
         parameters[param] = st.sidebar.slider(
             "Guidance Scale",
@@ -113,7 +149,7 @@ for param in preset["parameters"]:
             preset["guidance_scale_max"],
             preset["guidance_scale"],
             0.1,
-            disabled=st.session_state.txt2img_running,
+            disabled=st.session_state.running,
         )
     if param == "num_inference_steps":
         parameters[param] = st.sidebar.slider(
@@ -122,21 +158,19 @@ for param in preset["parameters"]:
             preset["num_inference_steps_max"],
             preset["num_inference_steps"],
             1,
-            disabled=st.session_state.txt2img_running,
+            disabled=st.session_state.running,
         )
-    if param == "seed":
-        parameters[param] = st.sidebar.number_input(
-            "Seed",
-            min_value=-1,
-            max_value=(1 << 53) - 1,
-            value=-1,
-            disabled=st.session_state.txt2img_running,
+    if param == "expand_prompt":
+        parameters[param] = st.sidebar.checkbox(
+            "Expand Prompt",
+            value=False,
+            disabled=st.session_state.running,
         )
-    if param == "negative_prompt":
-        parameters[param] = st.sidebar.text_area(
-            label="Negative Prompt",
-            value=Config.TXT2IMG_NEGATIVE_PROMPT,
-            disabled=st.session_state.txt2img_running,
+    if param == "prompt_expansion":
+        parameters[param] = st.sidebar.checkbox(
+            "Prompt Expansion",
+            value=False,
+            disabled=st.session_state.running,
         )
 
 # wrap the prompt in an expander to display additional parameters
@@ -154,8 +188,13 @@ 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
+                        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 message["parameters"].items()])
+                    md += "\n\n".join([f"`{k}`: {v}" for k, v in filtered_parameters.items()])
                     st.markdown(md)
 
             if role == "assistant":
@@ -192,7 +231,7 @@ if st.session_state.txt2img_messages:
         col1, col2 = st.columns(2)
         with col1:
             if (
-                st.button("❌", help="Delete last generation", disabled=st.session_state.txt2img_running)
+                st.button("❌", help="Delete last generation", disabled=st.session_state.running)
                 and len(st.session_state.txt2img_messages) >= 2
             ):
                 st.session_state.txt2img_messages.pop()
@@ -200,7 +239,7 @@ if st.session_state.txt2img_messages:
                 st.rerun()
 
         with col2:
-            if st.button("🗑️", help="Clear all generations", disabled=st.session_state.txt2img_running):
+            if st.button("🗑️", help="Clear all generations", disabled=st.session_state.running):
                 st.session_state.txt2img_messages = []
                 st.session_state.txt2img_seed = 0
                 st.rerun()
@@ -210,7 +249,7 @@ else:
 # show the prompt and spinner while loading then update state and re-render
 if prompt := st.chat_input(
     "What do you want to see?",
-    on_submit=lambda: setattr(st.session_state, "txt2img_running", True),
+    on_submit=lambda: setattr(st.session_state, "running", True),
 ):
     if "seed" in parameters and parameters["seed"] >= 0:
         st.session_state.txt2img_seed = parameters["seed"]
@@ -231,7 +270,7 @@ if prompt := st.chat_input(
                 parameters.update(preset["kwargs"])
             api_key = getattr(st.session_state, f"api_key_{service.lower()}", None)
             image = txt2img_generate(api_key, service, model, prompt, parameters)
-        st.session_state.txt2img_running = False
+        st.session_state.running = False
 
     model_name = PRESET_MODEL[model]["name"]
     st.session_state.txt2img_messages.append(