Update app.py

app.py

@@ -663,38 +663,50 @@ with gr.Blocks(title="Self-Forcing Streaming Demo") as demo:
         
         # Collect all frames from this generation
         collected_frames = []
+        last_frame = None
+        last_status = None
+        generation_complete = False
         
-        # Call the generator function and yield frames, status, and video updates
         try:
-            # Set save_frames=True to explicitly ensure frames are collected
-            for frame, status_html in video_generation_handler_streaming(p, s, f, save_frames=True):
+            # Handle frame generation
+            for output in video_generation_handler_streaming(p, s, f, save_frames=True):
+                # Unpack the output correctly
+                if isinstance(output, tuple):
+                    if len(output) == 2:
+                        frame, status_html = output
+                    else:
+                        # Handle any unexpected output format gracefully
+                        continue
+                else:
+                    # Skip if not a proper tuple
+                    continue
+                    
+                # Save the last valid frame and status
+                if frame is not None:
+                    last_frame = frame
+                if status_html is not None:
+                    last_status = status_html
+                
                 # Track frames for this specific session
                 if frame is not None and isinstance(frame, np.ndarray):
                     collected_frames.append(frame.copy())
-                    
-                # Show status update during generation
-                yield frame, status_html, gr.update(visible=False)
                 
                 # Check if this is the final frame
-                if "Complete" in str(status_html) or "100%" in str(status_html):
-                    # Create the final video
-                    if collected_frames:
-                        print(f"Creating final video from {len(collected_frames)} frames at {f} FPS")
-                        temp_file = save_frames_as_video(collected_frames, f)
-                        if temp_file:
-                            # Save these frames as the current set
-                            APP_STATE["current_frames"] = collected_frames
-                            yield frame, status_html, gr.update(visible=True, value=temp_file)
-                        else:
-                            yield frame, status_html, gr.update(visible=False)
+                if status_html and ("Complete" in str(status_html) or "100%" in str(status_html)):
+                    generation_complete = True
+                    
+                # Always keep final video hidden during streaming
+                yield frame, status_html, gr.update(visible=False)
             
-            # Ensure final frame is properly handled
-            if collected_frames and "Complete" not in str(status_html) and "100%" not in str(status_html):
-                print(f"Generation complete, creating final video from {len(collected_frames)} frames")
+            # After streaming is done, create the final video
+            if collected_frames:
+                print(f"Generation complete, creating final video from {len(collected_frames)} frames at {f} FPS")
                 temp_file = save_frames_as_video(collected_frames, f)
                 if temp_file:
-                    yield frame, status_html, gr.update(visible=True, value=temp_file)
-                    
+                    # Save these frames as the current set
+                    APP_STATE["current_frames"] = collected_frames
+                    # Use the last valid frame and status
+                    yield last_frame, last_status, gr.update(visible=True, value=temp_file)
         except Exception as e:
             import traceback
             traceback.print_exc()