Upload folder using huggingface_hub

app.py

@@ -5,7 +5,8 @@ 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
+from image_generators import FalImageGenerator, MultiModelFalImageGenerator, AVAILABLE_MODELS
+from multi_model_optimizer import MultiModelPromptOptimizer
 
 # Load environment variables from .env file
 load_dotenv()
@@ -17,19 +18,28 @@ st.set_page_config(
 )
 
 def main():
-    st.title("🎨 WeavePrompt: Iterative Prompt Optimization")
+    st.title("🎨 WeavePrompt: Multi-Model Prompt Optimization")
     st.markdown("""
-    Upload a target image and watch as WeavePrompt iteratively optimizes a text prompt to recreate it.
+    Upload a target image and watch as WeavePrompt optimizes prompts across multiple AI models to find the best result.
     """)
     
+    # Model selection state
+    if 'selected_models' not in st.session_state:
+        st.session_state.selected_models = ["FLUX.1 [pro]"]  # Default selection
+    
+    # Initialize max_iterations in session state if not exists
+    if 'max_iterations' not in st.session_state:
+        st.session_state.max_iterations = 2
+    
     # Initialize session state
     if 'optimizer' not in st.session_state:
-        st.session_state.optimizer = PromptOptimizer(
-            image_generator=FalImageGenerator(),
+        image_generator = MultiModelFalImageGenerator(st.session_state.selected_models)
+        st.session_state.optimizer = MultiModelPromptOptimizer(
+            image_generator=image_generator,
             evaluator=LlamaEvaluator(),
             refiner=LlamaPromptRefiner(),
             similarity_metric=LPIPSImageSimilarityMetric(),
-            max_iterations=10,
+            max_iterations=st.session_state.max_iterations,
             similarity_threshold=0.95
         )
     
@@ -50,6 +60,71 @@ def main():
     if 'auto_should_step' not in st.session_state:
         st.session_state.auto_should_step = False
 
+    # Model Selection UI
+    st.subheader("🤖 Model Selection")
+    st.markdown("Choose which AI models to optimize with:")
+    
+    # Organize models by category
+    flux_models = [k for k in AVAILABLE_MODELS.keys() if k.startswith("FLUX")]
+    google_models = [k for k in AVAILABLE_MODELS.keys() if k in ["Imagen 4", "Imagen 4 Ultra", "Gemini 2.5 Flash Image"]]
+    other_models = [k for k in AVAILABLE_MODELS.keys() if k not in flux_models and k not in google_models]
+    
+    # Track if selection changed
+    new_selection = []
+    
+    # FLUX Models Section
+    st.markdown("### 🔥 FLUX Models")
+    cols_flux = st.columns(2)
+    for i, model_name in enumerate(flux_models):
+        col_idx = i % 2
+        with cols_flux[col_idx]:
+            is_selected = model_name in st.session_state.selected_models
+            if st.checkbox(model_name, value=is_selected, key=f"model_{model_name}"):
+                new_selection.append(model_name)
+    
+    # Google Models Section  
+    st.markdown("### 🔍 Google Models")
+    cols_google = st.columns(2)
+    for i, model_name in enumerate(google_models):
+        col_idx = i % 2
+        with cols_google[col_idx]:
+            is_selected = model_name in st.session_state.selected_models
+            if st.checkbox(model_name, value=is_selected, key=f"model_{model_name}"):
+                new_selection.append(model_name)
+    
+    # Other Models Section
+    if other_models:
+        st.markdown("### 🎨 Other Models")
+        cols_other = st.columns(2)
+        for i, model_name in enumerate(other_models):
+            col_idx = i % 2
+            with cols_other[col_idx]:
+                is_selected = model_name in st.session_state.selected_models
+                if st.checkbox(model_name, value=is_selected, key=f"model_{model_name}"):
+                    new_selection.append(model_name)
+    
+    # Ensure at least one model is selected
+    if not new_selection:
+        st.error("Please select at least one model!")
+        new_selection = ["FLUX.1 [pro]"]  # Default fallback
+    
+    # Update selection if changed
+    if set(new_selection) != set(st.session_state.selected_models):
+        st.session_state.selected_models = new_selection
+        # Recreate optimizer with new models
+        image_generator = MultiModelFalImageGenerator(st.session_state.selected_models)
+        st.session_state.optimizer = MultiModelPromptOptimizer(
+            image_generator=image_generator,
+            evaluator=LlamaEvaluator(),
+            refiner=LlamaPromptRefiner(),
+            similarity_metric=LPIPSImageSimilarityMetric(),
+            max_iterations=st.session_state.max_iterations,
+            similarity_threshold=0.95
+        )
+        st.success(f"Updated to use {len(new_selection)} model(s): {', '.join(new_selection)}")
+    
+    st.markdown("---")
+
     # File uploader
     uploaded_file = st.file_uploader("Choose a target image", type=['png', 'jpg', 'jpeg'])
     
@@ -64,25 +139,106 @@ def main():
         
         # Start button
         if not st.session_state.optimization_started:
-            if st.button("Start Optimization"):
+            if st.button("🚀 Start Optimization", type="primary"):
+                # Set state first to ensure immediate UI update
                 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)
