Client improvements

lib/__init__.py

@@ -1,12 +1,11 @@
-from .api import HuggingFaceTxt2ImgAPI, HuggingFaceTxt2TxtAPI, PerplexityTxt2TxtAPI
+from .api import txt2img_generate, txt2txt_generate
 from .config import Config
 from .presets import ModelPresets, ServicePresets
 
 __all__ = [
     "Config",
-    "HuggingFaceTxt2ImgAPI",
-    "HuggingFaceTxt2TxtAPI",
     "ModelPresets",
-    "PerplexityTxt2TxtAPI",
     "ServicePresets",
+    "txt2img_generate",
+    "txt2txt_generate",
 ]

lib/api.py

@@ -1,87 +1,48 @@
 import io
-from abc import ABC, abstractmethod
 
 import requests
 import streamlit as st
 from openai import APIError, OpenAI
 from PIL import Image
 
-
-class Txt2TxtAPI(ABC):
-    @abstractmethod
-    def generate_text(self, model, parameters, **kwargs):
-        pass
-
-
-class Txt2ImgAPI(ABC):
-    @abstractmethod
-    def generate_image(self, model, prompt, parameters, **kwargs):
-        pass
-
-
-class HuggingFaceTxt2TxtAPI(Txt2TxtAPI):
-    def __init__(self, api_key):
-        self.api_key = api_key
-
-    def generate_text(self, model, parameters, **kwargs):
-        if not self.api_key:
-            return "API Key is required."
-        client = OpenAI(
-            api_key=self.api_key,
-            base_url=f"https://api-inference.huggingface.co/models/{model}/v1",
+from .config import Config
+
+
+def txt2txt_generate(api_key, service, model, parameters, **kwargs):
+    base_url = Config.SERVICES[service]
+    if service == "Huggingface":
+        base_url = f"{base_url}/{model}/v1"
+    client = OpenAI(api_key=api_key, base_url=base_url)
+
+    try:
+        stream = client.chat.completions.create(stream=True, model=model, **parameters, **kwargs)
+        return st.write_stream(stream)
+    except APIError as e:
+        return e.message
+    except Exception as e:
+        return str(e)
+
+
+def txt2img_generate(api_key, service, model, inputs, parameters, **kwargs):
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "X-Wait-For-Model": "true",
+        "X-Use-Cache": "false",
+    }
+    base_url = f"{Config.SERVICES[service]}/{model}"
+
+    try:
+        response = requests.post(
+            base_url,
+            headers=headers,
+            json={
+                "inputs": inputs,
+                "parameters": {**parameters, **kwargs},
+            },
         )
-        try:
-            stream = client.chat.completions.create(stream=True, model=model, **parameters, **kwargs)
-            return st.write_stream(stream)
-        except APIError as e:
-            return e.message
-        except Exception as e:
-            return str(e)
-
-
-class PerplexityTxt2TxtAPI(Txt2TxtAPI):
-    def __init__(self, api_key):
-        self.api_key = api_key
-
-    def generate_text(self, model, parameters, **kwargs):
-        if not self.api_key:
-            return "API Key is required."
-        client = OpenAI(
-            api_key=self.api_key,
-            base_url="https://api.perplexity.ai",
-        )
-        try:
-            stream = client.chat.completions.create(stream=True, model=model, **parameters, **kwargs)
-            return st.write_stream(stream)
-        except APIError as e:
-            return e.message
-        except Exception as e:
-            return str(e)
-
-
-# essentially the same as huggingface_hub's inference client
-class HuggingFaceTxt2ImgAPI(Txt2ImgAPI):
-    def __init__(self, token):
-        self.api_url = "https://api-inference.huggingface.co/models"
-        self.headers = {
-            "Authorization": f"Bearer {token}",
-            "X-Wait-For-Model": "true",
-            "X-Use-Cache": "false",
-        }
-
-    def generate_image(self, model, prompt, parameters, **kwargs):
-        try:
-            response = requests.post(
-                f"{self.api_url}/{model}",
-                headers=self.headers,
-                json={
-                    "inputs": prompt,
-                    "parameters": {**parameters, **kwargs},
-                },
-            )
-            if response.status_code == 200:
-                return Image.open(io.BytesIO(response.content))
-            else:
-                raise Exception(f"Error: {response.status_code} - {response.text}")
-        except Exception as e:
-            return str(e)
+        if response.status_code == 200:
+            return Image.open(io.BytesIO(response.content))
+        else:
+            return f"Error: {response.status_code} - {response.text}"
+    except Exception as e:
+        return str(e)

