Better environment variable handling

lib/api.py

@@ -11,7 +11,7 @@ from .config import Config
 
 def txt2txt_generate(api_key, service, model, parameters, **kwargs):
     base_url = Config.SERVICES[service]
-    if service == "Huggingface":
+    if service == "Hugging Face":
         base_url = f"{base_url}/{model}/v1"
     client = OpenAI(api_key=api_key, base_url=base_url)
 
@@ -19,14 +19,15 @@ def txt2txt_generate(api_key, service, model, parameters, **kwargs):
         stream = client.chat.completions.create(stream=True, model=model, **parameters, **kwargs)
         return st.write_stream(stream)
     except APIError as e:
-        return e.message
+        # OpenAI uses this message for streaming errors and attaches response.error to error.body
+        return e.body if e.message == "An error occurred during streaming" else e.message
     except Exception as e:
         return str(e)
 
 
 def txt2img_generate(api_key, service, model, inputs, parameters, **kwargs):
     headers = {}
-    if service == "Huggingface":
+    if service == "Hugging Face":
         headers["Authorization"] = f"Bearer {api_key}"
         headers["X-Wait-For-Model"] = "true"
         headers["X-Use-Cache"] = "false"
@@ -34,7 +35,7 @@ def txt2img_generate(api_key, service, model, inputs, parameters, **kwargs):
         headers["Authorization"] = f"Key {api_key}"
 
     json = {}
