Add Hugging Face txt2img client

lib/__init__.py

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

lib/api.py

@@ -0,0 +1,39 @@
+import io
+from abc import ABC, abstractmethod
+
+import requests
+from PIL import Image
+
+
+class Txt2ImgAPI(ABC):
+    @abstractmethod
+    def generate_image(self, model, prompt, parameters, **kwargs):
+        pass
+
+
+# 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)

lib/config.py

@@ -4,7 +4,7 @@ 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_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=[
         "black-forest-labs/flux.1-dev",
@@ -13,11 +13,11 @@ Config = SimpleNamespace(
     ],
     TXT2IMG_DEFAULT_AR="1:1",
     TXT2IMG_AR={
-        "9:7": (1152, 896),
         "7:4": (1344, 768),
+        "9:7": (1152, 896),
         "1:1": (1024, 1024),
-        "4:7": (768, 1344),
         "7:9": (896, 1152),
+        "4:7": (768, 1344),
     },
     TXT2TXT_DEFAULT_MODEL=4,
     TXT2TXT_MODELS=[

pages/2_🎨_Text_to_Image.py

@@ -1,12 +1,9 @@
-import io
 import os
 from datetime import datetime
 
-import requests
 import streamlit as st
-from PIL import Image
 
-from lib import Config, Presets
+from lib import Config, HuggingFaceTxt2ImgAPI, Presets
 
 # TODO: key input and store in cache_data
 # TODO: API dropdown; changes available models
@@ -20,22 +17,9 @@ PRESET_MODEL = {
 }
 
 
-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
+@st.cache_resource
+def get_txt2img_api():
+    return HuggingFaceTxt2ImgAPI(HF_TOKEN)
 
 
 # config
@@ -46,6 +30,9 @@ st.set_page_config(
 )
 
 # initialize state
+if "txt2img_running" not in st.session_state:
+    st.session_state.txt2img_running = False
+
 if "txt2img_messages" not in st.session_state:
     st.session_state.txt2img_messages = []
 
@@ -60,11 +47,13 @@ model = st.sidebar.selectbox(
     format_func=lambda x: x.split("/")[1],
     options=Config.TXT2IMG_MODELS,
     index=Config.TXT2IMG_DEFAULT_MODEL,
+    disabled=st.session_state.txt2img_running,
 )
 aspect_ratio = st.sidebar.select_slider(
     "Aspect Ratio",
     options=list(Config.TXT2IMG_AR.keys()),
     value=Config.TXT2IMG_DEFAULT_AR,
+    disabled=st.session_state.txt2img_running,
 )
 
 # heading
@@ -88,6 +77,7 @@ for param in preset["parameters"]:
             preset["guidance_scale_max"],
             preset["guidance_scale"],
             0.1,
+            disabled=st.session_state.txt2img_running,
         )
     if param == "num_inference_steps":
         parameters[param] = st.sidebar.slider(
@@ -96,6 +86,7 @@ for param in preset["parameters"]:
             preset["num_inference_steps_max"],
             preset["num_inference_steps"],
             1,
+            disabled=st.session_state.txt2img_running,
         )
     if param == "seed":
         parameters[param] = st.sidebar.number_input(
@@ -103,11 +94,13 @@ for param in preset["parameters"]:
             min_value=-1,
             max_value=(1 << 53) - 1,
             value=-1,
+            disabled=st.session_state.txt2img_running,
         )
     if param == "negative_prompt":
         parameters[param] = st.sidebar.text_area(
             label="Negative Prompt",
             value=Config.TXT2IMG_NEGATIVE_PROMPT,
+            disabled=st.session_state.txt2img_running,
         )
 
 # wrap the prompt in an expander to display additional parameters
@@ -142,7 +135,7 @@ for message in st.session_state.txt2img_messages:
                     }
                 </style>
                 """)
-                st.image(message["content"])
+                st.write(message["content"])  # success will be image, error will be text
 
 # button row
 if st.session_state.txt2img_messages:
@@ -162,13 +155,16 @@ if st.session_state.txt2img_messages:
         # retry
         col1, col2 = st.columns(2)
         with col1:
-            if st.button("❌", help="Delete last generation") and len(st.session_state.txt2img_messages) >= 2:
+            if (
+                st.button("❌", help="Delete last generation", disabled=st.session_state.txt2img_running)
+                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"):
+            if st.button("🗑️", help="Clear all generations", disabled=st.session_state.txt2img_running):
                 st.session_state.txt2img_messages = []
                 st.session_state.txt2img_seed = 0
                 st.rerun()
@@ -176,7 +172,10 @@ 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 prompt := st.chat_input(
+    "What do you want to see?",
+    on_submit=lambda: setattr(st.session_state, "txt2img_running", True),
+):
     if "seed" in parameters and parameters["seed"] >= 0:
         st.session_state.txt2img_seed = parameters["seed"]
     else:
@@ -195,7 +194,9 @@ if prompt := st.chat_input("What do you want to see?"):
             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)
+            api = get_txt2img_api()
+            image = api.generate_image(**generate_kwargs)
+        st.session_state.txt2img_running = False
 
     model_name = PRESET_MODEL[model]["name"]
     st.session_state.txt2img_messages.append(