# At the beginning of your app.py import streamlit as st import os import google.generativeai as genai from PIL import Image from dotenv import load_dotenv from model_loader import load_gemini_model, generate_caption_with_gemini, generate_detailed_description # Load environment variables - try to get from .env file first, then from secrets load_dotenv() api_key = os.getenv("GEMINI_API_KEY") # If running on Hugging Face Space, get API key from secrets if not api_key and 'GEMINI_API_KEY' in st.secrets: api_key = st.secrets["GEMINI_API_KEY"] # Configure Gemini with API Key genai.configure(api_key=api_key) # Setup a session state counter if "queries_done" not in st.session_state: st.session_state.queries_done = 0 if "user_api_key" not in st.session_state: st.session_state.user_api_key = None # Add a state variable for the dynamic prompt if "dynamic_prompt" not in st.session_state: st.session_state.dynamic_prompt = "" # Function to reset the dynamic prompt def reset_dynamic_prompt(): st.session_state.dynamic_prompt = "" # Streamlit App st.title("🖼️VisionCraft: AI Image Explorer") st.markdown(""" Welcome to VisionCraft: AI Image Explorer 🚀 Upload any image and let AI create detailed, creative, and emotional descriptions for you. - ✅ 2 free queries included - 🔑 After that, bring your own Gemini API key - 🎯 Supports both quick presets and custom prompts Get ready to explore your images in a whole new way! 🌟 """) st.write("Upload an image and generate creative descriptions based on your needs.") # If more than 2 queries, ask for user's API key if st.session_state.queries_done >= 2 and not st.session_state.user_api_key: user_api_key = st.text_input("Enter your Gemini API Key to continue:", type="password") if user_api_key: st.session_state.user_api_key = user_api_key genai.configure(api_key=user_api_key) else: st.warning("Please enter your Gemini API Key to continue using the app.") st.stop() else: # For first two requests, use your own key if st.session_state.queries_done < 2: genai.configure(api_key=api_key) elif st.session_state.user_api_key: genai.configure(api_key=st.session_state.user_api_key) # Upload image uploaded_image = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"]) if uploaded_image is not None: image = Image.open(uploaded_image).convert("RGB") # Use use_container_width instead of use_column_width st.image(image, caption="Uploaded Image", width=None) # Create tabs for different ways to interact tab1, tab2 = st.tabs(["Quick Presets", "Ask Anything"]) with tab1: # Prompt selection options prompt_options = { "basic": "Basic Description", "chain_of_thought": "Detailed Analysis", "story": "Creative Story", "emotional": "Emotional Analysis", "object": "Object Detection", "context": "Contextual Description", "action": "Action Description" } # Prompt selection prompt_type = st.selectbox( "Select description type", options=list(prompt_options.keys()), format_func=lambda x: prompt_options[x] ) if st.button("Generate Description", key="preset_button"): with st.spinner("Processing image..."): # Increment the counter BEFORE processing st.session_state.queries_done += 1 # Load model gemini_model = load_gemini_model() # Generate caption with st.spinner("Generating basic caption..."): caption = generate_caption_with_gemini(image, gemini_model) # Show caption st.subheader("✨ Basic Caption") st.write(caption) # Generate detailed description based on selected prompt type with st.spinner(f"Generating {prompt_options[prompt_type]}..."): detailed_description = generate_detailed_description( image, gemini_model, prompt_type ) # Show the result with appropriate header st.subheader(f"✨ {prompt_options[prompt_type]}") st.write(detailed_description) # Add download button for the generated text st.download_button( label="Download Description", data=detailed_description, file_name="image_description.txt", mime="text/plain", key="download1" ) # Show usage count st.info(f"You have used {st.session_state.queries_done} out of 2 free queries.") with tab2: # Dynamic prompt box for custom questions st.subheader("Ask anything about this image") dynamic_prompt = st.text_area( "Enter your question or what you'd like to know about the image", value=st.session_state.dynamic_prompt, placeholder="Examples:\n- Write a poem about this image\n- Explain what's happening in this scene\n- What emotion does this image evoke?\n- Describe this as if you were a detective", key="dynamic_prompt_input" ) if st.button("Get Answer", key="dynamic_button"): if dynamic_prompt: # Save the current prompt before processing current_prompt = dynamic_prompt # Clear the prompt for next use reset_dynamic_prompt() # Increment the counter BEFORE processing st.session_state.queries_done += 1 with st.spinner("Processing your request..."): # Load model gemini_model = load_gemini_model() # Generate response directly from the image and custom prompt response = generate_detailed_description( image, gemini_model, "custom", current_prompt ) # Display the response st.subheader("✨ Response") st.write(response) # Add download button st.download_button( label="Download Response", data=response, file_name="image_response.txt", mime="text/plain", key="download2" ) # Show usage count st.info(f"You have used {st.session_state.queries_done} out of 2 free queries.") else: st.warning("Please enter a question or prompt first")