323 lines
No EOL
12 KiB
Python
323 lines
No EOL
12 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
|
|
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
|
|
|
|
# --- 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)
|
|
|
|
# --- Cookie Manager Setup ---
|
|
cookies = EncryptedCookieManager(
|
|
password="your-secret-password-here", # Replace with a secure password
|
|
)
|
|
if not cookies.ready():
|
|
st.stop() # Wait for cookies to be ready
|
|
|
|
# --- Utility Functions ---
|
|
def create_slug(prompt, max_words=10):
|
|
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():
|
|
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'] = '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)
|
|
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":
|
|
pass
|
|
else:
|
|
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}
|
|
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 = False
|
|
try:
|
|
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:
|
|
api_key = os.environ.get("BLACK_FOREST_LABS_API_KEY")
|
|
|
|
if not api_key:
|
|
st.error("BLACK_FOREST_LABS_API_KEY is not set.")
|
|
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 Session() 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_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
|
|
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 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.session_state.namespace_id = namespace.id # Set the new album as active
|
|
st.success(f"Album '{new_namespace}' created!")
|
|
st.rerun() # Rerun to show the generation page
|
|
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:
|
|
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:
|
|
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.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 ---
|
|
st.subheader(f"Images in '{selected_namespace}' (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",
|
|
)
|
|
|
|
# --- 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."
|
|
)
|
|
if deployed:
|
|
st.write("Running in deployed environment")
|
|
else:
|
|
st.write("Running locally") |