Update app.py

app.py

@@ -1084,94 +1084,110 @@ class VideoGeneratorUI:
 
     def setup_ui(self):
         st.set_page_config(layout="wide")
+    
+        # Initialize session state
+        if 'selected_images' not in st.session_state:
+            st.session_state.selected_images = []
+        if 'processing_complete' not in st.session_state:
+            st.session_state.processing_complete = False
+        if 'current_step' not in st.session_state:
+            st.session_state.current_step = 'input'
         
-        # Custom CSS
-        st.markdown("""
-            <style>
-            .stApp {
-                max-width: 1200px;
-                margin: 0 auto;
-            }
-            .image-category {
-                margin-top: 2rem;
-                padding: 1rem;
-                border-radius: 0.5rem;
-                background: #f8f9fa;
-            }
-            .image-metadata {
-                font-size: 0.8rem;
-                color: #666;
-                margin-top: 0.5rem;
-            }
-            .submit-btn {
-                margin-top: 1rem;
-                padding: 0.5rem 1rem;
-            }
-            </style>
-        """, unsafe_allow_html=True)
-
         st.title("VaultGenix Video Generator")
         st.markdown("Create professional videos for your digital legacy management platform")
         
+        # Add sidebar for advanced settings
+        with st.sidebar:
+            st.subheader("Advanced Settings")
+            st.session_state.enable_transitions = st.checkbox("Enable Transitions", value=True)
+            st.session_state.enable_captions = st.checkbox("Enable Captions", value=True)
+            st.session_state.high_quality = st.checkbox("High Quality Export", value=True)
+        
+        # Main content area
         with st.container():
-            # Add form for prompt submission
-            with st.form(key='prompt_form'):
-                prompt = st.text_area("Enter your video script", height=200)
-                submit_button = st.form_submit_button(label='Analyze Script & Find Images')
+            if st.session_state.current_step == 'input':
+                self.show_input_form()
+            elif st.session_state.current_step == 'image_selection':
+                self.show_image_selection()
+            elif st.session_state.current_step == 'video_generation':
+                self.show_video_generation()
+
+    def show_input_form(self):
+        with st.form(key='prompt_form'):
+            prompt = st.text_area("Enter your video script", height=200)
+            col1, col2 = st.columns(2)
+            with col1:
+                st.session_state.video_style = st.selectbox(
+                    "Choose style",
+                    ["Professional", "Creative", "Educational"]
+                )
+            with col2:
+                st.session_state.voice_style = st.selectbox(
+                    "Choose voice style",
+                    ["Professional Male", "Professional Female", "Casual Male", "Casual Female"]
+                )
+            
+            submit_button = st.form_submit_button(label='Find Images')
             
             if submit_button and prompt:
-                # First show AI-selected images
-                with st.spinner("AI analyzing script and selecting relevant images..."):
-                    try:
-                        # Get AI-selected images first
-                        keywords = self.generator.image_scraper.extract_key_topics(prompt)
-                        st.write("🤖 AI-detected keywords:", ", ".join(keywords))
-                        
-                        image_categories = self.generator.image_scraper.get_images(prompt)
-                        
-                        # Store selections in session state
-                        if 'selected_images' not in st.session_state:
-                            st.session_state.selected_images = []
+                st.session_state.prompt = prompt
+                st.session_state.current_step = 'image_selection'
+                st.rerun()
+
+    def show_image_selection(self):
+        st.subheader("Select Images for Your Video")
+        
+        # Display back button
+        if st.button("← Back to Script"):
+            st.session_state.current_step = 'input'
+            st.rerun()
+        
+        try:
+            with st.spinner("AI analyzing script and selecting relevant images..."):
+                keywords = self.generator.image_scraper.extract_key_topics(st.session_state.prompt)
+                st.write("🤖 AI-detected keywords:", ", ".join(keywords))
+                
+                image_categories = self.generator.image_scraper.get_images(st.session_state.prompt)
+                
+                if image_categories and isinstance(image_categories, dict):
+                    # Create columns for image selection
+                    col1, col2 = st.columns([3, 1])
+                    
+                    with col1:
+                        # Display images with selection
+                        selected_count = 0
+                        st.session_state.selected_images = []  # Reset selected images
                         