+                # Force immediate rerun to show disabled state
                 st.rerun()
         else:
-            # Show disabled button or status when optimization has started
-            st.button("Start Optimization", disabled=True, help="Optimization in progress")
+            # Show disabled button when optimization has started
+            st.button("⏳ Optimization Running...", disabled=True, help="Optimization in progress", type="secondary")
+            st.info("💡 Optimization is running across selected models. Use the controls below to pause/resume or reset.")
+            
+            # Initialize optimization after state is set (only once)
+            if st.session_state.current_results is None:
+                try:
+                    is_completed, prompt, generated_image = st.session_state.optimizer.initialize(target_image)
+                    st.session_state.current_results = (is_completed, prompt, generated_image)
+                    st.rerun()
+                except Exception as e:
+                    st.error(f"Error initializing optimization: {str(e)}")
+                    st.session_state.optimization_started = False
+        
+        # Settings (always visible)
+        st.subheader("⚙️ Settings")
+        
+        # Check if optimization is actively running (only disable settings during auto mode)
+        is_actively_running = (st.session_state.optimization_started and 
+                              st.session_state.auto_mode and 
+                              not st.session_state.auto_paused)
+        
+        # Single row: Number of iterations, Auto-progress, and Pause/Resume controls
+        col_settings1, col_settings2, col_settings3 = st.columns(3)
+        
+        with col_settings1:
+            new_max_iterations = st.number_input(
+                "Number of Iterations", 
+                min_value=1, 
+                max_value=20, 
+                value=st.session_state.max_iterations,
+                help="Maximum number of optimization iterations per model",
+                disabled=is_actively_running
+            )
+            
+            # Update if changed (only when not running)
+            if new_max_iterations != st.session_state.max_iterations and not is_actively_running:
+                st.session_state.max_iterations = new_max_iterations
+                # Update the optimizer's max_iterations
+                st.session_state.optimizer.max_iterations = new_max_iterations
+                if hasattr(st.session_state.optimizer, 'current_optimizer') and st.session_state.optimizer.current_optimizer:
+                    st.session_state.optimizer.current_optimizer.max_iterations = new_max_iterations
+        
+        with col_settings2:
+            auto_mode = st.checkbox(
+                "Auto-progress steps", 
+                value=st.session_state.auto_mode,
+                disabled=is_actively_running
+            )
+            if auto_mode != st.session_state.auto_mode and not is_actively_running:
+                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_settings3:
+            # Pause/Resume controls (only when auto mode is enabled and optimization started)
+            if st.session_state.auto_mode and st.session_state.optimization_started:
+                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()
+            else:
+                # Show placeholder or empty space when not in auto mode
+                st.write("")
         
         # Display optimization progress
         if st.session_state.optimization_started:
             if st.session_state.current_results is not None:
+                is_completed, prompt, generated_image = st.session_state.current_results
+                
+                # Always show the actual current results (most up-to-date)
                 with col2:
-                    st.subheader("Generated Image")
-                    is_completed, prompt, generated_image = st.session_state.current_results
+                    st.subheader("Current Generated Image")
                     st.image(generated_image, width='stretch')
+                    # Show current step info
+                    current_model_name = st.session_state.optimizer.get_current_model_name()
+                    current_history = st.session_state.optimizer.history if hasattr(st.session_state.optimizer, 'history') else []
+                    
+                    if current_history:
+                        # Show info about the current state
+                        if is_completed:
+                            st.caption(f"🏁 {current_model_name} - Final Result")
+                        else:
+                            st.caption(f"🎯 {current_model_name}")
+                    else:
+                        st.caption(f"🚀 {current_model_name} - Initializing...")
                 
