.

main.py

@@ -29,7 +29,6 @@ NUM_MAX_FRAMES = 1
 SCREEN_WIDTH = 512
 SCREEN_HEIGHT = 384
 NUM_SAMPLING_STEPS = 32
-USE_RNN = True
 USE_RNN = False
 
 MODEL_NAME = "yuntian-deng/computer-model-ss005-cont-lr2e5-384k"
@@ -313,6 +312,19 @@ async def websocket_endpoint(websocket: WebSocket):
             # Send confirmation to client
             await websocket.send_json({"type": "steps_updated", "steps": steps})
         
+        # Add a function to update USE_RNN setting
+        async def update_use_rnn(use_rnn):
+            global USE_RNN
+            
+            # Update the global variable
+            old_setting = USE_RNN
+            USE_RNN = use_rnn
+            
+            print(f"[{time.perf_counter():.3f}] Updated USE_RNN from {old_setting} to {use_rnn}")
+            
+            # Send confirmation to client
+            await websocket.send_json({"type": "rnn_updated", "use_rnn": use_rnn})
+        
         async def process_input(data):
             nonlocal previous_frame, hidden_states, keys_down, frame_num, frame_count, is_processing
             
@@ -475,6 +487,12 @@ async def websocket_endpoint(websocket: WebSocket):
                     await update_sampling_steps(data.get("steps", 32))
                     continue
                 
+                # Handle USE_RNN update
+                if data.get("type") == "update_use_rnn":
+                    print(f"[{receive_time:.3f}] Received request to update USE_RNN")
+                    await update_use_rnn(data.get("use_rnn", False))
+                    continue
+                
                 # Add the input to our queue
                 await input_queue.put(data)
                 print(f"[{receive_time:.3f}] Received input. Queue size now: {input_queue.qsize()}")

static/index.html

@@ -104,6 +104,11 @@
                     <input type="number" id="samplingSteps" class="form-control" min="1" max="100" value="32">
                     <button id="updateStepsButton" class="btn btn-secondary control-button">Update</button>
                 </div>
+                
+                <div class="form-check form-switch">
+                    <input class="form-check-input" type="checkbox" role="switch" id="useRnnToggle">
+                    <label class="form-check-label" for="useRnnToggle">Use RNN</label>
+                </div>
             </div>
         </div>
         
@@ -118,6 +123,7 @@
                             <li>Right-click to perform context menu actions</li>
                             <li>Use your keyboard to type within the simulated environment</li>
                             <li>Adjust sampling steps to control quality/speed tradeoff</li>
+                            <li>Toggle "Use RNN" to switch between RNN and diffusion mode</li>
                         </ul>
                     </div>
                 </div>
@@ -188,6 +194,10 @@
                     console.log("Simulation reset confirmed by server");
                     // Clear the canvas and reset to initial state
                     drawInitialCanvas();
+                } else if (data.type === "rnn_updated") {
+                    console.log(`USE_RNN setting updated to: ${data.use_rnn}`);
+                    // Update the toggle to match the server state
+                    document.getElementById('useRnnToggle').checked = data.use_rnn;
                 }
             };
         }
@@ -375,6 +385,21 @@
                 console.error("WebSocket not connected, cannot update steps");
             }
         });
+        
+        // Add event listener for the USE_RNN toggle
+        document.getElementById('useRnnToggle').addEventListener('change', function() {
+            const useRnn = this.checked;
+            
+            if (socket && socket.readyState === WebSocket.OPEN) {
+                console.log(`Sending update to set USE_RNN to ${useRnn}`);
+                socket.send(JSON.stringify({
+                    type: "update_use_rnn",
+                    use_rnn: useRnn
+                }));
+            } else {
+                console.error("WebSocket not connected, cannot update USE_RNN");
+            }
+        });
     </script>
     
     <!-- Bootstrap JS (optional) -->