Convert config to dataclass

0_🏠_Home.py

@@ -1,11 +1,11 @@
 import streamlit as st
 
-from lib import Config
+from lib import config
 
 st.set_page_config(
-    page_title=Config.TITLE,
-    page_icon=Config.ICON,
-    layout=Config.LAYOUT,
+    page_title=config.title,
+    page_icon=config.icon,
+    layout=config.layout,
 )
 
 # sidebar

lib/__init__.py

@@ -1,9 +1,9 @@
 from .api import txt2img_generate, txt2txt_generate
-from .config import Config
+from .config import config
 from .presets import ModelPresets, ServicePresets
 
 __all__ = [
-    "Config",
+    "config",
     "ModelPresets",
     "ServicePresets",
     "txt2img_generate",

lib/api.py

@@ -7,11 +7,11 @@ import streamlit as st
 from openai import APIError, OpenAI
 from PIL import Image
 
-from .config import Config
+from .config import config
 
 
 def txt2txt_generate(api_key, service, model, parameters, **kwargs):
-    base_url = Config.SERVICES[service]
+    base_url = config.services[service]
     if service == "Hugging Face":
         base_url = f"{base_url}/{model}/v1"
     client = OpenAI(api_key=api_key, base_url=base_url)
@@ -62,23 +62,23 @@ def txt2img_generate(api_key, service, model, inputs, parameters, **kwargs):
         json = {**parameters, **kwargs}
         json["prompt"] = inputs
 
-    base_url = Config.SERVICES[service]
+    base_url = config.services[service]
 
     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)
+        response = httpx.post(base_url, headers=headers, json=json, timeout=config.txt2img.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 == "Black Forest Labs":
                 id = response.json()["id"]
-                url = f"{Config.SERVICES[service]}/get_result?id={id}"
+                url = f"{config.services[service]}/get_result?id={id}"
 
                 retries = 0
-                while retries < Config.TXT2IMG_TIMEOUT:
-                    response = httpx.get(url, timeout=Config.TXT2IMG_TIMEOUT)
+                while retries < config.txt2img.timeout:
+                    response = httpx.get(url, timeout=config.txt2img.timeout)
                     if response.status_code // 100 != 2:
                         return f"Error: {response.status_code} {response.text}"
 
@@ -86,7 +86,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=config.txt2img.timeout,
                         )
                         return Image.open(io.BytesIO(image.content))
 
@@ -102,7 +102,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=config.txt2img.timeout)
                     return Image.open(io.BytesIO(image.content))
 
             if service == "Hugging Face":
@@ -110,7 +110,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=config.txt2img.timeout)
                 return Image.open(io.BytesIO(image.content))
 
         else:

lib/config.py

@@ -1,115 +1,149 @@
-from types import SimpleNamespace
+from dataclasses import dataclass
+from typing import Dict, List
 
