Update app.py

app.py

@@ -146,7 +146,7 @@ pipeline.to(dtype=torch.float16).to(gpu)
 # --- Frame Streaming Video Generation Handler ---
 @torch.no_grad()
 @spaces.GPU
-def video_generation_handler(prompt, seed, fps, progress=gr.Progress()):
+def video_generation_handler(prompt, seed, fps):
     """
     Generator function that yields RGB frames for display in gr.Image.
     Includes timing delays for smooth playback.
@@ -230,8 +230,11 @@ def video_generation_handler(prompt, seed, fps, progress=gr.Progress()):
 
         print(f"📹 Decoded pixels shape: {pixels.shape}")
         
+        # Calculate actual frames that will be yielded for this block
+        actual_frames_this_block = pixels.shape[1]
+        
         # Yield individual frames with timing delays
-        for frame_idx in range(pixels.shape[1]):
+        for frame_idx in range(actual_frames_this_block):
             frame_tensor = pixels[0, frame_idx]  # Get single frame [C, H, W]
             
             # Normalize from [-1, 1] to [0, 255]
@@ -244,22 +247,47 @@ def video_generation_handler(prompt, seed, fps, progress=gr.Progress()):
             all_frames_for_video.append(frame_np)
             total_frames_yielded += 1
             
-            # Calculate progress
-            total_expected_frames = num_blocks * pipeline.num_frame_per_block
-            current_frame_count = (idx * pipeline.num_frame_per_block) + frame_idx + 1
-            frame_progress = current_frame_count / total_expected_frames
+            # Calculate progress based on blocks completed + current block progress
+            blocks_completed = idx
+            current_block_progress = (frame_idx + 1) / actual_frames_this_block
+            total_block_progress = (blocks_completed + current_block_progress) / num_blocks
+            frame_progress_percent = total_block_progress * 100
             
-            # Update progress
-            progress(frame_progress, desc=f"Frame {total_frames_yielded} | Block {idx+1}/{num_blocks}")
+            # Cap at 100% to avoid going over
+            frame_progress_percent = min(frame_progress_percent, 100.0)
             
             print(f"📺 Yielding frame {total_frames_yielded}: shape {frame_np.shape}")
             
-            # Yield frame with timing delay
-            yield gr.update(visible=True, frame_np), gr.update(visible=False)
+            # Create HTML status update
+            if frame_idx == actual_frames_this_block - 1 and idx + 1 == num_blocks:  # Last frame
+                status_html = (
+                    f"<div style='padding: 16px; border: 1px solid #198754; background-color: #d1e7dd; border-radius: 8px; font-family: sans-serif; text-align: center;'>"
+                    f"  <h4 style='margin: 0 0 8px 0; color: #0f5132; font-size: 18px;'>🎉 Generation Complete!</h4>"
+                    f"  <p style='margin: 0; color: #0f5132;'>"
+                    f"    Total frames: {total_frames_yielded}. The final video is now available."
+                    f"  </p>"
+                    f"</div>"
+                )
+            else:  # Regular frames
+                status_html = (
+                    f"<div style='padding: 10px; border: 1px solid #ddd; border-radius: 8px; font-family: sans-serif;'>"
+                    f"  <p style='margin: 0 0 8px 0; font-size: 16px; font-weight: bold;'>Generating Video...</p>"
+                    f"  <div style='background: #e9ecef; border-radius: 4px; width: 100%; overflow: hidden;'>"
+                    f"    <div style='width: {frame_progress_percent:.1f}%; height: 20px; background-color: #0d6efd; transition: width 0.2s;'></div>"
+                    f"  </div>"
+                    f"  <p style='margin: 8px 0 0 0; color: #555; font-size: 14px; text-align: right;'>"
+                    f"    Block {idx+1}/{num_blocks}   |   Frame {total_frames_yielded}   |   {frame_progress_percent:.1f}%"
+                    f"  </p>"
+                    f"</div>"
+                )
+            
+            # Yield frame with a small delay to ensure UI updates
+            yield gr.update(visible=True, value=frame_np), gr.update(visible=False), status_html
             
             # Sleep between frames for smooth playback (except for the last frame)
-            if not (frame_idx == pixels.shape[1] - 1 and idx + 1 == num_blocks):
-                time.sleep(frame_delay)
+            # Add minimum delay to ensure UI can update
+            if not (frame_idx == actual_frames_this_block - 1 and idx + 1 == num_blocks):
+                time.sleep(max(frame_delay, 0.1))  # Minimum 100ms delay
             
         current_start_frame += current_num_frames
     
@@ -270,10 +298,26 @@ def video_generation_handler(prompt, seed, fps, progress=gr.Progress()):
         video_path = f"gradio_tmp/{seed}_{hashlib.md5(prompt.encode()).hexdigest()}.mp4"
         imageio.mimwrite(video_path, all_frames_for_video, fps=fps, quality=8)
         print(f"✅ Video saved to {video_path}")
-        return gr.update(visible=False), gr.update(value=video_path, visible=True)
+        final_status_html = (
+            f"<div style='padding: 16px; border: 1px solid #198754; background-color: #d1e7dd; border-radius: 8px; font-family: sans-serif; text-align: center;'>"
+            f"  <h4 style='margin: 0 0 8px 0; color: #0f5132; font-size: 18px;'>🎉 Generation Complete!</h4>"
+            f"  <p style='margin: 0; color: #0f5132;'>"
+            f"    Video saved successfully with {total_frames_yielded} frames at {fps} FPS."
+            f"  </p>"
+            f"</div>"
+        )
+        yield gr.update(visible=False), gr.update(value=video_path, visible=True), final_status_html
     except Exception as e:
         print(f"⚠️ Could not save final video: {e}")
-        return None, None
+        error_status_html = (
+            f"<div style='padding: 16px; border: 1px solid #dc3545; background-color: #f8d7da; border-radius: 8px; font-family: sans-serif; text-align: center;'>"
+            f"  <h4 style='margin: 0 0 8px 0; color: #721c24; font-size: 18px;'>⚠️ Video Save Error</h4>"
+            f"  <p style='margin: 0; color: #721c24;'>"
+            f"    Could not save final video: {str(e)}"
+            f"  </p>"
+            f"</div>"
+        )
+        yield None, None, error_status_html
 
 # --- Gradio UI Layout ---
 with gr.Blocks(theme=gr.themes.Soft(), title="Self-Forcing Frame Streaming Demo") as demo:
@@ -326,17 +370,22 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Self-Forcing Frame Streaming Demo"
             
             final_video = gr.Video(
                 label="Final Rendered Video", 
-                visible=True, 
+                visible=False, 
                 interactive=False,
-                height=400
+                height=400,
+                autoplay=True
+            )
+            
+            status_html = gr.HTML(
+                value="<div style='text-align: center; padding: 20px; color: #666;'>Ready to start generation...</div>",
+                label="Generation Status"
             )
 
     # Connect the generator to the image display
     start_btn.click(
         fn=video_generation_handler,
         inputs=[prompt, seed, fps],
-        outputs=[frame_display, final_video],
-        show_progress="full"
+        outputs=[frame_display, final_video, status_html]
     )
 
 # --- Launch App ---