Update app.py

app.py

@@ -148,60 +148,57 @@ APP_STATE = {
     "current_vae_decoder": None,
 }
 
-# ONLY ADDITION: Store frames for download
+# Store frames for download
 DOWNLOAD_FRAMES = []
 
-def frames_to_ts_file(frames, filepath, fps = 15):
+def frames_to_mp4_chunk(frames, filepath, fps=15):
     """
-    Convert frames directly to .ts file using PyAV.
-    
-    Args:
-        frames: List of numpy arrays (HWC, RGB, uint8)
-        filepath: Output file path
-        fps: Frames per second
-    
-    Returns:
-        The filepath of the created file
+    Convert frames to MP4 chunk using imageio (more compatible than .ts for Gradio streaming)
     """
     if not frames:
         return filepath
     
-    height, width = frames[0].shape[:2]
-    
-    # Create container for MPEG-TS format
-    container = av.open(filepath, mode='w', format='mpegts')
-    
-    # Add video stream with optimized settings for streaming
-    stream = container.add_stream('h264', rate=fps)
-    stream.width = width
-    stream.height = height
-    stream.pix_fmt = 'yuv420p'
-    
-    # Optimize for low latency streaming
-    stream.options = {
-        'preset': 'ultrafast',
-        'tune': 'zerolatency', 
-        'crf': '23',
-        'profile': 'baseline',
-        'level': '3.0'
-    }
-    
     try:
-        for frame_np in frames:
-            frame = av.VideoFrame.from_ndarray(frame_np, format='rgb24')
-            frame = frame.reformat(format=stream.pix_fmt)
-            for packet in stream.encode(frame):
-                container.mux(packet)
+        # Use imageio to create MP4 chunk
+        with imageio.get_writer(filepath, fps=fps, codec='libx264', quality=6) as writer:
+            for frame_np in frames:
+                writer.append_data(frame_np)
         
-        for packet in stream.encode():
-            container.mux(packet)
+        return filepath
+        
+    except Exception as e:
+        print(f"❌ Error creating MP4 chunk: {e}")
+        # Fallback to PyAV if imageio fails
+        try:
+            height, width = frames[0].shape[:2]
+            container = av.open(filepath, mode='w', format='mp4')
             
-    finally:
-        container.close()
-    
-    return filepath
+            stream = container.add_stream('h264', rate=fps)
+            stream.width = width
+            stream.height = height
+            stream.pix_fmt = 'yuv420p'
+            stream.options = {
+                'preset': 'ultrafast',
+                'tune': 'zerolatency',
+                'crf': '28'
+            }
+            
+            for frame_np in frames:
+                frame = av.VideoFrame.from_ndarray(frame_np, format='rgb24')
+                frame = frame.reformat(format=stream.pix_fmt)
+                for packet in stream.encode(frame):
+                    container.mux(packet)
+            
+            for packet in stream.encode():
+                container.mux(packet)
+                
+            container.close()
+            return filepath
+            
+        except Exception as e2:
+            print(f"❌ Both imageio and PyAV failed: {e2}")
+            return filepath
 
-# ONLY ADDITION: Download function
 def create_download_mp4():
     global DOWNLOAD_FRAMES
     if not DOWNLOAD_FRAMES:
@@ -280,16 +277,15 @@ pipeline.to(dtype=torch.float16).to(gpu)
 @spaces.GPU  
 def video_generation_handler_streaming(prompt, seed=42, fps=15):
     """
-    Generator function that yields .ts video chunks using PyAV for streaming.
-    Now optimized for block-based processing.
+    Generator function that yields MP4 video chunks for streaming.
     """
     global DOWNLOAD_FRAMES
-    DOWNLOAD_FRAMES = []  # ONLY ADDITION: Reset frames
+    DOWNLOAD_FRAMES = []  # Reset frames
     
     if seed == -1: 
         seed = random.randint(0, 2**32 - 1)
     