-                        if image_categories and isinstance(image_categories, dict):
-                            # Display AI-selected primary matches first
-                            if 'primary' in image_categories and image_categories['primary']:
-                                st.subheader("🎯 AI-Selected Most Relevant Images")
-                                self.display_image_grid(image_categories['primary'])
-                            
-                            # Display secondary matches
-                            if 'secondary' in image_categories and image_categories['secondary']:
-                                st.subheader("🔄 AI-Selected Related Images")
-                                self.display_image_grid(image_categories['secondary'])
-                            
-                            # Collect selected images
-                            selected_images = []
-                            for category in image_categories.values():
-                                if isinstance(category, list):
-                                    for img in category:
-                                        key = f"img_{img['url']}"
-                                        if st.session_state.get(key, False):
-                                            selected_images.append(img['url'])
-                            
-                            st.session_state.selected_images = selected_images
-                            
-                            # Video generation section
-                            if selected_images:
-                                self.show_video_settings(prompt, selected_images)
-                            else:
-                                st.warning("Please select at least one image to generate the video.")
+                        for category, images in image_categories.items():
+                            if images:
+                                st.subheader(f"{category.title()} Images")
+                                selected = self.display_image_grid(images)
+                                selected_count += len(selected)
+                                st.session_state.selected_images.extend(selected)
+                    
+                    with col2:
+                        # Show selection summary
+                        st.subheader("Selection Summary")
+                        st.write(f"Selected: {selected_count} images")
                         
+                        if selected_count > 0:
+                            if st.button("Continue to Video Generation →", type="primary"):
+                                st.session_state.current_step = 'video_generation'
+                                st.rerun()
                         else:
-                            st.warning("No images found. Please try a different prompt.")
-                            
-                    except Exception as e:
-                        st.error(f"An error occurred: {str(e)}")
-                        print(f"Error in UI: {str(e)}")
-    
-    def display_image_grid(self, images: List[Dict[str, str]], cols: int = 3):
-        """Display images in a grid with metadata and confidence scores"""
-        if not images or not isinstance(images, list):
-            return
+                            st.warning("Please select at least one image")
+                
+                else:
+                    st.warning("No images found. Please try a different prompt.")
+                    
+        except Exception as e:
+            st.error(f"An error occurred: {str(e)}")
+            print(f"Error in image selection: {str(e)}")
+
+
+    def display_image_grid(self, images: List[Dict[str, str]], cols: int = 3) -> List[str]:
+        """Display images in a grid and return selected image URLs"""
+        selected_urls = []
         
         n_images = len(images)
         n_rows = (n_images + cols - 1) // cols
@@ -1187,40 +1203,48 @@ class VideoGeneratorUI:
                             try:
                                 st.image(img['url'], use_container_width=True)
                                 
-                                # Add confidence score to checkbox label
-                                confidence = img.get('relevance_score', 0) * 100
-                                checkbox_label = f"Select (AI Confidence: {confidence:.1f}%)"
+                                # Create unique key for each checkbox
+                                checkbox_key = f"img_select_{row}_{col}_{img['url'][-10:]}"
                                 
-                                st.checkbox(
-                                    checkbox_label,
-                                    key=f"img_{img['url']}",
-                                    help=f"Keywords: {img['keyword']}\nTags: {img['tags']}"
-                                )
+                                if st.checkbox(
+                                    f"Select (Confidence: {img.get('relevance_score', 0)*100:.1f}%)",
+                                    key=checkbox_key
+                                ):
+                                    selected_urls.append(img['url'])
                                 
-                                # Show relevance metadata
+                                # Show image metadata
                                 st.markdown(
-                                    f"<div class='image-metadata'>"
-                                    f"<b>AI Relevance:</b> {img['relevance']}<br>"
-                                    f"<b>Keywords:</b> {img['keyword']}<br>"
-                                    f"<b>Match Type:</b> {img.get('category', 'General')}"
-                                    f"</div>",
+                                    f"""<div style='font-size: 0.8em;'>
+                                    Keywords: {img['keyword']}<br>
+                                    Category: {img.get('category', 'General')}
+                                    </div>""",
                                     unsafe_allow_html=True
                                 )
                             except Exception as e:
-                                print(f"Error displaying image: {e}")
+                                st.error(f"Error displaying image: {str(e)}")
+        
+        return selected_urls
+
     
