Update app.py

app.py

@@ -264,15 +264,15 @@ pipeline.to(dtype=torch.float16).to(gpu)
 
 @torch.no_grad()
 @spaces.GPU  
-def video_generation_handler_streaming(prompt, seed=42, fps=15, save_frames=True):
+def video_generation_handler_frame_by_frame(prompt, seed=42, fps=15, save_frames=True):
     """
-    Generator function that yields .ts video chunks using PyAV for streaming.
-    Now optimized for block-based processing.
+    Generator function that yields individual frames and status updates.
+    No streaming - just frame by frame display.
     """
     if seed == -1: 
         seed = random.randint(0, 2**32 - 1)
     
-    print(f"🎬 Starting PyAV streaming: '{prompt}', seed: {seed}")
+    print(f"🎬 Starting frame-by-frame generation: '{prompt}', seed: {seed}")
     
     # Setup
     conditional_dict = text_encoder(text_prompts=[prompt])
@@ -381,55 +381,56 @@ def video_generation_handler_streaming(prompt, seed=42, fps=15, save_frames=True
                 f"</div>"
             )
             
-            # Yield None for video but update status (frame-by-frame tracking)
-            yield None, frame_status_html
+            # No streaming - show the current frame and update status
+            yield frame_np, frame_status_html
 
-        # Encode entire block as one chunk immediately
+        # Save frames for download without streaming
         if all_frames_from_block:
-            print(f"📹 Encoding block {idx} with {len(all_frames_from_block)} frames")
+            print(f"💹 Processed block {idx} with {len(all_frames_from_block)} frames")
             
-            try:
-                chunk_uuid = str(uuid.uuid4())[:8]
-                ts_filename = f"block_{idx:04d}_{chunk_uuid}.ts"
-                ts_path = os.path.join("gradio_tmp", ts_filename)
-                
-                frames_to_ts_file(all_frames_from_block, ts_path, fps)
-                
-                # Calculate final progress for this block
-                total_progress = (idx + 1) / num_blocks * 100
-                
-                # Yield the actual video chunk
-                yield ts_path, gr.update()
-                
-            except Exception as e:
-                print(f"⚠️ Error encoding block {idx}: {e}")
-                import traceback
-                traceback.print_exc()
+            # We already yielded each frame individually for display
+            # No need to encode video chunks for streaming anymore
                     
         current_start_frame += current_num_frames
     
-    # Final completion status
-    final_status_html = (
-        f"<div style='padding: 16px; border: 1px solid #198754; background: linear-gradient(135deg, #d1e7dd, #f8f9fa); border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.1);'>"
-        f"  <div style='display: flex; align-items: center; margin-bottom: 8px;'>"
-        f"    <span style='font-size: 24px; margin-right: 12px;'>🎉</span>"
-        f"    <h4 style='margin: 0; color: #0f5132; font-size: 18px;'>Stream Complete!</h4>"
-        f"  </div>"
-        f"  <div style='background: rgba(255,255,255,0.7); padding: 8px; border-radius: 4px;'>"
-        f"    <p style='margin: 0; color: #0f5132; font-weight: 500;'>"
-        f"      📊 Generated {total_frames_yielded} frames across {num_blocks} blocks"
-        f"    </p>"
-        f"    <p style='margin: 4px 0 0 0; color: #0f5132; font-size: 14px;'>"
-        f"      🎬 Playback: {fps} FPS • 📁 Format: MPEG-TS/H.264"
-        f"    </p>"
-        f"  </div>"
-        f"</div>"
-    )
-    yield None, final_status_html
-    print(f" PyAV streaming complete! {total_frames_yielded} frames across {num_blocks} blocks")
+    # Generate final video preview if we have frames
+    if APP_STATE["current_frames"]:
+        # Create a temporary preview file
+        preview_file = os.path.join("gradio_tmp", f"preview_{uuid.uuid4()}.mp4")
+        try:
+            # Save a preview video file
+            save_frames_as_video(APP_STATE["current_frames"], fps, preview_file)
+            
+            # Final completion status with success message
+            final_status_html = (
+                f"<div style='padding: 16px; border: 1px solid #198754; background: linear-gradient(135deg, #d1e7dd, #f8f9fa); border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.1);'>"
+                f"  <div style='display: flex; align-items: center; margin-bottom: 8px;'>"
+                f"    <span style='font-size: 24px; margin-right: 12px;'>🎉</span>"
+                f"    <h4 style='margin: 0; color: #0f5132; font-size: 18px;'>Generation Complete!</h4>"
+                f"  </div>"
+                f"  <div style='background: rgba(255,255,255,0.7); padding: 8px; border-radius: 4px;'>"
+                f"    <p style='margin: 0; color: #0f5132; font-weight: 500;'>"
+                f"      📈 Generated {total_frames_yielded} frames across {num_blocks} blocks"
+                f"    </p>"
+                f"    <p style='margin: 4px 0 0 0; color: #0f5132; font-size: 14px;'>"
+                f"      🎬 Preview available • Click Download to save as MP4"
+                f"    </p>"
+                f"  </div>"
+                f"</div>"
+            )
+            
+            # Return the last frame and completion message along with final video
+            yield APP_STATE["current_frames"][-1], final_status_html, gr.update(value=preview_file, visible=True)
+        except Exception as e:
+            print(f"Error creating preview: {e}")
+            # Just return the last frame and completion message
+            final_status_html = f"<div style='color: green; padding: 10px;'>Generation complete! {total_frames_yielded} frames generated. Ready to download.</div>"
+            yield APP_STATE["current_frames"][-1], final_status_html, gr.update(visible=False)
+    
+    print(f"✅ Generation complete! {total_frames_yielded} frames across {num_blocks} blocks")
 
 # Function to save frames as downloadable video
