Move presets to config

lib/__init__.py

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

lib/api.py

@@ -11,7 +11,7 @@ from .config import config
 
 
 def txt2txt_generate(api_key, service, model, parameters, **kwargs):
-    base_url = config.service[service].url
+    base_url = config.services[service].url
     if service == "hf":
         base_url = f"{base_url}/{model}/v1"
     client = OpenAI(api_key=api_key, base_url=base_url)
@@ -52,23 +52,24 @@ def txt2img_generate(api_key, service, model, inputs, parameters, **kwargs):
         headers["Authorization"] = f"Bearer {api_key}"
         json["prompt"] = inputs
 
-    base_url = config.service[service].url
+    base_url = config.services[service].url
 
     if service not in ["together"]:
         base_url = f"{base_url}/{model}"
 
     try:
-        response = httpx.post(base_url, headers=headers, json=json, timeout=config.txt2img.timeout)
+        timeout = config.timeout
+        response = httpx.post(base_url, headers=headers, json=json, timeout=timeout)
         if response.status_code // 100 == 2:  # 2xx
             # 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.service[service].url}/get_result?id={id}"
+                url = f"{config.services[service].url}/get_result?id={id}"
 
                 retries = 0
-                while retries < config.txt2img.timeout:
-                    response = httpx.get(url, timeout=config.txt2img.timeout)
+                while retries < timeout:
+                    response = httpx.get(url, timeout=timeout)
                     if response.status_code // 100 != 2:
                         return f"Error: {response.status_code} {response.text}"
 
@@ -76,7 +77,7 @@ def txt2img_generate(api_key, service, model, inputs, parameters, **kwargs):
                         image = httpx.get(
                             response.json()["result"]["sample"],
                             headers=headers,
-                            timeout=config.txt2img.timeout,
+                            timeout=timeout,
                         )
                         return Image.open(io.BytesIO(image.content))
 
@@ -92,7 +93,7 @@ def txt2img_generate(api_key, service, model, inputs, parameters, **kwargs):
                     return Image.open(io.BytesIO(bytes))
                 else:
                     url = response.json()["images"][0]["url"]
-                    image = httpx.get(url, headers=headers, timeout=config.txt2img.timeout)
+                    image = httpx.get(url, headers=headers, timeout=timeout)
                     return Image.open(io.BytesIO(image.content))
 
             if service == "hf":
@@ -100,7 +101,7 @@ def txt2img_generate(api_key, service, model, inputs, parameters, **kwargs):
 
             if service == "together":
                 url = response.json()["data"][0]["url"]
-                image = httpx.get(url, headers=headers, timeout=config.txt2img.timeout)
+                image = httpx.get(url, headers=headers, timeout=timeout)
                 return Image.open(io.BytesIO(image.content))
 
         else:

lib/config.py

@@ -1,123 +1,415 @@
 import os
-from dataclasses import dataclass
-from typing import Dict, List, Optional
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Union
+
+TEXT_SYSTEM_PROMPT = "You are a helpful assistant. Be precise and concise."
+
+IMAGE_NEGATIVE_PROMPT = "ugly, unattractive, disfigured, deformed, mutated, malformed, blurry, grainy, oversaturated, undersaturated, overexposed, underexposed, worst quality, low details, lowres, watermark, signature, sloppy, cluttered"
+
+IMAGE_IMAGE_SIZES = [
+    "landscape_16_9",
+    "landscape_4_3",
+    "square_hd",
+    "portrait_4_3",
+    "portrait_16_9",
+]
+
+IMAGE_ASPECT_RATIOS = [
+    "704x1408",  # 1:2
+    "704x1344",  # 11:21
+    "768x1344",  # 4:7
+    "768x1280",  # 3:5
+    "832x1216",  # 13:19
+    "832x1152",  # 13:18
+    "896x1152",  # 7:9
+    "896x1088",  # 14:17
+    "960x1088",  # 15:17
+    "960x1024",  # 15:16
+    "1024x1024",
+    "1024x960",  # 16:15
+    "1088x960",  # 17:15
+    "1088x896",  # 17:14
+    "1152x896",  # 9:7
+    "1152x832",  # 18:13
+    "1216x832",  # 19:13
+    "1280x768",  # 5:3
+    "1344x768",  # 7:4
+    "1344x704",  # 21:11
+    "1408x704",  # 2:1
+]
+
+IMAGE_RANGE = (256, 1408)
 
 
 @dataclass