lib/config.py

@@ -4,6 +4,10 @@ Config = SimpleNamespace(
     TITLE="API Inference",
     ICON="⚡",
     LAYOUT="wide",
+    SERVICES={
+        "Huggingface": "https://api-inference.huggingface.co/models",
+        "Perplexity": "https://api.perplexity.ai",
+    },
     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_DEFAULT_MODEL=2,
     TXT2IMG_MODELS=[

lib/presets.py

@@ -1,5 +1,6 @@
 from types import SimpleNamespace
 
+# txt2txt services
 ServicePresets = SimpleNamespace(
     Huggingface={
         # every service has model and system messages
@@ -16,6 +17,7 @@ ServicePresets = SimpleNamespace(
     },
 )
 
+# txt2img models
 ModelPresets = SimpleNamespace(
     FLUX_1_DEV={
         "name": "FLUX.1 Dev",

pages/1_💬_Text_Generation.py

@@ -3,21 +3,12 @@ from datetime import datetime
 
 import streamlit as st
 
-from lib import Config, HuggingFaceTxt2TxtAPI, PerplexityTxt2TxtAPI, ServicePresets
+from lib import Config, ServicePresets, txt2txt_generate
 
 HF_TOKEN = os.environ.get("HF_TOKEN") or None
 PPLX_API_KEY = os.environ.get("PPLX_API_KEY") or None
 
 
-@st.cache_resource
-def get_txt2txt_api(service="Huggingface", api_key=None):
-    if service == "Huggingface":
-        return HuggingFaceTxt2TxtAPI(api_key)
-    if service == "Perplexity":
-        return PerplexityTxt2TxtAPI(api_key)
-    return None
-
-
 # config
 st.set_page_config(
     page_title=f"{Config.TITLE} | Text Generation",
@@ -26,8 +17,11 @@ st.set_page_config(
 )
 
 # initialize state
-if "txt2txt_running" not in st.session_state:
-    st.session_state.txt2txt_running = False
+if "api_key_huggingface" not in st.session_state:
+    st.session_state.api_key_huggingface = ""
+
+if "api_key_perplexity" not in st.session_state:
+    st.session_state.api_key_perplexity = ""
 
 if "txt2txt_messages" not in st.session_state:
     st.session_state.txt2txt_messages = []
@@ -35,6 +29,9 @@ if "txt2txt_messages" not in st.session_state:
 if "txt2txt_prompt" not in st.session_state:
     st.session_state.txt2txt_prompt = ""
 
+if "txt2txt_running" not in st.session_state:
+    st.session_state.txt2txt_running = False
+
 # sidebar
 st.logo("logo.svg")
 st.sidebar.header("Settings")
@@ -45,22 +42,40 @@ service = st.sidebar.selectbox(
     disabled=st.session_state.txt2txt_running,
 )
 
-# hide key input if environment variables are set
-if (service == "Huggingface" and HF_TOKEN is None) or (service == "Perplexity" and PPLX_API_KEY is None):
-    api_key = st.sidebar.text_input(
+if service == "Huggingface" and HF_TOKEN is None:
+    st.session_state.api_key_huggingface = st.sidebar.text_input(
         "API Key",
-        value="",
         type="password",
         help="Cleared on page refresh",
         disabled=st.session_state.txt2txt_running,
+        value=st.session_state.api_key_huggingface,
     )
+else:
+    st.session_state.api_key_huggingface = None
+
+if service == "Perplexity" and PPLX_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.txt2txt_running,
+        value=st.session_state.api_key_perplexity,
+    )
+else:
+    st.session_state.api_key_perplexity = None
+
+if service == "Huggingface" and HF_TOKEN is not None:
+    st.session_state.api_key_huggingface = HF_TOKEN
+
+if service == "Perplexity" and PPLX_API_KEY is not None:
+    st.session_state.api_key_perplexity = PPLX_API_KEY
 
 model = st.sidebar.selectbox(
     "Model",
-    format_func=lambda x: x.split("/")[1] if service == "Huggingface" else x,
-    index=Config.TXT2TXT_DEFAULT_MODEL[service],
     options=Config.TXT2TXT_MODELS[service],
+    index=Config.TXT2TXT_DEFAULT_MODEL[service],
     disabled=st.session_state.txt2txt_running,
+    format_func=lambda x: x.split("/")[1] if service == "Huggingface" else x,
 )
 system = st.sidebar.text_area(
     "System Message",
@@ -70,52 +85,48 @@ system = st.sidebar.text_area(
 
 # build parameters from preset
 parameters = {}
-preset = getattr(ServicePresets, service)
+preset = getattr(ServicePresets, service, {})
 for param in preset["parameters"]:
     if param == "max_tokens":
         parameters[param] = st.sidebar.slider(
             "Max Tokens",
+            step=128,
+            value=512,
             min_value=512,
             max_value=4096,
-            value=512,
-            step=128,
-            help="Maximum number of tokens to generate (default: 512)",
             disabled=st.session_state.txt2txt_running,
+            help="Maximum number of tokens to generate (default: 512)",
         )
     if param == "temperature":
         parameters[param] = st.sidebar.slider(
             "Temperature",
+            step=0.1,
+            value=1.0,
             min_value=0.0,
             max_value=2.0,
-            value=1.0,
-            step=0.1,
-            help="Used to modulate the next token probabilities (default: 1.0)",
             disabled=st.session_state.txt2txt_running,
+            help="Used to modulate the next token probabilities (default: 1.0)",
         )
     if param == "frequency_penalty":
         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=preset["frequency_penalty"],
-            step=0.1,
-            help="Penalize new tokens based on their existing frequency in the text (default: 0.0)",
             disabled=st.session_state.txt2txt_running,
+            help="Penalize new tokens based on their existing frequency in the text (default: 0.0)",
         )
     if param == "seed":
         parameters[param] = st.sidebar.number_input(
             "Seed",
+            value=-1,
             min_value=-1,
             max_value=(1 << 53) - 1,
-            value=-1,
-            help="Make a best effort to sample deterministically (default: -1)",
             disabled=st.session_state.txt2txt_running,
+            help="Make a best effort to sample deterministically (default: -1)",
         )
 
-# random seed
-if parameters.get("seed", 0) < 0:
-    parameters["seed"] = int(datetime.now().timestamp() * 1e6) % (1 << 53)
-
 # heading
 st.html("""
     <h1 style="padding: 0; margin-bottom: 0.5rem">Text Generation</h1>
@@ -173,9 +184,8 @@ if prompt := st.chat_input(
 ):
     st.session_state.txt2txt_prompt = prompt
 
-if st.session_state.txt2txt_prompt:
-    with st.chat_message("user"):
-        st.markdown(st.session_state.txt2txt_prompt)
+    if parameters.get("seed", 0) < 0:
+        parameters["seed"] = int(datetime.now().timestamp() * 1e6) % (1 << 53)
 
     if button_container:
         button_container.empty()
@@ -185,16 +195,12 @@ if st.session_state.txt2txt_prompt:
     messages.append({"role": "user", "content": st.session_state.txt2txt_prompt})
     parameters["messages"] = messages
 
+    with st.chat_message("user"):
+        st.markdown(st.session_state.txt2txt_prompt)
+
     with st.chat_message("assistant"):
-        # allow environment variables in development for convenience
-        if service == "Huggingface" and HF_TOKEN is not None:
-            key = HF_TOKEN
-        elif service == "Perplexity" and PPLX_API_KEY is not None:
-            key = PPLX_API_KEY
-        else:
-            key = api_key
-        api = get_txt2txt_api(service, key)
-        response = api.generate_text(model, parameters)
+        api_key = getattr(st.session_state, f"api_key_{service.lower()}", None)
+        response = txt2txt_generate(api_key, service, model, parameters)
         st.session_state.txt2txt_running = False
 
     st.session_state.txt2txt_messages.append({"role": "user", "content": st.session_state.txt2txt_prompt})

pages/2_🎨_Text_to_Image.py

@@ -3,11 +3,10 @@ from datetime import datetime
 
 import streamlit as st
 
-from lib import Config, HuggingFaceTxt2ImgAPI, ModelPresets
+from lib import Config, ModelPresets, txt2img_generate
 
-# TODO: key input and store in cache_data
-# TODO: API dropdown; changes available models
 HF_TOKEN = os.environ.get("HF_TOKEN")
+FAL_KEY = os.environ.get("FAL_KEY")
 API_URL = "https://api-inference.huggingface.co/models"
 HEADERS = {"Authorization": f"Bearer {HF_TOKEN}", "X-Wait-For-Model": "true", "X-Use-Cache": "false"}
 PRESET_MODEL = {
@@ -16,12 +15,6 @@ PRESET_MODEL = {
     "stabilityai/stable-diffusion-xl-base-1.0": ModelPresets.STABLE_DIFFUSION_XL,
 }
 
-
-@st.cache_resource
-def get_txt2img_api():
-    return HuggingFaceTxt2ImgAPI(HF_TOKEN)
-
-
 # config
 st.set_page_config(
     page_title=f"{Config.TITLE} | Text to Image",
@@ -30,24 +23,65 @@ st.set_page_config(
 )
 
 # initialize state
-if "txt2img_running" not in st.session_state:
-    st.session_state.txt2img_running = False
+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 "txt2img_running" not in st.session_state:
+    st.session_state.txt2img_running = False
+
 if "txt2img_seed" not in st.session_state:
     st.session_state.txt2img_seed = 0
 
 # sidebar
 st.logo("logo.svg")
 st.sidebar.header("Settings")
+service = st.sidebar.selectbox(
+    "Service",
+    options=["Huggingface"],
+    index=0,
+    disabled=st.session_state.txt2img_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.txt2txt_running,
+    )
+else:
+    st.session_state.api_key_huggingface = None
+
+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.txt2txt_running,
+    )
+else:
+    st.session_state.api_key_fal = None
+
+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
+
 model = st.sidebar.selectbox(
     "Model",
-    format_func=lambda x: x.split("/")[1],
     options=Config.TXT2IMG_MODELS,
     index=Config.TXT2IMG_DEFAULT_MODEL,
     disabled=st.session_state.txt2img_running,
+    format_func=lambda x: x.split("/")[1],
 )
 aspect_ratio = st.sidebar.select_slider(
     "Aspect Ratio",
@@ -191,11 +225,10 @@ if prompt := st.chat_input(
 
     with st.chat_message("assistant"):
         with st.spinner("Running..."):
-            generate_kwargs = {"model": model, "prompt": prompt, "parameters": parameters}
             if preset.get("kwargs") is not None:
-                generate_kwargs.update(preset["kwargs"])
-            api = get_txt2img_api()
-            image = api.generate_image(**generate_kwargs)
+                parameters.update(preset["kwargs"])
+            api_key = getattr(st.session_state, f"api_key_{service.lower()}", None)
+            image = txt2img_generate(api_key, service, model, prompt, parameters)
         st.session_state.txt2img_running = False
 
     model_name = PRESET_MODEL[model]["name"]