-    def show_video_settings(self, prompt: str, selected_images: List[str]):
-        """Show video generation settings and controls with effects"""
-        st.subheader("Video Settings")
+    def show_video_generation(self):
+        st.subheader("Video Generation Settings")
+        
+        # Display back button
+        if st.button("← Back to Image Selection"):
+            st.session_state.current_step = 'image_selection'
+            st.rerun()
         
+        # Show selected images preview
+        st.write("Selected Images:")
+        cols = st.columns(4)
+        for idx, img_url in enumerate(st.session_state.selected_images):
+            with cols[idx % 4]:
+                st.image(img_url, width=150)
+        
+        # Video settings
         col1, col2 = st.columns(2)
+        
         with col1:
-            style = st.selectbox(
-                "Choose style",
-                options=["Professional", "Creative", "Educational"],
-                index=0
-            )
-        with col2:
             duration = st.slider(
                 "Video duration (seconds)",
                 min_value=30,
@@ -1228,148 +1252,145 @@ class VideoGeneratorUI:
                 value=60,
                 step=30
             )
-    
-        # Video effects controls
-        st.subheader("Video Effects")
-        effect_col1, effect_col2 = st.columns(2)
+            
+            background_music = st.selectbox(
+                "Background Music",
+                ["None", "Corporate", "Upbeat", "Inspirational", "Technology"]
+            )
         
