Merge pull request #1 from russellballestrini/settings
streamlit settings
This commit is contained in:
GitHub
commit
d50e2f3608
1 changed files with 143 additions and 66 deletions
|
|
@ -9,6 +9,8 @@ from datetime import datetime
|
|||
import http.client
|
||||
import time
|
||||
import re
|
||||
import random
|
||||
|
||||
|
||||
def create_slug(prompt, max_words=10):
|
||||
words = prompt.split()[:max_words]
|
||||
|
|
@ -17,33 +19,34 @@ def create_slug(prompt, max_words=10):
|
|||
timestamp = int(time.time())
|
||||
return f"{slug}-{timestamp}"
|
||||
|
||||
|
||||
def poll_for_result(conn, headers, request_id):
|
||||
"""Poll the BFL server for the generation result."""
|
||||
while True:
|
||||
conn.request("GET", f"/v1/get_result?id={request_id}", headers=headers)
|
||||
res = conn.getresponse()
|
||||
data = res.read()
|
||||
response = json.loads(data.decode("utf-8"))
|
||||
|
||||
if 'status' in response:
|
||||
if response['status'] == 'Ready':
|
||||
|
||||
if "status" in response:
|
||||
if response["status"] == "Ready":
|
||||
return response
|
||||
elif response['status'] == 'Failed':
|
||||
elif response["status"] == "Failed":
|
||||
raise Exception("Image generation failed")
|
||||
elif response['status'] == 'Pending':
|
||||
elif response["status"] == "Pending":
|
||||
st.text("Still processing...")
|
||||
else:
|
||||
st.text(f"Unknown status: {response['status']}")
|
||||
else:
|
||||
st.text("Unexpected response structure")
|
||||
|
||||
|
||||
time.sleep(5)
|
||||
|
||||
def generate_image(prompt, api_key):
|
||||
|
||||
def generate_image(prompt, api_key, endpoint, seed):
|
||||
"""Send the generation request to the chosen BFL endpoint with the provided seed."""
|
||||
conn = http.client.HTTPSConnection("api.bfl.ml")
|
||||
headers = {
|
||||
'Content-Type': "application/json",
|
||||
'X-Key': api_key
|
||||
}
|
||||
headers = {"Content-Type": "application/json", "X-Key": api_key}
|
||||
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
|
|
@ -53,104 +56,178 @@ def generate_image(prompt, api_key):
|
|||
"top_k": 40,
|
||||
"repetition_penalty": 1.1,
|
||||
"stop": ["\n\n"],
|
||||
"seed": 42
|
||||
"seed": seed, # Pass the user-chosen (or random) seed
|
||||
}
|
||||
|
||||
conn.request("POST", "/v1/flux-dev", body=json.dumps(payload), headers=headers)
|
||||
# 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"Image generation request ID: {request_id}")
|
||||
request_id = response_data["id"]
|
||||
st.text(f"Generation request ID: {request_id}")
|
||||
|
||||
result = poll_for_result(conn, headers, request_id)
|
||||
|
||||
if 'result' in result and 'sample' in result['result']:
|
||||
return result['result']['sample']
|
||||
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
|
||||
|
||||
st.title('Black Forest Labs Image Generation')
|
||||
|
||||
# -------------------- Streamlit App --------------------
|
||||
st.title("Black Forest Labs Image Generation")
|
||||
|
||||
# Track whether we're deployed or running locally
|
||||
deployed = False
|
||||
|
||||
try:
|
||||
if hasattr(st, 'secrets') and st.secrets:
|
||||
if hasattr(st, "secrets") and st.secrets:
|
||||
deployed = True
|
||||
api_key = st.secrets.get("BLACK_FOREST_LABS_API_KEY")
|
||||
else:
|
||||
api_key = os.environ.get("BLACK_FOREST_LABS_API_KEY")
|
||||
except FileNotFoundError:
|
||||
# If secrets file is not found, assume we're running locally
|
||||
api_key = os.environ.get("BLACK_FOREST_LABS_API_KEY")
|
||||
|
||||
if not api_key:
|
||||
st.error("BLACK_FOREST_LABS_API_KEY is not set in secrets or environment variables. Please set it before running the app.")
|
||||
st.error(
|
||||
"BLACK_FOREST_LABS_API_KEY is not set in secrets or environment variables. Please set it before running the app."
|
||||
)
|
||||
st.stop()
|
||||
|
||||
prompt = st.text_input('Enter your image prompt:')
|
||||
# -- Initialize session state --
|
||||
if "generated_images" not in st.session_state:
|
||||
st.session_state.generated_images = []
|
||||
|
||||
if st.button('Generate Image'):
|
||||
# -----------------------------------------------------------------
|
||||
# Load previously generated images from the local DB (if any).
|
||||
# We show newest images first, so we ORDER BY timestamp DESC in the query.
|
||||
# -----------------------------------------------------------------
|
||||
if not deployed:
|
||||
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)"""
|
||||
)
|
||||
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
|
||||
if len(st.session_state.generated_images) == 0:
|
||||
for slug, prompt, filename, base64_image, _ in rows:
|
||||
image_data = base64.b64decode(base64_image)
|
||||
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))
|
||||
|
||||
# ---------------------------------------------------------
|
||||
# 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
|
||||
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}")
|
||||
else:
|
||||
seed_value = st.sidebar.number_input(
|
||||
"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 = st.text_input("Enter your image prompt:")
|
||||
|
||||
if st.button("Generate Image"):
|
||||
if prompt:
|
||||
image_url = generate_image(prompt, api_key)
|
||||
image_url = generate_image(prompt, api_key, selected_model, seed_value)
|
||||
if image_url:
|
||||
response = requests.get(image_url)
|
||||
if response.status_code == 200:
|
||||
# Create a slug for the filename
|
||||
slug = create_slug(prompt)
|
||||
filename = f"{slug}.jpg"
|
||||
|
||||
if not deployed:
|
||||
# If local, save to 'images' directory and SQLite
|
||||
os.makedirs('images', exist_ok=True)
|
||||
with open(os.path.join('images', filename), 'wb') as f:
|
||||
# Local environment: save to 'images' directory and SQLite
|
||||
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
|
||||
conn = sqlite3.connect('image_metadata.db')
|
||||
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 (?, ?, ?, ?, ?)',
|
||||
(slug, prompt, filename, base64_image, datetime.now()))
|
||||
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 (?, ?, ?, ?, ?)",
|
||||
(slug, prompt, filename, base64_image, datetime.now()),
|
||||
)
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
st.image(os.path.join('images', filename), caption='Generated Image')
|
||||
st.success(f"Image saved locally as '{filename}' and metadata stored in SQLite.")
|
||||
else:
|
||||
# If deployed, use a temporary file
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as tmp_file:
|
||||
tmp_file.write(response.content)
|
||||
st.image(tmp_file.name, caption='Generated Image')
|
||||
|
||||
# Provide download button
|
||||
with open(tmp_file.name, "rb") as file:
|
||||
btn = st.download_button(
|
||||
label="Download Image",
|
||||
data=file,
|
||||
file_name=filename,
|
||||
mime="image/jpeg"
|
||||
)
|
||||
# Clean up the temporary file
|
||||
os.unlink(tmp_file.name)
|
||||
else:
|
||||
st.error(f"Failed to download image. Status code: {response.status_code}")
|
||||
else:
|
||||
st.error('Please enter a prompt')
|
||||
# Insert the new image at the top so it's displayed first
|
||||
st.session_state.generated_images.insert(
|
||||
0, (filepath, filename, prompt)
|
||||
)
|
||||
|
||||
# Display all generated images if running locally
|
||||
if not deployed:
|
||||
st.subheader('Generated Images')
|
||||
image_dir = 'images'
|
||||
if os.path.exists(image_dir) and os.listdir(image_dir):
|
||||
for image in sorted(os.listdir(image_dir), key=lambda f: os.path.getctime(os.path.join(image_dir, f)), reverse=True):
|
||||
st.image(os.path.join(image_dir, image), caption=image, use_column_width=True)
|
||||
st.success(
|
||||
f"Image saved locally as '{filename}' and metadata stored in SQLite."
|
||||
)
|
||||
else:
|
||||
# Deployed environment: use a temporary file
|
||||
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}"
|
||||
)
|
||||
else:
|
||||
st.info("No images have been generated yet.")
|
||||
st.error("Please enter a prompt")
|
||||
|
||||
# --------------------------------------------------
|
||||
# Display images: newest first
|
||||
# --------------------------------------------------
|
||||
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}",
|
||||
data=file,
|
||||
file_name=filename,
|
||||
mime="image/jpeg",
|
||||
)
|
||||
|
||||
# Add a note about setting the API key
|
||||
st.sidebar.info(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue