App

0_🏠_Home.py

@@ -0,0 +1,49 @@
+import streamlit as st
+
+from lib import Config
+
+st.set_page_config(
+    page_title=Config.TITLE,
+    page_icon=Config.ICON,
+    layout=Config.LAYOUT,
+)
+
+# sidebar
+st.logo("logo.svg")
+
+# title
+st.html("""
+<style>
+.pro-badge {
+    display: inline-block;
+    transform: skew(-12deg);
+    font-size: 0.875rem;
+    line-height: 1.25rem;
+    font-weight: 700;
+    padding: 0.125rem 0.625rem;
+    border-radius: 0.5rem;
+    color: rgb(0 0 0 / 1);
+    box-shadow: 0 0 #0000, 0 0 #0000, 0 10px 15px -3px rgb(16 185 129 / .1), 0 4px 6px -4px rgb(16 185 129 / .1);
+    background-image: linear-gradient(to bottom right, #f9a8d4, #a7f3d0, #fde68a);
+    border: 1px solid rgb(229 231 235 / 1);
+}
+@media (prefers-color-scheme: dark) {
+    .pro-badge {
+        box-shadow: 0 0 #0000, 0 0 #0000, 0 10px 15px -3px rgb(16 185 129 / .2), 0 4px 6px -4px rgb(16 185 129 / .2);
+        background-image: linear-gradient(to bottom right, #ec4899, #10b981, #f59e0b);
+        border: 1px solid rgb(20 28 46 / 1);
+    }
+}
+</style>
+<div style="display: flex; align-items: center; gap: 0.75rem">
+    <h1 style="padding: 0; margin-bottom: 0.5rem">API Inference</h1>
+    <span class="pro-badge">PRO</span>
+</div>
+<p>Run inference on API endpoints. Hugging Face for now; more coming soon!</p>
+""")
+
+# TODO: categorize tasks by service (SAI, FAL, etc)
+# content
+st.markdown("## Tasks")
+st.page_link("pages/1_💬_Text_Generation.py", label="Text Generation", icon="💬")
+st.page_link("pages/2_🎨_Text_to_Image.py", label="Text to Image", icon="🎨")

app.py

@@ -0,0 +1,10 @@
+#!/usr/bin/env python
+import subprocess
+import sys
+
+try:
+    subprocess.run([sys.executable, "-m", "streamlit", "run", "0_🏠_Home.py"], check=True)
+except KeyboardInterrupt:
+    sys.exit(0)
+except subprocess.CalledProcessError as e:
+    sys.exit(e.returncode)

lib/__init__.py

@@ -0,0 +1,4 @@
+from .config import Config
+from .presets import Presets
+
+__all__ = ["Config", "Presets"]

lib/config.py

@@ -0,0 +1,22 @@
+from types import SimpleNamespace
+
+Config = SimpleNamespace(
+    TITLE="API Inference",
+    ICON="⚡",
+    LAYOUT="wide",
+    TXT2IMG_NEGATIVE_PROMPT="ugly, bad, asymmetrical, malformed, mutated, disgusting, blurry, grainy, oversaturated, undersaturated, overexposed, underexposed, worst quality, low details, lowres, watermark, signature, autograph, trademark",
+    TXT2IMG_DEFAULT_INDEX=2,
+    TXT2IMG_MODELS=[
+        "black-forest-labs/flux.1-dev",
+        "black-forest-labs/flux.1-schnell",
+        "stabilityai/stable-diffusion-xl-base-1.0",
+    ],
+    TXT2TXT_DEFAULT_INDEX=4,
+    TXT2TXT_MODELS=[
+        "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",
+    ],
+)

lib/presets.py

@@ -0,0 +1,33 @@
+from types import SimpleNamespace
+
+Presets = SimpleNamespace(
+    FLUX_1_DEV={
+        "name": "FLUX.1 Dev",
+        "num_inference_steps": 30,
+        "num_inference_steps_min": 10,
+        "num_inference_steps_max": 40,
+        "guidance_scale": 3.5,
+        "guidance_scale_min": 1.0,
+        "guidance_scale_max": 7.0,
+        "parameters": ["width", "height", "guidance_scale", "num_inference_steps"],
+        "kwargs": {"max_sequence_length": 512},
+    },
+    FLUX_1_SCHNELL={
+        "name": "FLUX.1 Schnell",
+        "num_inference_steps": 4,
+        "num_inference_steps_min": 1,
+        "num_inference_steps_max": 8,
+        "parameters": ["width", "height", "num_inference_steps"],
+        "kwargs": {"guidance_scale": 0.0, "max_sequence_length": 256},
+    },
+    STABLE_DIFFUSION_XL={
+        "name": "SDXL",
+        "guidance_scale": 7.0,
+        "guidance_scale_min": 1.0,
+        "guidance_scale_max": 15.0,
+        "num_inference_steps": 40,
+        "num_inference_steps_min": 10,
+        "num_inference_steps_max": 50,
+        "parameters": ["width", "height", "guidance_scale", "num_inference_steps", "seed", "negative_prompt"],
+    },
+)

pages/1_💬_Text_Generation.py

@@ -0,0 +1,174 @@
+import os
+from datetime import datetime
+
+import streamlit as st
+from openai import APIError, OpenAI
+
+from lib import Config
+
+# TODO: key input and store in cache_data
+# api key
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+
+# TODO: eventually support different APIs like OpenAI and Perplexity
+@st.cache_resource
+def get_chat_client(api_key, model):
+    client = OpenAI(
+        api_key=api_key,
+        base_url=f"https://api-inference.huggingface.co/models/{model}/v1",
+    )
+    return client
+
+
+# config
+st.set_page_config(
+    page_title=f"{Config.TITLE} | Text Generation",
+    page_icon=Config.ICON,
+    layout=Config.LAYOUT,
+)
+
+# initialize 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 = ""
+
+# sidebar
+st.logo("logo.svg")
+st.sidebar.header("Settings")
+model = st.sidebar.selectbox(
+    "Model",
+    placeholder="Select a model",
+    format_func=lambda x: x.split("/")[1],
+    index=Config.TXT2TXT_DEFAULT_INDEX,
+    options=Config.TXT2TXT_MODELS,
+)
+max_tokens = st.sidebar.slider(
+    "Max Tokens",
+    min_value=512,
+    max_value=4096,
+    value=512,
+    step=128,
+    help="Maximum number of tokens to generate (default: 512)",
+)
+temperature = st.sidebar.slider(
+    "Temperature",
+    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)",
+)
+frequency_penalty = st.sidebar.slider(
+    "Frequency Penalty",
+    min_value=-2.0,
+    max_value=2.0,
+    value=0.0,
+    step=0.1,
+    help="Penalize new tokens based on their existing frequency in the text (default: 0.0)",
+)
+seed = st.sidebar.number_input(
+    "Seed",
+    min_value=-1,
+    max_value=(1 << 53) - 1,
+    value=-1,
+    help="Make a best effort to sample deterministically (default: -1)",
+)
+system = st.sidebar.text_area(
+    "System Message",
+    value="You are a helpful assistant. Be precise and concise.",
+)
+
+# random seed
+if seed < 0:
+    seed = int(datetime.now().timestamp() * 1e6) % (1 << 53)
+
+# heading
+st.html("""
+    <h1 style="padding: 0; margin-bottom: 0.5rem">Text Generation</h1>
+    <p>Chat with large language models.</p>
+""")
+
+# chat messages
+for message in st.session_state.txt2txt_messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+# button row
+if st.session_state.txt2txt_messages:
+    button_container = st.empty()
+    with button_container.container():
+        # https://discuss.streamlit.io/t/st-button-in-one-line/25966/6
+        st.html("""
+        <style>
+            div[data-testid="column"] {
+                width: fit-content;
+                min-width: 0;
+                flex: none;
+            }
+        </style>
+        """)
+
+        # remove last assistant message and resend prompt
+        col1, col2, col3 = st.columns(3)
+        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:
+            if st.button("🗑️", help="Clear all messages"):
+                st.session_state.txt2txt_messages = []
+                st.session_state.txt2txt_prompt = ""
+                st.rerun()
+else:
+    button_container = None
+
+# chat input
+if prompt := st.chat_input("What would you like to know?"):
+    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 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})
+
+    with st.chat_message("assistant"):
+        try:
+            client = get_chat_client(HF_TOKEN, model)
+            stream = client.chat.completions.create(
+                frequency_penalty=frequency_penalty,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                messages=messages,
+                model=model,
+                stream=True,
+                seed=seed,
+            )
+            response = st.write_stream(stream)
+        except APIError as e:
+            response = e.message
+        except Exception as e:
+            response = str(e)
+
+    st.session_state.txt2txt_messages.append({"role": "user", "content": st.session_state.txt2txt_prompt})
+    st.session_state.txt2txt_messages.append({"role": "assistant", "content": response})
+    st.session_state.txt2txt_prompt = ""
+    st.rerun()

