Update app.py robust

app.py

@@ -46,6 +46,27 @@ def create_temp_folder():
     return run_folder
 
 def main():
+    # Initialize session state for tracking processing status
+    if 'processing' not in st.session_state:
+        st.session_state.processing = False
+    
+    if 'process_complete' not in st.session_state:
+        st.session_state.process_complete = False
+    
+    if 'temp_dir' not in st.session_state:
+        st.session_state.temp_dir = None
+    
+    if 'final_video_path' not in st.session_state:
+        st.session_state.final_video_path = None
+        
+    # Restart function to reset session state
+    def restart():
+        st.session_state.processing = False
+        st.session_state.process_complete = False 
+        st.session_state.temp_dir = None
+        st.session_state.final_video_path = None
+        st.rerun()  # Use st.rerun() instead of experimental_rerun
+
     # ---------------- CUSTOM CSS FOR A PROFESSIONAL, DARK THEME ----------------
     st.markdown(
         """
@@ -118,6 +139,13 @@ def main():
             transform: scale(1.02);
             box-shadow: 0 0 20px rgba(142,68,173,0.8), 0 0 30px rgba(114,45,145,0.6);
         }
+        /* Disabled button styling */
+        div.stButton > button:disabled {
+            background-image: linear-gradient(45deg, #6b6b6b, #4a4a4a);
+            box-shadow: none;
+            cursor: not-allowed;
+            opacity: 0.7;
+        }
         /* File uploader label styling */
         .stFileUploader label {
             font-size: 1rem;
@@ -129,6 +157,22 @@ def main():
             padding-left: 1rem;
             margin-left: 1rem;
         }
+        /* Processing overlay */
+        .processing-overlay {
+            position: fixed;
+            top: 0;
+            left: 0;
+            right: 0;
+            bottom: 0;
+            background-color: rgba(0, 0, 0, 0.7);
+            z-index: 1000;
+            display: flex;
+            flex-direction: column;
+            justify-content: center;
+            align-items: center;
+            color: white;
+            font-size: 1.5rem;
+        }
         </style>
         """,
         unsafe_allow_html=True
@@ -171,23 +215,61 @@ def main():
         unsafe_allow_html=True
     )
 
+    # 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 Result 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
+                
+        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
+
+    # Show processing indicator if currently processing
+    if st.session_state.processing:
+        progress_container = st.container()
+        with progress_container:
+            st.info("Processing your morphing request. Please wait...")
+            progress_bar = st.progress(0)
+            status_text = st.empty()
+            status_text.text("Initializing pipeline. Please do not close this page.")
+        st.warning("⚠️ Please wait for the current process to complete. Do not press the run button again.")
+        
     # ---------------- SECTION 1: IMAGE & PROMPT INPUTS ----------------
     st.subheader("1. Upload Source Images & Prompts")
     col_imgA, col_imgB = st.columns(2)
     with col_imgA:
         st.markdown("#### Image A")
-        uploaded_image_A = st.file_uploader("Upload your first image", type=["png", "jpg", "jpeg"], key="imgA")
+        uploaded_image_A = st.file_uploader("Upload your first image", type=["png", "jpg", "jpeg"], key="imgA", 
+                                         disabled=st.session_state.processing)
         if uploaded_image_A is not None:
             st.image(uploaded_image_A, caption="Preview - Image A", use_container_width=True)
         prompt_A = st.text_input("Short Description for Image A (optional)", value="", key="promptA", 
-        help="For added interpolation between the two descriptions")
+                              help="For added interpolation between the two descriptions",
+                              disabled=st.session_state.processing)
     with col_imgB:
         st.markdown("#### Image B")
-        uploaded_image_B = st.file_uploader("Upload your second image", type=["png", "jpg", "jpeg"], key="imgB")
+        uploaded_image_B = st.file_uploader("Upload your second image", type=["png", "jpg", "jpeg"], key="imgB",
+                                         disabled=st.session_state.processing)
         if uploaded_image_B is not None:
             st.image(uploaded_image_B, caption="Preview - Image B", use_container_width=True)
         prompt_B = st.text_input("Short Description for Image B (optional)", value="", key="promptB",
-        help="For added interpolation between the two descriptions")
+                              help="For added interpolation between the two descriptions",
+                              disabled=st.session_state.processing)
 
     st.markdown("<hr>", unsafe_allow_html=True)
 