-def save_frames_as_video(frames, fps=15):
+def save_frames_as_video(frames, fps=15, output_path=None):
     """
     Convert frames to a downloadable MP4 video file.
     
@@ -444,8 +445,8 @@ def save_frames_as_video(frames, fps=15):
         print("No frames available to save")
         return None
     
-    # Create a temporary file with a unique name
-    temp_file = os.path.join("gradio_tmp", f"download_{uuid.uuid4()}.mp4")
+    # Create a temporary file with a unique name or use provided path
+    temp_file = output_path if output_path else os.path.join("gradio_tmp", f"download_{uuid.uuid4()}.mp4")
     
     # Use PyAV for better quality and reliability
     try:
@@ -549,16 +550,21 @@ with gr.Blocks(title="Self-Forcing Streaming Demo") as demo:
                 )
             
         with gr.Column(scale=3):
-            gr.Markdown("### 📺 Video Stream")
+            gr.Markdown("### 📺 Video Preview")
 
-            streaming_video = gr.Video(
-                label="Live Stream",
-                streaming=True,
-                loop=True,
+            # Replace streaming video with image display
+            preview_image = gr.Image(
+                label="Current Frame",
                 height=400,
-                autoplay=True,
                 show_label=False,
-                format="mp4"  # Use more stable mp4 format when possible
+            )
+            
+            # Add a non-streaming video component for final result preview
+            final_video = gr.Video(
+                label="Final Video Preview",
+                visible=False,
+                autoplay=True,
+                loop=True
             )
             
             status_display = gr.HTML(
@@ -574,29 +580,34 @@ with gr.Blocks(title="Self-Forcing Streaming Demo") as demo:
             download_btn = gr.Button("💾 Download Video", variant="secondary")
             download_output = gr.File(label="Download")
 
-    # Define a wrapper function to ensure we always return the expected outputs
-    def safe_video_generator(p, s, f):
+    # Define a wrapper function to ensure proper handling of outputs
+    def safe_frame_generator(p, s, f):
         # Clear frames from previous generation
         APP_STATE["current_frames"] = []
         
-        # Call the generator function and ensure it returns both video and status
+        # Reset the final video display
+        yield None, None, gr.update(visible=False)
+        
+        # Call the generator function and yield frames, status, and video updates
         try:
-            for video_path, status_html in video_generation_handler_streaming(p, s, f):
-                if video_path is not None:
-                    yield video_path, status_html
-                else:
-                    yield None, status_html
+            for frame, status_html, *video_update in video_generation_handler_frame_by_frame(p, s, f):
+                # Ensure we always have three outputs
+                if not video_update:
+                    video_update = [gr.update(visible=False)]  # Default - hide video
+                    
+                yield frame, status_html, video_update[0]
         except Exception as e:
             import traceback
             print(f"Error in generator: {e}")
             traceback.print_exc()
-            yield None, 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: {e}</div>"
+            yield None, error_html, gr.update(visible=False)
     
-    # Connect the generator to the streaming video
+    # Connect the generator to the image and status display
     start_btn.click(
-        fn=safe_video_generator,
+        fn=safe_frame_generator,
         inputs=[prompt, seed, fps],
-        outputs=[streaming_video, status_display]
+        outputs=[preview_image, status_display, final_video]
     )
     
     # Function to handle download button click