From 189ef6a01eba1b4a5ef73d479da436a8c3a6fa7a Mon Sep 17 00:00:00 2001 From: Russell Ballestrini Date: Tue, 17 Dec 2024 10:51:11 -0500 Subject: [PATCH] fix up --- black_forest_streamlit.py | 121 ++++++++++++++++++++------------------ 1 file changed, 63 insertions(+), 58 deletions(-) diff --git a/black_forest_streamlit.py b/black_forest_streamlit.py index ecc8b03..53a57f0 100644 --- a/black_forest_streamlit.py +++ b/black_forest_streamlit.py @@ -34,11 +34,12 @@ def poll_for_result(conn, headers, request_id): elif response["status"] == "Failed": raise Exception("Image generation failed") elif response["status"] == "Pending": - st.text("Still processing...") + # Instead of spamming top real estate, let's keep minimal text or logs here + pass else: - st.text(f"Unknown status: {response['status']}") + raise Exception(f"Unknown status: {response['status']}") else: - st.text("Unexpected response structure") + raise Exception("Unexpected response structure from poll_for_result.") time.sleep(5) @@ -56,31 +57,29 @@ def generate_image(prompt, api_key, endpoint, seed): "top_k": 40, "repetition_penalty": 1.1, "stop": ["\n\n"], - "seed": seed, # Pass the user-chosen (or random) seed + "seed": seed, } - # Use the chosen endpoint (e.g., "flux-dev", "flux-pro-1.1-ultra", etc.) conn.request("POST", f"/v1/{endpoint}", body=json.dumps(payload), headers=headers) res = conn.getresponse() data = res.read() response_data = json.loads(data.decode("utf-8")) - request_id = response_data["id"] - st.text(f"Generation request ID: {request_id}") + request_id = response_data.get("id") + if not request_id: + raise Exception(f"Failed to get request_id from response: {response_data}") result = poll_for_result(conn, headers, request_id) - if "result" in result and "sample" in result["result"]: return result["result"]["sample"] else: - st.error("No image URL found in the response") - return None + raise Exception("No image URL found in the final response.") # -------------------- Streamlit App -------------------- st.title("Black Forest Labs Image Generation") -# Track whether we're deployed or running locally +# Check environment vs. secrets for API key deployed = False try: if hasattr(st, "secrets") and st.secrets: @@ -102,8 +101,8 @@ if "generated_images" not in st.session_state: st.session_state.generated_images = [] # ----------------------------------------------------------------- -# Load previously generated images from the local DB (if any). -# We show newest images first, so we ORDER BY timestamp DESC in the query. +# Load previously generated images from the local DB (if any), only if running locally. +# We ORDER BY timestamp DESC to get newest first. # ----------------------------------------------------------------- if not deployed: conn = sqlite3.connect("image_metadata.db") @@ -114,34 +113,40 @@ if not deployed: ) conn.commit() - rows = cursor.execute( - "SELECT slug, prompt, filename, base64_image, timestamp FROM images ORDER BY timestamp DESC" - ).fetchall() - conn.close() - - # If the session just started (no images in session state), load images from DB + # Only load these once, if the session is fresh (no images in session state yet) if len(st.session_state.generated_images) == 0: - for slug, prompt, filename, base64_image, _ in rows: - image_data = base64.b64decode(base64_image) + rows = cursor.execute( + "SELECT slug, prompt, filename, base64_image, timestamp FROM images ORDER BY timestamp DESC" + ).fetchall() + conn.close() + + for slug, prompt_text, filename, base64_img_str, _ in rows: + image_data = base64.b64decode(base64_img_str) temp_dir = tempfile.mkdtemp() filepath = os.path.join(temp_dir, filename) with open(filepath, "wb") as f: f.write(image_data) - # Insert at the bottom of the list if we're iterating in DESC order - # But we want newest at top, so actually we can just append in the order they come in - st.session_state.generated_images.append((filepath, filename, prompt)) + # Insert at the *end* of the list if we want to preserve the DESC from DB, + # but we want the newest first in final display. We'll just append in the loop + # and reverse the final list or insert(0). Let's keep it simple: + st.session_state.generated_images.append((filepath, filename, prompt_text)) + +# We'll store any errors in a variable and show them later (at the bottom) to save vertical space +error_message = None # --------------------------------------------------------- -# UI: Model Selection, Seed selection (random or fixed) +# Sidebar UI: Model Selection, Seed selection (random or fixed) # --------------------------------------------------------- st.sidebar.subheader("BFL Model Settings") -# Available model endpoints -model_options = ["flux-dev", "flux-pro-1.1-ultra", "flux-pro-1.1"] # default +model_options = [ + "flux-dev", # default + "flux-pro-1.1-ultra", # more expensive + "flux-pro-1.1", # another variant +] selected_model = st.sidebar.selectbox("Choose a model endpoint", model_options, index=0) use_random_seed = st.sidebar.checkbox("Use random seed?", value=False) -seed_value = 42 # default if use_random_seed: seed_value = random.randint(1, 9999999) st.sidebar.write(f"Random seed chosen: {seed_value}") @@ -150,34 +155,34 @@ else: "Set a specific seed", value=42, min_value=0, max_value=99999999, step=1 ) -st.write("**Current model endpoint:**", selected_model) -st.write("**Current seed:**", seed_value) +# --------------------------------------------------------- +# Prompt + Generate (with Enter key to submit) +# --------------------------------------------------------- +with st.form("prompt_form", clear_on_submit=False): + prompt = st.text_input("Enter your image prompt:") + generate_submitted = st.form_submit_button( + "Generate Image" + ) # Pressing Enter or the button triggers this -prompt = st.text_input("Enter your image prompt:") - -if st.button("Generate Image"): - if prompt: - image_url = generate_image(prompt, api_key, selected_model, seed_value) - if image_url: +if generate_submitted: + if prompt.strip(): + try: + image_url = generate_image(prompt, api_key, selected_model, seed_value) response = requests.get(image_url) if response.status_code == 200: slug = create_slug(prompt) filename = f"{slug}.jpg" if not deployed: - # Local environment: save to 'images' directory and SQLite + # Local environment: save file + metadata os.makedirs("images", exist_ok=True) filepath = os.path.join("images", filename) with open(filepath, "wb") as f: f.write(response.content) - # Save to SQLite + # Store in SQLite conn = sqlite3.connect("image_metadata.db") cursor = conn.cursor() - cursor.execute( - """CREATE TABLE IF NOT EXISTS images - (id INTEGER PRIMARY KEY, slug TEXT, prompt TEXT, filename TEXT, base64_image TEXT, timestamp DATETIME)""" - ) base64_image = base64.b64encode(response.content).decode("utf-8") cursor.execute( "INSERT INTO images (slug, prompt, filename, base64_image, timestamp) VALUES (?, ?, ?, ?, ?)", @@ -186,33 +191,28 @@ if st.button("Generate Image"): conn.commit() conn.close() - # Insert the new image at the top so it's displayed first + # Insert at top of session state st.session_state.generated_images.insert( 0, (filepath, filename, prompt) ) - st.success( - f"Image saved locally as '{filename}' and metadata stored in SQLite." - ) else: - # Deployed environment: use a temporary file + # Deployed environment temp_dir = tempfile.mkdtemp() filepath = os.path.join(temp_dir, filename) with open(filepath, "wb") as f: f.write(response.content) - - # Insert at top st.session_state.generated_images.insert( 0, (filepath, filename, prompt) ) st.success("Image generated successfully!") else: - st.error( - f"Failed to download image. Status code: {response.status_code}" - ) + error_message = f"Failed to download image. HTTP status code: {response.status_code}" + except Exception as e: + error_message = str(e) else: - st.error("Please enter a prompt") + error_message = "Please enter a non-empty prompt." # -------------------------------------------------- # Display images: newest first @@ -220,7 +220,6 @@ if st.button("Generate Image"): st.subheader("Generated Images (Newest First)") for filepath, filename, prompt_text in st.session_state.generated_images: st.image(filepath, caption=f"Prompt: {prompt_text}", use_column_width=True) - with open(filepath, "rb") as file: st.download_button( label=f"Download {filename}", @@ -229,12 +228,18 @@ for filepath, filename, prompt_text in st.session_state.generated_images: mime="image/jpeg", ) -# Add a note about setting the API key +# -------------------------------------------------- +# Show any error messages at the bottom +# -------------------------------------------------- +if error_message: + st.error(error_message) + +# -------------------------------------------------- +# Sidebar info for local environment usage +# -------------------------------------------------- st.sidebar.info( - "Make sure to set the BLACK_FOREST_LABS_API_KEY environment variable before running this app locally. " - "You can do this by running:\n\n" + "To run locally, set the BLACK_FOREST_LABS_API_KEY environment variable:\n\n" "export BLACK_FOREST_LABS_API_KEY='your_api_key_here'\n\n" - "Replace 'your_api_key_here' with your actual API key.\n\n" "If deploying, set the API key in Streamlit secrets." ) -- 2.49.1