-Config = SimpleNamespace(
-    TITLE="API Inference",
-    ICON="⚡",
-    LAYOUT="wide",
-    SERVICES={
+
+@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
+    image_sizes: List[str]
+    default_aspect_ratio: str
+    aspect_ratios: List[str]
+    timeout: int = 60
+
+
+@dataclass
+class Config:
+    title: str
+    icon: str
+    layout: str
+    services: Dict[str, str]
+    txt2img: Txt2ImgConfig
+    txt2txt: Txt2TxtConfig
+
+
+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",
     },
-    TXT2IMG_TIMEOUT=60,
-    TXT2IMG_HIDDEN_PARAMETERS=[
-        # sent to API but not shown in generation parameters accordion
-        "enable_safety_checker",
-        "max_sequence_length",
-        "num_images",
-        "output_format",
-        "performance",
-        "safety_tolerance",
-        "scheduler",
-        "sharpness",
-        "style",
-        "styles",
-        "sync_mode",
-    ],
-    TXT2IMG_NEGATIVE_PROMPT="ugly, unattractive, disfigured, deformed, mutated, malformed, blurry, grainy, noisy, oversaturated, undersaturated, overexposed, underexposed, worst quality, low details, lowres, watermark, signature, autograph, trademark, sloppy, cluttered",
-    TXT2IMG_DEFAULT_MODEL={
-        # The index of model in below lists
-        "Black Forest Labs": 2,
-        "Fal": 0,
-        "Hugging Face": 2,
-        "Together": 0,
-    },
-    TXT2IMG_MODELS={
-        # Model IDs referenced in Text_to_Image.py
-        "Black Forest Labs": [
-            "flux-dev",
-            "flux-pro",
-            "flux-pro-1.1",
-        ],
-        "Fal": [
-            "fal-ai/aura-flow",
-            "fal-ai/flux/dev",
-            "fal-ai/flux/schnell",
-            "fal-ai/flux-pro",
-            "fal-ai/flux-pro/v1.1",
-            "fal-ai/fooocus",
-            "fal-ai/kolors",
-            "fal-ai/stable-diffusion-v3-medium",
+    txt2img=Txt2ImgConfig(
+        default_model={
+            "Black Forest Labs": 2,
+            "Fal": 0,
+            "Hugging Face": 2,
+            "Together": 0,
+        },
+        models={
+            # Model identifiers referenced in Text_to_Image.py
+            "Black Forest Labs": [
+                "flux-dev",
+                "flux-pro",
+                "flux-pro-1.1",
+            ],
+            "Fal": [
+                "fal-ai/aura-flow",
+                "fal-ai/flux/dev",
+                "fal-ai/flux/schnell",
+                "fal-ai/flux-pro",
+                "fal-ai/flux-pro/v1.1",
+                "fal-ai/fooocus",
+                "fal-ai/kolors",
+                "fal-ai/stable-diffusion-v3-medium",
+            ],
+            "Hugging Face": [
+                "black-forest-labs/flux.1-dev",
+                "black-forest-labs/flux.1-schnell",
+                "stabilityai/stable-diffusion-xl-base-1.0",
+            ],
+            "Together": [
+                "black-forest-labs/FLUX.1-schnell-Free",
+            ],
+        },
+        hidden_parameters=[
+            # sent to API but not shown in generation parameters accordion
+            "enable_safety_checker",
+            "max_sequence_length",
+            "num_images",
+            "output_format",
+            "performance",
+            "safety_tolerance",
+            "scheduler",
+            "sharpness",
+            "style",
+            "styles",
+            "sync_mode",
         ],
-        "Hugging Face": [
-            "black-forest-labs/flux.1-dev",
-            "black-forest-labs/flux.1-schnell",
-            "stabilityai/stable-diffusion-xl-base-1.0",
+        negative_prompt="ugly, unattractive, disfigured, deformed, mutated, malformed, blurry, grainy, noisy, oversaturated, undersaturated, overexposed, underexposed, worst quality, low details, lowres, watermark, signature, autograph, trademark, sloppy, cluttered",
+        default_image_size="square_hd",
+        image_sizes=[
+            "landscape_16_9",
+            "landscape_4_3",
+            "square_hd",
+            "portrait_4_3",
+            "portrait_16_9",
         ],
-        "Together": [
-            "black-forest-labs/FLUX.1-schnell-Free",
+        default_aspect_ratio="1024x1024",
+        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
         ],
-    },
-    TXT2IMG_DEFAULT_IMAGE_SIZE="square_hd",  # fal image sizes
-    TXT2IMG_IMAGE_SIZES=[
-        "landscape_16_9",
-        "landscape_4_3",
-        "square_hd",
-        "portrait_4_3",
-        "portrait_16_9",
-    ],
-    TXT2IMG_DEFAULT_ASPECT_RATIO="1024x1024",  # fooocus aspect ratios
-    TXT2IMG_ASPECT_RATIOS=[
-        "704x1408",  # 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_DEFAULT_SYSTEM="You are a helpful assistant. Be precise and concise.",
-    TXT2TXT_DEFAULT_MODEL={
-        "Hugging Face": 4,
-        "Perplexity": 3,
-    },
-    TXT2TXT_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",
-        ],
-    },
+    ),
+    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",
+            ],
+        },
+    ),
 )

pages/1_💬_Text_Generation.py

