Update app.py

app.py

@@ -653,31 +653,52 @@ with gr.Blocks(title="Self-Forcing Streaming Demo") as demo:
                 label="Generation Status"
             )
             
-            # Add download button and output
-            download_btn = gr.Button("💾 Download Video", variant="secondary")
-            download_output = gr.File(label="Download")
-
     # Define a wrapper function to ensure proper handling of outputs
     def safe_frame_generator(p, s, f):
         # Clear frames from previous generation
-        APP_STATE.update({"current_frames": []})
+        APP_STATE["current_frames"] = []
         
         # Reset the final video display
         yield None, None, gr.update(visible=False)
         
+        # Collect all frames from this generation
+        collected_frames = []
+        
         # Call the generator function and yield frames, status, and video updates
         try:
-            for frame, status_html, *video_update in video_generation_handler_streaming(p, s, f):
-                # Ensure we always have three outputs
-                if not video_update:
-                    video_update = [gr.update(visible=False)]  # Default - hide video
+            # Set save_frames=True to explicitly ensure frames are collected
+            for frame, status_html in video_generation_handler_streaming(p, s, f, save_frames=True):
+                # Track frames for this specific session
+                if frame is not None and isinstance(frame, np.ndarray):
+                    collected_frames.append(frame.copy())
+                    
+                # Show status update during generation
+                yield frame, status_html, gr.update(visible=False)
+                
+                # Check if this is the final frame
+                if "Complete" in str(status_html) or "100%" in str(status_html):
+                    # Create the final video
+                    if collected_frames:
+                        print(f"Creating final video from {len(collected_frames)} frames at {f} FPS")
+                        temp_file = save_frames_as_video(collected_frames, f)
+                        if temp_file:
+                            # Save these frames as the current set
+                            APP_STATE["current_frames"] = collected_frames
+                            yield frame, status_html, gr.update(visible=True, value=temp_file)
+                        else:
+                            yield frame, status_html, gr.update(visible=False)
+            
+            # Ensure final frame is properly handled
+            if collected_frames and "Complete" not in str(status_html) and "100%" not in str(status_html):
+                print(f"Generation complete, creating final video from {len(collected_frames)} frames")
+                temp_file = save_frames_as_video(collected_frames, f)
+                if temp_file:
+                    yield frame, status_html, gr.update(visible=True, value=temp_file)
                     
-                yield frame, status_html, video_update[0]
         except Exception as e:
             import traceback
-            print(f"Error in generator: {e}")
             traceback.print_exc()
-            error_html = f"<div style='color: red; padding: 10px; border: 1px solid #ffcccc; border-radius: 5px;'>Error: {e}</div>"
+            error_html = f"<div style='color: red; padding: 10px; border: 1px solid #ffcccc; border-radius: 5px;'>Error: {str(e)}</div>"
             yield None, error_html, gr.update(visible=False)
     
     # Connect the generator to the streaming video
@@ -687,23 +708,7 @@ with gr.Blocks(title="Self-Forcing Streaming Demo") as demo:
         outputs=[streaming_video, status_display, final_video]
     )
     
-    # Function to handle download button click
-    def download_video(fps):
-        if not APP_STATE.get("current_frames"):
-            return None
-        video_path = save_frames_as_video(APP_STATE["current_frames"], fps)
-        return video_path
-    
-    # Connect download button
-    download_btn.click(
-        fn=download_video,
-        inputs=[fps],
-        outputs=[download_output],
-        show_progress=True,
-        api_name="download_video"  # Make it accessible via API
-    )
-    
-    # Make the FPS slider visible for download quality control
+    # Make the FPS slider visible for video quality control
     fps.visible = True
     
     enhance_button.click(