Update app.py

app.py

@@ -176,13 +176,18 @@ def frames_to_ts_file(frames, filepath, fps = 15):
     stream.height = height
     stream.pix_fmt = 'yuv420p'
     
-    # Optimize for low latency streaming
+    # Optimize for low latency streaming with better buffering
     stream.options = {
-        'preset': 'ultrafast',
-        'tune': 'zerolatency', 
-        'crf': '23',
-        'profile': 'baseline',
-        'level': '3.0'
+        'preset': 'ultrafast',  # Speed over quality for real-time
+        'tune': 'zerolatency',  # Reduce latency
+        'crf': '28',            # Slightly lower quality (higher number) for better throughput
+        'profile': 'baseline',   # Simpler profile for better compatibility
+        'level': '3.0',         # Compatibility level
+        'g': '15',              # Keyframe interval matching fps for better seeking
+        'b:v': '2000k',         # Target bitrate - reducing for smoother playback
+        'maxrate': '2500k',     # Maximum bitrate
+        'bufsize': '5000k',     # Larger buffer size
+        'sc_threshold': '0'     # Disable scene detection for smoother streaming
     }
     
     try:
@@ -423,6 +428,7 @@ def video_generation_handler_streaming(prompt, seed=42, fps=15, save_frames=True
     yield None, final_status_html
     print(f" PyAV streaming complete! {total_frames_yielded} frames across {num_blocks} blocks")
 
+# Function to save frames as downloadable video
 def save_frames_as_video(frames, fps=15):
     """
     Convert frames to a downloadable MP4 video file.
@@ -435,21 +441,66 @@ def save_frames_as_video(frames, fps=15):
         Path to the saved video file
     """
     if not frames:
+        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")
     
-    # Use imageio to write the video file
+    # Use PyAV for better quality and reliability
     try:
-        writer = imageio.get_writer(temp_file, fps=fps, codec='h264', quality=9)
-        for frame in frames:
-            writer.append_data(frame)
-        writer.close()
-        return temp_file
+        # First try PyAV which has better compatibility
+        container = av.open(temp_file, mode='w')
+        stream = container.add_stream('h264', rate=fps)
+        
+        # Get dimensions from first frame
+        height, width = frames[0].shape[:2]
+        stream.width = width
+        stream.height = height
+        stream.pix_fmt = 'yuv420p'
+        
+        # Use higher quality for downloads
+        stream.options = {
+            'preset': 'medium',    # Better quality than ultrafast
+            'crf': '23',          # Better quality than streaming
+            'profile': 'high',    # Higher quality profile
+            'g': f'{fps*2}',      # GOP size
+            'b:v': '4000k',       # Higher bitrate for downloads
+            'refs': '3'           # Number of reference frames
+        }
+        
+        print(f"Saving video with {len(frames)} frames at {fps} FPS")
+        for frame_np in frames:
+            frame = av.VideoFrame.from_ndarray(frame_np, format='rgb24')
+            for packet in stream.encode(frame):
+                container.mux(packet)
+                
+        # Flush the stream
+        for packet in stream.encode():
+            container.mux(packet)
+            
+        container.close()
+        
+        # Verify the file exists and has content
+        if os.path.exists(temp_file) and os.path.getsize(temp_file) > 0:
+            print(f"Video saved successfully: {temp_file} ({os.path.getsize(temp_file)} bytes)")
+            return temp_file
+        else:
+            print("Video file is empty or missing, falling back to imageio")
+            raise RuntimeError("Empty file created")
+            
     except Exception as e:
-        print(f"Error saving video: {e}")
-        return None
+        # Fall back to imageio if PyAV fails
+        print(f"PyAV encoding failed: {e}, falling back to imageio")
+        try:
+            writer = imageio.get_writer(temp_file, fps=fps, codec='h264', quality=9, bitrate='4000k')
+            for frame in frames:
+                writer.append_data(frame)
+            writer.close()
+            return temp_file
+        except Exception as e2:
+            print(f"Error saving video with imageio: {e2}")
+            return None
 
 # --- Gradio UI Layout ---
 with gr.Blocks(title="Self-Forcing Streaming Demo") as demo:
@@ -506,7 +557,8 @@ with gr.Blocks(title="Self-Forcing Streaming Demo") as demo:
                 loop=True,
                 height=400,
                 autoplay=True,
-                show_label=False
+                show_label=False,
+                format="mp4"  # Use more stable mp4 format when possible
             )
             
             status_display = gr.HTML(
@@ -559,9 +611,13 @@ with gr.Blocks(title="Self-Forcing Streaming Demo") as demo:
         fn=download_video,
         inputs=[fps],
         outputs=[download_output],
-        show_progress=True
+        show_progress=True,
+        api_name="download_video"  # Make it accessible via API
     )
     
+    # Make the FPS slider visible for download quality control
+    fps.visible = True
+    
     enhance_button.click(
         fn=enhance_prompt,
         inputs=[prompt],