Update app.py

app.py

@@ -17,8 +17,12 @@ except ImportError:
 
 import base64
 import os
+import queue
 import random
-from dataclasses import dataclass
+import threading
+import time
+from collections import deque
+from dataclasses import dataclass, field
 from io import BytesIO
 from multiprocessing import Queue
 from pathlib import Path
@@ -54,6 +58,20 @@ MODEL_ID = os.environ.get("MODEL_PATH", "diffusers-internal-dev/world-engine-mod
 pipe = ModularPipeline.from_pretrained(MODEL_ID, trust_remote_code=True, revision="aot-compatible")
 pipe.load_components(["transformer", "vae"], trust_remote_code=True, revision="aot-compatible", torch_dtype=torch.bfloat16)
 pipe.load_components(["text_encoder", "tokenizer"], trust_remote_code=True, torch_dtype=torch.bfloat16)
+pipe.to("cuda")
+pipe.blocks.sub_blocks['before_denoise'].sub_blocks['setup_kv_cache']._setup_kv_cache(pipe.transformer, pipe.device, torch.bfloat16)
+aoti_load_(
+    pipe.transformer,
+    "diffusers-internal-dev/world-engine-aot",
+    "transformer-fp8.pt2",
+    "transformer-fp8-constants.pt"
+)
+aoti_load_(
+    pipe.vae.decoder,
+    "diffusers-internal-dev/world-engine-aot",
+    "decoder.pt2",
+    "decoder-constants.pt"
+)
 
 SEED_FRAME_URLS = [
     "https://gist.github.com/user-attachments/assets/5d91c49a-2ae9-418f-99c0-e93ae387e1de",
@@ -137,9 +155,65 @@ class StopCommand:
     pass
 
 
+# --- Session State ---
+@dataclass
+class GameSession:
+    """Per-user game session with background worker thread."""
+    command_queue: Queue
+    frame_queue: queue.Queue  # Thread-safe queue for output frames
+    worker_thread: threading.Thread
+    stop_event: threading.Event
+    generator: object = None  # The GPU generator
+    frame_times: deque = field(default_factory=lambda: deque(maxlen=30))  # Track last 30 frame times for FPS
+
+
+def gpu_worker_thread(gen, command_queue, frame_queue, stop_event, frame_times):
+    """
+    Worker thread that consumes the GPU generator and pushes frames to frame_queue.
+    Runs independently of Gradio's timer, allowing non-blocking frame reads.
+    """
+    try:
+        while not stop_event.is_set():
+            # Get next frame from generator (this blocks on GPU work)
+            try:
+                frame, frame_count = next(gen)
+
+                # Track frame generation time for FPS
+                now = time.time()
+                frame_times.append(now)
+
+                # Calculate FPS from generation times
+                if len(frame_times) >= 2:
+                    elapsed = frame_times[-1] - frame_times[0]
+                    fps = (len(frame_times) - 1) / elapsed if elapsed > 0 else 0.0
+                else:
+                    fps = 0.0
+
+                # Put frame in queue, replacing old frame if queue is full
+                try:
+                    # Clear old frame if present (keep only latest)
+                    while not frame_queue.empty():
+                        try:
+                            frame_queue.get_nowait()
+                        except queue.Empty:
+                            break
+                except:
+                    pass
+                frame_queue.put_nowait((frame, frame_count, round(fps, 1)))
+            except StopIteration:
+                print("Generator exhausted, worker thread ending")
+                break
+            except Exception as e:
+                print(f"Worker thread error: {e}")
+                break
+    finally:
+        print("Worker thread finished")
+
+
 # --- GPU Session Generator ---
 def create_gpu_game_loop(command_queue: Queue, initial_seed_image=None, initial_seed_url=None, initial_prompt="An explorable world"):
     """Create GPU game loop generator with closure over command_queue."""
+    print(f"create_gpu_game_loop: initial_seed_image={type(initial_seed_image)}, initial_seed_url={initial_seed_url}")
 
     @spaces.GPU(duration=90)
     def gpu_game_loop():
@@ -147,29 +221,20 @@ def create_gpu_game_loop(command_queue: Queue, initial_seed_image=None, initial_
         Generator that keeps GPU allocated and processes commands.
         Yields (frame, frame_count) tuples.
         """
-        pipe.to("cuda")
-        pipe.blocks.sub_blocks['before_denoise'].sub_blocks['setup_kv_cache']._setup_kv_cache(pipe.transformer, pipe.device, torch.bfloat16)
-        aoti_load_(
-            pipe.transformer,
-            "diffusers-internal-dev/world-engine-aot",
-            "transformer-fp8.pt2",
-            "transformer-fp8-constants.pt"
-        )
-        aoti_load_(
-            pipe.vae.decoder,
-            "diffusers-internal-dev/world-engine-aot",
-            "decoder.pt2",
-            "decoder-constants.pt"
-        )
+
         n_frames = pipe.transformer.config.n_frames
         print(f"Model loaded! (n_frames={n_frames})")
+        print(f"gpu_game_loop: initial_seed_image={type(initial_seed_image)}, initial_seed_url={initial_seed_url}")
 
         # Initialize state with provided seed or random
         if initial_seed_image is not None:
+            print(f"gpu_game_loop init: Using initial_seed_image {initial_seed_image.size if hasattr(initial_seed_image, 'size') else type(initial_seed_image)}")
             seed_image = initial_seed_image.resize((640, 360), Image.BILINEAR)
         elif initial_seed_url is not None:
+            print(f"gpu_game_loop init: Using initial_seed_url {initial_seed_url}")
             seed_image = load_seed_frame(initial_seed_url)
         else:
+            print("gpu_game_loop init: Using random seed")
             seed_image = load_seed_frame(random.choice(SEED_FRAME_URLS))
 
         state = pipe(
@@ -188,70 +253,89 @@ def create_gpu_game_loop(command_queue: Queue, initial_seed_image=None, initial_
         # Yield initial frame
         yield (frame, frame_count)
 
-        # Main loop - process commands and yield frames
-        while True:
-            try:
-                # Non-blocking get with short timeout for responsiveness
-                command = command_queue.get(timeout=0.005)
-            except:
-                # No command, continue idle
-                command = None
-
-            if command is None:
-                continue
+        # Track current input state (updated by commands)
+        current_buttons = set()
+        current_mouse = (0.0, 0.0)
+        current_prompt = initial_prompt
 
-            if isinstance(command, StopCommand):
+        # Main loop - generate frames continuously, sample latest input
+        while True:
+            # Drain command queue - get all pending commands (non-blocking)
+            stop_requested = False
+            reset_command = None
+            while True:
+                try:
+                    command = command_queue.get_nowait()
+                    if isinstance(command, StopCommand):
+                        stop_requested = True
+                        break
+                    elif isinstance(command, ResetCommand):
+                        reset_command = command
+                    elif isinstance(command, GenerateCommand):
+                        # Update current input state with latest command
+                        current_buttons = command.buttons
+                        current_mouse = command.mouse
+                        current_prompt = command.prompt
+                except:
+                    break  # Queue empty
+
+            if stop_requested:
                 print("Stop command received, ending GPU session")
                 break
 
-            elif isinstance(command, ResetCommand):
-                print("Reset command received")
-                if command.seed_image is not None:
-                    seed_img = command.seed_image.resize((640, 360), Image.BILINEAR)
-                elif command.seed_url:
-                    seed_img = load_seed_frame(command.seed_url)
+            # Handle reset if requested
+            if reset_command is not None:
+                print(f"Reset command received: seed_image={type(reset_command.seed_image)}, seed_url={reset_command.seed_url}")
+                if reset_command.seed_image is not None:
+                    print(f"Using seed_image from command: {reset_command.seed_image.size if hasattr(reset_command.seed_image, 'size') else 'unknown'}")
+                    seed_img = reset_command.seed_image.resize((640, 360), Image.BILINEAR)
+                elif reset_command.seed_url:
+                    print(f"Using seed_url from command: {reset_command.seed_url}")
+                    seed_img = load_seed_frame(reset_command.seed_url)
                 else:
+                    print("Using random seed from command")
                     seed_img = load_seed_frame(random.choice(SEED_FRAME_URLS))
 
                 state = pipe(
-                    prompt=command.prompt,
+                    prompt=reset_command.prompt,
                     image=seed_img,
                     button=set(),
                     mouse=(0.0, 0.0),
                     output_type="pil",
                 )
                 frame_count = 1
+                current_prompt = reset_command.prompt
                 frame = state.values.get("images")
                 yield (frame, frame_count)
+                continue
 
-            elif isinstance(command, GenerateCommand):
-                # Generate next frame
+            # Generate next frame with current input state (ALWAYS generates)
+            state = pipe(
+                state,
+                prompt=current_prompt,
+                button=current_buttons,
+                mouse=current_mouse,
+                image=None,
+                output_type="pil",
+            )
+            frame_count += 1
+            frame = state.values.get("images")
+
+            # Auto-reset near end of context
+            if frame_count >= n_frames - 2:
+                print(f"Auto-reset at frame {frame_count}")
+                seed_img = load_seed_frame(random.choice(SEED_FRAME_URLS))
                 state = pipe(
-                    state,
-                    prompt=command.prompt,
-                    button=command.buttons,
-                    mouse=command.mouse,
-                    image=None,
+                    prompt=current_prompt,
+                    image=seed_img,
+                    button=set(),
+                    mouse=(0.0, 0.0),
                     output_type="pil",
                 )
-                frame_count += 1
+                frame_count = 1
                 frame = state.values.get("images")
 
-                # Auto-reset near end of context
-                if frame_count >= n_frames - 2:
-                    print(f"Auto-reset at frame {frame_count}")
-                    seed_img = load_seed_frame(random.choice(SEED_FRAME_URLS))
-                    state = pipe(
-                        prompt=command.prompt,
-                        image=seed_img,
-                        button=set(),
-                        mouse=(0.0, 0.0),
-                        output_type="pil",
-                    )
-                    frame_count = 1
-                    frame = state.values.get("images")
-
-                yield (frame, frame_count)
+            yield (frame, frame_count)
 
         print("GPU session ended")
 
@@ -681,6 +765,9 @@ def create_app():
         
         # Selected seed URL from examples
         selected_seed_url = gr.State(None)
+        
+        # Store uploaded image in state (workaround for Gradio component value issues)
+        uploaded_image_state = gr.State(None)
 
         gr.Markdown("""
         # 🌍 Waypoint 1 Small
@@ -742,7 +829,9 @@ def create_app():
                     lines=2,
                 )
 
-                frame_display = gr.Number(label="Frame", value=0, interactive=False)
+                with gr.Row():
+                    frame_display = gr.Number(label="Frame", value=0, interactive=False)
+                    fps_display = gr.Number(label="FPS", value=0.0, interactive=False)
 
         # --- Event Handlers ---
         
@@ -754,26 +843,25 @@ def create_app():
 
         def on_gallery_start(state, evt: gr.SelectData, uploaded_image, prompt):
             """Handle gallery selection - start/restart game with selected world."""
+            print(f"on_gallery_start CALLED: evt.index={evt.index}", flush=True)
             if evt.index is None or evt.index >= len(SEED_FRAME_URLS):
                 # No valid selection, do nothing
                 yield (state, None, 0, None, 0, gr.update(), gr.update())
                 return
-            
+
             selected_url = SEED_FRAME_URLS[evt.index]
-            
+
             # If game is running, stop it first
-            if len(state) > 0:
-                gen, command_queue = state
-                command_queue.put(StopCommand())
-                try:
-                    while True:
-                        next(gen)
-                except StopIteration:
-                    pass
-            
+            if state and isinstance(state, GameSession):
+                session = state
+                session.stop_event.set()
+                session.command_queue.put(StopCommand())
+                if session.worker_thread.is_alive():
+                    session.worker_thread.join(timeout=2.0)
+
             # Show info about controls
             gr.Info("Controls locked! Press ESC to release mouse/keyboard capture.", duration=5)
-            
+
             # Show loading state
             loading_img = create_loading_image(text="Generating World ...")
             yield (
@@ -785,21 +873,42 @@ def create_app():
                 gr.update(interactive=False),
                 gr.update(interactive=False),
             )
-            
+
             # Start new game with selected world
             command_queue = Queue()
+            frame_queue = queue.Queue(maxsize=2)
+            stop_event = threading.Event()
+
             gen = create_gpu_game_loop(
                 command_queue,
                 initial_seed_image=None,
                 initial_seed_url=selected_url,
                 initial_prompt=prompt or "An explorable world"
             )
-            
+
             # Get initial frame
             frame, frame_count = next(gen)
-            
+
+            # Start worker thread
+            frame_times = deque(maxlen=30)
+            worker = threading.Thread(
+                target=gpu_worker_thread,
+                args=(gen, command_queue, frame_queue, stop_event, frame_times),
+                daemon=True
+            )
+            worker.start()
+
+            session = GameSession(
+                command_queue=command_queue,
+                frame_queue=frame_queue,
+                worker_thread=worker,
+                stop_event=stop_event,
+                generator=gen,
+                frame_times=frame_times,
+            )
+
             yield (
-                (gen, command_queue),
+                session,
                 frame,
                 frame_count,
                 frame,
@@ -809,10 +918,15 @@ def create_app():
             )
 
         def on_start(selected_url, uploaded_image, prompt):
-            """Start GPU session - generator that shows loading then first frame."""
+            """Start GPU session - creates background worker thread for non-blocking frames."""
+            print(f"on_start CALLED:", flush=True)
+            print(f"  uploaded_image (from state) type: {type(uploaded_image)}", flush=True)
+            print(f"  uploaded_image is PIL: {isinstance(uploaded_image, Image.Image) if uploaded_image else False}", flush=True)
+            print(f"  selected_url: {selected_url}", flush=True)
+            
             # Show info about controls
             gr.Info("Controls locked! Press ESC to release mouse/keyboard capture.", duration=5)
-            
+
             # Show loading state immediately
             loading_img = create_loading_image(text="Generating World ...")
             yield (
@@ -826,17 +940,32 @@ def create_app():
             )
 
             # Determine seed image/url
-            # Priority: uploaded image > selected from gallery > random
+            # Priority: uploaded image (from state) > selected from gallery > random
             seed_image = None
             seed_url = None
 
-            if uploaded_image is not None:
+            # Check if uploaded_image is a valid PIL Image
+            is_pil_image = isinstance(uploaded_image, Image.Image)
+            has_uploaded_image = uploaded_image is not None and is_pil_image
+            
+            print(f"on_start decision:", flush=True)
+            print(f"  is_pil_image: {is_pil_image}", flush=True)
+            print(f"  has_uploaded_image: {has_uploaded_image}", flush=True)
+
+            if has_uploaded_image:
                 seed_image = uploaded_image
+                print(f"on_start: Using uploaded image: {seed_image.size}", flush=True)
             elif selected_url is not None:
                 seed_url = selected_url
+                print(f"on_start: Using selected URL: {seed_url}", flush=True)
+            else:
+                print("on_start: Using random seed", flush=True)
             # else: random will be chosen in create_gpu_game_loop
 
             command_queue = Queue()
+            frame_queue = queue.Queue(maxsize=2)  # Thread-safe output queue
+            stop_event = threading.Event()
+
             gen = create_gpu_game_loop(
                 command_queue,
                 initial_seed_image=seed_image,
@@ -844,11 +973,29 @@ def create_app():
                 initial_prompt=prompt or "An explorable world"
             )
 
-            # Get initial frame
+            # Get initial frame synchronously (needed to show first frame)
             frame, frame_count = next(gen)
 
+            # Start worker thread to consume generator in background
+            frame_times = deque(maxlen=30)
+            worker = threading.Thread(
+                target=gpu_worker_thread,
+                args=(gen, command_queue, frame_queue, stop_event, frame_times),
+                daemon=True
+            )
+            worker.start()
+
+            session = GameSession(
+                command_queue=command_queue,
+                frame_queue=frame_queue,
+                worker_thread=worker,
+                stop_event=stop_event,
+                generator=gen,
+                frame_times=frame_times,
+            )
+
             yield (
-                (gen, command_queue),  # session_state
+                session,                # session_state
                 frame,                  # latest_frame
                 frame_count,            # latest_frame_count
                 frame,                  # video_output
@@ -858,20 +1005,19 @@ def create_app():
             )
 
         def on_stop(state):
-            """Stop GPU session."""
-            if len(state) == 0:
+            """Stop GPU session and cleanup worker thread."""
+            if not state or not isinstance(state, GameSession):
                 return ((), None, 0, None, 0,
                         gr.update(interactive=True), gr.update(interactive=False))
 
-            gen, command_queue = state
-            command_queue.put(StopCommand())
+            session = state
+            # Signal worker to stop
+            session.stop_event.set()
+            session.command_queue.put(StopCommand())
 
-            # Drain generator
-            try:
-                while True:
-                    next(gen)
-            except StopIteration:
-                pass
+            # Wait for worker thread to finish (with timeout)
+            if session.worker_thread.is_alive():
+                session.worker_thread.join(timeout=2.0)
 
             return (
                 (),
@@ -883,32 +1029,42 @@ def create_app():
                 gr.update(interactive=False),
             )
 
-        def on_generate_tick(state, controls, prompt, current_frame, current_count):
-            """Called by timer - send generate command and get next frame."""
-            if len(state) == 0:
-                return current_frame, current_count, current_frame, current_count
+        def on_generate_tick(state, controls, prompt, current_frame, current_count, current_fps):
+            """Called by timer - send generate command and get next frame (non-blocking)."""
+            if not state or not isinstance(state, GameSession):
+                return current_frame, current_count, current_frame, current_count, 0.0
 
-            gen, command_queue = state
+            session = state
 
-            # Send generate command
+            # Send generate command (non-blocking)
             buttons = set(controls.get("buttons", []))
             mouse = (controls.get("mouse_x", 0.0), controls.get("mouse_y", 0.0))
-            command_queue.put(GenerateCommand(buttons=buttons, mouse=mouse, prompt=prompt))
+            session.command_queue.put(GenerateCommand(buttons=buttons, mouse=mouse, prompt=prompt))
 
-            # Get next frame
+            # Non-blocking read from frame_queue - get latest frame if available
             try:
-                frame, frame_count = next(gen)
-                return frame, frame_count, frame, frame_count
-            except StopIteration:
-                return current_frame, current_count, current_frame, current_count
+                frame, frame_count, fps = session.frame_queue.get_nowait()
+                return frame, frame_count, frame, frame_count, fps
+            except queue.Empty:
+                # No new frame yet, show previous frame (never blocks!)
+                return current_frame, current_count, current_frame, current_count, current_fps
 
         def on_reset(state, selected_url, uploaded_image, prompt):
             """Reset world with new seed - starts game if not running."""
+            print(f"on_reset CALLED:", flush=True)
+            print(f"  uploaded_image (from state) type: {type(uploaded_image)}", flush=True)
+            print(f"  uploaded_image is PIL: {isinstance(uploaded_image, Image.Image) if uploaded_image else False}", flush=True)
+            print(f"  selected_url: {selected_url}", flush=True)
+            
+            # Check if uploaded_image is a valid PIL Image
+            is_pil_image = isinstance(uploaded_image, Image.Image)
+            has_uploaded_image = uploaded_image is not None and is_pil_image
+            
             # If game is not running, start it
-            if len(state) == 0:
+            if not state or not isinstance(state, GameSession):
                 # Show info about controls
                 gr.Info("Controls locked! Press ESC to release mouse/keyboard capture.", duration=5)
-                
+
                 # Show loading state
                 loading_img = create_loading_image(text="Generating World ...")
                 yield (
@@ -920,29 +1076,54 @@ def create_app():
                     gr.update(interactive=False),
                     gr.update(interactive=False),
                 )
-                
+
                 # Priority: uploaded image > selected from gallery > random
                 seed_image = None
                 seed_url = None
-                
-                if uploaded_image is not None:
+
+                if has_uploaded_image:
                     seed_image = uploaded_image
+                    print(f"on_reset (start): Using uploaded image: {seed_image.size}", flush=True)
                 elif selected_url is not None:
                     seed_url = selected_url
-                
+                    print(f"on_reset (start): Using selected URL: {seed_url}", flush=True)
+                else:
+                    print("on_reset (start): Using random seed", flush=True)
+
                 command_queue = Queue()
+                frame_queue = queue.Queue(maxsize=2)
+                stop_event = threading.Event()
+
                 gen = create_gpu_game_loop(
                     command_queue,
                     initial_seed_image=seed_image,
                     initial_seed_url=seed_url,
                     initial_prompt=prompt or "An explorable world"
                 )
-                
+
                 # Get initial frame
                 frame, frame_count = next(gen)
-                
+
+                # Start worker thread
+                frame_times = deque(maxlen=30)
+                worker = threading.Thread(
+                    target=gpu_worker_thread,
+                    args=(gen, command_queue, frame_queue, stop_event, frame_times),
+                    daemon=True
+                )
+                worker.start()
+
+                session = GameSession(
+                    command_queue=command_queue,
+                    frame_queue=frame_queue,
+                    worker_thread=worker,
+                    stop_event=stop_event,
+                    generator=gen,
+                    frame_times=frame_times,
+                )
+
                 yield (
-                    (gen, command_queue),
+                    session,
                     frame,
                     frame_count,
                     frame,
@@ -953,40 +1134,33 @@ def create_app():
                 return
 
             # Game is running - reset it
-            gen, command_queue = state
+            session = state
 
             # Priority: uploaded image > selected from gallery > random
             seed_image = None
             seed_url = None
 
-            if uploaded_image is not None:
+            if has_uploaded_image:
                 seed_image = uploaded_image
+                print(f"on_reset (running): Using uploaded image: {seed_image.size}", flush=True)
             elif selected_url is not None:
                 seed_url = selected_url
+                print(f"on_reset (running): Using selected URL: {seed_url}", flush=True)
+            else:
+                print("on_reset (running): Using random seed", flush=True)
 
-            command_queue.put(ResetCommand(seed_image=seed_image, seed_url=seed_url, prompt=prompt))
+            session.command_queue.put(ResetCommand(seed_image=seed_image, seed_url=seed_url, prompt=prompt))
 
-            try:
-                frame, frame_count = next(gen)
-                yield (
-                    state,
-                    frame,
-                    frame_count,
-                    frame,
-                    frame_count,
-                    gr.update(),
-                    gr.update(),
-                )
-            except StopIteration:
-                yield (
-                    state,
-                    None,
-                    0,
-                    None,
-                    0,
-                    gr.update(),
-                    gr.update(),
-                )
+            # Just return current state - next timer tick will pick up the reset frame
+            yield (
+                state,
+                None,
+                0,
+                None,
+                0,
+                gr.update(),
+                gr.update(),
+            )
 
         def on_controller_change(value):
             """Update current controls state."""
@@ -1014,7 +1188,7 @@ def create_app():
         
         start_btn.click(
             fn=on_start,
-            inputs=[selected_seed_url, seed_image_upload, prompt_input],
+            inputs=[selected_seed_url, uploaded_image_state, prompt_input],
             outputs=[session_state, latest_frame, latest_frame_count,
                     video_output, frame_display, start_btn, stop_btn],
             js="() => { setTimeout(() => { if (window.worldEngineRequestPointerLock) window.worldEngineRequestPointerLock(); }, 500); }",
@@ -1029,7 +1203,7 @@ def create_app():
 
         reset_btn.click(
             fn=on_reset,
-            inputs=[session_state, selected_seed_url, seed_image_upload, current_prompt],
+            inputs=[session_state, selected_seed_url, uploaded_image_state, current_prompt],
             outputs=[session_state, latest_frame, latest_frame_count,
                     video_output, frame_display, start_btn, stop_btn],
             js="() => { setTimeout(() => { if (window.worldEngineRequestPointerLock) window.worldEngineRequestPointerLock(); }, 500); }",
@@ -1037,13 +1211,30 @@ def create_app():
 
         control_input.change(fn=on_controller_change, inputs=[control_input], outputs=[current_controls])
         prompt_input.change(fn=on_prompt_change, inputs=[prompt_input], outputs=[current_prompt])
+        
+        # Store uploaded image in state and clear gallery selection
+        def on_image_upload(image):
+            """When user uploads an image, store it and clear the gallery selection."""
+            print(f"on_image_upload: image type={type(image)}, is PIL={isinstance(image, Image.Image) if image else False}", flush=True)
+            if image is not None and isinstance(image, Image.Image):
+                print(f"on_image_upload: Storing uploaded image {image.size}", flush=True)
+                return image, None  # Store image, clear selected_seed_url
+            else:
+                print(f"on_image_upload: Clearing stored image", flush=True)
+                return None, gr.update()  # Clear stored image, no change to URL
+        
+        seed_image_upload.change(
+            fn=on_image_upload,
+            inputs=[seed_image_upload],
+            outputs=[uploaded_image_state, selected_seed_url],
+        )
 
         # Timer for continuous generation
-        timer = gr.Timer(value=1/30)  # 15 FPS target
+        timer = gr.Timer(value=1/30)
         timer.tick(
             fn=on_generate_tick,
-            inputs=[session_state, current_controls, current_prompt, latest_frame, latest_frame_count],
-            outputs=[latest_frame, latest_frame_count, video_output, frame_display],
+            inputs=[session_state, current_controls, current_prompt, latest_frame, latest_frame_count, fps_display],
+            outputs=[latest_frame, latest_frame_count, video_output, frame_display, fps_display],
         )
 
         # Pointer lock JS - also allows clicking the game window