engineering/art.ai.unturf.com
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

modified: black_forest_streamlit.py

Browse source
This commit is contained in:
Russell Ballestrini 2025-03-20 20:08:43 -04:00
parent dd80435e0b
commit 62c7d43367
1 changed files with 62 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

83
black_forest_streamlit.py
View file

@ -12,12 +12,19 @@ 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 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)
@ -25,6 +32,7 @@ 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"
@ -34,6 +42,7 @@ class User(Base):
otp_expiry = Column(DateTime)
namespaces = relationship("Namespace", back_populates="user")
class Namespace(Base):
__tablename__ = "namespaces"
id = Column(Integer, primary_key=True)
@ -43,6 +52,7 @@ class Namespace(Base):
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)
@ -54,12 +64,14 @@ class Image(Base):
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]
@ -68,21 +80,25 @@ 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')
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:
msg["Subject"] = "Your Image Generation OTP"
msg["From"] = "art-no-reply@master.unturf.com"
msg["To"] = email
with smtplib.SMTP("localhost", 25) as smtp:
smtp.send_message(msg)
def poll_for_result(conn, headers, request_id):
while True:
conn.request("GET", f"/v1/get_result?id={request_id}", headers=headers)
@ -99,6 +115,7 @@ def poll_for_result(conn, headers, request_id):
raise Exception(f"Unknown status: {response.get('status', 'N/A')}")
time.sleep(5)
def generate_image(prompt, api_key, endpoint, seed):
conn = http.client.HTTPSConnection("api.bfl.ml")
headers = {"Content-Type": "application/json", "X-Key": api_key}
@ -123,6 +140,7 @@ def generate_image(prompt, api_key, endpoint, seed):
return result["result"]["sample"]
raise Exception("No image URL found.")
# --- Streamlit App ---
st.title("Black Forest Labs Image Generation")
@ -178,7 +196,9 @@ if not st.session_state.user_id:
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')):
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()
@ -194,16 +214,22 @@ 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()
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()))
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)
namespace = Namespace(
user_id=st.session_state.user_id, name=new_namespace
)
session.add(namespace)
try:
session.commit()
@ -218,25 +244,36 @@ else:
# 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()
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))
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)
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)
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:
@ -258,7 +295,7 @@ else:
slug=slug,
prompt=prompt,
filename=filename,
base64_image=base64_image
base64_image=base64_image,
)
session.add(image)
session.commit()
@ -266,10 +303,14 @@ else:
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}")
st.error(
f"Failed to download image. HTTP status code: {response.status_code}"
)
except Exception as e:
st.error(str(e))
elif generate_submitted:
@ -297,4 +338,4 @@ st.sidebar.info(
if deployed:
st.write("Running in deployed environment")
else:
st.write("Running locally")
st.write("Running locally")