Convert presets to dataclasses

lib/__init__.py

@@ -1,11 +1,10 @@
 from .api import txt2img_generate, txt2txt_generate
 from .config import config
-from .presets import ModelPresets, ServicePresets
+from .preset import preset
 
 __all__ = [
     "config",
-    "ModelPresets",
-    "ServicePresets",
+    "preset",
     "txt2img_generate",
     "txt2txt_generate",
 ]

lib/preset.py

@@ -0,0 +1,252 @@
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Union
+
+
+@dataclass
+class Txt2TxtPreset:
+    frequency_penalty: float
+    frequency_penalty_min: float
+    frequency_penalty_max: float
+    parameters: Optional[List[str]] = field(default_factory=list)
+
+
+@dataclass
+class Txt2ImgPreset:
+    # FLUX1.1 has no scale or steps
+    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 Txt2TxtPresets:
+    hugging_face: Txt2TxtPreset
+    perplexity: Txt2TxtPreset
+
+
+@dataclass
+class Txt2ImgPresets:
+    # bfl
+    flux_1_1_pro_bfl: Txt2ImgPreset
+    flux_dev_bfl: Txt2ImgPreset
+    flux_pro_bfl: Txt2ImgPreset
+    # fal
+    aura_flow: Txt2ImgPreset
+    flux_1_1_pro_fal: Txt2ImgPreset
+    flux_dev_fal: Txt2ImgPreset
+    flux_pro_fal: Txt2ImgPreset
+    flux_schnell_fal: Txt2ImgPreset
+    fooocus: Txt2ImgPreset
+    kolors: Txt2ImgPreset
+    stable_diffusion_3: Txt2ImgPreset
+    # hf
+    flux_dev_hf: Txt2ImgPreset
+    flux_schnell_hf: Txt2ImgPreset
+    stable_diffusion_xl: Txt2ImgPreset
+    # together
+    flux_schnell_free_together: Txt2ImgPreset
+
+
+@dataclass
+class Preset:
+    txt2txt: Txt2TxtPresets
+    txt2img: Txt2ImgPresets
+
+
+preset = Preset(
+    txt2txt=Txt2TxtPresets(
+        # Every service has model and system messages
+        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(
+        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},
+        ),
+        flux_1_1_pro_bfl=Txt2ImgPreset(
+            "FLUX1.1 Pro",
+            parameters=["seed", "width", "height", "prompt_upsampling"],
+            kwargs={"safety_tolerance": 6},
+        ),
+        flux_pro_bfl=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_bfl=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},
+        ),
+        flux_1_1_pro_fal=Txt2ImgPreset(
+            "FLUX1.1 Pro",
+            parameters=["seed", "image_size"],
+            kwargs={
+                "num_images": 1,
+                "sync_mode": False,
+                "safety_tolerance": 6,
+                "enable_safety_checker": False,
+            },
+        ),
+        flux_pro_fal=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},
+        ),
+        flux_dev_fal=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},
+        ),
+        flux_schnell_fal=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},
+        ),
+        flux_dev_hf=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},
+        ),
+        flux_schnell_hf=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},
+        ),
+        flux_schnell_free_together=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},
+        ),
+        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",
+            },
+        ),
+        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",
+            },
+        ),
+        stable_diffusion_3=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},
+        ),
+        stable_diffusion_xl=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",
+            ],
+        ),
+    ),
+)

lib/presets.py