@@ -3,7 +3,7 @@ from datetime import datetime
 
 import streamlit as st
 
-from lib import Config, ServicePresets, txt2txt_generate
+from lib import ServicePresets, config, txt2txt_generate
 
 SERVICE_SESSION = {
     "Hugging Face": "api_key_hugging_face",
@@ -17,9 +17,9 @@ SESSION_TOKEN = {
 
 # config
 st.set_page_config(
-    page_title=f"{Config.TITLE} | Text Generation",
-    page_icon=Config.ICON,
-    layout=Config.LAYOUT,
+    page_title=f"{config.title} | Text Generation",
+    page_icon=config.icon,
+    layout=config.layout,
 )
 
 # initialize state
@@ -61,14 +61,14 @@ for display_name, session_key in SERVICE_SESSION.items():
 
 model = st.sidebar.selectbox(
     "Model",
-    options=Config.TXT2TXT_MODELS[service],
-    index=Config.TXT2TXT_DEFAULT_MODEL[service],
+    options=config.txt2txt.models[service],
+    index=config.txt2txt.default_model[service],
     disabled=st.session_state.running,
     format_func=lambda x: x.split("/")[1] if service == "Hugging Face" else x,
 )
 system = st.sidebar.text_area(
     "System Message",
-    value=Config.TXT2TXT_DEFAULT_SYSTEM,
+    value=config.txt2txt.default_system,
     disabled=st.session_state.running,
 )
 

pages/2_🎨_Text_to_Image.py

@@ -3,7 +3,7 @@ from datetime import datetime
 
 import streamlit as st
 
-from lib import Config, ModelPresets, txt2img_generate
+from lib import ModelPresets, config, txt2img_generate
 
 # The token name is the service in lower_snake_case
 SERVICE_SESSION = {
@@ -45,9 +45,9 @@ PRESET_MODEL = {
 }
 
 st.set_page_config(
-    page_title=f"{Config.TITLE} | Text to Image",
-    page_icon=Config.ICON,
-    layout=Config.LAYOUT,
+    page_title=f"{config.title} | Text to Image",
+    page_icon=config.icon,
+    layout=config.layout,
 )
 
 # Initialize Streamlit session state
@@ -94,8 +94,8 @@ for display_name, session_key in SERVICE_SESSION.items():
 
 model = st.sidebar.selectbox(
     "Model",
-    options=Config.TXT2IMG_MODELS[service],
-    index=Config.TXT2IMG_DEFAULT_MODEL[service],
+    options=config.txt2img.models[service],
+    index=config.txt2img.default_model[service],
     disabled=st.session_state.running,
 )
 
@@ -122,7 +122,7 @@ for param in preset["parameters"]:
     if param == "negative_prompt":
         parameters[param] = st.sidebar.text_area(
             "Negative Prompt",
-            value=Config.TXT2IMG_NEGATIVE_PROMPT,
+            value=config.txt2img.negative_prompt,
             disabled=st.session_state.running,
         )
     if param == "width":
@@ -146,15 +146,15 @@ for param in preset["parameters"]:
     if param == "image_size":
         parameters[param] = st.sidebar.select_slider(
             "Image Size",
-            options=Config.TXT2IMG_IMAGE_SIZES,
-            value=Config.TXT2IMG_DEFAULT_IMAGE_SIZE,
+            options=config.txt2img.image_sizes,
+            value=config.txt2img.default_image_size,
             disabled=st.session_state.running,
         )
     if param == "aspect_ratio":
         parameters[param] = st.sidebar.select_slider(
             "Aspect Ratio",
-            options=Config.TXT2IMG_ASPECT_RATIOS,
-            value=Config.TXT2IMG_DEFAULT_ASPECT_RATIO,
+            options=config.txt2img.aspect_ratios,
+            value=config.txt2img.default_aspect_ratio,
             disabled=st.session_state.running,
         )
     if param in ["guidance_scale", "guidance"]:
@@ -206,7 +206,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.txt2img.hidden_parameters
                     ]
                     st.markdown(f"`model`: {message['model']}\n\n" + "\n\n".join(filtered_parameters))