.

main.py

@@ -23,6 +23,7 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 
 DEBUG_MODE = False
+DEBUG_MODE_2 = True
 
 SCREEN_WIDTH = 512
 SCREEN_HEIGHT = 384
@@ -101,6 +102,10 @@ def prepare_model_inputs(
     if DEBUG_MODE:
         print ('DEBUG MODE, SETTING TIME STEP TO 0')
         time_step = 0
+    if DEBUG_MODE_2:
+        if time_step > 1:
+            print ('DEBUG MODE_2, SETTING TIME STEP TO 0')
+            time_step = 0
     
     inputs = {
         'image_features': previous_frame.to(device),
@@ -121,6 +126,11 @@ def prepare_model_inputs(
         print ('DEBUG MODE, REMOVING INPUTS')
         if 'hidden_states' in inputs:
             del inputs['hidden_states']
+    if DEBUG_MODE_2:
+        if time_step > 1:
+            print ('DEBUG MODE_2, REMOVING HIDDEN STATES')
+            if 'hidden_states' in inputs:
+                del inputs['hidden_states']
     return inputs
 
 @torch.no_grad()
@@ -244,13 +254,17 @@ async def websocket_endpoint(websocket: WebSocket):
                 for key in keys_up_list:
                     if key in keys_down:  # Check if key exists to avoid KeyError
                         keys_down.remove(key)
-                
-                inputs = prepare_model_inputs(previous_frame, hidden_states, x, y, is_right_click, is_left_click, list(keys_down), stoi, itos, frame_num)
-                print(f"[{time.perf_counter():.3f}] Starting model inference...")
-                previous_frame, sample_img, hidden_states, timing_info = await process_frame(model, inputs)
                 if DEBUG_MODE:
                     print (f"DEBUG MODE, REMOVING HIDDEN STATES")
                     previous_frame = padding_image
+                if DEBUG_MODE_2:
+                    if frame_num > 1:
+                        print (f"DEBUG MODE_2, REMOVING HIDDEN STATES")
+                        previous_frame = padding_image
+                inputs = prepare_model_inputs(previous_frame, hidden_states, x, y, is_right_click, is_left_click, list(keys_down), stoi, itos, frame_num)
+                print(f"[{time.perf_counter():.3f}] Starting model inference...")
+                previous_frame, sample_img, hidden_states, timing_info = await process_frame(model, inputs)
+                
                 timing_info['full_frame'] = time.perf_counter() - process_start_time
                 
                 print(f"[{time.perf_counter():.3f}] Model inference complete. Queue size now: {input_queue.qsize()}")