Update app.py

app.py

@@ -492,12 +492,17 @@ def video_generation_handler_streaming(prompt, seed=42, fps=15):
         f"</div>"
     )
     
-    # Show download button if MP4 was created successfully
-    download_visible = final_mp4_path is not None
-    yield None, final_status_html, gr.update(visible=download_visible), final_mp4_path
+    # Show download button and complete video
+    yield None, final_status_html, final_mp4_path, gr.update(value=final_mp4_path)
     
     print(f"✅ PyAV streaming complete! {total_frames_yielded} frames across {num_blocks} blocks")
 
+def download_video(video_path):
+    """Handle video download."""
+    if video_path and os.path.exists(video_path):
+        return video_path
+    return None
+
 # --- Gradio UI Layout ---
 with gr.Blocks(title="Self-Forcing Streaming Demo") as demo:
     gr.Markdown("# 🚀 Self-Forcing Video Generation")
@@ -556,9 +561,19 @@ with gr.Blocks(title="Self-Forcing Streaming Demo") as demo:
                 show_label=False
             )
             
-            # Download button - initially hidden
+            gr.Markdown("### 🎬 Complete Video")
+            
+            # Complete video display
+            complete_video = gr.Video(
+                label="Complete Video",
+                height=400,
+                show_label=False,
+                visible=False
+            )
+            
+            # Download button
             with gr.Row():
-                download_button = gr.DownloadButton(
+                download_button = gr.Button(
                     "📥 Download Video",
                     visible=False,
                     variant="primary"
@@ -578,16 +593,26 @@ with gr.Blocks(title="Self-Forcing Streaming Demo") as demo:
     download_file = gr.File(visible=False)
 
     # Connect the generator to the streaming video
-    start_btn.click(
+    generation_event = start_btn.click(
         fn=video_generation_handler_streaming,
         inputs=[prompt, seed, fps],
-        outputs=[streaming_video, status_display, download_button, download_file]
+        outputs=[streaming_video, status_display, complete_video, download_button]
+    )
+    
+    # When generation completes, show the complete video and download button
+    generation_event.then(
+        fn=lambda x: (gr.update(visible=True), gr.update(visible=True)),
+        inputs=[complete_video],
+        outputs=[complete_video, download_button]
     )
     
-    # Set up download button to use the generated file
+    # Handle download button click
     download_button.click(
-        lambda x: x,
-        inputs=[download_file],
+        fn=download_video,
+        inputs=[complete_video],
+        outputs=[download_file]
+    ).then(
+        fn=lambda: gr.update(visible=True),
         outputs=[download_file]
     )