@@ -215,7 +297,8 @@ def main():
             "Custom ⚙️"
         ],
         index=0,
-        label_visibility="collapsed"  # Hide the label in the UI but keep it for accessibility
+        label_visibility="collapsed",
+        disabled=st.session_state.processing
     )
 
     # Determine preset defaults based on selection
@@ -252,34 +335,37 @@ def main():
         ]
         default_model = preset_model if preset_model is not None else "Base Stable Diffusion V1-5"
         default_index = options_list.index(default_model)
-        model_option = st.selectbox("Select Model Card", options=options_list, index=default_index)
+        model_option = st.selectbox("Select Model Card", options=options_list, index=default_index, disabled=st.session_state.processing)
         
         col_left, col_right = st.columns(2)
         # Left Column: Keyframe Generator Parameters
         with col_left:
             st.markdown("##### Keyframe Generator Parameters")
-            num_frames = st.number_input("Number of keyframes (2–50)", min_value=2, max_value=50, value=16)
+            num_frames = st.number_input("Number of keyframes (2–50)", min_value=2, max_value=50, value=16, disabled=st.session_state.processing)
             lcm_default = preset_lcm if preset_lcm is not None else False
             enable_lcm_lora = st.checkbox(
                 "Enable LCM-LoRA",
                 value=lcm_default,
-                help="Accelerates inference with slight quality decrease"
+                help="Accelerates inference with slight quality decrease",
+                disabled=st.session_state.processing
             )
-            use_adain = st.checkbox("Use AdaIN", value=True, help="Adaptive Instance Normalization for improved generation")
-            use_reschedule = st.checkbox("Use reschedule sampling", value=True, help="Better sampling strategy")
+            use_adain = st.checkbox("Use AdaIN", value=True, help="Adaptive Instance Normalization for improved generation", disabled=st.session_state.processing)
+            use_reschedule = st.checkbox("Use reschedule sampling", value=True, help="Better sampling strategy", disabled=st.session_state.processing)
         
         # Right Column: Inter-frame Interpolator Parameters (FILM)
         with col_right:
             st.markdown("<div class='right-column-divider'>", unsafe_allow_html=True)
             st.markdown("##### Inter-frame Interpolator Parameters")
             default_use_film = preset_film if preset_film is not None else True
-            use_film = st.checkbox("Use FILM interpolation", value=default_use_film, help="Frame Interpolation for Large Motion - creates smooth transitions")
+            use_film = st.checkbox("Use FILM interpolation", value=default_use_film, help="Frame Interpolation for Large Motion - creates smooth transitions", disabled=st.session_state.processing)
             film_recursions = st.number_input("FILM recursion passes (1–6)", min_value=1, max_value=6, value=3, 
-                                             help="Higher values create more intermediate frames (smoother but slower)")
+                                             help="Higher values create more intermediate frames (smoother but slower)",
+                                             disabled=st.session_state.processing)
             # Set default FPS based on whether FILM is enabled
             default_fps = 30 if use_film else 4
             output_fps = st.number_input("Output FPS (1–120)", min_value=1, max_value=120, value=default_fps, 
-                                    help="Output video frames per second")
+                                    help="Output video frames per second",
+                                    disabled=st.session_state.processing)
             st.markdown("</div>", unsafe_allow_html=True)
 
     st.markdown("<hr>", unsafe_allow_html=True)
@@ -289,122 +375,132 @@ def main():
     st.markdown("Once satisfied with your inputs, click below to start the process.")
 
     # New checkbox for SLAB execution toggle
-    # using_slab = st.checkbox("Using SLAB GPU Cluster?", value=False, help="If enabled, the pipeline command will be prefixed with SLAB cluster execution parameters.")
+    # using_slab = st.checkbox("Using SLAB GPU Cluster?", value=False, help="If enabled, the pipeline command will be prefixed with SLAB cluster execution parameters.", disabled=st.session_state.processing)
 
-    if st.button("Run Morphing Pipeline", key="run_pipeline"):
-        if not (uploaded_image_A and uploaded_image_B):
-            st.error("Please upload both images before running the morphing pipeline.")
-            return
-
-        # Instead of using /tmp, create a folder in the repo for temporary processing.
-        temp_dir = create_temp_folder()
+    # Create a container for the run button and its status messages
+    run_container = st.container()
+    
+    with run_container:
+        # Button text changes based on processing state
+        button_text = "Processing... Please Wait" if st.session_state.processing else "Run Morphing Pipeline"
         