-class ServiceConfig:
+class ModelConfig:
     name: str
-    url: str
-    api_key: Optional[str] = None
+    parameters: List[str]
+    kwargs: Optional[Dict[str, Union[str, int, float, bool]]] = field(default_factory=dict)
 
 
 @dataclass
-class Txt2TxtConfig:
-    default_system: str
+class TextModelConfig(ModelConfig):
+    system_prompt: Optional[str] = None
+    frequency_penalty: Optional[float] = None
+    frequency_penalty_range: Optional[tuple[float, float]] = None
+    max_tokens: Optional[int] = None
+    max_tokens_range: Optional[tuple[int, int]] = None
+    temperature: Optional[float] = None
+    temperature_range: Optional[tuple[float, float]] = None
 
 
 @dataclass
-class Txt2ImgConfig:
-    hidden_parameters: List[str]
-    negative_prompt: str
-    default_image_size: str
-    image_sizes: List[str]
-    default_aspect_ratio: str
-    aspect_ratios: List[str]
-    timeout: int = 60
+class ImageModelConfig(ModelConfig):
+    negative_prompt: Optional[str] = None
+    width: Optional[int] = None
+    width_range: Optional[tuple[int, int]] = None
+    height: Optional[int] = None
+    height_range: Optional[tuple[int, int]] = None
+    image_size: Optional[str] = None
+    image_sizes: Optional[List[str]] = field(default_factory=list)
+    aspect_ratio: Optional[str] = None
+    aspect_ratios: Optional[List[str]] = field(default_factory=list)
+    guidance_scale: Optional[float] = None
+    guidance_scale_range: Optional[tuple[float, float]] = None
+    num_inference_steps: Optional[int] = None
+    num_inference_steps_range: Optional[tuple[int, int]] = None
 
 
 @dataclass
-class Config:
+class ServiceConfig:
+    name: str
+    url: str
+    api_key: Optional[str]
+    text: Optional[Dict[str, TextModelConfig]] = field(default_factory=dict)
+    image: Optional[Dict[str, ImageModelConfig]] = field(default_factory=dict)
+
+
+@dataclass
+class AppConfig:
     title: str
     layout: str
     logo: str
-    service: Dict[str, ServiceConfig]
-    txt2img: Txt2ImgConfig
-    txt2txt: Txt2TxtConfig
+    timeout: int
+    hidden_parameters: List[str]
+    services: Dict[str, ServiceConfig]
+
+
+_hf_text_kwargs = {
+    "system_prompt": TEXT_SYSTEM_PROMPT,
+    "frequency_penalty": 0.0,
+    "frequency_penalty_range": (-2.0, 2.0),
+    "max_tokens": 512,
+    "max_tokens_range": (512, 4096),
+    "temperature": 1.0,
+    "temperature_range": (0.0, 2.0),
+    "parameters": ["max_tokens", "temperature", "frequency_penalty", "seed"],
+}
 
+_pplx_text_kwargs = {
+    "system_prompt": TEXT_SYSTEM_PROMPT,
+    "frequency_penalty": 1.0,
+    "frequency_penalty_range": (1.0, 2.0),
+    "max_tokens": 512,
+    "max_tokens_range": (512, 4096),
+    "temperature": 1.0,
+    "temperature_range": (0.0, 2.0),
+    "parameters": ["max_tokens", "temperature", "frequency_penalty", "seed"],
+}
 
