app.py

import streamlit as st
from PIL import Image
from dotenv import load_dotenv
from image_evaluators import LlamaEvaluator
from prompt_refiners import LlamaPromptRefiner
from weave_prompt import PromptOptimizer
from similarity_metrics import LPIPSImageSimilarityMetric
from image_generators import FalImageGenerator

# Load environment variables from .env file
load_dotenv()

st.set_page_config(
    page_title="WeavePrompt",
    page_icon="🎨",
    layout="wide"
)

def main():
    st.title("🎨 WeavePrompt: Iterative Prompt Optimization")
    st.markdown("""
    Upload a target image and watch as WeavePrompt iteratively optimizes a text prompt to recreate it.
    """)
    
    # Initialize session state
    if 'optimizer' not in st.session_state:
        st.session_state.optimizer = PromptOptimizer(
            image_generator=FalImageGenerator(),
            evaluator=LlamaEvaluator(),
            refiner=LlamaPromptRefiner(),
            similarity_metric=LPIPSImageSimilarityMetric(),
            max_iterations=10,
            similarity_threshold=0.95
        )
    
    if 'optimization_started' not in st.session_state:
        st.session_state.optimization_started = False
    
    if 'current_results' not in st.session_state:
        st.session_state.current_results = None
    
    # Auto mode state
    if 'auto_mode' not in st.session_state:
        st.session_state.auto_mode = False
    
    if 'auto_paused' not in st.session_state:
        st.session_state.auto_paused = False
    
    # Auto mode step control - use this to control when to step vs when to display
    if 'auto_should_step' not in st.session_state:
        st.session_state.auto_should_step = False

    # File uploader
    uploaded_file = st.file_uploader("Choose a target image", type=['png', 'jpg', 'jpeg'])
    
    if uploaded_file is not None:
        # Display target image
        target_image = Image.open(uploaded_file)
        
        col1, col2 = st.columns(2)
        with col1:
            st.subheader("Target Image")
            st.image(target_image, width='stretch')
        
        # Start button
        if not st.session_state.optimization_started:
            if st.button("Start Optimization"):
                st.session_state.optimization_started = True
                # Initialize optimization
                is_completed, prompt, generated_image = st.session_state.optimizer.initialize(target_image)
                st.session_state.current_results = (is_completed, prompt, generated_image)
                st.rerun()
        else:
            # Show disabled button or status when optimization has started
            st.button("Start Optimization", disabled=True, help="Optimization in progress")
        
        # Display optimization progress
        if st.session_state.optimization_started:
            if st.session_state.current_results is not None:
                with col2:
                    st.subheader("Generated Image")
                    is_completed, prompt, generated_image = st.session_state.current_results
                    st.image(generated_image, width='stretch')
                
                # Display prompt and controls
                st.text_area("Current Prompt", prompt, height=100)
            else:
                # Show loading state
                with col2:
                    st.subheader("Generated Image")
                    st.info("Initializing optimization...")
                st.text_area("Current Prompt", "Generating initial prompt...", height=100)
                is_completed = False
                prompt = ""
            
            # Auto mode controls
            st.subheader("Auto Mode Controls")
            col_auto1, col_auto2 = st.columns(2)
            
            with col_auto1:
                auto_mode = st.checkbox("Auto-progress steps", value=st.session_state.auto_mode)
                if auto_mode != st.session_state.auto_mode:
                    st.session_state.auto_mode = auto_mode
                    if auto_mode:
                        st.session_state.auto_paused = False
                        st.session_state.auto_should_step = True  # Start by stepping
                    st.rerun()
            
            with col_auto2:
                if st.session_state.auto_mode:
                    if st.session_state.auto_paused:
                        if st.button("▶️ Resume", key="resume_btn"):
                            st.session_state.auto_paused = False
                            st.rerun()
                    else:
                        if st.button("⏸️ Pause", key="pause_btn"):
                            st.session_state.auto_paused = True
                            st.rerun()
            
            # Progress metrics
            col1, col2, col3 = st.columns(3)
            with col1:
                # Show current iteration: completed steps + 1 if still in progress
                current_iteration = len(st.session_state.optimizer.history) + (0 if is_completed else 1)
                st.metric("Iteration", current_iteration)
            with col2:
                if len(st.session_state.optimizer.history) > 0:
                    similarity = st.session_state.optimizer.history[-1]['similarity']
                    st.metric("Similarity", f"{similarity:.2%}")
            with col3:
                status = "Completed" if is_completed else "Paused" if st.session_state.auto_paused else "In Progress"
                st.metric("Status", status)
            
            # Progress bar
            st.subheader("Progress")
            max_iterations = st.session_state.optimizer.max_iterations
            progress_value = min(current_iteration / max_iterations, 1.0) if max_iterations > 0 else 0.0
            st.progress(progress_value, text=f"Step {current_iteration} of {max_iterations}")
            
            # Next step logic
            if not is_completed:
                # Auto mode logic - mimic pause button behavior
                if st.session_state.auto_mode and not st.session_state.auto_paused:
                    if st.session_state.auto_should_step:
                        # Execute the step
                        is_completed, prompt, generated_image = st.session_state.optimizer.step()
                        st.session_state.current_results = (is_completed, prompt, generated_image)
                        # Set flag to NOT step on next render (let history display)
                        st.session_state.auto_should_step = False
                        st.rerun()
                    else:
                        # Don't step, just display current state and history, then set flag to step next time
                        st.session_state.auto_should_step = True
                        # Use a small delay then rerun to continue auto mode
                        import time
                        time.sleep(0.5)  # Give user time to see the history
                        st.rerun()
                
                # Manual mode
                elif not st.session_state.auto_mode:
                    if st.button("Next Step"):
                        is_completed, prompt, generated_image = st.session_state.optimizer.step()
                        st.session_state.current_results = (is_completed, prompt, generated_image)
                        st.rerun()
                
                # Show status when auto mode is paused
                elif st.session_state.auto_paused:
                    st.info("Auto mode is paused. Click Resume to continue or uncheck Auto-progress to use manual mode.")
            else:
                st.success("Optimization completed! Click 'Reset' to try another image.")
                # Turn off auto mode when completed
                if st.session_state.auto_mode:
                    st.session_state.auto_mode = False
                    st.session_state.auto_paused = False
            
            # Reset button
            if st.button("Reset"):
                st.session_state.optimization_started = False
                st.session_state.current_results = None
                st.session_state.auto_mode = False
                st.session_state.auto_paused = False
                st.rerun()
            
            # Display history - simple approach
            if len(st.session_state.optimizer.history) > 0:
                st.subheader(f"Optimization History ({len(st.session_state.optimizer.history)} steps)")
                
                for idx, hist_entry in enumerate(st.session_state.optimizer.history):
                    st.markdown(f"### Step {idx + 1}")
                    col1, col2 = st.columns([2, 3])
                    with col1:
                        st.image(hist_entry['image'], width='stretch')
                    with col2:
                        st.text(f"Similarity: {hist_entry['similarity']:.2%}")
                        st.text("Prompt:")
                        st.text(hist_entry['prompt'])
                        # Toggle analysis view per history entry
                        expand_key = f"expand_analysis_{idx}"
                        if 'analysis_expanded' not in st.session_state:
                            st.session_state['analysis_expanded'] = {}
                        if expand_key not in st.session_state['analysis_expanded']:
                            st.session_state['analysis_expanded'][expand_key] = False

                        if st.session_state['analysis_expanded'][expand_key]:
                            if st.button("Hide Analysis", key=f"hide_{expand_key}"):
                                st.session_state['analysis_expanded'][expand_key] = False
                                st.rerun()
                            st.text("Analysis:")
                            for key, value in hist_entry['analysis'].items():
                                st.text(f"{key}: {value}")
                        else:
                            if st.button("Expand Analysis", key=expand_key):
                                st.session_state['analysis_expanded'][expand_key] = True
                                st.rerun()

if __name__ == "__main__":
    main()