-    if service == "Huggingface":
+    if service == "Hugging Face":
         json = {
             "inputs": inputs,
             "parameters": {**parameters, **kwargs},
@@ -48,7 +49,7 @@ def txt2img_generate(api_key, service, model, inputs, parameters, **kwargs):
     try:
         response = requests.post(base_url, headers=headers, json=json)
         if response.status_code // 100 == 2:  # 2xx
-            if service == "Huggingface":
+            if service == "Hugging Face":
                 return Image.open(io.BytesIO(response.content))
             if service == "Fal":
                 bytes = base64.b64decode(response.json()["images"][0]["url"].split(",")[-1])

lib/config.py

@@ -5,11 +5,12 @@ Config = SimpleNamespace(
     ICON="⚡",
     LAYOUT="wide",
     SERVICES={
-        "Huggingface": "https://api-inference.huggingface.co/models",
+        "Hugging Face": "https://api-inference.huggingface.co/models",
         "Perplexity": "https://api.perplexity.ai",
         "Fal": "https://fal.run",
     },
     TXT2IMG_HIDDEN_PARAMETERS=[
+        # sent to API but not shown in generation parameters accordion
         "enable_safety_checker",
         "max_sequence_length",
         "num_images",
@@ -22,18 +23,20 @@ Config = SimpleNamespace(
         "styles",
         "sync_mode",
     ],
-    TXT2IMG_NEGATIVE_PROMPT="ugly, unattractive, malformed, mutated, disgusting, blurry, grainy, noisy, oversaturated, undersaturated, overexposed, underexposed, worst quality, low details, lowres, watermark, signature, autograph, trademark, sloppy, cluttered",
+    TXT2IMG_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={
-        "Huggingface": 2,
+        # index of model in below lists
+        "Hugging Face": 2,
         "Fal": 1,
     },
     TXT2IMG_MODELS={
-        "Huggingface": [
+        "Hugging Face": [
             "black-forest-labs/flux.1-dev",
             "black-forest-labs/flux.1-schnell",
             "stabilityai/stable-diffusion-xl-base-1.0",
         ],
         "Fal": [
+            # TODO: fix these models
             # "fal-ai/aura-flow",
             # "fal-ai/flux-pro",
             "fal-ai/fooocus",
@@ -52,40 +55,36 @@ Config = SimpleNamespace(
     ],
     TXT2IMG_DEFAULT_ASPECT_RATIO="1024x1024",  # fooocus aspect ratios
     TXT2IMG_ASPECT_RATIOS=[
-        "704x1408",
-        "704x1344",
-        "768x1344",
-        "768x1280",
-        "832x1216",
-        "832x1152",
-        "896x1152",
-        "896x1088",
-        "960x1088",
-        "960x1024",
+        "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",
-        "1088x960",
-        "1088x896",
-        "1152x896",
-        "1152x832",
-        "1216x832",
-        "1280x768",
-        "1344x768",
-        "1344x704",
-        "1408x704",
-        "1472x704",
-        "1536x640",
-        "1600x640",
-        "1664x576",
-        "1728x576",
+        "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
     ],
+    # TODO: txt2img fooocus styles like "Fooocus V2" and "Fooocus Enhance" (use multiselect in UI)
     TXT2TXT_DEFAULT_SYSTEM="You are a helpful assistant. Be precise and concise.",
     TXT2TXT_DEFAULT_MODEL={
-        "Huggingface": 4,
+        "Hugging Face": 4,
         "Perplexity": 3,
     },
     TXT2TXT_MODELS={
-        "Huggingface": [
+        "Hugging Face": [
             "codellama/codellama-34b-instruct-hf",
             "meta-llama/llama-2-13b-chat-hf",
             "meta-llama/meta-llama-3.1-405b-instruct-fp8",

lib/presets.py

@@ -2,14 +2,14 @@ from types import SimpleNamespace
 
 # txt2txt services
 ServicePresets = SimpleNamespace(
-    Huggingface={
+    HUGGING_FACE={
         # every service has model and system messages
         "frequency_penalty": 0.0,
         "frequency_penalty_min": -2.0,
         "frequency_penalty_max": 2.0,
         "parameters": ["max_tokens", "temperature", "frequency_penalty", "seed"],
     },
-    Perplexity={
+    PERPLEXITY={
         "frequency_penalty": 1.0,
         "frequency_penalty_min": 1.0,
         "frequency_penalty_max": 2.0,

pages/1_💬_Text_Generation.py

@@ -5,8 +5,15 @@ import streamlit as st
 
 from lib import Config, ServicePresets, txt2txt_generate
 
-HF_TOKEN = os.environ.get("HF_TOKEN") or None
-PERPLEXITY_API_KEY = os.environ.get("PERPLEXITY_API_KEY") or None
+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("PERPLEXITY_API_KEY") or None,
+}
 
 # config
 st.set_page_config(
@@ -16,8 +23,8 @@ st.set_page_config(
 )
 
 # initialize state
-if "api_key_huggingface" not in st.session_state:
-    st.session_state.api_key_huggingface = ""
+if "api_key_hugging_face" not in st.session_state:
+    st.session_state.api_key_hugging_face = ""
 
 if "api_key_perplexity" not in st.session_state:
     st.session_state.api_key_perplexity = ""
@@ -28,53 +35,36 @@ if "running" not in st.session_state:
 if "txt2txt_messages" not in st.session_state:
     st.session_state.txt2txt_messages = []
 
-if "txt2txt_prompt" not in st.session_state:
-    st.session_state.txt2txt_prompt = ""
+if "txt2txt_seed" not in st.session_state:
+    st.session_state.txt2txt_seed = 0
 
 # sidebar
 st.logo("logo.svg")
 st.sidebar.header("Settings")
 service = st.sidebar.selectbox(
     "Service",
-    options=["Huggingface", "Perplexity"],
+    options=["Hugging Face", "Perplexity"],
     index=0,
     disabled=st.session_state.running,
 )
 
-if service == "Huggingface" and HF_TOKEN is None:
-    st.session_state.api_key_huggingface = st.sidebar.text_input(
-        "API Key",
-        type="password",
-        help="Cleared on page refresh",
-        disabled=st.session_state.running,
-        value=st.session_state.api_key_huggingface,
-    )
-else:
-    st.session_state.api_key_huggingface = st.session_state.api_key_huggingface
-
-if service == "Perplexity" and PERPLEXITY_API_KEY is None:
-    st.session_state.api_key_perplexity = st.sidebar.text_input(
-        "API Key",
-        type="password",
-        help="Cleared on page refresh",
-        disabled=st.session_state.running,
-        value=st.session_state.api_key_perplexity,
-    )
-else:
-    st.session_state.api_key_perplexity = st.session_state.api_key_perplexity
-
-if service == "Huggingface" and HF_TOKEN is not None:
-    st.session_state.api_key_huggingface = HF_TOKEN
-
-if service == "Perplexity" and PERPLEXITY_API_KEY is not None:
-    st.session_state.api_key_perplexity = PERPLEXITY_API_KEY
+# 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:
+        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",
+        )
 
 model = st.sidebar.selectbox(
     "Model",
     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 == "Huggingface" else x,
+    format_func=lambda x: x.split("/")[1] if service == "Hugging Face" else x,
 )
 system = st.sidebar.text_area(
     "System Message",
@@ -84,7 +74,8 @@ system = st.sidebar.text_area(
 
 # build parameters from preset
 parameters = {}
-preset = getattr(ServicePresets, service, {})
+service_key = service.upper().replace(" ", "_")
+preset = getattr(ServicePresets, service_key, {})
 for param in preset["parameters"]:
     if param == "max_tokens":
         parameters[param] = st.sidebar.slider(
@@ -152,26 +143,15 @@ if st.session_state.txt2txt_messages:
         </style>
         """)
 
-        # remove last assistant message and resend prompt
-        col1, col2, col3 = st.columns(3)
+        col1, col2 = st.columns(2)
         with col1:
-            if st.button("🔄️", help="Retry last message") and len(st.session_state.txt2txt_messages) >= 2:
-                st.session_state.txt2txt_messages.pop()
-                st.session_state.txt2txt_prompt = st.session_state.txt2txt_messages.pop()["content"]
-                st.rerun()
-
-        # delete last message pair
-        with col2:
             if st.button("❌", help="Delete last message") and len(st.session_state.txt2txt_messages) >= 2:
                 st.session_state.txt2txt_messages.pop()
                 st.session_state.txt2txt_messages.pop()
                 st.rerun()
-
-        # reset app state
-        with col3:
+        with col2:
             if st.button("🗑️", help="Clear all messages"):
                 st.session_state.txt2txt_messages = []
-                st.session_state.txt2txt_prompt = ""
                 st.rerun()
 else:
     button_container = None
@@ -181,28 +161,30 @@ if prompt := st.chat_input(
     "What would you like to know?",
     on_submit=lambda: setattr(st.session_state, "running", True),
 ):
-    st.session_state.txt2txt_prompt = prompt
-
-    if parameters.get("seed", 0) < 0:
-        parameters["seed"] = int(datetime.now().timestamp() * 1e6) % (1 << 53)
+    if "seed" in parameters and parameters["seed"] >= 0:
+        st.session_state.txt2txt_seed = parameters["seed"]
+    else:
+        st.session_state.txt2txt_seed = int(datetime.now().timestamp() * 1e6) % (1 << 53)
+        if "seed" in parameters:
+            parameters["seed"] = st.session_state.txt2txt_seed
 
     if button_container:
         button_container.empty()
 
     messages = [{"role": "system", "content": system}]
     messages.extend([{"role": m["role"], "content": m["content"]} for m in st.session_state.txt2txt_messages])
-    messages.append({"role": "user", "content": st.session_state.txt2txt_prompt})
+    messages.append({"role": "user", "content": prompt})
     parameters["messages"] = messages
 
     with st.chat_message("user"):
-        st.markdown(st.session_state.txt2txt_prompt)
+        st.markdown(prompt)
 
     with st.chat_message("assistant"):
-        api_key = getattr(st.session_state, f"api_key_{service.lower()}", None)
+        session_key = f"api_key_{service.lower().replace(' ', '_')}"
+        api_key = st.session_state[session_key] or SESSION_TOKEN[session_key]
         response = txt2txt_generate(api_key, service, model, parameters)
         st.session_state.running = False
 
-    st.session_state.txt2txt_messages.append({"role": "user", "content": st.session_state.txt2txt_prompt})
+    st.session_state.txt2txt_messages.append({"role": "user", "content": prompt})
     st.session_state.txt2txt_messages.append({"role": "assistant", "content": response})
-    st.session_state.txt2txt_prompt = ""
     st.rerun()

pages/2_🎨_Text_to_Image.py

@@ -5,8 +5,16 @@ import streamlit as st
 
 from lib import Config, ModelPresets, txt2img_generate
 
-HF_TOKEN = os.environ.get("HF_TOKEN") or None
-FAL_KEY = os.environ.get("FAL_KEY") or None
+SERVICE_SESSION = {
+    "Fal": "api_key_fal",
+    "Hugging Face": "api_key_hugging_face",
+}
+
+SESSION_TOKEN = {
+    "api_key_fal": os.environ.get("FAL_KEY") or None,
+    "api_key_hugging_face": os.environ.get("HF_TOKEN") or None,
+}
+
 PRESET_MODEL = {
     "black-forest-labs/flux.1-dev": ModelPresets.FLUX_DEV,
     "black-forest-labs/flux.1-schnell": ModelPresets.FLUX_SCHNELL,
@@ -28,15 +36,15 @@ st.set_page_config(
 if "api_key_fal" not in st.session_state:
     st.session_state.api_key_fal = ""
 
-if "api_key_huggingface" not in st.session_state:
-    st.session_state.api_key_huggingface = ""
-
-if "txt2img_messages" not in st.session_state:
-    st.session_state.txt2img_messages = []
+if "api_key_hugging_face" not in st.session_state:
+    st.session_state.api_key_hugging_face = ""
 
 if "running" not in st.session_state:
     st.session_state.running = False
 
+if "txt2img_messages" not in st.session_state:
+    st.session_state.txt2img_messages = []
+
 if "txt2img_seed" not in st.session_state:
     st.session_state.txt2img_seed = 0
 
@@ -45,38 +53,21 @@ st.logo("logo.svg")
 st.sidebar.header("Settings")
 service = st.sidebar.selectbox(
     "Service",
-    options=["Fal", "Huggingface"],
+    options=list(SERVICE_SESSION.keys()),
     index=1,
     disabled=st.session_state.running,
 )
 
-if service == "Huggingface" and HF_TOKEN is None:
-    st.session_state.api_key_huggingface = st.sidebar.text_input(
-        "API Key",
-        type="password",
-        help="Cleared on page refresh",
-        value=st.session_state.api_key_huggingface,
-        disabled=st.session_state.running,
-    )
-else:
-    st.session_state.api_key_huggingface = st.session_state.api_key_huggingface
-
-if service == "Fal" and FAL_KEY is None:
-    st.session_state.api_key_fal = st.sidebar.text_input(
-        "API Key",
-        type="password",
-        help="Cleared on page refresh",
-        value=st.session_state.api_key_fal,
-        disabled=st.session_state.running,
-    )
-else:
-    st.session_state.api_key_fal = st.session_state.api_key_fal
-
-if service == "Huggingface" and HF_TOKEN is not None:
-    st.session_state.api_key_huggingface = HF_TOKEN
-
-if service == "Fal" and FAL_KEY is not None:
-    st.session_state.api_key_fal = FAL_KEY
+# 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:
+        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",
+        )
 
 model = st.sidebar.selectbox(
     "Model",
@@ -227,7 +218,6 @@ if st.session_state.txt2img_messages:
         </style>
         """)
 
-        # retry
         col1, col2 = st.columns(2)
         with col1:
             if (
@@ -268,7 +258,8 @@ if prompt := st.chat_input(
         with st.spinner("Running..."):
             if preset.get("kwargs") is not None:
                 parameters.update(preset["kwargs"])
-            api_key = getattr(st.session_state, f"api_key_{service.lower()}", None)
+            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