-        with effect_col1:
-            zoom = st.slider(
-                "Zoom effect",
-                min_value=1.0,
-                max_value=1.5,
-                value=1.0,
-                step=0.1,
-                help="Gradually zoom in during each image segment"
+        with col2:
+            transition_style = st.selectbox(
+                "Transition Style",
+                ["Fade", "Slide", "Zoom", "None"]
             )
             
-            brightness = st.slider(
-                "Brightness",
-                min_value=0.5,
-                max_value=1.5,
-                value=1.0,
+            music_volume = st.slider(
+                "Music Volume",
+                min_value=0.0,
+                max_value=1.0,
+                value=0.3,
                 step=0.1
             )
         
-        with effect_col2:
-            contrast = st.slider(
-                "Contrast",
-                min_value=0.5,
-                max_value=1.5,
-                value=1.0,
-                step=0.1
+        # Advanced video effects
+        with st.expander("Advanced Video Effects"):
+            effect_col1, effect_col2 = st.columns(2)
+            
+            with effect_col1:
+                zoom = st.slider(
+                    "Zoom Effect",
+                    min_value=1.0,
+                    max_value=1.5,
+                    value=1.0,
+                    step=0.1
+                )
+                
+                brightness = st.slider(
+                    "Brightness",
+                    min_value=0.5,
+                    max_value=1.5,
+                    value=1.0,
+                    step=0.1
+                )
+            
+            with effect_col2:
+                contrast = st.slider(
+                    "Contrast",
+                    min_value=0.5,
+                    max_value=1.5,
+                    value=1.0,
+                    step=0.1
+                )
+                
+                enable_blur = st.checkbox(
+                    "Enable Blur Effect",
+                    value=False
+                )
+        
+        # Text overlay settings
+        with st.expander("Text Overlay Settings"):
+            subtitle_style = st.selectbox(
+                "Subtitle Style",
+                ["Modern", "Classic", "Minimal", "None"]
             )
             
-            enable_blur = st.checkbox(
-                "Enable subtle blur effect",
-                value=False,
-                help="Add a subtle blur effect to transitions"
+            caption_position = st.selectbox(
+                "Caption Position",
+                ["Bottom", "Top", "Middle"]
             )
-    
-        video_effects = {
-            'zoom': zoom,
-            'brightness': brightness,
-            'contrast': contrast,
-            'blur': enable_blur
-        }
-    
+        
+        # Generate video button
         if st.button("🎬 Generate Video", type="primary"):
-            if not selected_images:
-                st.error("Please select at least one image before generating the video.")
+            if not st.session_state.selected_images:
+                st.error("No images selected. Please go back and select images.")
                 return
-                
+            
             try:
+                video_effects = {
+                    'zoom': zoom,
+                    'brightness': brightness,
+                    'contrast': contrast,
+                    'blur': enable_blur,
+                    'transition_style': transition_style,
+                    'subtitle_style': subtitle_style,
+                    'caption_position': caption_position,
+                    'background_music': background_music,
+                    'music_volume': music_volume
+                }
+                
                 output_dir = "temp_videos"
                 os.makedirs(output_dir, exist_ok=True)
                 output_path = os.path.join(output_dir, f"vaultgenix_video_{int(time.time())}.mp4")
                 
-                with st.spinner("🎥 Generating your video..."):
-                    video_path = self.generator.create_video(
-                        prompt,
-                        style,
-                        duration,
-                        output_path,
-                        selected_images,
-                        video_effects
-                    )
+                progress_text = st.empty()
+                progress_bar = st.progress(0)
                 
-                if os.path.exists(video_path):
-                    st.success("✨ Video generated successfully!")
-                    
-                    # Display video
-                    with open(video_path, 'rb') as video_file:
-                        video_bytes = video_file.read()
-                        st.video(video_bytes)
-                        
-                        # Download button
-                        st.download_button(
-                            label="⬇️ Download Video",
-                            data=video_bytes,
-                            file_name=os.path.basename(video_path),
-                            mime="video/mp4"
-                        )
-                else:
-                    st.error("Video generation failed. Please try again.")
-                    
-            except Exception as e:
-                st.error(f"Error generating video: {str(e)}")
-                print(f"Video generation error: {str(e)}")
-        
-    def generate_video(self, prompt: str, style: str, duration: int, selected_images: List[str]):
-        """Handle video generation with improved error handling"""
-        if not selected_images:
-            st.error("Please select at least one image before generating the video.")
-            return
-            
-        with st.spinner("🎥 Generating your video..."):
-            try:
-                # Create temp directory if it doesn't exist
-                temp_dir = Path("temp_videos")
-                temp_dir.mkdir(exist_ok=True)
+                # Generate video with progress updates
+                def progress_callback(progress):
+                    progress_bar.progress(progress)
+                    progress_text.text(f"Generating video: {progress}% complete")
                 
-                # Generate unique output path
-                output_path = str(temp_dir / f"vaultgenix_video_{int(time.time())}.mp4")
-                
-                # Generate video
                 video_path = self.generator.create_video(
-                    prompt,
-                    style,
+                    st.session_state.prompt,
+                    st.session_state.video_style,
                     duration,
                     output_path,
-                    selected_images
+                    st.session_state.selected_images,
+                    video_effects,
+                    progress_callback
                 )
                 
-                if video_path and os.path.exists(video_path):
+                if os.path.exists(video_path):
                     st.success("✨ Video generated successfully!")
                     
                     # Display video
-                    video_file = open(video_path, 'rb')
-                    video_bytes = video_file.read()
-                    st.video(video_bytes)
-                    
-                    # Download button
-                    st.download_button(
-                        label="⬇️ Download Video",
-                        data=video_bytes,
-                        file_name=os.path.basename(video_path),
-                        mime="video/mp4"
-                    )
-                    
-                    video_file.close()
+                    with open(video_path, 'rb') as video_file:
+                        video_bytes = video_file.read()
+                        st.video(video_bytes)
+                        
+                        # Download options
+                        col1, col2 = st.columns(2)
+                        with col1:
+                            st.download_button(
+                                label="⬇️ Download Video",
+                                data=video_bytes,
+                                file_name=os.path.basename(video_path),
+                                mime="video/mp4"
+                            )
+                        with col2:
+                            if st.session_state.high_quality:
+                                st.download_button(
+                                    label="⬇️ Download High Quality",
+                                    data=video_bytes,  # You'll need to implement high quality export
+                                    file_name=f"high_quality_{os.path.basename(video_path)}",
+                                    mime="video/mp4"
+                                )
                 else:
                     st.error("Video generation failed. Please try again.")
                     
             except Exception as e:
-                st.error(f"Failed to generate video: {str(e)}")
+                st.error(f"Error generating video: {str(e)}")
                 print(f"Video generation error: {str(e)}")
-                
-            finally:
-                # Cleanup temporary files
-                try:
-                    if 'video_file' in locals():
-                        video_file.close()
-                except Exception as e:
-                    print(f"Cleanup error: {e}")
-
+    
 if __name__ == "__main__":
     ui = VideoGeneratorUI()