pages/2_🎨_Text_to_Image.py

@@ -0,0 +1,212 @@
+import io
+import os
+from datetime import datetime
+
+import requests
+import streamlit as st
+from PIL import Image
+
+from lib import Config, Presets
+
+# TODO: key input and store in cache_data
+# TODO: API dropdown; changes available models
+HF_TOKEN = os.environ.get("HF_TOKEN")
+API_URL = "https://api-inference.huggingface.co/models"
+HEADERS = {"Authorization": f"Bearer {HF_TOKEN}", "X-Wait-For-Model": "true", "X-Use-Cache": "false"}
+SIZE_AR = {
+    "9:7": (1152, 896),
+    "7:4": (1344, 768),
+    "1:1": (1024, 1024),
+    "4:7": (768, 1344),
+    "7:9": (896, 1152),
+}
+PRESET_MODEL = {
+    "black-forest-labs/flux.1-dev": Presets.FLUX_1_DEV,
+    "black-forest-labs/flux.1-schnell": Presets.FLUX_1_SCHNELL,
+    "stabilityai/stable-diffusion-xl-base-1.0": Presets.STABLE_DIFFUSION_XL,
+}
+
+
+def generate_image(model, prompt, parameters, **kwargs):
+    response = requests.post(
+        f"{API_URL}/{model}",
+        headers=HEADERS,
+        json={
+            "inputs": prompt,
+            "parameters": {**parameters, **kwargs},
+        },
+    )
+
+    if response.status_code == 200:
+        image = Image.open(io.BytesIO(response.content))
+        return image
+    else:
+        st.error(f"Error: {response.status_code} - {response.text}")
+        return None
+
+
+# config
+st.set_page_config(
+    page_title=f"{Config.TITLE} | Text to Image",
+    page_icon=Config.ICON,
+    layout=Config.LAYOUT,
+)
+
+# initialize state
+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
+
+# sidebar
+st.logo("logo.svg")
+st.sidebar.header("Settings")
+model = st.sidebar.selectbox(
+    "Model",
+    format_func=lambda x: x.split("/")[1],
+    options=Config.TXT2IMG_MODELS,
+    index=Config.TXT2IMG_DEFAULT_INDEX,
+)
+aspect_ratio = st.sidebar.select_slider(
+    "Aspect Ratio",
+    options=list(SIZE_AR.keys()),
+    value=list(SIZE_AR.keys())[1],
+)
+
+# heading
+st.html("""
+    <h1 style="padding: 0; margin-bottom: 0.5rem">Text to Image</h1>
+    <p>Generate an image from a text prompt.</p>
+""")
+
+# build parameters from preset
+parameters = {}
+preset = PRESET_MODEL[model]
+for param in preset["parameters"]:
+    if param == "width":
+        parameters[param] = SIZE_AR[aspect_ratio][0]
+    if param == "height":
+        parameters[param] = SIZE_AR[aspect_ratio][1]
+    if param == "guidance_scale":
+        parameters[param] = st.sidebar.slider(
+            "Guidance Scale",
+            preset["guidance_scale_min"],
+            preset["guidance_scale_max"],
+            preset["guidance_scale"],
+            0.1,
+        )
+    if param == "num_inference_steps":
+        parameters[param] = st.sidebar.slider(
+            "Inference Steps",
+            preset["num_inference_steps_min"],
+            preset["num_inference_steps_max"],
+            preset["num_inference_steps"],
+            1,
+        )
+    if param == "seed":
+        parameters[param] = st.sidebar.number_input(
+            "Seed",
+            min_value=-1,
+            max_value=(1 << 53) - 1,
+            value=-1,
+        )
+    if param == "negative_prompt":
+        parameters[param] = st.sidebar.text_area(
+            label="Negative Prompt",
+            value=Config.TXT2IMG_NEGATIVE_PROMPT,
+        )
+
+# wrap the prompt in an expander to display additional parameters
+for message in st.session_state.txt2img_messages:
+    role = message["role"]
+    with st.chat_message(role):
+        image_container = st.empty()
+
+        with image_container.container():
+            if role == "user":
+                with st.expander(message["content"]):
+                    # build a markdown string for additional parameters
+                    st.html("""
+                    <style>
+                        div[data-testid="stMarkdownContainer"] p:not(:last-of-type) { margin-bottom: 0 }
+                    </style>
+                    """)
+                    md = f"`model`: {message['model']}\n\n"
+                    md += "\n\n".join([f"`{k}`: {v}" for k, v in message["parameters"].items()])
+                    st.markdown(md)
+
+            if role == "assistant":
+                # image is full width when _not_ in full-screen mode
+                st.html("""
+                <style>
+                    div[data-testid="stImage"]:has(img[style*="max-width: 100%"]) {
+                        height: auto;
+                        max-width: 512px;
+                    }
+                    div[data-testid="stImage"] img[style*="max-width: 100%"] {
+                        border-radius: 8px;
+                    }
+                </style>
+                """)
+                st.image(message["content"])
+
+# button row
+if st.session_state.txt2img_messages:
+    button_container = st.empty()
+    with button_container.container():
+        # https://discuss.streamlit.io/t/st-button-in-one-line/25966/6
+        st.html("""
+        <style>
+            div[data-testid="column"] {
+                width: fit-content;
+                min-width: 0;
+                flex: none;
+            }
+        </style>
+        """)
+
+        # retry
+        col1, col2 = st.columns(2)
+        with col1:
+            if st.button("❌", help="Delete last generation") and len(st.session_state.txt2img_messages) >= 2:
+                st.session_state.txt2img_messages.pop()
+                st.session_state.txt2img_messages.pop()
+                st.rerun()
+
+        with col2:
+            if st.button("🗑️", help="Clear all generations"):
+                st.session_state.txt2img_messages = []
+                st.session_state.txt2img_seed = 0
+                st.rerun()
+else:
+    button_container = None
+
+# show the prompt and spinner while loading then update state and re-render
+if prompt := st.chat_input("What do you want to see?"):
+    if "seed" in parameters and parameters["seed"] >= 0:
+        st.session_state.txt2img_seed = parameters["seed"]
+    else:
+        st.session_state.txt2img_seed = int(datetime.now().timestamp() * 1e6) % (1 << 53)
+        if "seed" in parameters:
+            parameters["seed"] = st.session_state.txt2img_seed
+
+    if button_container:
+        button_container.empty()
+
+    with st.chat_message("user"):
+        st.markdown(prompt)
+
+    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"])
+            image = generate_image(**generate_kwargs)
+
+    model_name = PRESET_MODEL[model]["name"]
+    st.session_state.txt2img_messages.append(
+        {"role": "user", "content": prompt, "parameters": parameters, "model": model_name}
+    )
+    st.session_state.txt2img_messages.append({"role": "assistant", "content": image})
+    st.rerun()