Spaces:
Running
Running
File size: 4,612 Bytes
d4acb0c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 |
# app.py
import streamlit as st
from cerebras.cloud.sdk import Cerebras
import openai
import os
from dotenv import load_dotenv
from together import Together
# --- Assuming config.py and utils.py exist ---
import config
import utils
try:
from prompt import BASE_PROMPT
except ImportError:
st.error(
"Error: 'prompt.py' not found or 'BASE_PROMPT' is not defined within it.")
st.stop()
# --- Import column rendering functions ---
from chat_column import render_chat_column
from image_column import render_image_column
load_dotenv()
st.set_page_config(page_icon="🤖", layout="wide",
page_title="Prompt & Image Generator")
utils.display_icon("🤖")
st.title("Prompt & Image Generator")
st.subheader("Generate text prompts (left) and edit/generate images (right)",
divider="orange", anchor=False)
api_key_from_env = os.getenv("CEREBRAS_API_KEY")
show_api_key_input = not bool(api_key_from_env)
cerebras_api_key = None
together_api_key = os.getenv("TOGETHER_API_KEY")
# --- サイドバーの設定 ---
with st.sidebar:
st.title("Settings")
if show_api_key_input:
st.markdown("### :red[Enter your Cerebras API Key below]")
api_key_input = st.text_input(
"Cerebras API Key:", type="password", key="cerebras_api_key_input_field")
if api_key_input:
cerebras_api_key = api_key_input
else:
cerebras_api_key = api_key_from_env
st.success("✓ Cerebras API Key loaded from environment")
# Together Key Status
if not together_api_key:
st.warning(
"TOGETHER_API_KEY environment variable not set. Image generation (right column) will not work.", icon="⚠️")
else:
st.success("✓ Together API Key loaded from environment")
# Model selection
model_option = st.selectbox(
"Choose a LLM model:",
options=list(config.MODELS.keys()),
format_func=lambda x: config.MODELS[x]["name"],
key="model_select"
)
# Max tokens slider
max_tokens_range = config.MODELS[model_option]["tokens"]
default_tokens = min(2048, max_tokens_range)
max_tokens = st.slider(
"Max Tokens (LLM):",
min_value=512,
max_value=max_tokens_range,
value=default_tokens,
step=512,
help="Max tokens for the LLM's text prompt response."
)
# Check if Cerebras API key is available
if not cerebras_api_key and show_api_key_input and 'cerebras_api_key_input_field' in st.session_state and st.session_state.cerebras_api_key_input_field:
cerebras_api_key = st.session_state.cerebras_api_key_input_field
if not cerebras_api_key:
st.error("Cerebras API Key is required. Please enter it in the sidebar or set the CEREBRAS_API_KEY environment variable.", icon="🚨")
st.stop()
llm_client = None
image_client = None
try:
llm_client = Cerebras(api_key=cerebras_api_key)
if together_api_key:
image_client = Together(api_key=together_api_key)
except Exception as e:
st.error(f"Failed to initialize API client(s): {str(e)}", icon="🚨")
st.stop()
# --- Session State Initialization ---
# Initialize state variables if they don't exist
if "messages" not in st.session_state:
st.session_state.messages = []
if "current_image_prompt_text" not in st.session_state:
st.session_state.current_image_prompt_text = ""
# --- MODIFICATION START ---
# Replace single image state with a list to store multiple images and their prompts
if "generated_images_list" not in st.session_state:
st.session_state.generated_images_list = [] # Initialize as empty list
# Remove old state variable if it exists (optional cleanup)
if "latest_generated_image" in st.session_state:
del st.session_state["latest_generated_image"]
# --- MODIFICATION END ---
if "selected_model" not in st.session_state:
st.session_state.selected_model = None
# --- Clear history if model changes ---
if st.session_state.selected_model != model_option:
st.session_state.messages = []
st.session_state.current_image_prompt_text = ""
# --- MODIFICATION START ---
# Clear the list of generated images when model changes
st.session_state.generated_images_list = []
# --- MODIFICATION END ---
st.session_state.selected_model = model_option
st.rerun()
# --- Define Main Columns ---
chat_col, image_col = st.columns([2, 1])
# --- Render Columns using imported functions ---
with chat_col:
render_chat_column(st, llm_client, model_option, max_tokens, BASE_PROMPT)
with image_col:
render_image_column(st, image_client) # Pass the client
|