Update app.py

app.py

@@ -148,7 +148,7 @@ APP_STATE = {
     "current_vae_decoder": None,
 }
 
-# Global variable to store frames for download
+# ONLY ADDITION: Store frames for download
 DOWNLOAD_FRAMES = []
 
 def frames_to_ts_file(frames, filepath, fps = 15):
@@ -201,59 +201,22 @@ def frames_to_ts_file(frames, filepath, fps = 15):
     
     return filepath
 
-def create_hls_playlist(ts_files, playlist_dir, fps=15):
-    """
-    Create HLS playlist (.m3u8) file for streaming.
-    """
-    playlist_path = os.path.join(playlist_dir, "playlist.m3u8")
-    segment_duration = 2.0  # Each segment duration in seconds
-    
-    playlist_content = [
-        "#EXTM3U",
-        "#EXT-X-VERSION:3",
-        f"#EXT-X-TARGETDURATION:{int(segment_duration) + 1}",
-        "#EXT-X-MEDIA-SEQUENCE:0",
-        "#EXT-X-PLAYLIST-TYPE:VOD"
-    ]
-    
-    for ts_file in ts_files:
-        ts_filename = os.path.basename(ts_file)
-        playlist_content.extend([
-            f"#EXTINF:{segment_duration:.1f},",
-            ts_filename
-        ])
-    
-    playlist_content.append("#EXT-X-ENDLIST")
-    
-    with open(playlist_path, 'w') as f:
-        f.write('\n'.join(playlist_content))
-    
-    return playlist_path
-
-def create_mp4_download():
-    """Create MP4 file from stored frames for download."""
+# ONLY ADDITION: Download function
+def create_download_mp4():
     global DOWNLOAD_FRAMES
-    
     if not DOWNLOAD_FRAMES:
         return None
-    
     try:
         os.makedirs("downloads", exist_ok=True)
-        
         timestamp = int(time.time())
-        mp4_filename = f"pixio_video_{timestamp}.mp4"
-        mp4_path = os.path.join("downloads", mp4_filename)
-        
-        # Use imageio to create MP4
+        mp4_path = f"downloads/video_{timestamp}.mp4"
         with imageio.get_writer(mp4_path, fps=args.fps, codec='libx264', quality=8) as writer:
             for frame in DOWNLOAD_FRAMES:
                 writer.append_data(frame)
-        
-        print(f"✅ MP4 created for download: {mp4_path}")
+        print(f"✅ Download MP4 created: {mp4_path}")
         return mp4_path
-        
     except Exception as e:
-        print(f"❌ Error creating MP4: {e}")
+        print(f"❌ Download error: {e}")
         return None
 
 def initialize_vae_decoder(use_taehv=False, use_trt=False):
@@ -317,20 +280,16 @@ pipeline.to(dtype=torch.float16).to(gpu)
 @spaces.GPU  
 def video_generation_handler_streaming(prompt, seed=42, fps=15):
     """
-    Generator function that creates HLS stream and stores frames for download.
+    Generator function that yields .ts video chunks using PyAV for streaming.
+    Now optimized for block-based processing.
     """
     global DOWNLOAD_FRAMES
-    DOWNLOAD_FRAMES = []  # Reset frames for new generation
+    DOWNLOAD_FRAMES = []  # ONLY ADDITION: Reset frames
     
     if seed == -1: 
         seed = random.randint(0, 2**32 - 1)
     
-    print(f"🎬 Starting HLS streaming: '{prompt}', seed: {seed}")
-    
-    # Create unique session directory for HLS files
-    session_id = str(uuid.uuid4())[:8]
-    session_dir = os.path.join("gradio_tmp", f"session_{session_id}")
-    os.makedirs(session_dir, exist_ok=True)
+    print(f"🎬 Starting PyAV streaming: '{prompt}', seed: {seed}")
     
     # Setup
     conditional_dict = text_encoder(text_prompts=[prompt])
@@ -351,7 +310,9 @@ def video_generation_handler_streaming(prompt, seed=42, fps=15):
     all_num_frames = [pipeline.num_frame_per_block] * num_blocks
     
     total_frames_yielded = 0
-    ts_files = []  # Store TS file paths for HLS playlist
+    
+    # Ensure temp directory exists
+    os.makedirs("gradio_tmp", exist_ok=True)
     
     # Generation loop
     for idx, current_num_frames in enumerate(all_num_frames):
@@ -415,7 +376,7 @@ 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)  # Store for download
+            DOWNLOAD_FRAMES.append(frame_np)  # ONLY ADDITION: Store for download
             total_frames_yielded += 1
             
             # Yield status update for each frame (cute tracking!)
@@ -441,25 +402,25 @@ 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
 
-        # Create TS segment for this block
+        # Encode entire block as one chunk immediately
         if all_frames_from_block:
             print(f"📹 Encoding block {idx} with {len(all_frames_from_block)} frames")
             
             try:
-                ts_filename = f"segment_{idx:04d}.ts"
-                ts_path = os.path.join(session_dir, ts_filename)
+                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)
-                ts_files.append(ts_path)
                 
-                # Create/update HLS playlist
-                playlist_path = create_hls_playlist(ts_files, session_dir, fps)
+                # Calculate final progress for this block
+                total_progress = (idx + 1) / num_blocks * 100
                 
-                # Yield the HLS playlist for streaming
-                yield playlist_path, gr.update()
+                # Yield the actual video chunk
+                yield ts_path, gr.update()
                 
             except Exception as e:
-                print(f"⚠️ Error creating HLS segment {idx}: {e}")
+                print(f"⚠️ Error encoding block {idx}: {e}")
                 import traceback
                 traceback.print_exc()
                     
@@ -477,13 +438,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: HLS/H.264 • 📥 Download ready!"
+        f"      🎬 Playback: {fps} FPS • 📁 Format: MPEG-TS/H.264 • 📥 Download ready!"
         f"    </p>"
         f"  </div>"
         f"</div>"
     )
     yield None, final_status_html
-    print(f"✅ HLS streaming complete! {total_frames_yielded} frames across {num_blocks} blocks")
+    print(f"✅ PyAV streaming complete! {total_frames_yielded} frames across {num_blocks} blocks")
 
 # --- Gradio UI Layout ---
 with gr.Blocks(title="Self-Forcing Streaming Demo") as demo:
@@ -553,10 +514,10 @@ with gr.Blocks(title="Self-Forcing Streaming Demo") as demo:
                 label="Generation Status"
             )
             
-            # Download button
+            # ONLY ADDITION: Download button
             download_btn = gr.DownloadButton(
                 label="📥 Download MP4",
-                value=create_mp4_download,
+                value=create_download_mp4,
                 variant="secondary"
             )
 
@@ -579,13 +540,13 @@ if __name__ == "__main__":
         import shutil
         shutil.rmtree("gradio_tmp")
     os.makedirs("gradio_tmp", exist_ok=True)
-    os.makedirs("downloads", exist_ok=True)
+    os.makedirs("downloads", exist_ok=True)  # ONLY ADDITION
     
     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/")
-    print(f"🎯 Streaming: HLS (.m3u8 + .ts segments)")
-    print(f"📱 Download: MP4 (imageio)")
+    print(f"📥 Download files will be stored in: downloads/")  # ONLY ADDITION
+    print(f"🎯 Chunk encoding: PyAV (MPEG-TS/H.264)")
+    print(f"⚡ GPU acceleration: {gpu}")
     
     demo.queue().launch(
         server_name=args.host,