Update app.py

app.py

@@ -432,7 +432,7 @@ def video_generation_handler_streaming(prompt, seed=42, fps=15):
             )
             
             # Yield None for video but update status (frame-by-frame tracking)
-            yield None, frame_status_html, gr.update(visible=False), None
+            yield None, frame_status_html, gr.update(visible=False), gr.update(visible=False)
 
         # Encode entire block as one chunk immediately
         if all_frames_from_block:
@@ -450,7 +450,7 @@ def video_generation_handler_streaming(prompt, seed=42, fps=15):
                 total_progress = (idx + 1) / num_blocks * 100
                 
                 # Yield the actual video chunk
-                yield ts_path, gr.update(), gr.update(visible=False), None
+                yield ts_path, gr.update(), gr.update(visible=False), gr.update(visible=False)
                 
             except Exception as e:
                 print(f"⚠️ Error encoding block {idx}: {e}")
@@ -492,17 +492,11 @@ def video_generation_handler_streaming(prompt, seed=42, fps=15):
         f"</div>"
     )
     
-    # Show download button and complete video
-    yield None, final_status_html, final_mp4_path, gr.update(value=final_mp4_path)
+    # Show complete video and file download
+    yield None, final_status_html, final_mp4_path, gr.update(value=final_mp4_path, visible=True)
     
     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")
@@ -563,21 +557,20 @@ with gr.Blocks(title="Self-Forcing Streaming Demo") as demo:
             
             gr.Markdown("### 🎬 Complete Video")
             
-            # Complete video display
+            # Complete video display with download enabled
             complete_video = gr.Video(
                 label="Complete Video",
                 height=400,
                 show_label=False,
-                visible=False
+                visible=False,
+                show_download_button=True  # Enable download button in video control
             )
             
-            # Download button
-            with gr.Row():
-                download_button = gr.Button(
-                    "📥 Download Video",
-                    visible=False,
-                    variant="primary"
-                )
+            # File component for download
+            download_file = gr.File(
+                label="📥 Download Video File",
+                visible=False
+            )
             
             status_display = gr.HTML(
                 value=(
@@ -589,31 +582,18 @@ with gr.Blocks(title="Self-Forcing Streaming Demo") as demo:
                 label="Generation Status"
             )
 
-    # Hidden component to store the download file path
-    download_file = gr.File(visible=False)
-
     # Connect the generator to the streaming video
     generation_event = start_btn.click(
         fn=video_generation_handler_streaming,
         inputs=[prompt, seed, fps],
-        outputs=[streaming_video, status_display, complete_video, download_button]
+        outputs=[streaming_video, status_display, complete_video, download_file]
     )
     
-    # When generation completes, show the complete video and download button
+    # When generation completes, show the complete video
     generation_event.then(
-        fn=lambda x: (gr.update(visible=True), gr.update(visible=True)),
-        inputs=[complete_video],
-        outputs=[complete_video, download_button]
-    )
-    
-    # Handle download button click
-    download_button.click(
-        fn=download_video,
+        fn=lambda x: gr.update(visible=True),
         inputs=[complete_video],
-        outputs=[download_file]
-    ).then(
-        fn=lambda: gr.update(visible=True),
-        outputs=[download_file]
+        outputs=[complete_video]
     )
     
     enhance_button.click(