Update app.py

app.py

@@ -201,6 +201,35 @@ 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."""
     global DOWNLOAD_FRAMES
@@ -288,8 +317,7 @@ 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 creates HLS stream and stores frames for download.
     """
     global DOWNLOAD_FRAMES
     DOWNLOAD_FRAMES = []  # Reset frames for new generation
@@ -297,7 +325,12 @@ def video_generation_handler_streaming(prompt, seed=42, fps=15):
     if seed == -1: 
         seed = random.randint(0, 2**32 - 1)
     
-    print(f"🎬 Starting PyAV streaming: '{prompt}', seed: {seed}")
+    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)
     
     # Setup
     conditional_dict = text_encoder(text_prompts=[prompt])
@@ -318,9 +351,7 @@ def video_generation_handler_streaming(prompt, seed=42, fps=15):
     all_num_frames = [pipeline.num_frame_per_block] * num_blocks
     
     total_frames_yielded = 0
-    
-    # Ensure temp directory exists
-    os.makedirs("gradio_tmp", exist_ok=True)
+    ts_files = []  # Store TS file paths for HLS playlist
     
     # Generation loop
     for idx, current_num_frames in enumerate(all_num_frames):
@@ -410,25 +441,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
 
-        # Encode entire block as one chunk immediately
+        # Create TS segment 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)
+                ts_filename = f"segment_{idx:04d}.ts"
+                ts_path = os.path.join(session_dir, ts_filename)
                 
                 frames_to_ts_file(all_frames_from_block, ts_path, fps)
+                ts_files.append(ts_path)
                 
-                # Calculate final progress for this block
-                total_progress = (idx + 1) / num_blocks * 100
+                # Create/update HLS playlist
+                playlist_path = create_hls_playlist(ts_files, session_dir, fps)
                 
-                # Yield the actual video chunk
-                yield ts_path, gr.update()
+                # Yield the HLS playlist for streaming
+                yield playlist_path, gr.update()
                 
             except Exception as e:
-                print(f"⚠️ Error encoding block {idx}: {e}")
+                print(f"⚠️ Error creating HLS segment {idx}: {e}")
                 import traceback
                 traceback.print_exc()
                     
@@ -446,13 +477,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: HLS/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"✅ HLS streaming complete! {total_frames_yielded} frames across {num_blocks} blocks")
 
 # --- Gradio UI Layout ---
 with gr.Blocks(title="Self-Forcing Streaming Demo") as demo:
@@ -553,8 +584,8 @@ if __name__ == "__main__":
     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"🎯 Chunk encoding: PyAV (MPEG-TS/H.264)")
-    print(f"⚡ GPU acceleration: {gpu}")
+    print(f"🎯 Streaming: HLS (.m3u8 + .ts segments)")
+    print(f"📱 Download: MP4 (imageio)")
     
     demo.queue().launch(
         server_name=args.host, 
@@ -565,6 +596,7 @@ if __name__ == "__main__":
         mcp_server=True
     )
 
+
 # import subprocess
 # subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)