-                # Display prompt and controls
+                # Display current prompt
                 st.text_area("Current Prompt", prompt, height=100)
             else:
                 # Show loading state
@@ -93,49 +249,81 @@ def main():
                 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()
+            # Multi-model progress info
+            progress_info = st.session_state.optimizer.get_progress_info()
+            current_model = progress_info['current_model_name']
             
             # Progress metrics
-            col1, col2, col3 = st.columns(3)
+            col1, col2, col3, col4 = st.columns(4)
             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)
+                st.metric("Current Model", current_model)
             with col2:
+                # Show current model number vs total models (1-indexed for display)
+                current_model_num = progress_info['current_model_index'] + 1
+                total_models = progress_info['total_models']
+                
+                # Fix: Don't let current model number exceed total models
+                if current_model_num > total_models:
+                    current_model_num = total_models
+                
+                st.metric("Models Progress", f"{current_model_num}/{total_models}")
+            with col3:
+                if 'current_iteration' in progress_info:
+                    # Show completed iterations vs max iterations
+                    current_iteration = progress_info['current_iteration']
+                    max_iterations = progress_info['max_iterations']
+                    # If not completed and we haven't reached max, show next iteration number
+                    if not is_completed and current_iteration < max_iterations:
+                        display_iteration = current_iteration + 1
+                        st.metric("Current Iteration", f"{display_iteration}/{max_iterations}")
+                    else:
+                        st.metric("Current Iteration", f"{current_iteration}/{max_iterations}")
+            with col4:
                 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)
+                else:
+                    status = "Completed" if is_completed else "Paused" if st.session_state.auto_paused else "In Progress"
+                    st.metric("Status", status)
             
-            # Progress bar
+            # Progress bars
             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}")
+            
+            # Overall progress across all models
+            overall_progress = progress_info['overall_progress']
+            if 'current_iteration' in progress_info and not is_completed:
+                # Add current model progress to overall
+                model_progress_fraction = progress_info['model_progress']
+                overall_progress = (progress_info['models_completed'] + model_progress_fraction) / progress_info['total_models']
+            
+            # Create more descriptive progress labels
+            models_completed = progress_info['models_completed']
+            total_models = progress_info['total_models']
+            
+            if total_models > 1:
+                # Multi-model scenario
+                if models_completed == total_models:
+                    overall_label = f"🏁 All {total_models} models completed!"
+                else:
+                    overall_label = f"🔄 Processing {current_model} (Model {models_completed + 1} of {total_models})"
+            else:
+                # Single model scenario
+                overall_label = f"🎯 Optimizing with {current_model}"
+            
+            st.progress(overall_progress, text=overall_label)
+            
+            # Current model progress
+            if 'current_iteration' in progress_info and not is_completed:
+                model_progress_value = progress_info['model_progress']
+                current_iter = progress_info['current_iteration']
+                max_iter = progress_info['max_iterations']
+                
+                if current_iter == max_iter:
+                    iteration_label = f"✅ {current_model}: Completed all {max_iter} iterations"
+                else:
+                    iteration_label = f"⚡ {current_model}: Step {current_iter + 1} of {max_iter}"
+                
+                st.progress(model_progress_value, text=iteration_label)
             
             # Next step logic
             if not is_completed:
@@ -181,37 +369,136 @@ def main():
                 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)")
+            # Display multi-model history with tabs
+            _display_multi_model_history()
+            
+            # Model Results Comparison (show when optimization is complete)
+            if is_completed and hasattr(st.session_state.optimizer, 'get_all_results'):
+                all_results = st.session_state.optimizer.get_all_results()
+                best_result = st.session_state.optimizer.get_best_result()
                 