@@ -1,193 +0,0 @@
-from types import SimpleNamespace
-
-# txt2txt
-ServicePresets = SimpleNamespace(
-    # Every service has model and system messages
-    HUGGING_FACE={
-        "frequency_penalty": 0.0,
-        "frequency_penalty_min": -2.0,
-        "frequency_penalty_max": 2.0,
-        "parameters": ["max_tokens", "temperature", "frequency_penalty", "seed"],
-    },
-    PERPLEXITY={
-        "frequency_penalty": 1.0,
-        "frequency_penalty_min": 1.0,
-        "frequency_penalty_max": 2.0,
-        "parameters": ["max_tokens", "temperature", "frequency_penalty"],
-    },
-)
-
-# txt2img
-ModelPresets = SimpleNamespace(
-    AURA_FLOW={
-        "name": "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},
-    },
-    FLUX_1_1_PRO_BFL={
-        "name": "FLUX1.1 Pro",
-        "parameters": ["seed", "width", "height", "prompt_upsampling"],
-        "kwargs": {"safety_tolerance": 6},
-    },
-    FLUX_PRO_BFL={
-        "name": "FLUX.1 Pro",
-        "guidance_scale": 2.5,
-        "guidance_scale_min": 1.5,
-        "guidance_scale_max": 5.0,
-        "num_inference_steps": 40,
-        "num_inference_steps_min": 10,
-        "num_inference_steps_max": 50,
-        "parameters": ["seed", "width", "height", "steps", "guidance", "prompt_upsampling"],
-        "kwargs": {"safety_tolerance": 6, "interval": 1},
-    },
-    FLUX_DEV_BFL={
-        "name": "FLUX.1 Dev",
-        "num_inference_steps": 28,
-        "num_inference_steps_min": 10,
-        "num_inference_steps_max": 50,
-        "guidance_scale": 3.0,
-        "guidance_scale_min": 1.5,
-        "guidance_scale_max": 5.0,
-        "parameters": ["seed", "width", "height", "steps", "guidance", "prompt_upsampling"],
-        "kwargs": {"safety_tolerance": 6},
-    },
-    FLUX_1_1_PRO_FAL={
-        "name": "FLUX1.1 Pro",
-        "parameters": ["seed", "image_size"],
-        "kwargs": {
-            "num_images": 1,
-            "sync_mode": False,
-            "safety_tolerance": 6,
-            "enable_safety_checker": False,
-        },
-    },
-    FLUX_PRO_FAL={
-        "name": "FLUX.1 Pro",
-        "guidance_scale": 2.5,
-        "guidance_scale_min": 1.5,
-        "guidance_scale_max": 5.0,
-        "num_inference_steps": 40,
-        "num_inference_steps_min": 10,
-        "num_inference_steps_max": 50,
-        "parameters": ["seed", "image_size", "num_inference_steps", "guidance_scale"],
-        "kwargs": {"num_images": 1, "sync_mode": False, "safety_tolerance": 6},
-    },
-    FLUX_DEV_FAL={
-        "name": "FLUX.1 Dev",
-        "num_inference_steps": 28,
-        "num_inference_steps_min": 10,
-        "num_inference_steps_max": 50,
-        "guidance_scale": 3.0,
-        "guidance_scale_min": 1.5,
-        "guidance_scale_max": 5.0,
-        "parameters": ["seed", "image_size", "num_inference_steps", "guidance_scale"],
-        "kwargs": {"num_images": 1, "sync_mode": False, "safety_tolerance": 6},
-    },
-    FLUX_SCHNELL_FAL={
-        "name": "FLUX.1 Schnell",
-        "num_inference_steps": 4,
-        "num_inference_steps_min": 1,
-        "num_inference_steps_max": 12,
-        "parameters": ["seed", "image_size", "num_inference_steps"],
-        "kwargs": {"num_images": 1, "sync_mode": False, "enable_safety_checker": False},
-    },
-    FLUX_DEV_HF={
-        "name": "FLUX.1 Dev",
-        "num_inference_steps": 28,
-        "num_inference_steps_min": 10,
-        "num_inference_steps_max": 50,
-        "guidance_scale": 3.0,
-        "guidance_scale_min": 1.5,
-        "guidance_scale_max": 5.0,
-        "parameters": ["width", "height", "guidance_scale", "num_inference_steps"],
-        "kwargs": {"max_sequence_length": 512},
-    },
-    FLUX_SCHNELL_HF={
-        "name": "FLUX.1 Schnell",
-        "num_inference_steps": 4,
-        "num_inference_steps_min": 1,
-        "num_inference_steps_max": 12,
-        "parameters": ["width", "height", "num_inference_steps"],
-        "kwargs": {"guidance_scale": 0.0, "max_sequence_length": 256},
-    },
-    FLUX_SCHNELL_FREE_TOGETHER={
-        "name": "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},
-    },
-    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",
-        },
-    },
-    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": 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"],
-    },
-)

