Refactor preset

0_🏠_Home.py

@@ -4,12 +4,12 @@ from lib import config
 
 st.set_page_config(
     page_title=config.title,
-    page_icon=config.icon,
+    page_icon=config.logo,
     layout=config.layout,
 )
 
 # sidebar
-st.logo("logo.png")
+st.logo(config.logo)
 
 # title
 st.html("""
@@ -39,6 +39,7 @@ st.html("""
     <h1>API Inference</h1>
     <span class="pro-badge">PRO</span>
 </div>
+<p>Explore popular AI endpoints in one place.</p>
 """)
 
 st.markdown("## Tasks")
@@ -58,7 +59,7 @@ st.markdown("""
 st.markdown("""
 ## Usage
 
-Choose a task from the sidebar. Enter your API key for the service you want to use. Refresh your browser to remove it.
+Choose a task. Select a service. Enter your API key (refresh browser to clear).
 
 I recommend [duplicating this space](https://huggingface.co/spaces/adamelliotfields/api-inference?duplicate=true) **privately** and persisting your keys as secrets. See [`README.md`](https://huggingface.co/spaces/adamelliotfields/api-inference/blob/main/README.md).
 """)

lib/__init__.py

@@ -1,8 +1,10 @@
 from .api import txt2img_generate, txt2txt_generate
 from .config import config
-from .preset import preset
+from .preset import Txt2ImgPreset, Txt2TxtPreset, preset
 
 __all__ = [
+    "Txt2ImgPreset",
+    "Txt2TxtPreset",
     "config",
     "preset",
     "txt2img_generate",

lib/api.py

@@ -11,8 +11,8 @@ from .config import config
 
 
 def txt2txt_generate(api_key, service, model, parameters, **kwargs):
-    base_url = config.services[service]
-    if service == "Hugging Face":
+    base_url = config.service[service].url
+    if service == "hf":
         base_url = f"{base_url}/{model}/v1"
     client = OpenAI(api_key=api_key, base_url=base_url)
 
@@ -29,42 +29,32 @@ def txt2txt_generate(api_key, service, model, parameters, **kwargs):
 
 def txt2img_generate(api_key, service, model, inputs, parameters, **kwargs):
     headers = {}
-    if service == "Black Forest Labs":
+    json = {**parameters, **kwargs}
+
+    if service == "bfl":
         headers["x-key"] = api_key
+        json["prompt"] = inputs
 
-    if service == "Fal":
+    if service == "fal":
         headers["Authorization"] = f"Key {api_key}"
+        json["prompt"] = inputs
 
-    if service == "Hugging Face":
+    if service == "hf":
         headers["Authorization"] = f"Bearer {api_key}"
         headers["X-Wait-For-Model"] = "true"
         headers["X-Use-Cache"] = "false"
-
-    if service == "Together":
-        headers["Authorization"] = f"Bearer {api_key}"
-
-    json = {}
-    if service == "Black Forest Labs":
-        json = {**parameters, **kwargs}
-        json["prompt"] = inputs
-
-    if service == "Fal":
-        json = {**parameters, **kwargs}
-        json["prompt"] = inputs
-
-    if service == "Hugging Face":
         json = {
             "inputs": inputs,
             "parameters": {**parameters, **kwargs},
         }
 
-    if service == "Together":
-        json = {**parameters, **kwargs}
+    if service == "together":
+        headers["Authorization"] = f"Bearer {api_key}"
         json["prompt"] = inputs
 
-    base_url = config.services[service]
+    base_url = config.service[service].url
 
-    if service not in ["Together"]:
+    if service not in ["together"]:
         base_url = f"{base_url}/{model}"
 
     try:
@@ -72,9 +62,9 @@ def txt2img_generate(api_key, service, model, inputs, parameters, **kwargs):
         if response.status_code // 100 == 2:  # 2xx
             # BFL is async so we need to poll for result
             # https://api.bfl.ml/docs
-            if service == "Black Forest Labs":
+            if service == "bfl":
                 id = response.json()["id"]
-                url = f"{config.services[service]}/get_result?id={id}"
+                url = f"{config.service[service].url}/get_result?id={id}"
 
                 retries = 0
                 while retries < config.txt2img.timeout:
@@ -95,7 +85,7 @@ def txt2img_generate(api_key, service, model, inputs, parameters, **kwargs):
 
                 return "Error: API timeout"
 
-            if service == "Fal":
+            if service == "fal":
                 # 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])
@@ -105,10 +95,10 @@ def txt2img_generate(api_key, service, model, inputs, parameters, **kwargs):
                     image = httpx.get(url, headers=headers, timeout=config.txt2img.timeout)
                     return Image.open(io.BytesIO(image.content))
 
-            if service == "Hugging Face":
+            if service == "hf":
                 return Image.open(io.BytesIO(response.content))
 
-            if service == "Together":
+            if service == "together":
                 url = response.json()["data"][0]["url"]
                 image = httpx.get(url, headers=headers, timeout=config.txt2img.timeout)
                 return Image.open(io.BytesIO(image.content))

lib/config.py

@@ -1,31 +1,22 @@
+import os
 from dataclasses import dataclass
-from typing import Dict, List
+from typing import Dict, List, Optional
 
-from .preset import preset
 
-
-def txt2img_models_from_presets(presets):
-    models = {}
-    for p in presets:
-        service = p.service
-        model_id = p.model_id
-        if service not in models:
-            models[service] = []
-        models[service].append(model_id)
-    return models
+@dataclass
+class ServiceConfig:
+    name: str
+    url: str
+    api_key: Optional[str] = None
 
 
 @dataclass
 class Txt2TxtConfig:
     default_system: str
-    default_model: Dict[str, int]
-    models: Dict[str, List[str]]
 
 
 @dataclass
 class Txt2ImgConfig:
-    default_model: Dict[str, int]
-    models: Dict[str, List[str]]
     hidden_parameters: List[str]
     negative_prompt: str
     default_image_size: str
@@ -38,37 +29,50 @@ class Txt2ImgConfig:
 @dataclass
 class Config:
     title: str
-    icon: str
     layout: str
-    services: Dict[str, str]
+    logo: str
+    service: Dict[str, ServiceConfig]
     txt2img: Txt2ImgConfig
     txt2txt: Txt2TxtConfig
 
 
-# TODO: API keys should be with services (make a dataclass)
 config = Config(
     title="API Inference",
-    icon="⚡",
     layout="wide",
-    services={
-        "Black Forest Labs": "https://api.bfl.ml/v1",
-        "Fal": "https://fal.run",
-        "Hugging Face": "https://api-inference.huggingface.co/models",
-        "Perplexity": "https://api.perplexity.ai",
-        "Together": "https://api.together.xyz/v1/images/generations",
+    logo="logo.png",
+    service={
+        "bfl": ServiceConfig(
+            "Black Forest Labs",
+            "https://api.bfl.ml/v1",
+            os.environ.get("BFL_API_KEY"),
+        ),
+        "fal": ServiceConfig(
+            "Fal",
+            "https://fal.run",
+            os.environ.get("FAL_KEY"),
+        ),
+        "hf": ServiceConfig(
+            "Hugging Face",
+            "https://api-inference.huggingface.co/models",
+            os.environ.get("HF_TOKEN"),
+        ),
+        "pplx": ServiceConfig(
+            "Perplexity",
+            "https://api.perplexity.ai",
+            os.environ.get("PPLX_API_KEY"),
+        ),
+        "together": ServiceConfig(
+            "Together",
+            "https://api.together.xyz/v1/images/generations",
+            os.environ.get("TOGETHER_API_KEY"),
+        ),
     },
     txt2img=Txt2ImgConfig(
-        default_model={
-            "Black Forest Labs": 2,
-            "Fal": 0,
-            "Hugging Face": 2,
-            "Together": 0,
-        },
-        models=txt2img_models_from_presets(preset.txt2img.presets),
         hidden_parameters=[
             # Sent to API but not shown in generation parameters accordion
             "enable_safety_checker",
             "max_sequence_length",
+            "n",
             "num_images",
             "output_format",
             "performance",
@@ -115,25 +119,5 @@ config = Config(
     ),
     txt2txt=Txt2TxtConfig(
         default_system="You are a helpful assistant. Be precise and concise.",
-        default_model={
-            "Hugging Face": 4,
-            "Perplexity": 3,
-        },
-        models={
-            "Hugging Face": [
-                "codellama/codellama-34b-instruct-hf",
-                "meta-llama/llama-2-13b-chat-hf",
-                "meta-llama/meta-llama-3.1-405b-instruct-fp8",
-                "mistralai/mistral-7b-instruct-v0.2",
-                "nousresearch/nous-hermes-2-mixtral-8x7b-dpo",
-            ],
-            "Perplexity": [
-                "llama-3.1-sonar-small-128k-chat",
-                "llama-3.1-sonar-large-128k-chat",
-                "llama-3.1-sonar-small-128k-online",
-                "llama-3.1-sonar-large-128k-online",
-                "llama-3.1-sonar-huge-128k-online",
-            ],
-        },
     ),
 )

lib/preset.py

@@ -4,6 +4,7 @@ from typing import Dict, List, Optional, Union
 
 @dataclass
 class Txt2TxtPreset:
+    name: str
     frequency_penalty: float
     frequency_penalty_min: float
     frequency_penalty_max: float
@@ -12,10 +13,7 @@ class Txt2TxtPreset:
 
 @dataclass
 class Txt2ImgPreset:
-    # FLUX1.1 has no scale or steps
     name: str
-    service: str
-    model_id: str
     guidance_scale: Optional[float] = None
     guidance_scale_min: Optional[float] = None
     guidance_scale_max: Optional[float] = None
@@ -27,66 +25,55 @@ class Txt2ImgPreset:
 
 
 @dataclass
-class Txt2TxtPresets:
-    hugging_face: Txt2TxtPreset
-    perplexity: Txt2TxtPreset
-
-
-@dataclass
-class Txt2ImgPresets:
-    presets: List[Txt2ImgPreset] = field(default_factory=list)
+class Preset:
+    txt2txt: Dict[str, Txt2TxtPreset]
+    txt2img: Dict[str, Txt2ImgPreset]
 
-    def __iter__(self):
-        return iter(self.presets)
 
+hf_txt2txt_kwargs = {
+    "frequency_penalty": 0.0,
+    "frequency_penalty_min": -2.0,
+    "frequency_penalty_max": 2.0,
+    "parameters": ["max_tokens", "temperature", "frequency_penalty", "seed"],
+}
 
-@dataclass
-class Preset:
-    txt2txt: Txt2TxtPresets
-    txt2img: Txt2ImgPresets
+pplx_txt2txt_kwargs = {
+    "frequency_penalty": 1.0,
+    "frequency_penalty_min": 1.0,
+    "frequency_penalty_max": 2.0,
+    "parameters": ["max_tokens", "temperature", "frequency_penalty"],
+}
 
 
 preset = Preset(
-    txt2txt=Txt2TxtPresets(
-        hugging_face=Txt2TxtPreset(
-            frequency_penalty=0.0,
-            frequency_penalty_min=-2.0,
-            frequency_penalty_max=2.0,
-            parameters=["max_tokens", "temperature", "frequency_penalty", "seed"],
-        ),
-        perplexity=Txt2TxtPreset(
-            frequency_penalty=1.0,
-            frequency_penalty_min=1.0,
-            frequency_penalty_max=2.0,
-            parameters=["max_tokens", "temperature", "frequency_penalty"],
-        ),
-    ),
-    txt2img=Txt2ImgPresets(
-        presets=[
-            Txt2ImgPreset(
-                "AuraFlow",
-                "Fal",
-                "fal-ai/aura-flow",
-                guidance_scale=3.5,
-                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", "num_inference_steps", "guidance_scale", "expand_prompt"],
-                kwargs={"num_images": 1, "sync_mode": False},
+    txt2txt={
+        "hf": {
+            # TODO: update models
+            "codellama/codellama-34b-instruct-hf": Txt2TxtPreset("Code Llama 34B", **hf_txt2txt_kwargs),
+            "meta-llama/llama-2-13b-chat-hf": Txt2TxtPreset("Llama 2 13B", **hf_txt2txt_kwargs),
+            "mistralai/mistral-7b-instruct-v0.2": Txt2TxtPreset("Mistral v0.2 7B", **hf_txt2txt_kwargs),
+            "nousresearch/nous-hermes-2-mixtral-8x7b-dpo": Txt2TxtPreset(
+                "Nous Hermes 2 Mixtral 8x7B",
+                **hf_txt2txt_kwargs,
             ),
-            Txt2ImgPreset(
+        },
+        "pplx": {
+            "llama-3.1-sonar-small-128k-chat": Txt2TxtPreset("Sonar Small (Offline)", **pplx_txt2txt_kwargs),
+            "llama-3.1-sonar-large-128k-chat": Txt2TxtPreset("Sonar Large (Offline)", **pplx_txt2txt_kwargs),
+            "llama-3.1-sonar-small-128k-online": Txt2TxtPreset("Sonar Small (Online)", **pplx_txt2txt_kwargs),
+            "llama-3.1-sonar-large-128k-online": Txt2TxtPreset("Sonar Large (Online)", **pplx_txt2txt_kwargs),
+            "llama-3.1-sonar-huge-128k-online": Txt2TxtPreset("Sonar Huge (Online)", **pplx_txt2txt_kwargs),
+        },
+    },
+    txt2img={
+        "bfl": {
+            "flux-pro-1.1": Txt2ImgPreset(
                 "FLUX1.1 Pro",
-                "Black Forest Labs",
-                "flux-pro-1.1",
                 parameters=["seed", "width", "height", "prompt_upsampling"],
                 kwargs={"safety_tolerance": 6},
             ),
-            Txt2ImgPreset(
+            "flux-pro": Txt2ImgPreset(
                 "FLUX.1 Pro",
-                "Black Forest Labs",
-                "flux-pro",
                 guidance_scale=2.5,
                 guidance_scale_min=1.5,
                 guidance_scale_max=5.0,
@@ -96,10 +83,8 @@ preset = Preset(
                 parameters=["seed", "width", "height", "steps", "guidance", "prompt_upsampling"],
                 kwargs={"safety_tolerance": 6, "interval": 1},
             ),
-            Txt2ImgPreset(
+            "flux-dev": Txt2ImgPreset(
                 "FLUX.1 Dev",
-                "Black Forest Labs",
-                "flux-dev",
                 num_inference_steps=28,
                 num_inference_steps_min=10,
                 num_inference_steps_max=50,
@@ -109,10 +94,21 @@ preset = Preset(
                 parameters=["seed", "width", "height", "steps", "guidance", "prompt_upsampling"],
                 kwargs={"safety_tolerance": 6},
             ),
-            Txt2ImgPreset(
+        },
+        "fal": {
+            "fal-ai/aura-flow": Txt2ImgPreset(
+                "AuraFlow",
+                guidance_scale=3.5,
+                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", "num_inference_steps", "guidance_scale", "expand_prompt"],
+                kwargs={"num_images": 1, "sync_mode": False},
+            ),
+            "fal-ai/flux-pro/v1.1": Txt2ImgPreset(
                 "FLUX1.1 Pro",
-                "Fal",
-                "fal-ai/flux-pro/v1.1",
                 parameters=["seed", "image_size"],
                 kwargs={
                     "num_images": 1,
@@ -121,10 +117,8 @@ preset = Preset(
                     "enable_safety_checker": False,
                 },
             ),
-            Txt2ImgPreset(
+            "fal-ai/flux-pro": Txt2ImgPreset(
                 "FLUX.1 Pro",
-                "Fal",
-                "fal-ai/flux-pro",
                 guidance_scale=2.5,
                 guidance_scale_min=1.5,
                 guidance_scale_max=5.0,
@@ -134,10 +128,8 @@ preset = Preset(
                 parameters=["seed", "image_size", "num_inference_steps", "guidance_scale"],
                 kwargs={"num_images": 1, "sync_mode": False, "safety_tolerance": 6},
             ),
-            Txt2ImgPreset(
+            "fal-ai/flux/dev": Txt2ImgPreset(
                 "FLUX.1 Dev",
-                "Fal",
-                "fal-ai/flux/dev",
                 num_inference_steps=28,
                 num_inference_steps_min=10,
                 num_inference_steps_max=50,
@@ -147,53 +139,16 @@ preset = Preset(
                 parameters=["seed", "image_size", "num_inference_steps", "guidance_scale"],
                 kwargs={"num_images": 1, "sync_mode": False, "safety_tolerance": 6},
             ),
-            Txt2ImgPreset(
+            "fal-ai/flux/schnell": Txt2ImgPreset(
                 "FLUX.1 Schnell",
-                "Fal",
-                "fal-ai/flux/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},
             ),
-            Txt2ImgPreset(
-                "FLUX.1 Dev",
-                "Hugging Face",
-                "black-forest-labs/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},
-            ),
-            Txt2ImgPreset(
-                "FLUX.1 Schnell",
-                "Hugging Face",
-                "black-forest-labs/flux.1-schnell",
-                num_inference_steps=4,
-                num_inference_steps_min=1,
-                num_inference_steps_max=12,
-                parameters=["width", "height", "num_inference_steps"],
-                kwargs={"guidance_scale": 0.0, "max_sequence_length": 256},
-            ),
-            Txt2ImgPreset(
-                "FLUX.1 Schnell Free",
-                "Together",
-                "black-forest-labs/FLUX.1-schnell-Free",
-                num_inference_steps=4,
-                num_inference_steps_min=1,
-                num_inference_steps_max=12,
-                parameters=["model", "seed", "width", "height", "steps"],
-                kwargs={"n": 1},
-            ),
-            Txt2ImgPreset(
+            "fal-ai/fooocus": Txt2ImgPreset(
                 "Fooocus",
-                "Fal",
-                "fal-ai/fooocus",
                 guidance_scale=4.0,
                 guidance_scale_min=1.0,
                 guidance_scale_max=10.0,
@@ -208,10 +163,8 @@ preset = Preset(
                     "performance": "Quality",
                 },
             ),
-            Txt2ImgPreset(
+            "fal-ai/kolors": Txt2ImgPreset(
                 "Kolors",
-                "Fal",
-                "fal-ai/kolors",
                 guidance_scale=5.0,
                 guidance_scale_min=1.0,
                 guidance_scale_max=10.0,
@@ -226,10 +179,8 @@ preset = Preset(
                     "scheduler": "EulerDiscreteScheduler",
                 },
             ),
-            Txt2ImgPreset(
+            "fal-ai/stable-diffusion-v3-medium": Txt2ImgPreset(
                 "SD3",
-                "Fal",
-                "fal-ai/stable-diffusion-v3-medium",
                 guidance_scale=5.0,
                 guidance_scale_min=1.0,
                 guidance_scale_max=10.0,
@@ -246,10 +197,29 @@ preset = Preset(
                 ],
                 kwargs={"num_images": 1, "sync_mode": True, "enable_safety_checker": False},
             ),
-            Txt2ImgPreset(
+        },
+        "hf": {
+            "black-forest-labs/flux.1-dev": Txt2ImgPreset(
+                "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},
+            ),
+            "black-forest-labs/flux.1-schnell": Txt2ImgPreset(
+                "FLUX.1 Schnell",
+                num_inference_steps=4,
+                num_inference_steps_min=1,
+                num_inference_steps_max=12,
+                parameters=["width", "height", "num_inference_steps"],
+                kwargs={"guidance_scale": 0.0, "max_sequence_length": 256},
+            ),
+            "stabilityai/stable-diffusion-xl-base-1.0": Txt2ImgPreset(
                 "SDXL",
-                "Hugging Face",
-                "stabilityai/stable-diffusion-xl-base-1.0",
                 guidance_scale=7.0,
                 guidance_scale_min=1.0,
                 guidance_scale_max=10.0,
@@ -265,6 +235,16 @@ preset = Preset(
                     "num_inference_steps",
                 ],
             ),
-        ]
-    ),
+        },
+        "together": {
+            "black-forest-labs/FLUX.1-schnell-Free": Txt2ImgPreset(
+                "FLUX.1 Schnell Free",
+                num_inference_steps=4,
+                num_inference_steps_min=1,
+                num_inference_steps_max=12,
+                parameters=["model", "seed", "width", "height", "steps"],
+                kwargs={"n": 1},
+            ),
+        },
+    },
 )

pages/1_💬_Text_Generation.py

@@ -1,33 +1,23 @@
-import os
 from datetime import datetime
+from typing import Dict
 
 import streamlit as st
 
-from lib import config, preset, txt2txt_generate
-
-SERVICE_SESSION = {
-    "Hugging Face": "api_key_hugging_face",
-    "Perplexity": "api_key_perplexity",
-}
-
-SESSION_TOKEN = {
-    "api_key_hugging_face": os.environ.get("HF_TOKEN") or None,
-    "api_key_perplexity": os.environ.get("PPLX_API_KEY") or None,
-}
+from lib import Txt2TxtPreset, config, preset, txt2txt_generate
 
 # config
 st.set_page_config(
     page_title=f"{config.title} | Text Generation",
-    page_icon=config.icon,
+    page_icon=config.logo,
     layout=config.layout,
 )
 
 # initialize state
-if "api_key_hugging_face" not in st.session_state:
-    st.session_state.api_key_hugging_face = ""
+if "api_key_hf" not in st.session_state:
+    st.session_state.api_key_hf = ""
 
-if "api_key_perplexity" not in st.session_state:
-    st.session_state.api_key_perplexity = ""
+if "api_key_pplx" not in st.session_state:
+    st.session_state.api_key_pplx = ""
 
 if "running" not in st.session_state:
     st.session_state.running = False
@@ -39,33 +29,41 @@ if "txt2txt_seed" not in st.session_state:
     st.session_state.txt2txt_seed = 0
 
 # sidebar
-st.logo("logo.png")
+st.logo(config.logo)
 st.sidebar.header("Settings")
+
 service = st.sidebar.selectbox(
     "Service",
-    options=SERVICE_SESSION.keys(),
-    index=0,
+    options=preset.txt2txt.keys(),
+    format_func=lambda x: config.service[x].name,
     disabled=st.session_state.running,
 )
 
 # 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:
+# for display_name, session_key in SERVICE_SESSION.items():
+for service_id, service_config in config.service.items():
+    if service == service_id:
+        session_key = f"api_key_{service}"
+        api_key = config.service[service].api_key
         st.session_state[session_key] = st.sidebar.text_input(
             "API Key",
             type="password",
-            value="" if SESSION_TOKEN[session_key] else st.session_state[session_key],
-            disabled=bool(SESSION_TOKEN[session_key]) or st.session_state.running,
-            help="Set by environment variable" if SESSION_TOKEN[session_key] else "Cleared on page refresh",
+            value="" if api_key else st.session_state[session_key],
+            disabled=bool(api_key) or st.session_state.running,
+            help="Set by environment variable" if api_key else "Cleared on page refresh",
         )
 
+service_preset: Dict[str, Txt2TxtPreset] = preset.txt2txt[service]
+
 model = st.sidebar.selectbox(
     "Model",
-    options=config.txt2txt.models[service],
-    index=config.txt2txt.default_model[service],
+    options=service_preset.keys(),
+    format_func=lambda x: service_preset[x].name,
     disabled=st.session_state.running,
-    format_func=lambda x: x.split("/")[1] if service == "Hugging Face" else x,
 )
+
+model_preset = service_preset[model]
+
 system = st.sidebar.text_area(
     "System Message",
     value=config.txt2txt.default_system,
@@ -74,9 +72,7 @@ system = st.sidebar.text_area(
 
 # build parameters from preset
 parameters = {}
-service_key = service.lower().replace(" ", "_")
-service_preset = getattr(preset.txt2txt, service_key)
-for param in service_preset.parameters:
+for param in model_preset.parameters:
     if param == "max_tokens":
         parameters[param] = st.sidebar.slider(
             "Max Tokens",
@@ -101,9 +97,9 @@ for param in service_preset.parameters:
         parameters[param] = st.sidebar.slider(
             "Frequency Penalty",
             step=0.1,
-            value=service_preset.frequency_penalty,
-            min_value=service_preset.frequency_penalty_min,
-            max_value=service_preset.frequency_penalty_max,
+            value=model_preset.frequency_penalty,
+            min_value=model_preset.frequency_penalty_min,
+            max_value=model_preset.frequency_penalty_max,
             disabled=st.session_state.running,
             help="Penalize new tokens based on their existing frequency in the text (default: 0.0)",
         )
@@ -180,8 +176,8 @@ if prompt := st.chat_input(
         st.markdown(prompt)
 
     with st.chat_message("assistant"):
-        session_key = f"api_key_{service.lower().replace(' ', '_')}"
-        api_key = st.session_state[session_key] or SESSION_TOKEN[session_key]
+        session_key = f"api_key_{service}"
+        api_key = st.session_state[session_key] or config.service[service].api_key
         response = txt2txt_generate(api_key, service, model, parameters)
         st.session_state.running = False
 

pages/2_🎨_Text_to_Image.py

@@ -1,44 +1,25 @@
-import os
 from datetime import datetime
+from typing import Dict
 
 import streamlit as st
 
-from lib import config, preset, txt2img_generate
-
-# The token name is the service in lower_snake_case
-SERVICE_SESSION = {
-    "Black Forest Labs": "api_key_black_forest_labs",
-    "Fal": "api_key_fal",
-    "Hugging Face": "api_key_hugging_face",
-    "Together": "api_key_together",
-}
-
-SESSION_TOKEN = {
-    "api_key_black_forest_labs": 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,
-    "api_key_together": os.environ.get("TOGETHER_API_KEY") or None,
-}
-
-PRESET_MODEL = {}
-for p in preset.txt2img.presets:
-    PRESET_MODEL[p.model_id] = p
+from lib import Txt2ImgPreset, config, preset, txt2img_generate
 
 st.set_page_config(
     page_title=f"{config.title} | Text to Image",
-    page_icon=config.icon,
+    page_icon=config.logo,
     layout=config.layout,
 )
 
 # Initialize Streamlit session state
-if "api_key_black_forest_labs" not in st.session_state:
-    st.session_state.api_key_black_forest_labs = ""
+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 = ""
 
-if "api_key_hugging_face" not in st.session_state:
-    st.session_state.api_key_hugging_face = ""
+if "api_key_hf" not in st.session_state:
+    st.session_state.api_key_hf = ""
 
 if "api_key_together" not in st.session_state:
     st.session_state.api_key_together = ""
@@ -52,34 +33,42 @@ if "txt2img_messages" not in st.session_state:
 if "txt2img_seed" not in st.session_state:
     st.session_state.txt2img_seed = 0
 
-st.logo("logo.png")
+st.logo(config.logo)
 st.sidebar.header("Settings")
+
 service = st.sidebar.selectbox(
     "Service",
-    options=list(SERVICE_SESSION.keys()),
+    options=preset.txt2img.keys(),
+    format_func=lambda x: config.service[x].name,
     disabled=st.session_state.running,
-    index=2,  # Hugging Face
 )
 
 # Show the API key input for the selected service.
 # Disable 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:
+# for display_name, session_key in SERVICE_SESSION.items():
+for service_id in config.service.keys():
+    if service == service_id:
+        session_key = f"api_key_{service}"
+        api_key = config.service[service].api_key
         st.session_state[session_key] = st.sidebar.text_input(
             "API Key",
             type="password",
-            value="" if SESSION_TOKEN[session_key] else st.session_state[session_key],
-            disabled=bool(SESSION_TOKEN[session_key]) or st.session_state.running,
-            help="Set by environment variable" if SESSION_TOKEN[session_key] else "Cleared on page refresh",
+            value="" if api_key else st.session_state[session_key],
+            disabled=bool(api_key) or st.session_state.running,
+            help="Set by environment variable" if api_key else "Cleared on page refresh",
         )
 
+service_preset: Dict[str, Txt2ImgPreset] = preset.txt2img[service]
+
 model = st.sidebar.selectbox(
     "Model",
-    options=config.txt2img.models[service],
-    index=config.txt2img.default_model[service],
+    options=service_preset.keys(),
+    format_func=lambda x: service_preset[x].name,
     disabled=st.session_state.running,
 )
 
+model_preset = service_preset[model]
+
 # heading
 st.html("""
     <h1>Text to Image</h1>
@@ -88,7 +77,6 @@ st.html("""
 
 # Build parameters from preset by rendering the appropriate input widgets
 parameters = {}
-model_preset = PRESET_MODEL[model]
 for param in model_preset.parameters:
     if param == "model":
         parameters[param] = model
@@ -262,13 +250,13 @@ if prompt := st.chat_input(
         with st.spinner("Running..."):
             if model_preset.kwargs:
                 parameters.update(model_preset.kwargs)
-            session_key = f"api_key_{service.lower().replace(' ', '_')}"
-            api_key = st.session_state[session_key] or SESSION_TOKEN[session_key]
+            session_key = f"api_key_{service}"
+            api_key = st.session_state[session_key] or config.service[service].api_key
             image = txt2img_generate(api_key, service, model, prompt, parameters)
         st.session_state.running = False
 
     st.session_state.txt2img_messages.append(
-        {"role": "user", "content": prompt, "parameters": parameters, "model": PRESET_MODEL[model].name}
+        {"role": "user", "content": prompt, "parameters": parameters, "model": model_preset.name}
     )
     st.session_state.txt2img_messages.append({"role": "assistant", "content": image})
     st.rerun()