5th ver, last

app.py

@@ -59,8 +59,9 @@ def main():
     if 'final_video_path' not in st.session_state:
         st.session_state.final_video_path = None
     
-    if 'progress_stage' not in st.session_state:
-        st.session_state.progress_stage = 0
+    # Initialize the download container in session state
+    if 'show_download' not in st.session_state:
+        st.session_state.show_download = False
         
     # Restart function to reset session state
     def restart():
@@ -68,7 +69,7 @@ def main():
         st.session_state.process_complete = False 
         st.session_state.temp_dir = None
         st.session_state.final_video_path = None
-        st.session_state.progress_stage = 0
+        st.session_state.show_download = False
         st.rerun()  # Use st.rerun() instead of experimental_rerun
 
     # ---------------- CUSTOM CSS FOR A PROFESSIONAL, DARK THEME ----------------
@@ -161,61 +162,13 @@ def main():
             padding-left: 1rem;
             margin-left: 1rem;
         }
-        /* Processing status box */
-        .processing-box {
-            background: rgba(0, 0, 0, 0.6);
-            border-radius: 10px;
-            border: 1px solid rgba(142, 68, 173, 0.6);
-            padding: 20px;
-            margin: 20px 0;
-            box-shadow: 0 0 30px rgba(142, 68, 173, 0.4);
-            text-align: center;
-        }
-        .processing-title {
-            font-size: 1.5rem;
-            color: #ffffff;
-            margin-bottom: 15px;
-            font-weight: bold;
-        }
-        .processing-stage {
-            margin: 15px 0;
-            font-size: 1.1rem;
-            color: #f1f1f1;
-        }
-        .pulse {
-            animation: pulse 2s infinite;
-        }
-        @keyframes pulse {
-            0% { opacity: 1; }
-            50% { opacity: 0.6; }
-            100% { opacity: 1; }
-        }
-        .progress-bar-container {
-            width: 100%;
-            background-color: rgba(255, 255, 255, 0.2);
-            border-radius: 5px;
-            margin: 15px 0;
-            overflow: hidden;
-        }
-        .progress-bar {
-            height: 10px;
-            background-image: linear-gradient(45deg, #8e44ad, #732d91);
-            border-radius: 5px;
-            transition: width 0.5s ease;
-        }
-        /* Rainbow spinner animation */
-        .spinner {
-            width: 40px;
-            height: 40px;
-            margin: 15px auto;
-            border: 4px solid transparent;
-            border-radius: 50%;
-            border-image: linear-gradient(45deg, purple, blue, cyan, green, yellow, orange, red) 1;
-            animation: spin 1s linear infinite;
-        }
-        @keyframes spin {
-            0% { transform: rotate(0deg); }
-            100% { transform: rotate(360deg); }
+        /* Results section styling */
+        .results-section {
+            margin-top: 2rem;
+            padding: 1.5rem;
+            background-color: rgba(0, 0, 0, 0.2);
+            border-radius: 8px;
+            border: 1px solid rgba(255, 255, 255, 0.1);
         }
         </style>
         """,
@@ -259,60 +212,35 @@ def main():
         unsafe_allow_html=True
     )
 
+    # Create a persistent container for the results section
+    results_container = st.container()
+
     # Show results if processing is complete
     if st.session_state.process_complete and st.session_state.final_video_path is not None:
-        st.success("Morphing complete! 🎉")
-        # st.video(st.session_state.final_video_path)
-        try:
-            with open(st.session_state.final_video_path, "rb") as f:
-                video_bytes = f.read()
-                st.download_button(
-                    "Download Morphing Video",
-                    data=video_bytes,
-                    file_name="metamorph_result.mp4",
-                    mime="video/mp4"
-                )
-            
-            # Add a restart button to clear the session state and start over
-            if st.button("Start New Morphing Project", key="restart"):
-                restart()  # Use the restart function
+        with results_container:
+            st.markdown("<div class='results-section'>", unsafe_allow_html=True)
+            st.success("Morphing complete! 🎉")
+            try:
+                with open(st.session_state.final_video_path, "rb") as f:
+                    video_bytes = f.read()
+                    st.session_state.video_bytes = video_bytes  # Store video bytes in session state
+                    st.video(video_bytes)
+                    st.download_button(
+                        "Download Morphing Video",
+                        data=video_bytes,
+                        file_name="metamorph_result.mp4",
+                        mime="video/mp4"
+                    )
+                st.session_state.show_download = True
                 
-        except Exception as e:
-            st.error(f"Error preparing video for download: {e}")
-        
-        # Early return to not show the input forms again when we have a result
-        return
+                # Add a restart button to clear the session state and start over
+                if st.button("Start New Morphing Project", key="restart"):
+                    restart()
+                    
+            except Exception as e:
+                st.error(f"Error preparing video for download: {e}")
+            st.markdown("</div>", unsafe_allow_html=True)
 
-    # Show enhanced processing status if in processing mode
-    if st.session_state.processing:
-        # Define processing stages
-        stages = [
-            "Initializing the morphing pipeline...",
-            "Loading and processing images...",
-            "Generating keyframes with diffusion model...",
-            "Creating intermediate frames...",
-            "Finalizing your morphing video...",
-        ]
-        
-        # Calculate progress percentage (artifically progressing based on time since start)
-        progress = min(st.session_state.progress_stage / len(stages), 1.0)
-        
-        # Show  processing status
-        st.markdown(
-            f"""
-            <div class="processing-box">
-                <div class="processing-title">🔮 Creating Your Morphing Video</div>
-                <div class="spinner"></div>
-                <div class="processing-stage pulse">{stages[st.session_state.progress_stage]}</div>
-                <div class="progress-bar-container">
-                    <div class="progress-bar" style="width: {int(progress * 100)}%"></div>
-                </div>
-                <div class="processing-stage">Please wait while we work our magic...</div>
-            </div>
-            """,
-            unsafe_allow_html=True
-        )
-    
     # ---------------- SECTION 1: IMAGE & PROMPT INPUTS ----------------
     st.subheader("1. Upload Source Images & Prompts")
     st.markdown("**Note:** Your uploaded images must be of similar topology and same size to achieve the best results.")
@@ -442,132 +370,132 @@ def main():
     # Create a container for the run button and its status messages
     run_container = st.container()
     
+    # Create a separate container for displaying the download button
+    download_container = st.container()
+    
+    # Check if we should display the download button even when the form is disabled
+    if st.session_state.show_download and st.session_state.final_video_path is not None:
+        with download_container:
+            st.markdown("<div class='results-section'>", unsafe_allow_html=True)
+            st.success("Morphing complete! 🎉")
+            try:
+                with open(st.session_state.final_video_path, "rb") as f:
+                    video_bytes = f.read()
+                    st.video(video_bytes)
+                    st.download_button(
+                        "Download Morphing Video",
+                        data=video_bytes,
+                        file_name="metamorph_result.mp4",
+                        mime="video/mp4"
+                    )
+                
+                # Add a restart button to clear the session state and start over
+                if st.button("Start New Morphing Project", key="restart_bottom"):
+                    restart()
+                    
+            except Exception as e:
+                st.error(f"Error preparing video for download: {e}")
+            st.markdown("</div>", unsafe_allow_html=True)
+    
     with run_container:
         # Button text changes based on processing state
         button_text = "Processing... Please Wait" if st.session_state.processing else "Run Morphing Pipeline"
         
         # Button is disabled during processing
         if st.button(button_text, key="run_pipeline", disabled=st.session_state.processing):
-            # Set processing state to True immediately
-            st.session_state.processing = True
-            st.rerun()  # Force a refresh to apply the disabled state to all UI elements
-        else:
-            # This will run after the refresh
-            if st.session_state.processing:
-                # Check if both images are uploaded
-                if not (uploaded_image_A and uploaded_image_B):
-                    st.error("Please upload both images before running the morphing pipeline.")
-                    st.session_state.processing = False  # Reset processing state
-                    st.rerun()  # Refresh again to enable UI elements
-                else:
-                    # Increment progress stage periodically
-                    # In a real implementation, these stage updates would be triggered by actual progress
-                    # Here we simulate progress for demonstration
+            if not (uploaded_image_A and uploaded_image_B):
+                st.error("Please upload both images before running the morphing pipeline.")
+            else:
+                # Set the processing state to True to prevent multiple runs
+                st.session_state.processing = True
+                st.info("Processing your morphing request. Please wait...")
+                
+                # Create a temporary folder for processing
+                temp_dir = create_temp_folder()
+                st.session_state.temp_dir = temp_dir
+                
+                try:
+                    # Save uploaded images
+                    imgA_path = os.path.join(temp_dir, "imageA.png")
+                    imgB_path = os.path.join(temp_dir, "imageB.png")
+                    save_uploaded_file(uploaded_image_A, imgA_path)
+                    save_uploaded_file(uploaded_image_B, imgB_path)
+                    
+                    # Create output directories
+                    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+                    output_dir = os.path.join(temp_dir, f"morph_results_{timestamp}")
+                    film_output_dir = os.path.join(temp_dir, f"film_output_{timestamp}")
+                    os.makedirs(output_dir, exist_ok=True)
+                    os.makedirs(film_output_dir, exist_ok=True)
                     
-                    # Create a temporary folder for processing
+                    actual_model_path = (
+                        "lykon/dreamshaper-7" if model_option == "Dreamshaper-7 (fine-tuned SD V1-5)" 
+                        else "stabilityai/stable-diffusion-2-1-base" if model_option == "Base Stable Diffusion V2-1"
+                        else "sd-legacy/stable-diffusion-v1-5"
+                    )
+                    
+                    # Build the command for run_morphing.py
+                    cmd = [
+                        sys.executable, "run_morphing.py",
+                        "--model_path", actual_model_path,
+                        "--image_path_0", imgA_path,
+                        "--image_path_1", imgB_path,
+                        "--prompt_0", prompt_A,
+                        "--prompt_1", prompt_B,
+                        "--output_path", output_dir,
+                        "--film_output_folder", film_output_dir,
+                        "--num_frames", str(num_frames),
+                        "--fps", str(output_fps)
+                    ]
+                    
+                    if enable_lcm_lora:
+                        cmd.append("--use_lcm")
+                    if use_adain:
+                        cmd.append("--use_adain")
+                    if use_reschedule:
+                        cmd.append("--use_reschedule")
+                    if use_film:
+                        cmd.append("--use_film")
+                    
+                    # Add film recursion parameter
+                    cmd.extend(["--film_num_recursions", str(film_recursions)])
+                    
+                    # Run the morphing process
                     try:
-                        temp_dir = create_temp_folder()
-                        st.session_state.temp_dir = temp_dir
-                        
-                        # Save uploaded images
-                        imgA_path = os.path.join(temp_dir, "imageA.png")
-                        imgB_path = os.path.join(temp_dir, "imageB.png")
-                        save_uploaded_file(uploaded_image_A, imgA_path)
-                        save_uploaded_file(uploaded_image_B, imgB_path)
-                        
-                        # Update progress stage
-                        st.session_state.progress_stage = 1
-                        st.rerun()  # Refresh to update progress
-                        
-                        # Create output directories
-                        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-                        output_dir = os.path.join(temp_dir, f"morph_results_{timestamp}")
-                        film_output_dir = os.path.join(temp_dir, f"film_output_{timestamp}")
-                        os.makedirs(output_dir, exist_ok=True)
-                        os.makedirs(film_output_dir, exist_ok=True)
+                        subprocess.run(cmd, check=True)
                         
-                        actual_model_path = (
-                            "lykon/dreamshaper-7" if model_option == "Dreamshaper-7 (fine-tuned SD V1-5)" 
-                            else "stabilityai/stable-diffusion-2-1-base" if model_option == "Base Stable Diffusion V2-1"
-                            else "sd-legacy/stable-diffusion-v1-5"
-                        )
+                        # Check for output video
+                        video_found = False
+                        possible_outputs = [f for f in os.listdir(film_output_dir) if f.endswith(".mp4")]
+                        if possible_outputs:
+                            final_video_path = os.path.join(film_output_dir, possible_outputs[0])
+                            video_found = True
                         
-                        # Build the command for run_morphing.py
-                        cmd = [
-                            sys.executable, "run_morphing.py",
-                            "--model_path", actual_model_path,
-                            "--image_path_0", imgA_path,
-                            "--image_path_1", imgB_path,
-                            "--prompt_0", prompt_A,
-                            "--prompt_1", prompt_B,
-                            "--output_path", output_dir,
-                            "--film_output_folder", film_output_dir,
-                            "--num_frames", str(num_frames),
-                            "--fps", str(output_fps)
-                        ]
-                        
-                        if enable_lcm_lora:
-                            cmd.append("--use_lcm")
-                        if use_adain:
-                            cmd.append("--use_adain")
-                        if use_reschedule:
-                            cmd.append("--use_reschedule")
-                        if use_film:
-                            cmd.append("--use_film")
-                        
-                        # Add film recursion parameter
-                        cmd.extend(["--film_num_recursions", str(film_recursions)])
-                        
-                        # Update progress stage - Start generating keyframes
-                        st.session_state.progress_stage = 2
-                        st.rerun()  # Refresh to update progress
-                        
-                        # Run the morphing process
-                        try:
-                            subprocess.run(cmd, check=True)
-                            
-                            # Update progress stage - Creating intermediate frames
-                            st.session_state.progress_stage = 3
-                            st.rerun()  # Refresh to update progress
-                            
-                            # Check for output video
-                            video_found = False
-                            possible_outputs = [f for f in os.listdir(film_output_dir) if f.endswith(".mp4")]
+                        if not video_found:
+                            possible_outputs = [f for f in os.listdir(output_dir) if f.endswith(".mp4")]
                             if possible_outputs:
-                                final_video_path = os.path.join(film_output_dir, possible_outputs[0])
+                                final_video_path = os.path.join(output_dir, possible_outputs[0])
                                 video_found = True
-                            
-                            if not video_found:
-                                possible_outputs = [f for f in os.listdir(output_dir) if f.endswith(".mp4")]
-                                if possible_outputs:
-                                    final_video_path = os.path.join(output_dir, possible_outputs[0])
-                                    video_found = True
-                            
-                            # Update progress stage - Finalizing video
-                            st.session_state.progress_stage = 4
-                            st.rerun()  # Refresh to update progress
-                            
-                            if video_found:
-                                st.session_state.final_video_path = final_video_path
-                                st.session_state.process_complete = True
-                                st.rerun()  # Show the results page
-                            else:
-                                st.error("No output video was generated. Check logs for details.")
-                                st.session_state.processing = False
-                                st.session_state.process_complete = False
-                                st.rerun()  # Reset the interface
-                                
-                        except subprocess.CalledProcessError as e:
-                            st.error(f"Error running morphing pipeline: {e}")
+                        
+                        if video_found:
+                            st.session_state.final_video_path = final_video_path
+                            st.session_state.process_complete = True
+                            st.session_state.show_download = True
+                            st.rerun()  # This is crucial for showing the results page
+                        else:
+                            st.error("No output video was generated. Check logs for details.")
                             st.session_state.processing = False
                             st.session_state.process_complete = False
-                            st.rerun()  # Reset the interface
-                        
-                    except Exception as e:
-                        st.error(f"An error occurred during processing: {e}")
+                            
+                    except subprocess.CalledProcessError as e:
+                        st.error(f"Error running morphing pipeline: {e}")
                         st.session_state.processing = False
                         st.session_state.process_complete = False
-                        st.rerun()  # Reset the interface
+                    
+                except Exception as e:
+                    st.error(f"An error occurred during processing: {e}")
+                    st.session_state.processing = False
+                    st.session_state.process_complete = False
 
 if __name__ == "__main__":
     main()