-config = Config(
+config = AppConfig(
     title="API Inference",
     layout="wide",
     logo="logo.png",
-    service={
+    timeout=60,
+    hidden_parameters=[
+        # Sent to API but not shown in generation parameters accordion
+        "enable_safety_checker",
+        "max_sequence_length",
+        "n",
+        "num_images",
+        "output_format",
+        "performance",
+        "safety_tolerance",
+        "scheduler",
+        "sharpness",
+        "style",
+        "styles",
+        "sync_mode",
+    ],
+    services={
         "bfl": ServiceConfig(
-            "Black Forest Labs",
-            "https://api.bfl.ml/v1",
-            os.environ.get("BFL_API_KEY"),
+            name="Black Forest Labs",
+            url="https://api.bfl.ml/v1",
+            api_key=os.environ.get("BFL_API_KEY"),
+            image={
+                "flux-pro-1.1": ImageModelConfig(
+                    "FLUX1.1 Pro",
+                    width=1024,
+                    width_range=IMAGE_RANGE,
+                    height=1024,
+                    height_range=IMAGE_RANGE,
+                    parameters=["seed", "width", "height", "prompt_upsampling"],
+                    kwargs={"safety_tolerance": 6},
+                ),
+                "flux-pro": ImageModelConfig(
+                    "FLUX.1 Pro",
+                    width=1024,
+                    width_range=IMAGE_RANGE,
+                    height=1024,
+                    height_range=IMAGE_RANGE,
+                    guidance_scale=2.5,
+                    guidance_scale_range=(1.5, 5.0),
+                    num_inference_steps=50,
+                    num_inference_steps_range=(10, 50),
+                    parameters=["seed", "width", "height", "steps", "guidance", "prompt_upsampling"],
+                    kwargs={"safety_tolerance": 6, "interval": 1},
+                ),
+                "flux-dev": ImageModelConfig(
+                    "FLUX.1 Dev",
+                    width=1024,
+                    width_range=IMAGE_RANGE,
+                    height=1024,
+                    height_range=IMAGE_RANGE,
+                    num_inference_steps=28,
+                    num_inference_steps_range=(10, 50),
+                    guidance_scale=3.0,
+                    guidance_scale_range=(1.5, 5.0),
+                    parameters=["seed", "width", "height", "steps", "guidance", "prompt_upsampling"],
+                    kwargs={"safety_tolerance": 6},
+                ),
+            },
         ),
         "fal": ServiceConfig(
-            "Fal",
-            "https://fal.run",
-            os.environ.get("FAL_KEY"),
+            name="Fal",
+            url="https://fal.run",
+            api_key=os.environ.get("FAL_KEY"),
+            image={
+                "fal-ai/aura-flow": ImageModelConfig(
+                    "AuraFlow",
+                    guidance_scale=3.5,
+                    guidance_scale_range=(1.0, 10.0),
+                    num_inference_steps=28,
+                    num_inference_steps_range=(10, 50),
+                    parameters=["seed", "num_inference_steps", "guidance_scale", "expand_prompt"],
+                    kwargs={"num_images": 1, "sync_mode": False},
+                ),
+                "fal-ai/flux-pro/v1.1": ImageModelConfig(
+                    "FLUX1.1 Pro",
+                    parameters=["seed", "image_size"],
+                    image_size="square_hd",
+                    image_sizes=IMAGE_IMAGE_SIZES,
+                    kwargs={
+                        "num_images": 1,
+                        "sync_mode": False,
+                        "safety_tolerance": 6,
+                        "enable_safety_checker": False,
+                    },
+                ),
+                "fal-ai/flux-pro": ImageModelConfig(
+                    "FLUX.1 Pro",
+                    image_size="square_hd",
+                    image_sizes=IMAGE_IMAGE_SIZES,
+                    guidance_scale=2.5,
+                    guidance_scale_range=(1.5, 5.0),
+                    num_inference_steps=40,
+                    num_inference_steps_range=(10, 50),
+                    parameters=["seed", "image_size", "num_inference_steps", "guidance_scale"],
+                    kwargs={"num_images": 1, "sync_mode": False, "safety_tolerance": 6},
+                ),
+                "fal-ai/flux/dev": ImageModelConfig(
+                    "FLUX.1 Dev",
+                    image_size="square_hd",
+                    image_sizes=IMAGE_IMAGE_SIZES,
+                    num_inference_steps=28,
+                    num_inference_steps_range=(10, 50),
+                    guidance_scale=3.0,
+                    guidance_scale_range=(1.5, 5.0),
+                    parameters=["seed", "image_size", "num_inference_steps", "guidance_scale"],
+                    kwargs={"num_images": 1, "sync_mode": False, "safety_tolerance": 6},
+                ),
+                "fal-ai/flux/schnell": ImageModelConfig(
+                    "FLUX.1 Schnell",
+                    image_size="square_hd",
+                    image_sizes=IMAGE_IMAGE_SIZES,
+                    num_inference_steps=4,
+                    num_inference_steps_range=(1, 12),
+                    parameters=["seed", "image_size", "num_inference_steps"],
+                    kwargs={"num_images": 1, "sync_mode": False, "enable_safety_checker": False},
+                ),
+                "fal-ai/fooocus": ImageModelConfig(
+                    "Fooocus",
+                    aspect_ratio="1024x1024",
+                    aspect_ratios=IMAGE_ASPECT_RATIOS,
+                    guidance_scale=4.0,
+                    guidance_scale_range=(1.0, 15.0),
+                    parameters=["seed", "negative_prompt", "aspect_ratio", "guidance_scale"],
+                    # TODO: more of these can be params
+                    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",
+                    },
+                ),
+                "fal-ai/kolors": ImageModelConfig(
+                    "Kolors",
+                    image_size="square_hd",
+                    image_sizes=IMAGE_IMAGE_SIZES,
+                    guidance_scale=5.0,
+                    guidance_scale_range=(1.0, 10.0),
+                    num_inference_steps=50,
+                    num_inference_steps_range=(10, 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",
+                    },
+                ),
+                "fal-ai/stable-diffusion-v3-medium": ImageModelConfig(
+                    "SD3",
+                    image_size="square_hd",
+                    image_sizes=IMAGE_IMAGE_SIZES,
+                    guidance_scale=5.0,
+                    guidance_scale_range=(1.0, 10.0),
+                    num_inference_steps=28,
+                    num_inference_steps_range=(10, 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},
+                ),
+            },
         ),
         "hf": ServiceConfig(
-            "Hugging Face",
-            "https://api-inference.huggingface.co/models",
-            os.environ.get("HF_TOKEN"),
+            name="Hugging Face",
+            url="https://api-inference.huggingface.co/models",
+            api_key=os.environ.get("HF_TOKEN"),
+            text={
+                "codellama/codellama-34b-instruct-hf": TextModelConfig(
+                    "Code Llama 34B",
+                    **_hf_text_kwargs,
+                ),
+                "meta-llama/llama-2-13b-chat-hf": TextModelConfig(
+                    "Meta Llama 2 13B",
+                    **_hf_text_kwargs,
+                ),
+                "mistralai/mistral-7b-instruct-v0.2": TextModelConfig(
+                    "Mistral 0.2 7B",
+                    **_hf_text_kwargs,
+                ),
+                "nousresearch/nous-hermes-2-mixtral-8x7b-dpo": TextModelConfig(
+                    "Nous Hermes 2 Mixtral 8x7B",
+                    **_hf_text_kwargs,
+                ),
+            },
+            image={
+                "black-forest-labs/flux.1-dev": ImageModelConfig(
+                    "FLUX.1 Dev",
+                    width=1024,
+                    width_range=IMAGE_RANGE,
+                    height=1024,
+                    height_range=IMAGE_RANGE,
+                    guidance_scale=3.0,
+                    guidance_scale_range=(1.5, 5.0),
+                    num_inference_steps=28,
+                    num_inference_steps_range=(10, 50),
+                    parameters=["width", "height", "guidance_scale", "num_inference_steps"],
+                ),
+                "black-forest-labs/flux.1-schnell": ImageModelConfig(
+                    "FLUX.1 Schnell",
+                    width=1024,
+                    width_range=IMAGE_RANGE,
+                    height=1024,
+                    height_range=IMAGE_RANGE,
+                    num_inference_steps=4,
+                    num_inference_steps_range=(1, 12),
+                    parameters=["width", "height", "num_inference_steps"],
+                    kwargs={"guidance_scale": 0.0, "max_sequence_length": 256},
+                ),
+                "stabilityai/stable-diffusion-xl-base-1.0": ImageModelConfig(
+                    "Stable Diffusion XL 1.0",
+                    negative_prompt=IMAGE_NEGATIVE_PROMPT,
+                    width=1024,
+                    width_range=IMAGE_RANGE,
+                    height=1024,
+                    height_range=IMAGE_RANGE,
+                    guidance_scale=7.0,
+                    guidance_scale_range=(1.0, 15.0),
+                    num_inference_steps=40,
+                    num_inference_steps_range=(10, 50),
+                    parameters=[
+                        "seed",
+                        "negative_prompt",
+                        "width",
+                        "height",
+                        "guidance_scale",
+                        "num_inference_steps",
+                    ],
+                ),
+            },
         ),
         "pplx": ServiceConfig(
-            "Perplexity",
-            "https://api.perplexity.ai",
-            os.environ.get("PPLX_API_KEY"),
+            name="Perplexity",
+            url="https://api.perplexity.ai",
+            api_key=os.environ.get("PPLX_API_KEY"),
+            text={
+                "llama-3.1-sonar-small-128k-chat": TextModelConfig(
+                    "Sonar Small (Offline)",
+                    **_pplx_text_kwargs,
+                ),
+                "llama-3.1-sonar-large-128k-chat": TextModelConfig(
+                    "Sonar Large (Offline)",
+                    **_pplx_text_kwargs,
+                ),
+                "llama-3.1-sonar-small-128k-online": TextModelConfig(
+                    "Sonar Small (Online)",
+                    **_pplx_text_kwargs,
+                ),
+                "llama-3.1-sonar-large-128k-online": TextModelConfig(
+                    "Sonar Large (Online)",
+                    **_pplx_text_kwargs,
+                ),
+                "llama-3.1-sonar-huge-128k-online": TextModelConfig(
+                    "Sonar Huge (Online)",
+                    **_pplx_text_kwargs,
+                ),
+            },
         ),
+        # TODO: text models
         "together": ServiceConfig(
-            "Together",
-            "https://api.together.xyz/v1/images/generations",
-            os.environ.get("TOGETHER_API_KEY"),
+            name="Together",
+            url="https://api.together.xyz/v1/images/generations",
+            api_key=os.environ.get("TOGETHER_API_KEY"),
+            image={
+                "black-forest-labs/FLUX.1-schnell-Free": ImageModelConfig(
+                    "FLUX.1 Schnell Free",
+                    width=1024,
+                    width_range=IMAGE_RANGE,
+                    height=1024,
+                    height_range=IMAGE_RANGE,
+                    num_inference_steps=4,
+                    num_inference_steps_range=(1, 12),
+                    parameters=["model", "seed", "width", "height", "steps"],
+                    kwargs={"n": 1},
+                ),
+            },
         ),
     },
-    txt2img=Txt2ImgConfig(
-        hidden_parameters=[
-            # Sent to API but not shown in generation parameters accordion
-            "enable_safety_checker",
-            "max_sequence_length",
-            "n",
-            "num_images",
-            "output_format",
-            "performance",
-            "safety_tolerance",
-            "scheduler",
-            "sharpness",
-            "style",
-            "styles",
-            "sync_mode",
-        ],
-        negative_prompt="ugly, unattractive, disfigured, deformed, mutated, malformed, blurry, grainy, oversaturated, undersaturated, overexposed, underexposed, worst quality, low details, lowres, watermark, signature, sloppy, cluttered",
-        default_image_size="square_hd",
-        image_sizes=[
-            "landscape_16_9",
-            "landscape_4_3",
-            "square_hd",
-            "portrait_4_3",
-            "portrait_16_9",
-        ],
-        default_aspect_ratio="1024x1024",  # fooocus aspect ratios
-        aspect_ratios=[
-            "704x1408",  # 1:2
-            "704x1344",  # 11:21
-            "768x1344",  # 4:7
-            "768x1280",  # 3:5
-            "832x1216",  # 13:19
-            "832x1152",  # 13:18
-            "896x1152",  # 7:9
-            "896x1088",  # 14:17
-            "960x1088",  # 15:17
-            "960x1024",  # 15:16
-            "1024x1024",
-            "1024x960",  # 16:15
-            "1088x960",  # 17:15
-            "1088x896",  # 17:14
-            "1152x896",  # 9:7
-            "1152x832",  # 18:13
-            "1216x832",  # 19:13
-            "1280x768",  # 5:3
-            "1344x768",  # 7:4
-            "1344x704",  # 21:11
-            "1408x704",  # 2:1
-        ],
-    ),
-    txt2txt=Txt2TxtConfig(
-        default_system="You are a helpful assistant. Be precise and concise.",
-    ),
 )

lib/preset.py

@@ -1,250 +0,0 @@
-from dataclasses import dataclass, field
-from typing import Dict, List, Optional, Union
-
-
-@dataclass
-class Txt2TxtPreset:
-    name: str
-    frequency_penalty: float
-    frequency_penalty_min: float
-    frequency_penalty_max: float
-    parameters: Optional[List[str]] = field(default_factory=list)
-
-
-@dataclass
-class Txt2ImgPreset:
-    name: str
-    guidance_scale: Optional[float] = None
-    guidance_scale_min: Optional[float] = None
-    guidance_scale_max: Optional[float] = None
-    num_inference_steps: Optional[int] = None
-    num_inference_steps_min: Optional[int] = None
-    num_inference_steps_max: Optional[int] = None
-    parameters: Optional[List[str]] = field(default_factory=list)
-    kwargs: Optional[Dict[str, Union[str, int, float, bool]]] = field(default_factory=dict)
-
-
-@dataclass
-class Preset:
-    txt2txt: Dict[str, Txt2TxtPreset]
-    txt2img: Dict[str, Txt2ImgPreset]
-
-
-hf_txt2txt_kwargs = {
-    "frequency_penalty": 0.0,
-    "frequency_penalty_min": -2.0,
-    "frequency_penalty_max": 2.0,
-    "parameters": ["max_tokens", "temperature", "frequency_penalty", "seed"],
-}
-
-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={
-        "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,
-            ),
-        },
-        "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",
-                parameters=["seed", "width", "height", "prompt_upsampling"],
-                kwargs={"safety_tolerance": 6},
-            ),
-            "flux-pro": Txt2ImgPreset(
-                "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": 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=["seed", "width", "height", "steps", "guidance", "prompt_upsampling"],
-                kwargs={"safety_tolerance": 6},
-            ),
-        },
-        "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",
-                parameters=["seed", "image_size"],
-                kwargs={
-                    "num_images": 1,
-                    "sync_mode": False,
-                    "safety_tolerance": 6,
-                    "enable_safety_checker": False,
-                },
-            ),
-            "fal-ai/flux-pro": Txt2ImgPreset(
-                "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},
-            ),
-            "fal-ai/flux/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=["seed", "image_size", "num_inference_steps", "guidance_scale"],
-                kwargs={"num_images": 1, "sync_mode": False, "safety_tolerance": 6},
-            ),
-            "fal-ai/flux/schnell": Txt2ImgPreset(
-                "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},
-            ),
-            "fal-ai/fooocus": Txt2ImgPreset(
-                "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",
-                },
-            ),
-            "fal-ai/kolors": Txt2ImgPreset(
-                "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",
-                },
-            ),
-            "fal-ai/stable-diffusion-v3-medium": Txt2ImgPreset(
-                "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},
-            ),
-        },
-        "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",
-                guidance_scale=7.0,
-                guidance_scale_min=1.0,
-                guidance_scale_max=10.0,
-                num_inference_steps=40,
-                num_inference_steps_min=10,
-                num_inference_steps_max=50,
-                parameters=[
-                    "seed",
-                    "negative_prompt",
-                    "width",
-                    "height",
-                    "guidance_scale",
-                    "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,9 +1,8 @@
 from datetime import datetime
-from typing import Dict
 
 import streamlit as st
 
-from lib import Txt2TxtPreset, config, preset, txt2txt_generate
+from lib import config, txt2txt_generate
 
 # config
 st.set_page_config(
@@ -32,19 +31,24 @@ if "txt2txt_seed" not in st.session_state:
 st.logo(config.logo)
 st.sidebar.header("Settings")
 
+text_services = {
+    service_id: service_config
+    for service_id, service_config in config.services.items()
+    if getattr(service_config, "text", None)
+}
+
 service = st.sidebar.selectbox(
     "Service",
-    options=preset.txt2txt.keys(),
-    format_func=lambda x: config.service[x].name,
+    options=text_services.keys(),
+    format_func=lambda x: text_services[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():
-for service_id, service_config in config.service.items():
+# Show the API key input for the selected service.
+for service_id, service_preset in text_services.items():
     if service == service_id:
         session_key = f"api_key_{service}"
-        api_key = config.service[service].api_key
+        api_key = service_preset.api_key
         st.session_state[session_key] = st.sidebar.text_input(
             "API Key",
             type="password",
@@ -53,33 +57,33 @@ for service_id, service_config in config.service.items():
             help="Set by environment variable" if api_key else "Cleared on page refresh",
         )
 
-service_preset: Dict[str, Txt2TxtPreset] = preset.txt2txt[service]
+service_config = text_services[service]
 
 model = st.sidebar.selectbox(
     "Model",
-    options=service_preset.keys(),
-    format_func=lambda x: service_preset[x].name,
+    options=service_config.text.keys(),
+    format_func=lambda x: service_config.text[x].name,
     disabled=st.session_state.running,
 )
 
-model_preset = service_preset[model]
+model_config = service_config.text[model]
 
 system = st.sidebar.text_area(
     "System Message",
-    value=config.txt2txt.default_system,
+    value=model_config.system_prompt,
     disabled=st.session_state.running,
 )
 
 # build parameters from preset
 parameters = {}
-for param in model_preset.parameters:
+for param in model_config.parameters:
     if param == "max_tokens":
         parameters[param] = st.sidebar.slider(
             "Max Tokens",
             step=512,
-            value=512,
-            min_value=512,
-            max_value=4096,
+            value=model_config.max_tokens,
+            min_value=model_config.max_tokens_range[0],
+            max_value=model_config.max_tokens_range[1],
             disabled=st.session_state.running,
             help="Maximum number of tokens to generate (default: 512)",
         )
@@ -87,9 +91,9 @@ for param in model_preset.parameters:
         parameters[param] = st.sidebar.slider(
             "Temperature",
             step=0.1,
-            value=1.0,
-            min_value=0.0,
-            max_value=2.0,
+            value=model_config.temperature,
+            min_value=model_config.temperature_range[0],
+            max_value=model_config.temperature_range[1],
             disabled=st.session_state.running,
             help="Used to modulate the next token probabilities (default: 1.0)",
         )
@@ -97,9 +101,9 @@ for param in model_preset.parameters:
         parameters[param] = st.sidebar.slider(
             "Frequency Penalty",
             step=0.1,
-            value=model_preset.frequency_penalty,
-            min_value=model_preset.frequency_penalty_min,
-            max_value=model_preset.frequency_penalty_max,
+            value=model_config.frequency_penalty,
+            min_value=model_config.frequency_penalty_range[0],
+            max_value=model_config.frequency_penalty_range[1],
             disabled=st.session_state.running,
             help="Penalize new tokens based on their existing frequency in the text (default: 0.0)",
         )
@@ -177,7 +181,7 @@ if prompt := st.chat_input(
 
     with st.chat_message("assistant"):
         session_key = f"api_key_{service}"
-        api_key = st.session_state[session_key] or config.service[service].api_key
+        api_key = st.session_state[session_key] or text_services[service].api_key
         response = txt2txt_generate(api_key, service, model, parameters)
         st.session_state.running = False
 

pages/2_🎨_Text_to_Image.py

@@ -1,9 +1,8 @@
 from datetime import datetime
-from typing import Dict
 
 import streamlit as st
 
-from lib import Txt2ImgPreset, config, preset, txt2img_generate
+from lib import config, txt2img_generate
 
 st.set_page_config(
     page_title=f"{config.title} | Text to Image",
@@ -36,20 +35,24 @@ if "txt2img_seed" not in st.session_state:
 st.logo(config.logo)
 st.sidebar.header("Settings")
 
+image_services = {
+    service_id: service_config
+    for service_id, service_config in config.services.items()
+    if getattr(service_config, "image", None)
+}
+
 service = st.sidebar.selectbox(
     "Service",
-    options=preset.txt2img.keys(),
-    format_func=lambda x: config.service[x].name,
+    options=image_services.keys(),
+    format_func=lambda x: image_services[x].name,
     disabled=st.session_state.running,
 )
 
 # 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():
-for service_id in config.service.keys():
+for service_id, service_config in image_services.items():
     if service == service_id:
         session_key = f"api_key_{service}"
-        api_key = config.service[service].api_key
+        api_key = service_config.api_key
         st.session_state[session_key] = st.sidebar.text_input(
             "API Key",
             type="password",
@@ -58,16 +61,16 @@ for service_id in config.service.keys():
             help="Set by environment variable" if api_key else "Cleared on page refresh",
         )
 
-service_preset: Dict[str, Txt2ImgPreset] = preset.txt2img[service]
+service_config = image_services[service]
 
 model = st.sidebar.selectbox(
     "Model",
-    options=service_preset.keys(),
-    format_func=lambda x: service_preset[x].name,
+    options=service_config.image.keys(),
+    format_func=lambda x: service_config.image[x].name,
     disabled=st.session_state.running,
 )
 
-model_preset = service_preset[model]
+model_config = service_config.image[model]
 
 # heading
 st.html("""
@@ -77,7 +80,7 @@ st.html("""
 
 # Build parameters from preset by rendering the appropriate input widgets
 parameters = {}
-for param in model_preset.parameters:
+for param in model_config.parameters:
     if param == "model":
         parameters[param] = model
     if param == "seed":
@@ -91,56 +94,56 @@ for param in model_preset.parameters:
     if param == "negative_prompt":
         parameters[param] = st.sidebar.text_area(
             "Negative Prompt",
-            value=config.txt2img.negative_prompt,
+            value=model_config.negative_prompt,
             disabled=st.session_state.running,
         )
     if param == "width":
         parameters[param] = st.sidebar.slider(
             "Width",
             step=64,
-            value=1024,
-            min_value=512,
-            max_value=2048,
+            value=model_config.width,
+            min_value=model_config.width_range[0],
+            max_value=model_config.width_range[1],
             disabled=st.session_state.running,
         )
     if param == "height":
         parameters[param] = st.sidebar.slider(
             "Height",
             step=64,
-            value=1024,
-            min_value=512,
-            max_value=2048,
+            value=model_config.height,
+            min_value=model_config.height_range[0],
+            max_value=model_config.height_range[1],
             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,
+            options=model_config.image_sizes,
+            value=model_config.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,
+            options=model_config.aspect_ratios,
+            value=model_config.aspect_ratio,
             disabled=st.session_state.running,
         )
     if param in ["guidance_scale", "guidance"]:
         parameters[param] = st.sidebar.slider(
             "Guidance Scale",
-            model_preset.guidance_scale_min,
-            model_preset.guidance_scale_max,
-            model_preset.guidance_scale,
+            model_config.guidance_scale_range[0],
+            model_config.guidance_scale_range[1],
+            model_config.guidance_scale,
             0.1,
             disabled=st.session_state.running,
         )
     if param in ["num_inference_steps", "steps"]:
         parameters[param] = st.sidebar.slider(
             "Inference Steps",
-            model_preset.num_inference_steps_min,
-            model_preset.num_inference_steps_max,
-            model_preset.num_inference_steps,
+            model_config.num_inference_steps_range[0],
+            model_config.num_inference_steps_range[1],
+            model_config.num_inference_steps,
             1,
             disabled=st.session_state.running,
         )
@@ -175,7 +178,7 @@ for message in st.session_state.txt2img_messages:
                     filtered_parameters = [
                         f"`{k}`: {v}"
                         for k, v in message["parameters"].items()
-                        if k not in config.txt2img.hidden_parameters
+                        if k not in config.hidden_parameters
                     ]
                     st.markdown(f"`model`: {message['model']}\n\n" + "\n\n".join(filtered_parameters))
 
@@ -248,15 +251,15 @@ if prompt := st.chat_input(
 
     with st.chat_message("assistant"):
         with st.spinner("Running..."):
-            if model_preset.kwargs:
-                parameters.update(model_preset.kwargs)
+            if model_config.kwargs:
+                parameters.update(model_config.kwargs)
             session_key = f"api_key_{service}"
-            api_key = st.session_state[session_key] or config.service[service].api_key
+            api_key = st.session_state[session_key] or image_services[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": model_preset.name}
+        {"role": "user", "content": prompt, "parameters": parameters, "model": model_config.name}
     )
     st.session_state.txt2img_messages.append({"role": "assistant", "content": image})
     st.rerun()