-    print(f"🎬 Starting PyAV streaming: '{prompt}', seed: {seed}")
+    print(f"🎬 Starting MP4 streaming: '{prompt}', seed: {seed}")
     
     # Setup
     conditional_dict = text_encoder(text_prompts=[prompt])
@@ -376,15 +372,13 @@ def video_generation_handler_streaming(prompt, seed=42, fps=15):
             frame_np = np.transpose(frame_np, (1, 2, 0))  # CHW -> HWC
             
             all_frames_from_block.append(frame_np)
-            DOWNLOAD_FRAMES.append(frame_np)  # ONLY ADDITION: Store for download
+            DOWNLOAD_FRAMES.append(frame_np)  # Store for download
             total_frames_yielded += 1
             
-            # Yield status update for each frame (cute tracking!)
+            # Yield status update for each frame
             blocks_completed = idx
             current_block_progress = (frame_idx + 1) / pixels.shape[1]
             total_progress = (blocks_completed + current_block_progress) / num_blocks * 100
-            
-            # Cap at 100% to avoid going over
             total_progress = min(total_progress, 100.0)
             
             frame_status_html = (
@@ -399,25 +393,21 @@ def video_generation_handler_streaming(prompt, seed=42, fps=15):
                 f"</div>"
             )
             
-            # Yield None for video but update status (frame-by-frame tracking)
             yield None, frame_status_html
 
-        # Encode entire block as one chunk immediately
+        # Create MP4 chunk for this block
         if all_frames_from_block:
             print(f"📹 Encoding 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)
+                mp4_filename = f"block_{idx:04d}_{chunk_uuid}.mp4"
+                mp4_path = os.path.join("gradio_tmp", mp4_filename)
                 
-                # Calculate final progress for this block
-                total_progress = (idx + 1) / num_blocks * 100
+                frames_to_mp4_chunk(all_frames_from_block, mp4_path, fps)
                 
-                # Yield the actual video chunk
-                yield ts_path, gr.update()
+                # Yield the MP4 chunk
+                yield mp4_path, gr.update()
                 
             except Exception as e:
                 print(f"⚠️ Error encoding block {idx}: {e}")
@@ -438,13 +428,13 @@ def video_generation_handler_streaming(prompt, seed=42, fps=15):
         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 • 📥 Download ready!"
+        f"      🎬 Playback: {fps} FPS • 📁 Format: MP4/H.264 • 📥 Download ready!"
         f"    </p>"
         f"  </div>"
         f"</div>"
     )
     yield None, final_status_html
-    print(f"✅ PyAV streaming complete! {total_frames_yielded} frames across {num_blocks} blocks")
+    print(f"✅ MP4 streaming complete! {total_frames_yielded} frames across {num_blocks} blocks")
 
 # --- Gradio UI Layout ---
 with gr.Blocks(title="Self-Forcing Streaming Demo") as demo:
@@ -514,7 +504,7 @@ with gr.Blocks(title="Self-Forcing Streaming Demo") as demo:
                 label="Generation Status"
             )
             
-            # ONLY ADDITION: Download button
+            # Download button
             download_btn = gr.DownloadButton(
                 label="📥 Download MP4",
                 value=create_download_mp4,
@@ -540,12 +530,12 @@ if __name__ == "__main__":
         import shutil
         shutil.rmtree("gradio_tmp")
     os.makedirs("gradio_tmp", exist_ok=True)
-    os.makedirs("downloads", exist_ok=True)  # ONLY ADDITION
+    os.makedirs("downloads", exist_ok=True)
     
     print("🚀 Starting Self-Forcing Streaming Demo")
     print(f"📁 Temporary files will be stored in: gradio_tmp/")
-    print(f"📥 Download files will be stored in: downloads/")  # ONLY ADDITION
-    print(f"🎯 Chunk encoding: PyAV (MPEG-TS/H.264)")
+    print(f"📥 Download files will be stored in: downloads/")
+    print(f"🎯 Chunk encoding: MP4/H.264 (more compatible)")
     print(f"⚡ GPU acceleration: {gpu}")
     
     demo.queue().launch(