app.py

# 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")