-                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 all_results:
+                    st.subheader("🏆 Model Comparison Results")
+                    
+                    # Show best result prominently
+                    if best_result:
+                        st.success(f"🥇 **Best Result**: {best_result['model_name']} with {best_result['similarity']:.2%} similarity")
+                    
+                    # Create columns for each model result
+                    num_models = len(all_results)
+                    if num_models > 0:
+                        cols = st.columns(min(num_models, 3))  # Max 3 columns
+                        
+                        # Sort results by similarity (best first)
+                        sorted_results = sorted(all_results.items(), 
+                                              key=lambda x: x[1]['final_similarity'], 
+                                              reverse=True)
+                        
+                        for i, (model_name, result) in enumerate(sorted_results):
+                            col_idx = i % 3
+                            with cols[col_idx]:
+                                # Add medal emoji for top 3
+                                medal = "🥇" if i == 0 else "🥈" if i == 1 else "🥉" if i == 2 else "🏅"
+                                st.markdown(f"### {medal} {model_name}")
+                                
+                                # Show final image
+                                st.image(result['final_image'], width='stretch')
+                                
+                                # Show metrics
+                                st.metric("Final Similarity", f"{result['final_similarity']:.2%}")
+                                st.metric("Iterations", result['iterations'])
+                                
+                                # Show final prompt in expander
+                                with st.expander("View Final Prompt"):
+                                    st.text(result['final_prompt'])
+
+
+def _display_multi_model_history():
+    """Display optimization history with tabs for multi-model scenarios."""
+    
+    # Get all completed model results
+    all_model_results = {}
+    if hasattr(st.session_state.optimizer, 'get_all_results'):
+        all_model_results = st.session_state.optimizer.get_all_results()
+    
+    # Build separate histories for each model
+    all_histories = {}
+    
+    # Add completed model histories - each model gets its own stored history
+    for model_name, result in all_model_results.items():
+        if 'history' in result and result['history']:
+            # Make a deep copy to ensure complete separation
+            all_histories[model_name] = [step.copy() for step in result['history']]
+    
+    # Add current model's in-progress history (if any)
+    if hasattr(st.session_state.optimizer, 'current_optimizer') and st.session_state.optimizer.current_optimizer:
+        current_history = st.session_state.optimizer.current_optimizer.history
+        if current_history:
+            current_model_name = st.session_state.optimizer.get_current_model_name()
+            # Only add if this model doesn't already have a completed history
+            if current_model_name not in all_histories:
+                all_histories[current_model_name] = [step.copy() for step in current_history]
+    
+    # Display histories
+    if all_histories:
+        st.subheader("📊 Optimization History")
+        
+        if len(all_histories) == 1:
+            # Single model - no tabs needed
+            model_name = list(all_histories.keys())[0]
+            history = all_histories[model_name]
+            st.markdown(f"**{model_name}** ({len(history)} steps)")
+            _display_model_history(history, model_name)
+        else:
+            # Multi-model - use tabs
+            tab_names = []
+            for model_name, history in all_histories.items():
+                tab_names.append(f"{model_name} ({len(history)} steps)")
+            
+            tabs = st.tabs(tab_names)
+            
+            for i, (model_name, history) in enumerate(all_histories.items()):
+                with tabs[i]:
+                    _display_model_history(history, model_name)
+
+
+def _display_model_history(history, model_name):
+    """Display history for a single model."""
+    if not history:
+        st.info(f"No optimization steps yet for {model_name}")
+        return
+    
+    for idx, hist_entry in enumerate(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_{model_name}_{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 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()

image_generators.py

@@ -5,10 +5,29 @@ import requests
 from io import BytesIO
 
 from weave_prompt import ImageGenerator
+from typing import List, Tuple
 
 from dotenv import load_dotenv
 load_dotenv()
 
+# Available fal.ai models for text-to-image generation
+AVAILABLE_MODELS = {
+    # FLUX Models
+    "FLUX.1 [pro]": "fal-ai/flux-pro",
+    "FLUX.1 [dev]": "fal-ai/flux/dev", 
+    "FLUX.1 [schnell]": "fal-ai/flux/schnell",
+    "FLUX.1 with LoRAs": "fal-ai/flux-lora",
+    
+    # Google Models
+    "Imagen 4": "fal-ai/imagen4/preview",
+    "Imagen 4 Ultra": "fal-ai/imagen4/preview/ultra", 
+    "Gemini 2.5 Flash Image": "fal-ai/gemini-25-flash-image",
+    
+    # Other Models
+    "Stable Diffusion 3.5 Large": "fal-ai/stable-diffusion-v35-large",
+    "Qwen Image": "fal-ai/qwen-image"
+}
+
 class FalImageGenerator(ImageGenerator):
     """Handles image generation using fal_client."""
     
@@ -45,4 +64,61 @@ class FalImageGenerator(ImageGenerator):
             image = Image.open(BytesIO(response.content))
             return image
         else:
-            raise ValueError("No image found in the result")
+            raise ValueError("No image found in the result")
+
+
+class MultiModelFalImageGenerator(ImageGenerator):
+    """Handles image generation using multiple fal.ai models."""
+    
+    def __init__(self, selected_models: List[str] = None):
+        """Initialize with selected model names.
+        
+        Args:
+            selected_models: List of model display names from AVAILABLE_MODELS keys
+        """
+        if selected_models is None:
+            selected_models = ["FLUX.1 [pro]"]  # Default to single model
+        
+        self.selected_models = selected_models
+        self.current_model_index = 0
+        self.generators = {}
+        
+        # Create individual generators for each selected model
+        for model_name in selected_models:
+            if model_name in AVAILABLE_MODELS:
+                model_id = AVAILABLE_MODELS[model_name]
+                self.generators[model_name] = FalImageGenerator(model_id)
+    
+    def get_current_model_name(self) -> str:
+        """Get the name of the currently active model."""
+        if self.current_model_index < len(self.selected_models):
+            return self.selected_models[self.current_model_index]
+        return self.selected_models[0] if self.selected_models else "Unknown"
+    
+    def switch_to_next_model(self) -> bool:
+        """Switch to the next model in the sequence.
+        
+        Returns:
+            True if switched to next model, False if no more models
+        """
+        self.current_model_index += 1
+        return self.current_model_index < len(self.selected_models)
+    
+    def reset_to_first_model(self):
+        """Reset to the first model in the sequence."""
+        self.current_model_index = 0
+    
+    def generate(self, prompt: str, **kwargs) -> Image.Image:
+        """Generate an image using the current model."""
+        current_model = self.get_current_model_name()
+        if current_model in self.generators:
+            return self.generators[current_model].generate(prompt, **kwargs)
+        else:
+            raise ValueError(f"Model {current_model} not available")
+    
+    def generate_with_model(self, model_name: str, prompt: str, **kwargs) -> Image.Image:
+        """Generate an image using a specific model."""
+        if model_name in self.generators:
+            return self.generators[model_name].generate(prompt, **kwargs)
+        else:
+            raise ValueError(f"Model {model_name} not available")

multi_model_optimizer.py

@@ -0,0 +1,179 @@
+from typing import List, Dict, Any, Tuple
+from PIL import Image
+from weave_prompt import PromptOptimizer, ImageEvaluator, PromptRefiner, ImageSimilarityMetric
+from image_generators import MultiModelFalImageGenerator
+
+
+class MultiModelPromptOptimizer:
+    """Sequential multi-model prompt optimizer that finds the best model-prompt combination."""
+    
+    def __init__(self,
+                 image_generator: MultiModelFalImageGenerator,
+                 evaluator: ImageEvaluator,
+                 refiner: PromptRefiner,
+                 similarity_metric: ImageSimilarityMetric,
+                 max_iterations: int = 10,
+                 similarity_threshold: float = 0.95):
+        """Initialize the multi-model optimizer.
+        
+        Args:
+            image_generator: Multi-model image generator
+            evaluator: Image evaluator for generating initial prompt and analysis
+            refiner: Prompt refinement strategy
+            similarity_metric: Image similarity metric
+            max_iterations: Maximum number of optimization iterations per model
+            similarity_threshold: Target similarity threshold for early stopping
+        """
+        self.image_generator = image_generator
+        self.evaluator = evaluator
+        self.refiner = refiner
+        self.similarity_metric = similarity_metric
+        self.max_iterations = max_iterations
+        self.similarity_threshold = similarity_threshold
+        
+        # Multi-model state
+        self.target_img = None
+        self.current_model_index = 0
+        self.model_results = {}  # Results per model
+        self.current_optimizer = None
+        self.best_result = None
+        
+        # Initialize individual optimizers for each model
+        self._create_current_optimizer()
+    
+    def _create_current_optimizer(self):
+        """Create optimizer for the current model."""
+        if self.current_model_index < len(self.image_generator.selected_models):
+            # Set the image generator to current model
+            self.image_generator.current_model_index = self.current_model_index
+            
+            # Create individual optimizer for current model
+            self.current_optimizer = PromptOptimizer(
+                image_generator=self.image_generator,
+                evaluator=self.evaluator,
+                refiner=self.refiner,
+                similarity_metric=self.similarity_metric,
+                max_iterations=self.max_iterations,
+                similarity_threshold=self.similarity_threshold
+            )
+    
+    def get_current_model_name(self) -> str:
+        """Get the name of the currently active model."""
+        # Ensure the image generator index is synchronized
+        self.image_generator.current_model_index = self.current_model_index
+        return self.image_generator.get_current_model_name()
+    
+    def get_progress_info(self) -> Dict[str, Any]:
+        """Get current progress information."""
+        total_models = len(self.image_generator.selected_models)
+        current_model = self.current_model_index + 1
+        
+        info = {
+            'current_model_index': self.current_model_index,
+            'current_model_name': self.get_current_model_name(),
+            'total_models': total_models,
+            'models_completed': self.current_model_index,
+            'overall_progress': self.current_model_index / total_models if total_models > 0 else 0,
+            'is_last_model': self.current_model_index >= total_models - 1
+        }
+        
+        if self.current_optimizer:
+            info['current_iteration'] = len(self.current_optimizer.history)
+            info['max_iterations'] = self.max_iterations
+            info['model_progress'] = len(self.current_optimizer.history) / self.max_iterations
+        
+        return info
+    
+    def initialize(self, target_img: Image.Image) -> Tuple[bool, str, Image.Image]:
+        """Initialize the multi-model optimization process.
+        
+        Args:
+            target_img: Target image to optimize towards
+        Returns:
+            Tuple of (is_completed, current_prompt, current_generated_image)
+        """
+        self.target_img = target_img
+        self.current_model_index = 0
+        self.model_results = {}
+        self.best_result = None
+        
+        # Reset image generator to first model
+        self.image_generator.reset_to_first_model()
+        self._create_current_optimizer()
+        
+        # Initialize first model
+        return self.current_optimizer.initialize(target_img)
+    
+    def step(self) -> Tuple[bool, str, Image.Image]:
+        """Perform one optimization step.
+        
+        Returns:
+            Tuple of (is_completed, current_prompt, current_generated_image)
+        """
+        if not self.current_optimizer:
+            raise RuntimeError("Must call initialize() before step()")
+        
+        # Step the current model optimizer
+        is_model_completed, prompt, generated_image = self.current_optimizer.step()
+        
+        if is_model_completed:
+            # Store results for current model - use data from history to ensure consistency
+            model_name = self.get_current_model_name()
+            
+            if len(self.current_optimizer.history) > 0:
+                # Use the last step from history as the final result (ensures consistency)
+                last_step = self.current_optimizer.history[-1]
+                final_prompt = last_step['prompt']
+                final_image = last_step['image']
+                final_similarity = last_step['similarity']
+            else:
+                # Fallback to step results if no history (shouldn't happen)
+                final_prompt = prompt
+                final_image = generated_image
+                final_similarity = 0.0
+            
+            self.model_results[model_name] = {
+                'final_prompt': final_prompt,
+                'final_image': final_image,
+                'final_similarity': final_similarity,
+                'history': self.current_optimizer.history.copy(),
+                'iterations': len(self.current_optimizer.history)
+            }
+            
+            # Update best result if this is better
+            if self.best_result is None or final_similarity > self.best_result['similarity']:
+                self.best_result = {
+                    'model_name': model_name,
+                    'prompt': final_prompt,
+                    'image': final_image,
+                    'similarity': final_similarity
+                }
+            
+            # Move to next model
+            self.current_model_index += 1
+            
+            if self.current_model_index < len(self.image_generator.selected_models):
+                # Initialize next model - ensure both indices are synchronized
+                self.image_generator.current_model_index = self.current_model_index
+                self._create_current_optimizer()
+                return self.current_optimizer.initialize(self.target_img)
+            else:
+                # All models completed - return best result
+                return True, self.best_result['prompt'], self.best_result['image']
+        
+        return is_model_completed, prompt, generated_image
+    
+    def get_all_results(self) -> Dict[str, Dict[str, Any]]:
+        """Get results from all completed models."""
+        return self.model_results.copy()
+    
+    def get_best_result(self) -> Dict[str, Any]:
+        """Get the best result across all models."""
+        return self.best_result.copy() if self.best_result else None
+    
+    @property
+    def history(self):
+        """Get history from current optimizer for compatibility."""
+        if self.current_optimizer:
+            return self.current_optimizer.history
+        return []