Update black_forest_streamlit.py
This commit is contained in:
parent
e6613977d8
commit
86ca4e878a
1 changed files with 217 additions and 167 deletions
|
|
@ -3,15 +3,63 @@ import os
|
|||
import requests
|
||||
import json
|
||||
import base64
|
||||
import sqlite3
|
||||
import tempfile
|
||||
from datetime import datetime
|
||||
from datetime import datetime, timedelta
|
||||
import http.client
|
||||
import time
|
||||
import re
|
||||
import random
|
||||
import smtplib
|
||||
from email.mime.text import MIMEText
|
||||
import bcrypt
|
||||
from sqlalchemy import create_engine, Column, Integer, String, ForeignKey, DateTime, UniqueConstraint
|
||||
from sqlalchemy.ext.declarative import declarative_base
|
||||
from sqlalchemy.orm import sessionmaker, relationship
|
||||
from sqlalchemy.sql import func
|
||||
|
||||
# --- Database Setup ---
|
||||
DATA_DIR = "data"
|
||||
os.makedirs(DATA_DIR, exist_ok=True)
|
||||
DB_PATH = f"sqlite:///{os.path.join(DATA_DIR, 'image_metadata.db')}"
|
||||
engine = create_engine(DB_PATH, echo=False)
|
||||
Base = declarative_base()
|
||||
|
||||
# Define Models
|
||||
class User(Base):
|
||||
__tablename__ = "users"
|
||||
id = Column(Integer, primary_key=True)
|
||||
email = Column(String, unique=True, nullable=False)
|
||||
otp_hash = Column(String)
|
||||
otp_expiry = Column(DateTime)
|
||||
namespaces = relationship("Namespace", back_populates="user")
|
||||
|
||||
class Namespace(Base):
|
||||
__tablename__ = "namespaces"
|
||||
id = Column(Integer, primary_key=True)
|
||||
user_id = Column(Integer, ForeignKey("users.id"), nullable=False)
|
||||
name = Column(String, nullable=False)
|
||||
user = relationship("User", back_populates="namespaces")
|
||||
images = relationship("Image", back_populates="namespace")
|
||||
__table_args__ = (UniqueConstraint("user_id", "name", name="_user_namespace_uc"),)
|
||||
|
||||
class Image(Base):
|
||||
__tablename__ = "images"
|
||||
id = Column(Integer, primary_key=True)
|
||||
namespace_id = Column(Integer, ForeignKey("namespaces.id"), nullable=False)
|
||||
slug = Column(String)
|
||||
prompt = Column(String)
|
||||
filename = Column(String)
|
||||
base64_image = Column(String)
|
||||
timestamp = Column(DateTime, default=func.now())
|
||||
namespace = relationship("Namespace", back_populates="images")
|
||||
|
||||
# Create tables
|
||||
Base.metadata.create_all(engine)
|
||||
|
||||
# Session factory
|
||||
Session = sessionmaker(bind=engine)
|
||||
|
||||
# --- Utility Functions ---
|
||||
def create_slug(prompt, max_words=10):
|
||||
words = prompt.split()[:max_words]
|
||||
slug = "-".join(words)
|
||||
|
|
@ -19,36 +67,40 @@ def create_slug(prompt, max_words=10):
|
|||
timestamp = int(time.time())
|
||||
return f"{slug}-{timestamp}"
|
||||
|
||||
def generate_otp():
|
||||
return str(random.randint(100000, 999999))
|
||||
|
||||
def hash_otp(otp):
|
||||
return bcrypt.hashpw(otp.encode('utf-8'), bcrypt.gensalt()).decode('utf-8')
|
||||
|
||||
def send_otp_email(email, otp):
|
||||
msg = MIMEText(f"Your OTP is: {otp}\nIt expires in 10 minutes.")
|
||||
msg['Subject'] = 'Your Image Generation OTP'
|
||||
msg['From'] = 'no-reply@art.ai.unturf.com'
|
||||
msg['To'] = email
|
||||
|
||||
with smtplib.SMTP('localhost', 25) as smtp:
|
||||
smtp.send_message(msg)
|
||||
|
||||
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":
|
||||
return response
|
||||
elif response["status"] == "Failed":
|
||||
raise Exception("Image generation failed")
|
||||
elif response["status"] == "Pending":
|
||||
# Instead of spamming top real estate, let's keep minimal text or logs here
|
||||
pass
|
||||
else:
|
||||
raise Exception(f"Unknown status: {response['status']}")
|
||||
if response.get("status") == "Ready":
|
||||
return response
|
||||
elif response.get("status") == "Failed":
|
||||
raise Exception("Image generation failed")
|
||||
elif response.get("status") == "Pending":
|
||||
pass
|
||||
else:
|
||||
raise Exception("Unexpected response structure from poll_for_result.")
|
||||
|
||||
raise Exception(f"Unknown status: {response.get('status', 'N/A')}")
|
||||
time.sleep(5)
|
||||
|
||||
|
||||
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}
|
||||
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"max_tokens": 512,
|
||||
|
|
@ -59,27 +111,21 @@ def generate_image(prompt, api_key, endpoint, seed):
|
|||
"stop": ["\n\n"],
|
||||
"seed": seed,
|
||||
}
|
||||
|
||||
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.get("id")
|
||||
data = json.loads(res.read().decode("utf-8"))
|
||||
request_id = data.get("id")
|
||||
if not request_id:
|
||||
raise Exception(f"Failed to get request_id from response: {response_data}")
|
||||
|
||||
raise Exception(f"Failed to get request_id: {data}")
|
||||
result = poll_for_result(conn, headers, request_id)
|
||||
if "result" in result and "sample" in result["result"]:
|
||||
return result["result"]["sample"]
|
||||
else:
|
||||
raise Exception("No image URL found in the final response.")
|
||||
raise Exception("No image URL found.")
|
||||
|
||||
|
||||
# -------------------- Streamlit App --------------------
|
||||
# --- Streamlit App ---
|
||||
st.title("Black Forest Labs Image Generation")
|
||||
|
||||
# Check environment vs. secrets for API key
|
||||
# API Key Setup
|
||||
deployed = False
|
||||
try:
|
||||
if hasattr(st, "secrets") and st.secrets:
|
||||
|
|
@ -91,159 +137,163 @@ except FileNotFoundError:
|
|||
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.")
|
||||
st.stop()
|
||||
|
||||
# -- Initialize session state --
|
||||
# --- Session State Initialization ---
|
||||
if "user_id" not in st.session_state:
|
||||
st.session_state.user_id = None
|
||||
if "namespace_id" not in st.session_state:
|
||||
st.session_state.namespace_id = None
|
||||
if "generated_images" not in st.session_state:
|
||||
st.session_state.generated_images = []
|
||||
|
||||
# -----------------------------------------------------------------
|
||||
# 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")
|
||||
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()
|
||||
|
||||
# Only load these once, if the session is fresh (no images in session state yet)
|
||||
if len(st.session_state.generated_images) == 0:
|
||||
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 *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
|
||||
|
||||
# ---------------------------------------------------------
|
||||
# Sidebar UI: Model Selection, Seed selection (random or fixed)
|
||||
# ---------------------------------------------------------
|
||||
st.sidebar.subheader("BFL Model Settings")
|
||||
|
||||
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)
|
||||
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
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------
|
||||
# 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
|
||||
|
||||
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 file + metadata
|
||||
os.makedirs("images", exist_ok=True)
|
||||
filepath = os.path.join("images", filename)
|
||||
with open(filepath, "wb") as f:
|
||||
f.write(response.content)
|
||||
|
||||
# Store in SQLite
|
||||
conn = sqlite3.connect("image_metadata.db")
|
||||
cursor = conn.cursor()
|
||||
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()
|
||||
|
||||
# Insert at top of session state
|
||||
st.session_state.generated_images.insert(
|
||||
0, (filepath, filename, prompt)
|
||||
)
|
||||
|
||||
# --- Authentication ---
|
||||
if not st.session_state.user_id:
|
||||
st.subheader("Login with OTP")
|
||||
email = st.text_input("Enter your email:")
|
||||
if st.button("Send OTP"):
|
||||
if email:
|
||||
otp = generate_otp()
|
||||
otp_hash = hash_otp(otp)
|
||||
expiry = datetime.now().replace(microsecond=0) + timedelta(minutes=10)
|
||||
with Session() as session:
|
||||
user = session.query(User).filter_by(email=email).first()
|
||||
if user:
|
||||
user.otp_hash = otp_hash
|
||||
user.otp_expiry = expiry
|
||||
else:
|
||||
# Deployed environment
|
||||
user = User(email=email, otp_hash=otp_hash, otp_expiry=expiry)
|
||||
session.add(user)
|
||||
session.commit()
|
||||
send_otp_email(email, otp)
|
||||
st.success("OTP sent to your email!")
|
||||
else:
|
||||
st.error("Please enter an email.")
|
||||
|
||||
otp_input = st.text_input("Enter OTP:", type="password")
|
||||
if st.button("Verify OTP"):
|
||||
if otp_input and email:
|
||||
with Session() as session:
|
||||
user = session.query(User).filter_by(email=email).first()
|
||||
if user and user.otp_expiry and datetime.now() < user.otp_expiry:
|
||||
if bcrypt.checkpw(otp_input.encode('utf-8'), user.otp_hash.encode('utf-8')):
|
||||
st.session_state.user_id = user.id
|
||||
st.success("Logged in successfully!")
|
||||
st.rerun()
|
||||
else:
|
||||
st.error("Invalid OTP.")
|
||||
elif user and user.otp_expiry:
|
||||
st.error("OTP expired.")
|
||||
else:
|
||||
st.error("User not found.")
|
||||
else:
|
||||
st.error("Please enter both email and OTP.")
|
||||
else:
|
||||
# --- Namespace Management ---
|
||||
st.subheader("Your Albums")
|
||||
with Session() as session:
|
||||
namespaces = session.query(Namespace).filter_by(user_id=st.session_state.user_id).all()
|
||||
namespace_options = {ns.name: ns.id for ns in namespaces}
|
||||
namespace_options["Create New Album"] = None
|
||||
|
||||
selected_namespace = st.selectbox("Select or create an album:", list(namespace_options.keys()))
|
||||
if selected_namespace == "Create New Album":
|
||||
new_namespace = st.text_input("New album name:")
|
||||
if st.button("Create Album"):
|
||||
if new_namespace:
|
||||
namespace = Namespace(user_id=st.session_state.user_id, name=new_namespace)
|
||||
session.add(namespace)
|
||||
try:
|
||||
session.commit()
|
||||
st.success(f"Album '{new_namespace}' created!")
|
||||
st.rerun()
|
||||
except:
|
||||
session.rollback()
|
||||
st.error("Album name already exists for this user.")
|
||||
else:
|
||||
st.session_state.namespace_id = namespace_options[selected_namespace]
|
||||
|
||||
# Load images for the selected namespace
|
||||
if st.session_state.namespace_id:
|
||||
st.session_state.generated_images = []
|
||||
images = session.query(Image).filter_by(namespace_id=st.session_state.namespace_id).order_by(Image.timestamp.desc()).all()
|
||||
for img in images:
|
||||
image_data = base64.b64decode(img.base64_image)
|
||||
temp_dir = tempfile.mkdtemp()
|
||||
filepath = os.path.join(temp_dir, img.filename)
|
||||
with open(filepath, "wb") as f:
|
||||
f.write(image_data)
|
||||
st.session_state.generated_images.append((filepath, img.filename, img.prompt))
|
||||
|
||||
# --- Sidebar Settings ---
|
||||
st.sidebar.subheader("BFL Model Settings")
|
||||
model_options = ["flux-dev", "flux-pro-1.1-ultra", "flux-pro-1.1"]
|
||||
selected_model = st.sidebar.selectbox("Choose a model endpoint", model_options, index=0)
|
||||
use_random_seed = st.sidebar.checkbox("Use random seed?", value=False)
|
||||
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)
|
||||
|
||||
# --- Prompt + Generate ---
|
||||
if st.session_state.namespace_id:
|
||||
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")
|
||||
|
||||
if generate_submitted and 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"
|
||||
base64_image = base64.b64encode(response.content).decode("utf-8")
|
||||
with Session() as session:
|
||||
image = Image(
|
||||
namespace_id=st.session_state.namespace_id,
|
||||
slug=slug,
|
||||
prompt=prompt,
|
||||
filename=filename,
|
||||
base64_image=base64_image
|
||||
)
|
||||
session.add(image)
|
||||
session.commit()
|
||||
temp_dir = tempfile.mkdtemp()
|
||||
filepath = os.path.join(temp_dir, filename)
|
||||
with open(filepath, "wb") as f:
|
||||
f.write(response.content)
|
||||
st.session_state.generated_images.insert(
|
||||
0, (filepath, filename, prompt)
|
||||
)
|
||||
st.session_state.generated_images.insert(0, (filepath, filename, prompt))
|
||||
st.success("Image generated successfully!")
|
||||
else:
|
||||
st.error(f"Failed to download image. HTTP status code: {response.status_code}")
|
||||
except Exception as e:
|
||||
st.error(str(e))
|
||||
elif generate_submitted:
|
||||
st.error("Please enter a non-empty prompt.")
|
||||
|
||||
st.success("Image generated successfully!")
|
||||
else:
|
||||
error_message = f"Failed to download image. HTTP status code: {response.status_code}"
|
||||
except Exception as e:
|
||||
error_message = str(e)
|
||||
else:
|
||||
error_message = "Please enter a non-empty prompt."
|
||||
# --- Display Images ---
|
||||
if st.session_state.namespace_id:
|
||||
st.subheader(f"Images in Album (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",
|
||||
)
|
||||
|
||||
# --------------------------------------------------
|
||||
# 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",
|
||||
)
|
||||
|
||||
# --------------------------------------------------
|
||||
# Show any error messages at the bottom
|
||||
# --------------------------------------------------
|
||||
if error_message:
|
||||
st.error(error_message)
|
||||
|
||||
# --------------------------------------------------
|
||||
# Sidebar info for local environment usage
|
||||
# --------------------------------------------------
|
||||
# --- Sidebar Info ---
|
||||
st.sidebar.info(
|
||||
"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"
|
||||
"If deploying, set the API key in Streamlit secrets."
|
||||
)
|
||||
|
||||
if deployed:
|
||||
st.write("Running in deployed environment")
|
||||
else:
|
||||
st.write("Running locally")
|
||||
st.write("Running locally")
|
||||
Loading…
Add table
Add a link
Reference in a new issue