552 lines
20 KiB
Python
552 lines
20 KiB
Python
import streamlit as st
|
||
import os
|
||
import requests
|
||
import json
|
||
import base64
|
||
import tempfile
|
||
from datetime import datetime, timedelta
|
||
import http.client
|
||
import time
|
||
import re
|
||
import random
|
||
import smtplib
|
||
from email.mime.text import MIMEText
|
||
import bcrypt
|
||
import socket # For socket error handling
|
||
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
|
||
from streamlit_cookies_manager import EncryptedCookieManager
|
||
import secrets
|
||
|
||
# --- Database and Secret Setup ---
|
||
DATA_DIR = "data"
|
||
os.makedirs(DATA_DIR, exist_ok=True) # Ensure data directory exists
|
||
MAIN_DB_PATH = (
|
||
f"sqlite:///{os.path.join(DATA_DIR, 'main.db')}" # Main DB for users and namespaces
|
||
)
|
||
SECRET_FILE = os.path.join(
|
||
DATA_DIR, "cookie_secret.txt"
|
||
) # File for cookie secret persistence
|
||
|
||
|
||
def get_or_create_cookie_secret():
|
||
"""Generate or load a secure cookie secret for session encryption."""
|
||
if os.path.exists(SECRET_FILE):
|
||
with open(SECRET_FILE, "r") as f:
|
||
return f.read().strip()
|
||
else:
|
||
secret = secrets.token_hex(32) # 32 bytes = 64 hex chars
|
||
with open(SECRET_FILE, "w") as f:
|
||
f.write(secret)
|
||
return secret
|
||
|
||
|
||
COOKIE_SECRET = get_or_create_cookie_secret()
|
||
|
||
# SQLAlchemy setup for main database (users and namespaces)
|
||
engine = create_engine(MAIN_DB_PATH, echo=False)
|
||
Base = declarative_base()
|
||
|
||
|
||
# Define Models for main.db
|
||
class User(Base):
|
||
"""User model for authentication and namespace association."""
|
||
|
||
__tablename__ = "users"
|
||
id = Column(Integer, primary_key=True)
|
||
email = Column(String, unique=True, nullable=False)
|
||
otp_hash = Column(String) # Hashed OTP for login
|
||
otp_expiry = Column(DateTime) # OTP expiration time
|
||
namespaces = relationship(
|
||
"Namespace", secondary="namespace_users", back_populates="users"
|
||
)
|
||
|
||
|
||
class Namespace(Base):
|
||
"""Namespace (album) model, shared by multiple users."""
|
||
|
||
__tablename__ = "namespaces"
|
||
id = Column(Integer, primary_key=True)
|
||
name = Column(String, nullable=False)
|
||
users = relationship(
|
||
"User", secondary="namespace_users", back_populates="namespaces"
|
||
)
|
||
|
||
|
||
class NamespaceUser(Base):
|
||
"""Many-to-many relationship between users and namespaces."""
|
||
|
||
__tablename__ = "namespace_users"
|
||
id = Column(Integer, primary_key=True)
|
||
user_id = Column(Integer, ForeignKey("users.id"), nullable=False)
|
||
namespace_id = Column(Integer, ForeignKey("namespaces.id"), nullable=False)
|
||
__table_args__ = (
|
||
UniqueConstraint("user_id", "namespace_id", name="_user_namespace_uc"),
|
||
)
|
||
|
||
|
||
# Define Image model for per-namespace DBs
|
||
class Image(Base):
|
||
"""Image model stored in per-namespace SQLite files."""
|
||
|
||
__tablename__ = "images"
|
||
id = Column(Integer, primary_key=True)
|
||
slug = Column(String) # Unique slug based on prompt
|
||
prompt = Column(String) # Original prompt text
|
||
filename = Column(String) # Generated filename
|
||
base64_image = Column(String) # Base64-encoded image data
|
||
generation_time = Column(DateTime, default=func.now()) # Time of generation
|
||
seed = Column(Integer) # Seed used for generation
|
||
model_endpoint = Column(String) # BFL model used (e.g., "flux-dev")
|
||
|
||
|
||
# Create tables in main.db
|
||
Base.metadata.create_all(engine)
|
||
|
||
# Session factory for main database
|
||
MainSession = sessionmaker(bind=engine)
|
||
|
||
|
||
# Function to get or create a namespace-specific database
|
||
def get_namespace_db(namespace_id):
|
||
"""Create or connect to a namespace-specific SQLite database."""
|
||
db_path = f"sqlite:///{os.path.join(DATA_DIR, f'namespace_{namespace_id}.db')}"
|
||
engine = create_engine(db_path, echo=False)
|
||
Image.__table__.create(
|
||
bind=engine, checkfirst=True
|
||
) # Create images table if not exists
|
||
return sessionmaker(bind=engine)
|
||
|
||
|
||
# --- Cookie Manager Setup ---
|
||
cookies = EncryptedCookieManager(password=COOKIE_SECRET)
|
||
if not cookies.ready():
|
||
st.stop() # Wait for cookies to initialize
|
||
|
||
|
||
# --- Utility Functions ---
|
||
def create_slug(prompt, max_words=10):
|
||
"""Generate a slug from the prompt with a timestamp."""
|
||
words = prompt.split()[:max_words]
|
||
slug = "-".join(words)
|
||
slug = re.sub(r"[^\w\-]", "", slug.lower())
|
||
timestamp = int(time.time())
|
||
return f"{slug}-{timestamp}"
|
||
|
||
|
||
def generate_otp():
|
||
"""Generate a 6-digit one-time password."""
|
||
return str(random.randint(100000, 999999))
|
||
|
||
|
||
def hash_otp(otp):
|
||
"""Hash an OTP using bcrypt."""
|
||
return bcrypt.hashpw(otp.encode("utf-8"), bcrypt.gensalt()).decode("utf-8")
|
||
|
||
|
||
def send_email(to_email, subject, body):
|
||
"""Send an email via SMTP, falling back to STDOUT logging on socket error."""
|
||
msg = MIMEText(body)
|
||
msg["Subject"] = subject
|
||
msg["From"] = "art-no-reply@master.unturf.com"
|
||
msg["To"] = to_email
|
||
try:
|
||
with smtplib.SMTP("localhost", 25) as smtp:
|
||
smtp.send_message(msg)
|
||
except (socket.error, smtplib.SMTPException) as e:
|
||
# Log to STDOUT if SMTP fails
|
||
print(f"Failed to send email to {to_email}: {str(e)}")
|
||
print(f"Subject: {subject}")
|
||
print(f"Body: {body}")
|
||
|
||
|
||
def poll_for_result(conn, headers, request_id):
|
||
"""Poll the BFL API for image 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 response.get("status") == "Ready":
|
||
return response
|
||
elif response.get("status") == "Failed":
|
||
raise Exception("Image generation failed")
|
||
elif response.get("status") == "Pending":
|
||
time.sleep(5)
|
||
else:
|
||
raise Exception(f"Unknown status: {response.get('status', 'N/A')}")
|
||
|
||
|
||
def generate_image(prompt, api_key, endpoint, seed):
|
||
"""Generate an image using the Black Forest Labs API."""
|
||
conn = http.client.HTTPSConnection("api.bfl.ai")
|
||
headers = {"Content-Type": "application/json", "X-Key": api_key}
|
||
payload = {
|
||
"prompt": prompt,
|
||
"max_tokens": 512,
|
||
"temperature": 0.7,
|
||
"top_p": 0.9,
|
||
"top_k": 40,
|
||
"repetition_penalty": 1.1,
|
||
"stop": ["\n\n"],
|
||
"seed": seed,
|
||
}
|
||
conn.request("POST", f"/v1/{endpoint}", body=json.dumps(payload), headers=headers)
|
||
res = conn.getresponse()
|
||
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: {data}")
|
||
result = poll_for_result(conn, headers, request_id)
|
||
if "result" in result and "sample" in result["result"]:
|
||
return result["result"]["sample"]
|
||
raise Exception("No image URL found.")
|
||
|
||
|
||
# --- Streamlit App ---
|
||
st.title("Black Forest Labs Image Generation")
|
||
|
||
# API Key Setup
|
||
deployed = (
|
||
os.getenv("STREAMLIT_CLOUD", "false").lower() == "true"
|
||
) # Check for Streamlit Cloud deployment
|
||
api_key = os.environ.get("BLACK_FOREST_LABS_API_KEY") # Default to environment variable
|
||
|
||
if deployed and not api_key: # Fallback to secrets if deployed and env var is missing
|
||
try:
|
||
api_key = st.secrets.get("BLACK_FOREST_LABS_API_KEY")
|
||
except Exception:
|
||
pass # Ignore if secrets aren't available
|
||
|
||
if not api_key:
|
||
st.error(
|
||
"BLACK_FOREST_LABS_API_KEY is not set in environment variables or secrets."
|
||
)
|
||
st.stop()
|
||
|
||
# --- Session State Initialization with Cookies ---
|
||
user_id_cookie = cookies.get("user_id")
|
||
if user_id_cookie:
|
||
st.session_state.user_id = int(user_id_cookie)
|
||
else:
|
||
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 = []
|
||
|
||
# --- 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 MainSession() as session:
|
||
user = session.query(User).filter_by(email=email).first()
|
||
if user:
|
||
user.otp_hash = otp_hash
|
||
user.otp_expiry = expiry
|
||
else:
|
||
user = User(email=email, otp_hash=otp_hash, otp_expiry=expiry)
|
||
session.add(user)
|
||
session.commit()
|
||
send_email(
|
||
email,
|
||
"Your Image Generation OTP",
|
||
f"Your OTP is: {otp}\nIt expires in 10 minutes.",
|
||
)
|
||
st.success(
|
||
"OTP sent to your email (or logged to console if SMTP failed)!"
|
||
)
|
||
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 MainSession() 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
|
||
cookies["user_id"] = str(user.id)
|
||
cookies.save()
|
||
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 MainSession() as session:
|
||
namespaces = (
|
||
session.query(Namespace)
|
||
.join(NamespaceUser)
|
||
.filter(NamespaceUser.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(name=new_namespace)
|
||
session.add(namespace)
|
||
session.flush() # Get namespace.id before committing
|
||
namespace_user = NamespaceUser(
|
||
user_id=st.session_state.user_id, namespace_id=namespace.id
|
||
)
|
||
session.add(namespace_user)
|
||
try:
|
||
session.commit()
|
||
st.session_state.namespace_id = namespace.id
|
||
st.success(f"Album '{new_namespace}' created!")
|
||
st.rerun()
|
||
except:
|
||
session.rollback()
|
||
st.error("Album creation failed.")
|
||
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 = []
|
||
NamespaceSession = get_namespace_db(st.session_state.namespace_id)
|
||
with NamespaceSession() as session:
|
||
images = (
|
||
session.query(Image).order_by(Image.generation_time.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,
|
||
img.seed,
|
||
img.model_endpoint,
|
||
img.generation_time,
|
||
)
|
||
)
|
||
|
||
# --- 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
|
||
)
|
||
|
||
# --- Namespace Settings in Sidebar ---
|
||
if st.session_state.namespace_id:
|
||
st.sidebar.subheader(f"Settings for '{selected_namespace}'")
|
||
with MainSession() as session:
|
||
current_users = (
|
||
session.query(User)
|
||
.join(NamespaceUser)
|
||
.filter(NamespaceUser.namespace_id == st.session_state.namespace_id)
|
||
.all()
|
||
)
|
||
|
||
# Display current users with removal option (except self)
|
||
st.sidebar.write("Current Collaborators:")
|
||
for user in current_users:
|
||
if user.id != st.session_state.user_id: # Prevent self-removal
|
||
if st.sidebar.button(
|
||
f"Remove {user.email}", key=f"remove_{user.id}"
|
||
):
|
||
session.delete(
|
||
session.query(NamespaceUser)
|
||
.filter_by(
|
||
user_id=user.id,
|
||
namespace_id=st.session_state.namespace_id,
|
||
)
|
||
.first()
|
||
)
|
||
session.commit()
|
||
st.sidebar.success(
|
||
f"Removed {user.email} from '{selected_namespace}'"
|
||
)
|
||
st.rerun()
|
||
else:
|
||
st.sidebar.write(f"- {user.email} (You)")
|
||
|
||
# Invite new user
|
||
invite_email = st.sidebar.text_input(
|
||
"Invite a collaborator by email:", key="invite_email"
|
||
)
|
||
if st.sidebar.button("Send Invite", key="send_invite"):
|
||
if invite_email:
|
||
invited_user = (
|
||
session.query(User).filter_by(email=invite_email).first()
|
||
)
|
||
if not invited_user:
|
||
invited_user = User(
|
||
email=invite_email, otp_hash="", otp_expiry=datetime.now()
|
||
)
|
||
session.add(invited_user)
|
||
session.flush()
|
||
if (
|
||
session.query(NamespaceUser)
|
||
.filter_by(
|
||
user_id=invited_user.id,
|
||
namespace_id=st.session_state.namespace_id,
|
||
)
|
||
.first()
|
||
):
|
||
st.sidebar.error(f"{invite_email} is already a collaborator.")
|
||
else:
|
||
namespace_user = NamespaceUser(
|
||
user_id=invited_user.id,
|
||
namespace_id=st.session_state.namespace_id,
|
||
)
|
||
session.add(namespace_user)
|
||
session.commit()
|
||
send_email(
|
||
invite_email,
|
||
"Album Collaboration Invite",
|
||
f"You’ve been invited to collaborate on '{selected_namespace}' at https://art.ai.unturf.com!",
|
||
)
|
||
st.sidebar.success(
|
||
f"Invited {invite_email} to '{selected_namespace}' (check console if email failed)!"
|
||
)
|
||
else:
|
||
st.sidebar.error("Please enter an email to invite.")
|
||
|
||
# --- Logout Option ---
|
||
if st.sidebar.button("Logout"):
|
||
st.session_state.user_id = None
|
||
st.session_state.namespace_id = None
|
||
st.session_state.generated_images = []
|
||
cookies["user_id"] = ""
|
||
cookies.save()
|
||
st.rerun()
|
||
|
||
# --- 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."
|
||
)
|
||
|
||
# --- Prompt + Generate ---
|
||
if st.session_state.namespace_id:
|
||
st.subheader(f"Generate Images in '{selected_namespace}'")
|
||
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:
|
||
generation_time = datetime.now()
|
||
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")
|
||
NamespaceSession = get_namespace_db(st.session_state.namespace_id)
|
||
with NamespaceSession() as session:
|
||
image = Image(
|
||
slug=slug,
|
||
prompt=prompt,
|
||
filename=filename,
|
||
base64_image=base64_image,
|
||
generation_time=generation_time,
|
||
seed=seed_value,
|
||
model_endpoint=selected_model,
|
||
)
|
||
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,
|
||
seed_value,
|
||
selected_model,
|
||
generation_time,
|
||
),
|
||
)
|
||
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.")
|
||
|
||
# --- Display Images with Metadata ---
|
||
st.subheader(f"Images in '{selected_namespace}' (Newest First)")
|
||
for (
|
||
filepath,
|
||
filename,
|
||
prompt_text,
|
||
seed,
|
||
model,
|
||
gen_time,
|
||
) in st.session_state.generated_images:
|
||
st.image(
|
||
filepath, caption=f"Prompt: {prompt_text}", use_container_width=True
|
||
)
|
||
st.write(f"**Seed**: {seed}")
|
||
st.write(f"**Model**: {model}")
|
||
st.write(f"**Generated**: {gen_time.strftime('%Y-%m-%d %H:%M:%S')}")
|
||
with open(filepath, "rb") as file:
|
||
st.download_button(
|
||
label=f"Download {filename}",
|
||
data=file,
|
||
file_name=filename,
|
||
mime="image/jpeg",
|
||
)
|
||
st.write("---") # Separator between images
|
||
|
||
if deployed:
|
||
st.write("Running in deployed environment")
|
||
else:
|
||
st.write("Running locally")
|