pages/1_💬_Text_Generation.py

@@ -3,7 +3,7 @@ from datetime import datetime
 
 import streamlit as st
 
-from lib import ServicePresets, config, txt2txt_generate
+from lib import config, preset, txt2txt_generate
 
 SERVICE_SESSION = {
     "Hugging Face": "api_key_hugging_face",
@@ -74,9 +74,9 @@ system = st.sidebar.text_area(
 
 # build parameters from preset
 parameters = {}
-service_key = service.upper().replace(" ", "_")
-preset = getattr(ServicePresets, service_key, {})
-for param in preset["parameters"]:
+service_key = service.lower().replace(" ", "_")
+service_preset = getattr(preset.txt2txt, service_key)
+for param in service_preset.parameters:
     if param == "max_tokens":
         parameters[param] = st.sidebar.slider(
             "Max Tokens",
@@ -101,9 +101,9 @@ for param in preset["parameters"]:
         parameters[param] = st.sidebar.slider(
             "Frequency Penalty",
             step=0.1,
-            value=preset["frequency_penalty"],
-            min_value=preset["frequency_penalty_min"],
-            max_value=preset["frequency_penalty_max"],
+            value=service_preset.frequency_penalty,
+            min_value=service_preset.frequency_penalty_min,
+            max_value=service_preset.frequency_penalty_max,
             disabled=st.session_state.running,
             help="Penalize new tokens based on their existing frequency in the text (default: 0.0)",
         )

pages/2_🎨_Text_to_Image.py

@@ -3,7 +3,7 @@ from datetime import datetime
 
 import streamlit as st
 
-from lib import ModelPresets, config, txt2img_generate
+from lib import config, preset, txt2img_generate
 
 # The token name is the service in lower_snake_case
 SERVICE_SESSION = {
@@ -24,24 +24,24 @@ SESSION_TOKEN = {
 # Model IDs in lib/config.py
 PRESET_MODEL = {
     # bfl
-    "flux-pro-1.1": ModelPresets.FLUX_1_1_PRO_BFL,
-    "flux-pro": ModelPresets.FLUX_PRO_BFL,
-    "flux-dev": ModelPresets.FLUX_DEV_BFL,
+    "flux-pro-1.1": preset.txt2img.flux_1_1_pro_bfl,
+    "flux-pro": preset.txt2img.flux_pro_bfl,
+    "flux-dev": preset.txt2img.flux_dev_bfl,
     # fal
-    "fal-ai/aura-flow": ModelPresets.AURA_FLOW,
-    "fal-ai/flux/dev": ModelPresets.FLUX_DEV_FAL,
-    "fal-ai/flux/schnell": ModelPresets.FLUX_SCHNELL_FAL,
-    "fal-ai/flux-pro": ModelPresets.FLUX_PRO_FAL,
-    "fal-ai/flux-pro/v1.1": ModelPresets.FLUX_1_1_PRO_FAL,
-    "fal-ai/fooocus": ModelPresets.FOOOCUS,
-    "fal-ai/kolors": ModelPresets.KOLORS,
-    "fal-ai/stable-diffusion-v3-medium": ModelPresets.STABLE_DIFFUSION_3,
+    "fal-ai/aura-flow": preset.txt2img.aura_flow,
+    "fal-ai/flux/dev": preset.txt2img.flux_dev_fal,
+    "fal-ai/flux/schnell": preset.txt2img.flux_schnell_fal,
+    "fal-ai/flux-pro": preset.txt2img.flux_pro_fal,
+    "fal-ai/flux-pro/v1.1": preset.txt2img.flux_1_1_pro_fal,
+    "fal-ai/fooocus": preset.txt2img.fooocus,
+    "fal-ai/kolors": preset.txt2img.kolors,
+    "fal-ai/stable-diffusion-v3-medium": preset.txt2img.stable_diffusion_3,
     # hf
-    "black-forest-labs/flux.1-dev": ModelPresets.FLUX_DEV_HF,
-    "black-forest-labs/flux.1-schnell": ModelPresets.FLUX_SCHNELL_HF,
-    "stabilityai/stable-diffusion-xl-base-1.0": ModelPresets.STABLE_DIFFUSION_XL,
+    "black-forest-labs/flux.1-dev": preset.txt2img.flux_dev_hf,
+    "black-forest-labs/flux.1-schnell": preset.txt2img.flux_schnell_hf,
+    "stabilityai/stable-diffusion-xl-base-1.0": preset.txt2img.stable_diffusion_xl,
     # together
-    "black-forest-labs/FLUX.1-schnell-Free": ModelPresets.FLUX_SCHNELL_FREE_TOGETHER,
+    "black-forest-labs/FLUX.1-schnell-Free": preset.txt2img.flux_schnell_free_together,
 }
 
 st.set_page_config(
@@ -81,7 +81,8 @@ service = st.sidebar.selectbox(
     index=2,  # Hugging Face
 )
 
-# Disable API key input and hide value if set by environment variable; handle empty string value later.
+# 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:
         st.session_state[session_key] = st.sidebar.text_input(
@@ -108,7 +109,7 @@ st.html("""
 # Build parameters from preset by rendering the appropriate input widgets
 parameters = {}
 preset = PRESET_MODEL[model]
-for param in preset["parameters"]:
+for param in preset.parameters:
     if param == "model":
         parameters[param] = model
     if param == "seed":
@@ -160,18 +161,18 @@ for param in preset["parameters"]:
     if param in ["guidance_scale", "guidance"]:
         parameters[param] = st.sidebar.slider(
             "Guidance Scale",
-            preset["guidance_scale_min"],
-            preset["guidance_scale_max"],
-            preset["guidance_scale"],
+            preset.guidance_scale_min,
+            preset.guidance_scale_max,
+            preset.guidance_scale,
             0.1,
             disabled=st.session_state.running,
         )
     if param in ["num_inference_steps", "steps"]:
         parameters[param] = st.sidebar.slider(
             "Inference Steps",
-            preset["num_inference_steps_min"],
-            preset["num_inference_steps_max"],
-            preset["num_inference_steps"],
+            preset.num_inference_steps_min,
+            preset.num_inference_steps_max,
+            preset.num_inference_steps,
             1,
             disabled=st.session_state.running,
         )
@@ -279,16 +280,15 @@ if prompt := st.chat_input(
 
     with st.chat_message("assistant"):
         with st.spinner("Running..."):
-            if preset.get("kwargs") is not None:
-                parameters.update(preset["kwargs"])
+            if preset.kwargs:
+                parameters.update(preset.kwargs)
             session_key = f"api_key_{service.lower().replace(' ', '_')}"
             api_key = st.session_state[session_key] or SESSION_TOKEN[session_key]
             image = txt2img_generate(api_key, service, model, prompt, parameters)
         st.session_state.running = False
 
-    model_name = PRESET_MODEL[model]["name"]
     st.session_state.txt2img_messages.append(
-        {"role": "user", "content": prompt, "parameters": parameters, "model": model_name}
+        {"role": "user", "content": prompt, "parameters": parameters, "model": PRESET_MODEL[model].name}
     )
     st.session_state.txt2img_messages.append({"role": "assistant", "content": image})
     st.rerun()