diff --git a/black_forest_streamlit.py b/black_forest_streamlit.py index 3cd7f77..88b5580 100644 --- a/black_forest_streamlit.py +++ b/black_forest_streamlit.py @@ -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") \ No newline at end of file + st.write("Running locally")