neural-os

Runtime error

App Files Files Community

da03 commited on May 10

Commit

9766588

1 Parent(s): 274b099

.

Browse files

Files changed (2) hide show

main.py +26 -1
static/index.html +53 -2

main.py CHANGED Viewed

@@ -166,7 +166,7 @@ def _process_frame_sync(model, inputs):
     # UNet sampling
     start = time.perf_counter()
     use_rnn = False
-    print (f"use_rnn: {use_rnn}")
     if use_rnn:
         sample_latent = output_from_rnn[:, :16]
     else:
@@ -259,6 +259,25 @@ async def websocket_endpoint(websocket: WebSocket):
             # Send confirmation to client
             await websocket.send_json({"type": "reset_confirmed"})
         async def process_input(data):
             nonlocal previous_frame, hidden_states, keys_down, frame_num, frame_count, is_processing
@@ -407,6 +426,12 @@ async def websocket_endpoint(websocket: WebSocket):
                     await reset_simulation()
                     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()}")

     # UNet sampling
     start = time.perf_counter()
     use_rnn = False
+    print (f"use_rnn: {use_rnn}, NUM_SAMPLING_STEPS: {NUM_SAMPLING_STEPS}")
     if use_rnn:
         sample_latent = output_from_rnn[:, :16]
     else:
             # Send confirmation to client
             await websocket.send_json({"type": "reset_confirmed"})
+        # Add a function to update sampling steps
+        async def update_sampling_steps(steps):
+            global NUM_SAMPLING_STEPS
+            # Validate the input
+            if steps < 1:
+                print(f"[{time.perf_counter():.3f}] Invalid sampling steps value: {steps}")
+                await websocket.send_json({"type": "error", "message": "Invalid sampling steps value"})
+                return
+            # Update the global variable
+            old_steps = NUM_SAMPLING_STEPS
+            NUM_SAMPLING_STEPS = steps
+            print(f"[{time.perf_counter():.3f}] Updated NUM_SAMPLING_STEPS from {old_steps} to {steps}")
+            # Send confirmation to client
+            await websocket.send_json({"type": "steps_updated", "steps": steps})
         async def process_input(data):
             nonlocal previous_frame, hidden_states, keys_down, frame_num, frame_count, is_processing
                     await reset_simulation()
                     continue
+                # Handle sampling steps update
+                if data.get("type") == "update_sampling_steps":
+                    print(f"[{receive_time:.3f}] Received request to update sampling steps")
+                    await update_sampling_steps(data.get("steps", 32))
+                    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 CHANGED Viewed

@@ -9,13 +9,43 @@
         #displayCanvas {
             cursor: none;
         }
     </style>
 </head>
 <body>
     <canvas id="displayCanvas" width="512" height="384"></canvas>
-    <!-- Add this button somewhere in your UI, perhaps near other controls -->
-    <button id="resetButton" class="control-button">Reset Simulation</button>
     <script>
         const canvas = document.getElementById('displayCanvas');
@@ -230,6 +260,27 @@
                 console.error("WebSocket not connected, cannot reset");
             }
         });
     </script>
 </body>
 </html>

         #displayCanvas {
             cursor: none;
         }
+        /* Add some styling for the controls */
+        .controls {
+            margin-top: 10px;
+            display: flex;
+            gap: 10px;
+            align-items: center;
+        }
+        .control-button {
+            padding: 5px 10px;
+            cursor: pointer;
+        }
+        .step-control {
+            display: flex;
+            align-items: center;
+            gap: 5px;
+        }
+        #samplingSteps {
+            width: 60px;
+        }
     </style>
 </head>
 <body>
     <canvas id="displayCanvas" width="512" height="384"></canvas>
+    <div class="controls">
+        <button id="resetButton" class="control-button">Reset Simulation</button>
+        <div class="step-control">
+            <label for="samplingSteps">Sampling Steps:</label>
+            <input type="number" id="samplingSteps" min="1" max="100" value="32">
+            <button id="updateStepsButton" class="control-button">Update</button>
+        </div>
+    </div>
     <script>
         const canvas = document.getElementById('displayCanvas');
                 console.error("WebSocket not connected, cannot reset");
             }
         });
+        // Add event listener for updating sampling steps
+        document.getElementById('updateStepsButton').addEventListener('click', function() {
+            const stepsInput = document.getElementById('samplingSteps');
+            const newSteps = parseInt(stepsInput.value, 10);
+            if (isNaN(newSteps) || newSteps < 1) {
+                alert("Please enter a valid number of steps (minimum 1)");
+                return;
+            }
+            if (socket && socket.readyState === WebSocket.OPEN) {
+                console.log(`Sending update to set sampling steps to ${newSteps}`);
+                socket.send(JSON.stringify({
+                    type: "update_sampling_steps",
+                    steps: newSteps
+                }));
+            } else {
+                console.error("WebSocket not connected, cannot update steps");
+            }
+        });
     </script>
 </body>
 </html>