-        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
-            output_dir = os.path.join(temp_dir, "morph_results")
-            film_output_dir = os.path.join(temp_dir, "film_output")
-            os.makedirs(output_dir, exist_ok=True)
-            os.makedirs(film_output_dir, exist_ok=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"
-            )
-            
-            # 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)])
-            
-            # If SLAB execution is enabled, prepend the srun command prefix.
-            # if using_slab:
-            #     slab_prefix = [
-            #         "srun", "-p", "rtx3090_slab", "-w", "slabgpu05", "--gres=gpu:1",
-            #         "--job-name=test", "--kill-on-bad-exit=1"
-            #     ]
-            #     cmd = slab_prefix + cmd
-            
-            st.info("Initializing pipeline. This may take a few minutes...")
-            progress_bar = st.progress(0)
-            status_text = st.empty()
-
-            # Update progress status
-            for i in range(1, 11):
-                status_text.text(f"Step {i}/10: {'Preparing images' if i <= 2 else 'Generating keyframes' if i <= 6 else 'Interpolating frames' if i <= 9 else 'Finalizing video'}")
-                progress_bar.progress(i * 10)
+        # Button is disabled during processing
+        if st.button(button_text, key="run_pipeline", disabled=st.session_state.processing):
+            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
+                
+                # Create progress indicators
+                progress_bar = st.progress(0)
+                status_text = st.empty()
+                status_text.text("Initializing pipeline. This may take a few minutes...")
                 
-                if i == 3: # Start processing!
+                # Instead of using /tmp, create a folder in the repo for temporary 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
+                    output_dir = os.path.join(temp_dir, "morph_results")
+                    film_output_dir = os.path.join(temp_dir, "film_output")
+                    os.makedirs(output_dir, exist_ok=True)
+                    os.makedirs(film_output_dir, exist_ok=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"
+                    )
+                    
+                    # 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)])
+                    
+                    # If SLAB execution is enabled, prepend the srun command prefix.
+                    # if using_slab:
+                    #     slab_prefix = [
+                    #         "srun", "-p", "rtx3090_slab", "-w", "slabgpu05", "--gres=gpu:1",
+                    #         "--job-name=test", "--kill-on-bad-exit=1"
+                    #     ]
+                    #     cmd = slab_prefix + cmd
+                    
+                    # Update progress status for preparing phase
+                    progress_bar.progress(10)
+                    status_text.text("Preparing images and configuration...")
+                    
+                    # Run the morphing process
                     try:
                         subprocess.run(cmd, check=True)
+                        
+                        # Update progress
+                        progress_bar.progress(90)
+                        status_text.text("Processing complete! Preparing results...")
+                        
+                        # 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
+                        
+                        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
+                        
+                        if video_found:
+                            st.session_state.final_video_path = final_video_path
+                            st.session_state.process_complete = True
+                            
+                            # Set progress to 100%
+                            progress_bar.progress(100)
+                            
+                            st.rerun()
+                        else:
+                            status_text.warning("No output video was generated. Check logs for details.")
+                            st.session_state.processing = False
+                            st.session_state.process_complete = False
+                            
                     except subprocess.CalledProcessError as e:
                         st.error(f"Error running morphing pipeline: {e}")
-                        return
-                    break
-
-            # Set to 100% when done
-            progress_bar.progress(100)
-            status_text.text("Processing complete!")
-            
-            # 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
-            
-            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
-            
-            if video_found:
-                st.success("Morphing complete! 🎉")
-                # st.video(final_video_path) # Comment to remove video preview, buggy
-                try:
-                    with open(final_video_path, "rb") as f:
-                        video_bytes = f.read()
-                        st.download_button(
-                            "Download Result Video",
-                            data=video_bytes,
-                            file_name="metamorph_result.mp4",
-                            mime="video/mp4"
-                        )
-                except Exception as e:
-                    st.error(f"Error preparing video for download: {e}")
-            else:
-                st.warning("No output video was generated. Check logs for details.")
+                        st.session_state.processing = False
+                        st.session_state.process_complete = False
                     
-        except Exception as e:
-            st.error(f"An error occurred during processing: {e}")
+                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()