Spaces:
Sleeping
Sleeping
| # 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") |