fix up #2
1 changed files with 63 additions and 58 deletions
|
|
@ -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."
|
||||
)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue